feat(proxy): add connection timeout with default 120s and proper error reporting

- Add configurable connection_timeout (default 120s, 0 to disable) to DatabaseConfig
- Fix ConnectionTimeout error message to display milliseconds (was showing seconds)
- Map ConnectionTimeout to PostgreSQL error code 57P05 instead of generic 58000
- Send ErrorResponse to client on timeout via channel writer
- Add pull_policy: never to docker-compose proxy services to ensure local builds
- Add connect_timeout to integration test client config
- Add unit tests for config defaults, error message format, and error code mapping
- Add integration tests for connection isolation under load
- Document local build requirement in DEVELOPMENT.md

Author: tobyhede

DEVELOPMENT.md

@@ -321,6 +321,9 @@ mise run proxy:down
 Running Proxy in a container cross-compiles a binary for Linux and the current architecture (`amd64`, `arm64`), then copies the binary into the container.
 We cross-compile binary outside the container because it's generally faster, due to packages already being cached, and slower network and disk IO in Docker.
 
+> [!IMPORTANT]
+> **Proxy must always be built from source for testing.** The `proxy:up` task builds a binary from source (`build:binary`), packages it into a Docker image tagged `cipherstash/proxy:latest` (`build:docker`), then starts it via `docker compose up`. The `tests/docker-compose.yml` file uses `pull_policy: never` on the proxy services to ensure Docker never pulls the released image from Docker Hub. If you see an error like `pull access denied` or `image not found`, run `mise run build` first to build the local image.
+
 ### Building
 
 Build a binary and Docker image:
@@ -460,6 +463,8 @@ This project uses `docker compose` to manage containers and networking.
 
 The configuration for those containers is in `tests/docker-compose.yml`.
 
+The proxy services in `tests/docker-compose.yml` use `pull_policy: never` to ensure Docker never pulls the released `cipherstash/proxy:latest` image from Docker Hub. The image must be built locally from source via `mise run proxy:up` (or `mise run build`). This guarantees integration tests always run against the current source code.
+
 The integration tests use the `proxy:up` and `proxy:down` commands documented above to run containers in different configurations.
 
 #### Configuration: configuring PostgreSQL containers in integration tests

mise.toml

