fix(proxy): close two plaintext paths found by the Quint routing review

CHANGELOG.md

@@ -47,6 +47,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Fixed
 
+- **Encrypted tables off Proxy's search path**: Proxy now counts encrypted columns in schemas that are not on its database user's search path. When every encrypted table was in such a schema, Proxy saw no encrypted columns, so a write to a table it could not resolve was forwarded as written instead of refused. Tables in those schemas are still not loaded.
+
+- **SQL `EXECUTE` after `UNSAFE_DISABLE_MAPPING`**: Proxy now refuses `EXECUTE` of a statement that writes data and was prepared with SQL `PREPARE` while `CIPHERSTASH.UNSAFE_DISABLE_MAPPING` was on, once mapping is enabled again. It was forwarded with its arguments as written, so only PostgreSQL's type check stopped plaintext reaching an encrypted column. `DEALLOCATE` the statement and prepare it again. This matches the check Proxy already makes for statements prepared with the extended protocol.
+
 - **`--db-user`**: the flag now sets the database username. It previously set an unused variable, so it had no effect.
 
 - **Unknown column error message**: the type-checking error for a column a table does not have now names the column and the table the right way round, for example `Column: email not found for table: users`. It previously swapped them.

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

@@ -306,6 +306,56 @@ mod tests {
         }
     }
 
