fix(proxy): reset keyset and mapping settings on DISCARD ALL and RESET (2.2.x, CIP-5021)

CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 - **Memory growth on long-lived connections**: Proxy now releases the prepared-statement, portal and statement-metrics state it keeps for each client connection once PostgreSQL has finished with it. Previously some common client patterns left state behind until the connection closed, so with a connection pool that keeps connections open, memory grew steadily until Proxy was restarted. The affected patterns were statements prepared and described but never executed, prepared statements removed with SQL `DEALLOCATE`, `DEALLOCATE ALL` or `DISCARD ALL` (used by psycopg 3, Rails and connection-pool reset queries, whether sent as a simple query or through the extended protocol), uniquely named portals, pipelined batches that fail part-way, statements rejected by Proxy, and connections that use `COPY ... FROM STDIN`. No configuration change is required.
 
+- **Connection settings after a pool reset**: `DISCARD ALL` and `RESET ALL` now return Proxy's per-connection settings to their startup values once PostgreSQL confirms the command. The keyset chosen with `SET CIPHERSTASH.KEYSET_ID` or `KEYSET_NAME` is cleared, so Proxy uses `default_keyset_id` (or requires a keyset again when there is none), and `SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true` is undone. `RESET CIPHERSTASH.KEYSET_ID`, `RESET CIPHERSTASH.KEYSET_NAME`, `RESET CIPHERSTASH.UNSAFE_DISABLE_MAPPING` and the equivalent `SET CIPHERSTASH.<setting> TO DEFAULT` reset just that setting. Previously a connection pool that reset a connection between clients handed the next client the previous client's keyset, or left mapping disabled so the next client's queries were sent unencrypted; `SET ... TO DEFAULT` left the setting unchanged or, for `KEYSET_NAME`, selected a keyset named "DEFAULT". A reset that PostgreSQL rejects, for example inside a failed transaction, changes nothing. A `SET CIPHERSTASH.<setting>` prepared as a named statement now takes effect each time it runs, so running it again after a reset selects the keyset again, as it does in PostgreSQL. A `SET CIPHERSTASH.KEYSET_ID` or `KEYSET_NAME` sent while mapping is disabled is now kept when no `default_keyset_id` is configured, so the keyset is in place when mapping is turned back on. As a result, a `SET CIPHERSTASH.KEYSET_ID` that is not a valid UUID is now refused while mapping is disabled, as it is when mapping is enabled; previously it was forwarded to PostgreSQL. A prepared `SET CIPHERSTASH.<setting>` now takes effect only when it is executed, not when it is prepared. `RESET` statements, which Proxy previously rejected as invalid SQL, are now forwarded to PostgreSQL when a query contains only `RESET` statements.
+
 ## [2.2.4] - 2026-06-18
 
 ### Fixed

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

@@ -154,4 +154,53 @@ mod tests {
             assert_eq!(expected, actual);
         }
     }
+
+    ///
+    /// Tests that a connection pool's reset turns mapping back on, so that a disabled mapping
+    /// is not carried over to the next client of a pooled connection.
+    /// Test both extended and simple query protocols
+    ///
+    #[tokio::test]
+    async fn pool_reset_turns_mapping_back_on() {
+        trace();
+
+        clear().await;
+
+        let id = random_id();
+        let encrypted_text = "hello".to_string();
+        let expected = vec![encrypted_text.to_owned()];
+
+        let sql = "INSERT INTO encrypted (id, encrypted_text) VALUES ($1, $2)";
+        execute_query(sql, &[&id, &encrypted_text]).await;
+
+        let select_sql = "SELECT encrypted_text FROM encrypted";
+        let client = connect_with_tls(PROXY).await;
+
+        for reset in [
+            "DISCARD ALL",
+            "RESET ALL",
+            "RESET CIPHERSTASH.UNSAFE_DISABLE_MAPPING",
+            "SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING TO DEFAULT",
+        ] {
+            // Simple query
+            client
+                .simple_query("SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true")
+                .await
+                .unwrap();
+            client.simple_query(reset).await.unwrap();
+
+            let actual = query_with_client::<String>(select_sql, &client).await;
+            assert_eq!(expected, actual, "{reset}");
+
+            // Extended query
+            client
+                .execute("SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true", &[])
+                .await
+                .unwrap();
+            client.execute(reset, &[]).await.unwrap();
+
+            let actual = simple_query_with_client::<String>(select_sql, &client).await;
+            assert_eq!(expected, actual, "{reset}");
+        }
+    }
 }

packages/cipherstash-proxy-integration/src/multitenant/set_keyset_id.rs

@@ -369,4 +369,48 @@ mod tests {
         let result = client.query(&sql, &[]).await;
         assert!(result.is_ok());
     }