@@ -663,6 +663,9 @@ cp -v {{config_root}}/target/{{ target }}/release/cipherstash-proxy {{config_roo
 
 [tasks."build:docker"]
 depends = ["eql:download"]
+# Tags the image as cipherstash/proxy:latest locally.
+# tests/docker-compose.yml uses pull_policy: never to ensure this local image
+# is always used instead of the released image on Docker Hub.
 description = "Build a Docker image for cipherstash-proxy"
 run = """
 {% set default_platform = "linux/" ~ arch() | replace(from="x86_64", to="amd64") %}

packages/cipherstash-proxy-integration/src/common.rs

@@ -128,7 +128,8 @@ pub fn connection_config(port: u16) -> tokio_postgres::Config {
         .port(port)
         .user(&username)
         .password(&password)
-        .dbname(&name);
+        .dbname(&name)
+        .connect_timeout(std::time::Duration::from_secs(10));
 
     db_config
 }

packages/cipherstash-proxy-integration/src/connection_resilience.rs

@@ -0,0 +1,181 @@
+/// Tests that validate proxy connection isolation under load.
+///
+/// These tests verify that:
+/// - Slow queries on one connection don't block other connections
+/// - The proxy accepts new connections after client disconnect
+/// - Concurrent connections under load remain responsive
+/// - Blocked backend connections don't affect other proxy connections
+#[cfg(test)]
+mod tests {
+    use crate::common::{connect_with_tls, PROXY, PG_PORT};
+    use std::sync::Arc;
+    use std::time::Instant;
+    use tokio::sync::Notify;
+    use tokio::task::JoinSet;
+    use tokio::time::{timeout, Duration};
+
+    /// A slow query on one connection does not block other connections through the proxy.
+    #[tokio::test]
+    async fn slow_query_does_not_block_other_connections() {
+        let result = timeout(Duration::from_secs(30), async {
+            let client_a = connect_with_tls(PROXY).await;
+            let client_b = connect_with_tls(PROXY).await;
+
+            // Connection A: run a slow query
+            let a_handle = tokio::spawn(async move {
+                client_a
+                    .simple_query("SELECT pg_sleep(5)")
+                    .await
+                    .unwrap();
+            });
+
+            // Brief pause to ensure A's query is in flight
+            tokio::time::sleep(Duration::from_millis(200)).await;
+
+            // Connection B: run a fast query, should complete promptly
+            let start = Instant::now();
+            let rows = client_b.simple_query("SELECT 1").await.unwrap();
+            let elapsed = start.elapsed();
+
+            assert!(!rows.is_empty(), "Expected result from SELECT 1");
+            assert!(
+                elapsed < Duration::from_secs(2),
+                "Fast query took {elapsed:?}, expected < 2s — proxy may be blocking"
+            );
+
+            a_handle.await.unwrap();
+        })
+        .await;
+
+        result.expect("Test timed out after 30s");
+    }
+
+    /// Proxy accepts new connections after a client disconnects.
+    #[tokio::test]
+    async fn proxy_accepts_new_connections_after_client_disconnect() {
+        let result = timeout(Duration::from_secs(10), async {
+            // First connection: query, then drop
+            {
+                let client = connect_with_tls(PROXY).await;
+                let rows = client.simple_query("SELECT 1").await.unwrap();
+                assert!(!rows.is_empty());
+            }
+            // Client dropped here
+
+            // Brief pause
+            tokio::time::sleep(Duration::from_millis(100)).await;
+
+            // Second connection: should work fine
+            let client = connect_with_tls(PROXY).await;
+            let rows = client.simple_query("SELECT 1").await.unwrap();
+            assert!(!rows.is_empty());
+        })
+        .await;
+
+        result.expect("Test timed out after 10s");
+    }
+
+    /// Concurrent slow and fast connections: fast queries complete promptly under slow load.
+    #[tokio::test]
+    async fn concurrent_connections_under_slow_load() {
+        let result = timeout(Duration::from_secs(30), async {
+            let mut join_set = JoinSet::new();
+
+            // 5 slow connections
+            for _ in 0..5 {
+                join_set.spawn(async {
+                    let client = connect_with_tls(PROXY).await;
+                    client
+                        .simple_query("SELECT pg_sleep(3)")
+                        .await
+                        .unwrap();
+                });
+            }
+
+            // Brief pause to let slow queries start
+            tokio::time::sleep(Duration::from_millis(300)).await;
+
+            // 5 fast connections, each should complete promptly
+            for _ in 0..5 {
+                join_set.spawn(async {
+                    let start = Instant::now();
+                    let client = connect_with_tls(PROXY).await;
+                    let rows = client.simple_query("SELECT 1").await.unwrap();
+                    let elapsed = start.elapsed();
+
+                    assert!(!rows.is_empty());
+                    assert!(
+                        elapsed < Duration::from_secs(5),
+                        "Fast query took {elapsed:?} under slow load, expected < 5s"
+                    );
+                });
+            }
+
+            while let Some(result) = join_set.join_next().await {
+                result.unwrap();
+            }
+        })
+        .await;
+
+        result.expect("Test timed out after 30s");
+    }
+
+    /// An advisory-lock-blocked connection through the proxy does not block other proxy connections.
+    #[tokio::test]
+    async fn advisory_lock_blocked_connection_does_not_block_proxy() {
+        let result = timeout(Duration::from_secs(30), async {
+            // Connection A: hold an advisory lock (connect directly to PG to avoid proxy interference)
+            let client_a = connect_with_tls(PG_PORT).await;
+            client_a
+                .simple_query("SELECT pg_advisory_lock(12345)")
+                .await
+                .unwrap();
+
+            let a_ready = Arc::new(Notify::new());
+            let a_ready_tx = a_ready.clone();
+
+            // Connection B: through proxy, attempt to acquire the same lock (will block)
+            let b_handle = tokio::spawn(async move {
+                let client_b = connect_with_tls(PROXY).await;
+                a_ready_tx.notify_one();
+                // This will block until A releases the lock
+                client_b
+                    .simple_query("SELECT pg_advisory_lock(12345)")
+                    .await
+                    .unwrap();
+                // Release after acquiring
+                client_b
+                    .simple_query("SELECT pg_advisory_unlock(12345)")
+                    .await
+                    .unwrap();
+            });
+
+            // Wait for B to be connected and attempting the lock
+            a_ready.notified().await;
+            tokio::time::sleep(Duration::from_millis(500)).await;
+
+            // Connection C: through proxy, should complete immediately despite B being blocked
+            let start = Instant::now();
+            let client_c = connect_with_tls(PROXY).await;
+            let rows = client_c.simple_query("SELECT 1").await.unwrap();
+            let elapsed = start.elapsed();
+
+            assert!(!rows.is_empty());
+            assert!(
+                elapsed < Duration::from_secs(2),
+                "Connection C took {elapsed:?}, expected < 2s — blocked connection may be affecting proxy"
+            );
+
+            // Release the lock so B can complete
+            client_a
+                .simple_query("SELECT pg_advisory_unlock(12345)")
+                .await
+                .unwrap();
+
+            b_handle.await.unwrap();
+        })
+        .await;
+
+        result.expect("Test timed out after 30s");
+    }
+}

packages/cipherstash-proxy-integration/src/lib.rs

@@ -1,4 +1,5 @@
 mod common;
+mod connection_resilience;
 mod decrypt;
 mod diagnostics;
 mod disable_mapping;

packages/cipherstash-proxy/src/config/database.rs

@@ -75,8 +75,14 @@ impl DatabaseConfig {
         self.password.to_owned().risky_unwrap()
     }
 
+    const DEFAULT_CONNECTION_TIMEOUT_MS: u64 = 120_000;
+
     pub fn connection_timeout(&self) -> Option<Duration> {
-        self.connection_timeout.map(Duration::from_millis)
+        match self.connection_timeout {
+            Some(0) => None,
+            Some(ms) => Some(Duration::from_millis(ms)),
+            None => Some(Duration::from_millis(Self::DEFAULT_CONNECTION_TIMEOUT_MS)),
+        }
     }
 
     pub fn server_name(&self) -> Result<ServerName<'_>, Error> {
@@ -104,6 +110,37 @@ impl DatabaseConfig {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn connection_timeout_defaults_to_120_seconds() {
+        let config = DatabaseConfig::for_testing();
+        assert_eq!(
+            config.connection_timeout(),
+            Some(Duration::from_secs(120))
+        );
+    }
+
+    #[test]
+    fn connection_timeout_zero_disables_timeout() {
+        let mut config = DatabaseConfig::for_testing();
+        config.connection_timeout = Some(0);
+        assert_eq!(config.connection_timeout(), None);
+    }
+
+    #[test]
+    fn connection_timeout_custom_value_in_millis() {
+        let mut config = DatabaseConfig::for_testing();
+        config.connection_timeout = Some(5000);
+        assert_eq!(
+            config.connection_timeout(),
+            Some(Duration::from_millis(5000))
+        );
+    }
+}
+
 ///
 /// Password is NEVER EVER displayed
 ///

packages/cipherstash-proxy/src/error.rs

@@ -27,7 +27,7 @@ pub enum Error {
     #[error("Connection closed by client")]
     ConnectionClosed,
 
-    #[error("Connection timed out after {} ms", duration.as_secs())]
+    #[error("Connection timed out after {} ms", duration.as_millis())]
     ConnectionTimeout { duration: Duration },
 
     #[error("Error creating connection")]
@@ -522,4 +522,12 @@ mod tests {
 
         assert_eq!(format!("Statement encountered an internal error. This may be a bug in the statement mapping module of CipherStash Proxy. Please visit {ERROR_DOC_BASE_URL}#mapping-internal-error for more information."), message);
     }
+
+    #[test]
+    fn connection_timeout_message_shows_millis() {
+        let error = Error::ConnectionTimeout {
+            duration: Duration::from_millis(5000),
+        };
+        assert_eq!(error.to_string(), "Connection timed out after 5000 ms");
+    }
 }

packages/cipherstash-proxy/src/postgresql/error_handler.rs

@@ -53,6 +53,7 @@ pub trait PostgreSqlErrorHandler {
             Error::Encrypt(EncryptError::UnknownKeysetIdentifier { .. }) => {
                 ErrorResponse::system_error(err.to_string())
             }
+            Error::ConnectionTimeout { .. } => ErrorResponse::connection_timeout(),
             _ => ErrorResponse::system_error(err.to_string()),
         }
     }
@@ -67,3 +68,55 @@ pub trait PostgreSqlErrorHandler {
     /// * `error_response` - The ErrorResponse to send to the client
     fn send_error_response(&mut self, err: Error) -> Result<(), Error>;
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::postgresql::messages::error_response::{
+        ErrorResponseCode, CODE_IDLE_SESSION_TIMEOUT, CODE_SYSTEM_ERROR,
+    };
+    use std::time::Duration;
+
+    /// Minimal implementation of PostgreSqlErrorHandler for testing the default method.
+    struct TestHandler;
+
+    impl PostgreSqlErrorHandler for TestHandler {
+        fn client_sender(&mut self) -> &mut Sender {
+            unimplemented!("not needed for error_to_response tests")
+        }
+
+        fn client_id(&self) -> i32 {
+            0
+        }
+
+        fn send_error_response(&mut self, _err: Error) -> Result<(), Error> {
+            unimplemented!("not needed for error_to_response tests")
+        }
+    }
+
+    fn error_code(response: &ErrorResponse) -> Option<&str> {
+        response
+            .fields
+            .iter()
+            .find(|f| f.code == ErrorResponseCode::Code)
+            .map(|f| f.value.as_str())
+    }
+
+    #[test]
+    fn connection_timeout_maps_to_57p05() {
+        let handler = TestHandler;
+        let err = Error::ConnectionTimeout {
+            duration: Duration::from_millis(5000),
+        };
+        let response = handler.error_to_response(err);
+        assert_eq!(error_code(&response), Some(CODE_IDLE_SESSION_TIMEOUT));
+    }
+
+    #[test]
+    fn unknown_error_maps_to_system_error() {
+        let handler = TestHandler;
+        let err = Error::Unknown;
+        let response = handler.error_to_response(err);
+        assert_eq!(error_code(&response), Some(CODE_SYSTEM_ERROR));
+    }
+}

packages/cipherstash-proxy/src/postgresql/handler.rs

@@ -232,6 +232,7 @@ pub async fn handler(client_stream: AsyncStream, context: Context<ZeroKms>) -> R
         }
     }
 
+    let timeout_sender = channel_writer.sender();
     tokio::spawn(channel_writer.receive());
 
     let client_to_server = async {
@@ -258,7 +259,18 @@ pub async fn handler(client_stream: AsyncStream, context: Context<ZeroKms>) -> R
         Ok::<(), Error>(())
     };
 
-    tokio::try_join!(client_to_server, server_to_client)?;
+    let result = tokio::try_join!(client_to_server, server_to_client);
+
+    if let Err(Error::ConnectionTimeout { .. }) = &result {
+        let error_response = ErrorResponse::connection_timeout();
+        if let Ok(bytes) = BytesMut::try_from(error_response) {
+            let _ = timeout_sender.send(bytes);
+        }
+        // Brief yield to allow ChannelWriter to flush
+        tokio::task::yield_now().await;
+    }
+
+    result?;
 
     Ok(())
 }