+    /// A statement prepared with SQL `PREPARE` while mapping is disabled was
+    /// never mapped. Executing it after mapping is enabled again must be
+    /// refused by Proxy, rather than forwarded with its arguments as written.
+    #[tokio::test]
+    async fn sql_execute_of_a_statement_prepared_without_mapping_is_refused() {
+        trace();
+
+        clear().await;
+
+        let client = connect_with_tls(*PROXY).await;
+        let id = random_id();
+
+        client
+            .simple_query("SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true")
+            .await
+            .unwrap();
+        client
+            .simple_query(
+                "PREPARE insert_unmapped AS INSERT INTO encrypted (id, encrypted_text) VALUES ($1, $2)",
+            )
+            .await
+            .unwrap();
+        client
+            .simple_query("SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = false")
+            .await
+            .unwrap();
+
+        // Valid JSON, so only PostgreSQL's domain check would stop it otherwise
+        let err = client
+            .simple_query(&format!(
+                r#"EXECUTE insert_unmapped({id}, '{{"plaintext": "hello"}}')"#
+            ))
+            .await
+            .err()
+            .unwrap();
+        let message = err.as_db_error().unwrap().message().to_owned();
+        assert!(
+            message.starts_with(
+                "The prepared statement or bound values no longer match confirmed connection state"
+            ),
+            "{message}"
+        );
+
+        let rows = client
+            .query("SELECT id FROM encrypted WHERE id = $1", &[&id])
+            .await
+            .unwrap();
+        assert!(rows.is_empty());
+    }
+
     /// Writes a row through Proxy, which must encrypt it with the default keyset. With
     /// mapping disabled, the plaintext parameter would be refused for the encrypted column.
     async fn assert_written_with_default_keyset(client: &tokio_postgres::Client) {

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

@@ -33,6 +33,7 @@ mod pooler;
 mod postgres_rejects_plaintext;
 /// Database-backed transaction-aware schema middleware regressions.
 mod schema_change;
+mod schema_loader;
 mod select;
 mod set_keyset_error;
 mod simple_protocol;

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

@@ -0,0 +1,63 @@
+//! Proxy's schema loader, run directly against PostgreSQL.
+
+#[cfg(test)]
+mod tests {
+    use crate::common::{connect_with_tls, get_database_port, trace};
+    use cipherstash_proxy::{proxy::SchemaManager, DatabaseConfig};
+
+    const ROLE: &str = "schema_loader_narrow_path";
+
+    /// The loader reads only the schemas on its own search path. An encrypted
+    /// column outside them must still count, or Proxy forwards a write to a
+    /// table it cannot find as if no column were encrypted.
+    #[tokio::test]
+    async fn encrypted_columns_off_the_loaders_search_path_are_counted() {
+        trace();
+
+        let postgres = connect_with_tls(get_database_port()).await;
+        postgres
+            .batch_execute(&format!(
+                "
+                DROP SCHEMA IF EXISTS schema_loader_visible CASCADE;
+                DROP SCHEMA IF EXISTS schema_loader_hidden CASCADE;
+                DROP ROLE IF EXISTS {ROLE};
+                CREATE SCHEMA schema_loader_visible;
+                CREATE SCHEMA schema_loader_hidden;
+                CREATE TABLE schema_loader_visible.notes (id bigint, body text);
+                CREATE TABLE schema_loader_hidden.secrets (id bigint, secret eql_v3_text_search);
+                CREATE ROLE {ROLE} LOGIN PASSWORD 'p@ssword';
+                ALTER ROLE {ROLE} SET search_path = schema_loader_visible;
+                GRANT USAGE ON SCHEMA schema_loader_visible, schema_loader_hidden TO {ROLE};
+                GRANT SELECT ON ALL TABLES IN SCHEMA schema_loader_visible TO {ROLE};
+                GRANT SELECT ON ALL TABLES IN SCHEMA schema_loader_hidden TO {ROLE};
+                "
+            ))
+            .await
+            .unwrap();
+
+        let config: DatabaseConfig = serde_json::from_value(serde_json::json!({
+            "host": "localhost",
+            "port": get_database_port(),
+            "name": "cipherstash",
+            "username": ROLE,
+            "password": "p@ssword",
+        }))
+        .unwrap();
+        let snapshot = SchemaManager::init(&config).await.unwrap().load();
+        let has_encrypted_columns =
+            !snapshot.encrypt_config().is_empty() || snapshot.schema().has_encrypted_columns();
+
+        postgres
+            .batch_execute(&format!(
+                "
+                DROP SCHEMA schema_loader_visible CASCADE;
+                DROP SCHEMA schema_loader_hidden CASCADE;
+                DROP ROLE {ROLE};
+                "
+            ))
+            .await
+            .unwrap();
+
+        assert!(has_encrypted_columns);
+    }
+}

packages/cipherstash-proxy/src/postgresql/context/mod.rs

@@ -160,13 +160,21 @@ struct ConnectionProtocolState<K = OperationId> {
     prepared_statements: PreparedStatements<K>,
     portals: HashMap<Name, Arc<Portal>>,
     portal_operations: HashMap<Name, PortalOwner<K>>,
+    /// SQL `PREPARE` statements that write data, forwarded while mapping was
+    /// disabled, with the latest operation that ran each one. They were never
+    /// mapped, so `EXECUTE` of one is refused once mapping is enabled again. A
+    /// name leaves only when PostgreSQL confirms a `DEALLOCATE` or
+    /// `DISCARD ALL` that removes it, run by a later operation.
+    unmapped_sql_statements: HashMap<Name, K>,
 }
 
 /// The Bind operation that created a portal, and the SQL deallocation it runs when executed
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct PortalOwner<K = OperationId> {
     operation: K,
     deallocation: Option<Deallocation>,
+    /// The unmapped SQL prepared statement it prepares when executed
+    preparation: Option<Name>,
 }
 
 impl<K> Default for ConnectionProtocolState<K> {
@@ -180,6 +188,7 @@ impl<K> Default for ConnectionProtocolState<K> {
             prepared_statements: PreparedStatements::default(),
             portals: HashMap::new(),
             portal_operations: HashMap::new(),
+            unmapped_sql_statements: HashMap::new(),
         }
     }
 }
@@ -207,6 +216,7 @@ impl ConnectionProtocolState {
             PortalOwner {
                 operation,
                 deallocation: None,
+                preparation: None,
             },
         );
         self.portals.insert(name, Arc::new(portal));
@@ -267,6 +277,16 @@ impl ConnectionProtocolState {
 }
 
 impl ConnectionProtocolState {
+    fn set_unmapped_sql_statement(&mut self, name: Name, operation: OperationId) {
+        let prepared = self
+            .unmapped_sql_statements
+            .entry(name)
+            .or_insert(operation);
+        if *prepared < operation {
+            *prepared = operation;
+        }
+    }
+
     /// Everything Proxy keeps for the named statement, for a rollback to restore
     fn statement_snapshot(&self, name: &Name) -> settings_protocol::MetadataSnapshot {
         settings_protocol::MetadataSnapshot {
@@ -695,10 +715,14 @@ where
                 state.statement_metrics.insert(execution_id, metrics);
                 ExecuteContext::new(name.clone(), portal, Some(execution_id))
             });
-        let deallocation = state
+        let (deallocation, preparation) = state
             .portal_operations
             .get(&name)
-            .and_then(|owner| owner.deallocation.clone());
+            .map(|owner| (owner.deallocation.clone(), owner.preparation.clone()))
+            .unwrap_or_default();
+        if let Some(prepares) = preparation {
+            state.set_unmapped_sql_statement(prepares, operation);
+        }
         let context = state.operations.entry(operation).or_default();
         context.execute = Some(execute);
         context.deallocations.extend(deallocation);
@@ -810,15 +834,55 @@ where
         Ok(())
     }
 
