Merge pull request #441 from cipherstash/james/cip-3678-protocol-desync-transformation-error-after-a-mapped

fix(proxy): protocol desync on transformation error after a mapped statement (CIP-3678)

Author: freshtonic

CHANGELOG.md

@@ -24,6 +24,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Fixed
 
+- **Statement errors no longer desync the connection**: when a statement failed inside the proxy (an unsupported operation on an encrypted column, for instance), the error was written straight to the client and could overtake responses still in flight from the server — with connection pools and prepared-statement caching, the client then saw a protocol error (`unexpected message from server` in tokio_postgres) instead of the proxy's message, typically right after an encrypted statement had run on the same connection. The proxy now delivers statement errors through the server, so clients always receive the proxy's actual error message, in order, and the connection remains usable.
+
 - **A param bound as both a stored value and a query operand**: `UPDATE t SET enc = $1 WHERE enc = $1` failed with a domain CHECK violation. The two occurrences need different payloads — the stored one carries the ciphertext, the query one only search terms — but the role was tracked per input param, so marking the param as a query operand stripped the ciphertext from the value being stored. The role is now taken from the rewritten statement, per occurrence.
 
 - **JSON selector params when the client declares its own types**: a client that sends param OIDs in Parse (pgx in `cache_describe` mode, for example) got `function eql_v3.jsonb_path_exists(eql_v3_json_search, jsonb) does not exist`. A JSON field selector is passed to the rewritten function as bare text, but was being declared as `jsonb` like every other encrypted operand. Affects `->`, `->>`, `jsonb_path_exists`, `jsonb_path_query` and `jsonb_path_query_first`.

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

@@ -47,6 +47,17 @@ use tracing::info;
 use tracing_subscriber::{filter::Directive, EnvFilter, FmtSubscriber};
 
 pub const PROXY: u16 = 6432;
+
+/// Proxy port for tests: `CS_PROXY__PORT` if set, otherwise [`PROXY`].
+///
+/// Lets a local run target a proxy on a non-default port without patching the
+/// test source (mirrors [`get_database_port`] for `CS_DATABASE__PORT`).
+pub fn proxy_port() -> u16 {
+    std::env::var("CS_PROXY__PORT")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(PROXY)
+}
 pub const PROXY_METRICS_PORT: u16 = 9930;
 pub const PG_PORT: u16 = 5532;
 pub const PG_TLS_PORT: u16 = 5617;

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

@@ -2,7 +2,21 @@
 mod tests {
     use tracing::{debug, info};
 
-    use crate::common::{clear, connect_with_tls, random_id, reset_schema, trace, PROXY};
+    use crate::common::{
+        clear, connect_with_tls, proxy_port, random_id, reset_schema, trace, PROXY,
+    };
+
+    /// A statement that always fails inside the proxy, at Parse, in every
+    /// configuration: the proxy's SQL parser rejects it before it reaches the
+    /// server (same shape as [`invalid_sql_statement`]).
+    ///
+    /// A transformation failure (e.g. equality on the storage-only
+    /// `eql_v3_boolean`) cannot be used here: type-check errors only surface
+    /// when `CS_DEVELOPMENT__ENABLE_MAPPING_ERRORS` is on, and the CI proxy
+    /// (like production) runs with it off, silently passing such statements
+    /// through. A parse error takes the same failure path in the proxy
+    /// (`handle_statement_error`) regardless of that flag.
+    const FAILS_IN_PROXY: &str = "INSERT INTO encrypted id, encrypted_text VALUES ($1, $2)";
 
     struct Reset;
 
@@ -108,6 +122,78 @@ mod tests {
         }
     }
 
+    /// CIP-3678 regression: a statement that fails inside the proxy on a
+    /// connection that has already run a MAPPED (encrypted) statement must
+    /// surface the proxy's own error as a clean `db error` — not desync the
+    /// extended-protocol stream into a client-side protocol error
+    /// (`unexpected message from server`) — and the connection must remain
+    /// usable afterwards.
+    #[tokio::test]
+    async fn proxy_error_after_mapped_statement() {
+        trace();
+
+        let client = connect_with_tls(proxy_port()).await;
+
+        // Mapped warm-up: an encrypted statement that parses, binds and
+        // executes successfully.
+        client
+            .query(
+                "SELECT id FROM encrypted WHERE encrypted_text = $1",
+                &[&"cip-3678"],
+            )
+            .await
+            .unwrap();
+
+        // A statement that fails inside the proxy must return the proxy's
+        // error, delivered as a database error.
+        let err = client
+            .query(FAILS_IN_PROXY, &[&random_id(), &"cip-3678"])
+            .await
+            .unwrap_err();
+        let db_err = err.as_db_error().unwrap_or_else(|| {
+            panic!("expected a db error carrying the proxy's message, got: {err:?}")
+        });
+        assert!(
+            db_err.message().contains("sql parser error"),
+            "expected the proxy's parse error, got: {db_err:?}"
+        );
+
+        // The connection must remain usable.
+        let rows = client.query("SELECT 1::int4", &[]).await.unwrap();
+        let one: i32 = rows[0].get(0);
+        assert_eq!(one, 1);
+    }
+
+    /// Companion to [`proxy_error_after_mapped_statement`]: the same failing
+    /// statement on a connection that has only run passthrough statements.
+    /// This path already worked; keep it covered.
+    #[tokio::test]
+    async fn proxy_error_after_passthrough_statement() {
+        trace();
+
+        let client = connect_with_tls(proxy_port()).await;
+
+        // Passthrough warm-up.
+        client.query("SELECT 1::int4", &[]).await.unwrap();
+
+        let err = client
+            .query(FAILS_IN_PROXY, &[&random_id(), &"cip-3678"])
+            .await
+            .unwrap_err();
+        let db_err = err.as_db_error().unwrap_or_else(|| {
+            panic!("expected a db error carrying the proxy's message, got: {err:?}")
+        });
+        assert!(
+            db_err.message().contains("sql parser error"),
+            "expected the proxy's parse error, got: {db_err:?}"
+        );
+
+        // The connection must remain usable.
+        let rows = client.query("SELECT 1::int4", &[]).await.unwrap();
+        let one: i32 = rows[0].get(0);
+        assert_eq!(one, 1);
+    }
+
     #[tokio::test]
     async fn invalid_sql_statement() {
         trace();