+
+    ///
+    /// Tests that a connection pool's reset clears the keyset, so that one client's keyset is
+    /// not carried over to the next client of a pooled connection.
+    /// Test both extended and simple query protocols
+    ///
+    #[tokio::test]
+    async fn pool_reset_clears_keyset_id() {
+        trace();
+        clear().await;
+
+        let tenant_keyset_id_1 = std::env::var("CS_TENANT_KEYSET_ID_1")
+            .map(|s| Uuid::parse_str(&s).unwrap())
+            .unwrap();
+
+        let insert_sql = "INSERT INTO encrypted (id, encrypted_int4) VALUES ($1, $2)";
+        let set_sql = format!("SET CIPHERSTASH.KEYSET_ID = '{tenant_keyset_id_1}'");
+
+        let client = connect_with_tls(PROXY).await;
+
+        for reset in [
+            "DISCARD ALL",
+            "RESET ALL",
+            "RESET CIPHERSTASH.KEYSET_ID",
+            "SET CIPHERSTASH.KEYSET_ID TO DEFAULT",
+        ] {
+            // Simple query
+            client.simple_query(&set_sql).await.unwrap();
+            client.query(insert_sql, &[&random_id(), &1]).await.unwrap();
+            client.simple_query(reset).await.unwrap();
+
+            // CS_DEFAULT_KEYSET_ID is not set for these tests, so a keyset is required again
+            let result = client.query(insert_sql, &[&random_id(), &1]).await;
+            assert!(result.is_err(), "{reset}");
+
+            // Extended query
+            client.execute(&set_sql, &[]).await.unwrap();
+            client.query(insert_sql, &[&random_id(), &1]).await.unwrap();
+            client.execute(reset, &[]).await.unwrap();
+
+            let result = client.query(insert_sql, &[&random_id(), &1]).await;
+            assert!(result.is_err(), "{reset}");
+        }
+    }
 }

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

@@ -182,7 +182,7 @@ where
                 client_id = self.context.client_id,
                 msg = "Passthrough enabled"
             );
-            let deallocation = is_deallocation(&bytes);
+            let completed = CompletedCommand::from(&bytes);
             if code == u8::from(BackendCode::ReadyForQuery) {
                 self.context.ready_for_query(transaction_status(&bytes));
             }
@@ -204,9 +204,7 @@ where
                 | BackendCode::ErrorResponse => {
                     self.context.complete_execution();
                     self.context.finish_session();
-                    if deallocation {
-                        self.context.complete_deallocation();
-                    }
+                    self.command_completed(completed);
                 }
                 _ => {}
             }
@@ -243,9 +241,7 @@ where
 
                 self.context.complete_execution();
                 self.context.finish_session();
-                if is_deallocation(&bytes) {
-                    self.context.complete_deallocation();
-                }
+                self.command_completed(CompletedCommand::from(&bytes));
             }
             BackendCode::ErrorResponse => {
                 if let Some(b) = self.error_response_handler(&bytes)? {
@@ -374,6 +370,16 @@ where
         Ok(())
     }
 
+    /// Apply the SQL deallocation or settings reset that a `CommandComplete` confirms
+    fn command_completed(&mut self, completed: CompletedCommand) {
+        if completed.deallocation {
+            self.context.complete_deallocation();
+        }
+        if completed.settings_reset {
+            self.context.complete_settings_reset();
+        }
+    }
+
     ///
     /// Write a message to the client
     /// Flushes all messages in the buffer before writing the message
@@ -761,6 +767,30 @@ fn is_deallocation(bytes: &BytesMut) -> bool {
     )
 }
 
+/// What a `CommandComplete` confirms, read before the message is forwarded
+#[derive(Clone, Copy, Debug)]
+struct CompletedCommand {
+    deallocation: bool,
+    settings_reset: bool,
+}
+
+impl From<&BytesMut> for CompletedCommand {
+    fn from(bytes: &BytesMut) -> Self {
+        CompletedCommand {
+            deallocation: is_deallocation(bytes),
+            settings_reset: is_settings_reset(bytes),
+        }
+    }
+}
+
+/// Whether a `CommandComplete` confirms a `RESET`, `SET` or `DISCARD ALL`
+fn is_settings_reset(bytes: &BytesMut) -> bool {
+    if bytes.first() != Some(&b'C') {
+        return false;
+    }
+    matches!(command_tag(bytes), b"RESET" | b"SET" | b"DISCARD ALL")
+}
+
 /// The transaction status (`I`, `T` or `E`) of a `ReadyForQuery` message
 fn transaction_status(bytes: &BytesMut) -> u8 {
     bytes.get(5).copied().unwrap_or(b'I')
@@ -1014,6 +1044,18 @@ mod tests {
         assert!(!super::is_deallocation(&error_response_bytes()));
     }
 
+    #[test]
+    fn settings_reset_command_tags() {
+        assert!(super::is_settings_reset(&command_complete("RESET")));
+        assert!(super::is_settings_reset(&command_complete("SET")));
+        assert!(super::is_settings_reset(&command_complete("DISCARD ALL")));
+        assert!(!super::is_settings_reset(&command_complete(
+            "DISCARD PLANS"
+        )));
+        assert!(!super::is_settings_reset(&command_complete("SELECT 1")));
+        assert!(!super::is_settings_reset(&error_response_bytes()));
+    }
+
     #[test]
     fn ready_for_query_transaction_status() {
         assert_eq!(