-    /// Carry the deallocation of `statement`, if any, onto the portal just bound from it.
+    /// Record that `operation` forwards the SQL prepared statement `name`,
+    /// which writes data, without mapping.
+    pub fn set_unmapped_sql_statement(
+        &self,
+        name: Name,
+        operation: OperationId,
+    ) -> Result<(), Error> {
+        self.protocol_state
+            .write()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .set_unmapped_sql_statement(name, operation);
+        Ok(())
+    }
+
+    /// Record that the SQL of statement `name` is a `PREPARE` of `prepares`, an
+    /// unmapped statement that writes data, so that executing it through the
+    /// extended protocol records `prepares` in execution order.
+    pub fn set_statement_preparation(&self, name: Name, prepares: Name) -> Result<(), Error> {
+        self.protocol_state
+            .write()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .prepared_statements
+            .set_preparation(name, prepares);
+        Ok(())
+    }
+
+    /// Whether the SQL prepared statement `name` writes data and was forwarded
+    /// without mapping.
+    pub fn is_unmapped_sql_statement(&self, name: &Name) -> Result<bool, Error> {
+        Ok(self
+            .protocol_state
+            .read()
+            .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?
+            .unmapped_sql_statements
+            .contains_key(name))
+    }
+
+    /// Carry the deallocation and preparation of `statement`, if any, onto the
+    /// portal just bound from it.
     pub fn bind_statement_effects(&self, portal: &Name, statement: &Name) -> Result<(), Error> {
         let mut state = self
             .protocol_state
             .write()
             .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?;
         let deallocation = state.prepared_statements.deallocation(statement);
+        let preparation = state.prepared_statements.preparation(statement);
         if let Some(owner) = state.portal_operations.get_mut(portal) {
             owner.deallocation = deallocation;
+            owner.preparation = preparation;
         }
         Ok(())
     }
@@ -888,6 +952,22 @@ where
 
             debug!(target: CONTEXT, client_id = self.client_id, msg = "Deallocate", ?deallocation);
 
+            // Not a statement a later operation prepared again
+            match &deallocation {
+                Deallocation::Statement(name) => {
+                    if state
+                        .unmapped_sql_statements
+                        .get(name)
+                        .is_some_and(|prepared| *prepared < operation)
+                    {
+                        state.unmapped_sql_statements.remove(name);
+                    }
+                }
+                Deallocation::AllStatements => state
+                    .unmapped_sql_statements
+                    .retain(|_, prepared| *prepared >= operation),
+            }
+
             let names = state.prepared_statements.deallocation_names(deallocation);
             for name in &names {
                 // Also when a later Parse prepared the name again, so that Parse's rollback

packages/cipherstash-proxy/src/postgresql/context/prepared_statements.rs

@@ -12,6 +12,9 @@ struct PreparedStatement<K> {
     parse: Option<K>,
     /// Set when the statement's SQL is a `DEALLOCATE` or `DISCARD ALL`
     deallocation: Option<Deallocation>,
+    /// Set when the statement's SQL is a `PREPARE` of an unmapped statement
+    /// that writes data: the name it prepares
+    preparation: Option<Name>,
 }
 
 impl<K> Default for PreparedStatement<K> {
@@ -21,6 +24,7 @@ impl<K> Default for PreparedStatement<K> {
             metrics_scope: None,
             parse: None,
             deallocation: None,
+            preparation: None,
         }
     }
 }
@@ -82,6 +86,14 @@ where
         self.statements.entry(name).or_default().deallocation = Some(deallocation);
     }
 
+    pub fn preparation(&self, name: &Name) -> Option<Name> {
+        self.statements.get(name)?.preparation.clone()
+    }
+
+    pub fn set_preparation(&mut self, name: Name, prepares: Name) {
+        self.statements.entry(name).or_default().preparation = Some(prepares);
+    }
+
     /// Remove the statement, returning its metrics scope.
     pub fn remove(&mut self, name: &Name) -> Option<SessionId> {
         self.statements.remove(name)?.metrics_scope