fix(proxy): address code review feedback for connection timeout

- Pass actual timeout duration in ErrorResponse message to client
- Document rationale for using 57P05 over 08006 error code
- Document yield_now() as best-effort flush with known limitation
- Document advisory lock test race condition timing margin
- Replace magic number 12345 with named ADVISORY_LOCK_ID constant

Author: tobyhede

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

@@ -14,6 +14,10 @@ mod tests {
     use tokio::task::JoinSet;
     use tokio::time::{timeout, Duration};
 
+    /// Advisory lock ID used in isolation tests. Arbitrary value — just needs to be
+    /// unique across concurrently running test suites against the same database.
+    const ADVISORY_LOCK_ID: i64 = 99_001;
+
     /// 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() {
@@ -121,31 +125,40 @@ mod tests {
     }
 
     /// An advisory-lock-blocked connection through the proxy does not block other proxy connections.
+    ///
+    /// Note: Connection B notifies readiness before `pg_advisory_lock` reaches PostgreSQL.
+    /// The 500ms sleep provides a generous margin for the lock attempt to reach PG, but is
+    /// not strictly guaranteed. In practice this has not caused flakiness.
     #[tokio::test]
     async fn advisory_lock_blocked_connection_does_not_block_proxy() {
+        let lock_query = format!("SELECT pg_advisory_lock({ADVISORY_LOCK_ID})");
+        let unlock_query = format!("SELECT pg_advisory_unlock({ADVISORY_LOCK_ID})");
+
         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)")
+                .simple_query(&lock_query)
                 .await
                 .unwrap();
 
             let a_ready = Arc::new(Notify::new());
             let a_ready_tx = a_ready.clone();
+            let b_lock_query = lock_query.clone();
+            let b_unlock_query = unlock_query.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)")
+                    .simple_query(&b_lock_query)
                     .await
                     .unwrap();
                 // Release after acquiring
                 client_b
-                    .simple_query("SELECT pg_advisory_unlock(12345)")
+                    .simple_query(&b_unlock_query)
                     .await
                     .unwrap();
             });
@@ -168,7 +181,7 @@ mod tests {
 
             // Release the lock so B can complete
             client_a
-                .simple_query("SELECT pg_advisory_unlock(12345)")
+                .simple_query(&unlock_query)
                 .await
                 .unwrap();
 

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

@@ -53,7 +53,9 @@ pub trait PostgreSqlErrorHandler {
             Error::Encrypt(EncryptError::UnknownKeysetIdentifier { .. }) => {
                 ErrorResponse::system_error(err.to_string())
             }
-            Error::ConnectionTimeout { .. } => ErrorResponse::connection_timeout(),
+            Error::ConnectionTimeout { .. } => {
+                ErrorResponse::connection_timeout(err.to_string())
+            }
             _ => ErrorResponse::system_error(err.to_string()),
         }
     }
@@ -102,6 +104,14 @@ mod tests {
             .map(|f| f.value.as_str())
     }
 
+    fn error_message(response: &ErrorResponse) -> Option<&str> {
+        response
+            .fields
+            .iter()
+            .find(|f| f.code == ErrorResponseCode::Message)
+            .map(|f| f.value.as_str())
+    }
+
     #[test]
     fn connection_timeout_maps_to_57p05() {
         let handler = TestHandler;
@@ -110,6 +120,10 @@ mod tests {
         };
         let response = handler.error_to_response(err);
         assert_eq!(error_code(&response), Some(CODE_IDLE_SESSION_TIMEOUT));
+        assert_eq!(
+            error_message(&response),
+            Some("Connection timed out after 5000 ms")
+        );
     }
 
     #[test]

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

@@ -261,12 +261,15 @@ pub async fn handler(client_stream: AsyncStream, context: Context<ZeroKms>) -> R
 
     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 Err(ref err @ Error::ConnectionTimeout { .. }) = &result {
+        let error_response = ErrorResponse::connection_timeout(err.to_string());
         if let Ok(bytes) = BytesMut::try_from(error_response) {
             let _ = timeout_sender.send(bytes);
         }
-        // Brief yield to allow ChannelWriter to flush
+        // Best-effort yield to allow ChannelWriter to flush the error response
+        // before the connection tears down. Not guaranteed — if the runtime doesn't
+        // schedule the writer task before teardown, the client may see a connection
+        // reset instead of the ErrorResponse.
         tokio::task::yield_now().await;
     }
 

packages/cipherstash-proxy/src/postgresql/messages/error_response.rs

@@ -60,7 +60,14 @@ pub enum ErrorResponseCode {
 }
 
 impl ErrorResponse {
-    pub fn connection_timeout() -> Self {
+    /// Create a FATAL error response for connection timeout.
+    ///
+    /// Uses PostgreSQL error code 57P05 (idle_session_timeout). While this code
+    /// is technically for idle session timeouts, it is the closest match for a
+    /// proxy-enforced connection timeout. The alternative 08006 (connection_failure)
+    /// implies a network-level failure, which is misleading — the proxy is
+    /// deliberately terminating a connection that exceeded its time limit.
+    pub fn connection_timeout(message: String) -> Self {
         Self {
             fields: vec![
                 Field {
@@ -77,7 +84,7 @@ impl ErrorResponse {
                 },
                 Field {
                     code: ErrorResponseCode::Message,
-                    value: "Connection timeout".to_string(),
+                    value: message,
                 },
             ],
         }