From 7023ad04b44cc9ceff1a84add18c6ce8848a5c66 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:06:23 +0000 Subject: [PATCH 1/3] test(proxy): reproduce a simple Query keeping the unnamed statement PostgreSQL runs every simple Query in the unnamed statement, so it drops the client's unnamed prepared statement first. Proxy keeps its copy. One test shows the copy surviving a simple Query. Another pipelines a failing batch whose skipped unnamed Parse is rolled back after Proxy has read the following Query, which must not bring the old copy back. The DEALLOCATE ALL test ran the statement as a simple Query, so it only held because of this gap. It now runs DEALLOCATE ALL through the extended protocol, which is where the unnamed statement survives it. Co-Authored-By: Claude Opus 5.5 --- .../src/postgresql/middleware/mod.rs | 69 ++++++++++++++++++- 1 file changed, 66 insertions(+), 3 deletions(-) diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs index f7565d66..8425fbd8 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs @@ -103,6 +103,8 @@ mod tests { const ITERATIONS: usize = 200; + mod connection_model; + #[derive(Clone)] struct TestService; @@ -643,6 +645,53 @@ mod tests { } } + both_configs! { + async fn simple_query_releases_the_unnamed_statement(connect) { + let mut connection = connect(); + prepare(&mut connection, &["", "a1"]).await; + let select = connection.select; + + // PostgreSQL runs a simple Query in the unnamed statement, dropping the client's + let query = connection.query(select).await; + connection + .server(query, command_complete(b"SELECT 1")) + .await; + connection + .server(query, ready(TransactionStatus::Idle)) + .await; + + assert!(!connection.holds_statement(b"")); + assert_eq!(statement_names(&connection), vec!["a1"]); + } + } + + both_configs! { + async fn earlier_pipelined_failure_does_not_restore_the_unnamed_statement_a_query_released(connect) { + let mut connection = connect(); + prepare(&mut connection, &[""]).await; + let select = connection.select; + + let failed = connection.parse_named("a1".to_owned()).await; + // PostgreSQL skips this Parse after the failure, so it never answers it + connection.parse(b"").await; + let sync = connection.client(FrontendMessage::Sync).await; + let query = connection.query(select).await; + connection.server(failed, error()).await; + connection + .server(sync, ready(TransactionStatus::Idle)) + .await; + connection + .server(query, command_complete(b"SELECT 1")) + .await; + connection + .server(query, ready(TransactionStatus::Idle)) + .await; + + assert!(!connection.holds_statement(b"")); + assert!(statement_names(&connection).is_empty()); + } + } + both_configs! { async fn query_failing_part_way_keeps_the_statements_it_did_not_deallocate(connect) { let mut connection = connect(); @@ -835,16 +884,30 @@ mod tests { .server(sync, ready(TransactionStatus::Idle)) .await; - let query = connection.query("DEALLOCATE ALL").await; + // Through the extended protocol, because a simple Query releases the unnamed + // statement itself + let parse = connection + .parse_sql(b"deallocate", "DEALLOCATE ALL".to_owned()) + .await; + let bind = connection + .bind(Bytes::from_static(b"deallocate"), Bytes::new()) + .await; + let execute = connection.execute(Bytes::new()).await; + let sync = connection.client(FrontendMessage::Sync).await; connection - .server(query, command_complete(b"DEALLOCATE ALL")) + .server(parse, BackendMessage::ParseComplete) .await; + connection.server(bind, BackendMessage::BindComplete).await; connection - .server(query, ready(TransactionStatus::Idle)) + .server(execute, command_complete(b"DEALLOCATE ALL")) + .await; + connection + .server(sync, ready(TransactionStatus::Idle)) .await; assert!(connection.holds_statement(b"")); assert!(!connection.holds_statement(b"a1")); + assert!(!connection.holds_statement(b"deallocate")); } } From 70f85aa5cf109625187db3840a0301a621841425 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:06:23 +0000 Subject: [PATCH 2/3] fix(proxy): release the unnamed statement on a simple Query Context now releases the unnamed statement when it reads a simple Query, mapped, refused or with mapping disabled, as PostgreSQL does before it runs the Query. PostgreSQL releases it whatever happens to earlier operations, so the rollback snapshots of earlier operations that changed the unnamed statement forget it too, as they do for a SQL DEALLOCATE. Co-Authored-By: Claude Opus 5.5 --- .../src/postgresql/context/lifecycle.rs | 3 +++ .../src/postgresql/context/mod.rs | 18 ++++++++++++++++++ .../postgresql/context/settings_protocol.rs | 16 ++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs b/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs index ad0b18b9..97e220dc 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/lifecycle.rs @@ -56,6 +56,7 @@ impl Context { } pub fn unmapped_query(&mut self, operation: OperationId) -> Result<(), Error> { + self.release_unnamed_statement()?; let metrics = self.start_metrics_scope()?; self.add_portal(operation, Name::new(), Portal::passthrough(Some(metrics)))?; self.set_simple_query_execute_until_ready(operation, Name::new(), Some(metrics))?; @@ -373,6 +374,7 @@ impl Context { portal: Portal, statements: &[ast::Statement], ) -> Result<(), Error> { + self.release_unnamed_statement()?; self.add_portal(operation, Name::new(), portal)?; self.set_simple_query_execute_until_ready(operation, Name::new(), Some(metrics))?; self.execute_simple_schema_statements(statements); @@ -386,6 +388,7 @@ impl Context { metrics: SessionId, response: DiagnosticResponse, ) -> Result<(), Error> { + self.release_unnamed_statement()?; self.finish_metrics_scope(Some(metrics))?; self.set_operation_error(operation, response)?; self.set_execute(operation, Name::new(), None)?; diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index c69951c4..7268bc0a 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -803,6 +803,24 @@ where Ok(()) } + /// PostgreSQL runs a simple Query in the unnamed statement, releasing the client's. + fn release_unnamed_statement(&self) -> Result<(), Error> { + let finished = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let name = Name::new(); + let mut candidates = state.settings_protocol.released(&name); + candidates.extend(state.prepared_statements.remove(&name)); + state.take_unreferenced_metrics(candidates) + }; + for metrics in finished { + self.record_finished_metrics_scope(Some(metrics)); + } + Ok(()) + } + /// Record that the SQL of statement `name` is a `DEALLOCATE` or `DISCARD ALL`, so that /// executing it through the extended protocol releases what it removes. pub fn set_statement_deallocation( diff --git a/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs b/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs index 32d71744..ad0a5768 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/settings_protocol.rs @@ -509,6 +509,22 @@ impl SettingsProtocol { mutation.metadata = MetadataSnapshot::default(); } } + /// Forgets a statement that PostgreSQL releases unconditionally, as a simple Query + /// releases the unnamed statement. A rollback of an earlier operation must not restore + /// it, so returns the metrics scopes its rollback snapshots referenced. + pub fn released(&mut self, name: &Name) -> Vec { + self.prepared.remove(name); + self.journal + .values_mut() + .filter(|mutation| { + mutation.name == *name && matches!(mutation.previous, Previous::Statement(_)) + }) + .flat_map(|mutation| { + mutation.previous = Previous::Statement(None); + std::mem::take(&mut mutation.metadata).metrics() + }) + .collect() + } #[cfg(test)] pub fn statement_count(&self) -> usize { self.prepared.len() From 0e35fc1937c8eab624dd550aab63a580d77308a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:06:23 +0000 Subject: [PATCH 3/3] test(proxy): check a connection's mirrored state against a PostgreSQL model A property test drives one connection with random extended-protocol and simple Query traffic, and answers each message Proxy forwards from a small model of PostgreSQL: implicit and explicit transactions, failed transactions, skipping until Sync, duplicate names, the unnamed statement and portal, DISCARD ALL, SET LOCAL and temporary tables. After every ReadyForQuery it checks that Proxy holds no statement PostgreSQL released, holds no operation or portal state at idle, and agrees with PostgreSQL on UNSAFE_DISABLE_MAPPING and on whether the temporary table exists. It runs on a plaintext connection and on one with an encrypted column. The module lists what the model leaves out. Co-Authored-By: Claude Opus 5.5 --- .../middleware/tests/connection_model.rs | 741 ++++++++++++++++++ 1 file changed, 741 insertions(+) create mode 100644 packages/cipherstash-proxy/src/postgresql/middleware/tests/connection_model.rs diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/tests/connection_model.rs b/packages/cipherstash-proxy/src/postgresql/middleware/tests/connection_model.rs new file mode 100644 index 00000000..ad8fbf29 --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/middleware/tests/connection_model.rs @@ -0,0 +1,741 @@ +//! Drives one connection with random client traffic against a model of the PostgreSQL state +//! that traffic produces, and checks that Proxy's mirror of that state stays in step. +//! +//! The model answers each message Proxy forwards as PostgreSQL would, so refusals Proxy turns +//! into failing messages fail here too. After every `ReadyForQuery` it checks that Proxy: +//! - holds no prepared statement PostgreSQL has released, +//! - holds no operation, portal or suspended execution once the connection is idle, +//! - agrees on `CIPHERSTASH.UNSAFE_DISABLE_MAPPING`, and +//! - agrees on whether the connection's temporary table exists. +//! +//! The model covers less than PostgreSQL does: +//! - each message gets its reply before the client sends the next, so pipelined batches +//! are left to the harness tests, +//! - Execute always runs a portal to completion, and there is no Describe or Flush, +//! - each simple Query holds one statement, +//! - there are no savepoints, and +//! - only the statements a step marks as failing fail, besides the errors PostgreSQL +//! raises for missing, duplicate or aborted-transaction messages. + +use super::*; +use proptest::prelude::*; +use sqltk::parser::ast::{Ident, ObjectName, ObjectNamePart}; +use std::collections::{HashMap, HashSet}; + +/// SQL sent with Parse +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ExtendedSql { + Select, + SetMapping(bool), + Deallocate(u8), + DeallocateAll, +} + +/// SQL sent with a simple Query +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SimpleSql { + Select, + SetMapping(bool), + SetLocalMapping(bool), + Begin, + Commit, + Rollback, + DiscardAll, + CreateTemporaryTable, + DropTemporaryTable, + Deallocate(u8), +} + +/// What PostgreSQL runs for a statement or portal +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Sql { + Extended(ExtendedSql), + Simple(SimpleSql), + /// A statement Proxy rewrote so that PostgreSQL fails it + Refusal, +} + +#[derive(Clone, Debug)] +enum Step { + Parse { statement: u8, sql: ExtendedSql }, + Bind { statement: u8, portal: u8 }, + Execute { portal: u8, fails: bool }, + CloseStatement(u8), + ClosePortal(u8), + Sync, + Query { sql: SimpleSql, fails: bool }, +} + +fn statement_name(n: u8) -> String { + if n == 0 { + String::new() + } else { + format!("s{n}") + } +} + +fn portal_name(n: u8) -> String { + if n == 0 { + String::new() + } else { + format!("p{n}") + } +} + +fn mapping(value: bool) -> &'static str { + if value { + "true" + } else { + "false" + } +} + +fn steps() -> impl Strategy> { + let extended_sql = prop_oneof![ + 3 => Just(ExtendedSql::Select), + 2 => any::().prop_map(ExtendedSql::SetMapping), + 1 => (1_u8..3).prop_map(ExtendedSql::Deallocate), + 1 => Just(ExtendedSql::DeallocateAll), + ]; + let simple_sql = prop_oneof![ + 2 => Just(SimpleSql::Select), + 2 => any::().prop_map(SimpleSql::SetMapping), + 1 => any::().prop_map(SimpleSql::SetLocalMapping), + 2 => Just(SimpleSql::Begin), + 1 => Just(SimpleSql::Commit), + 1 => Just(SimpleSql::Rollback), + 1 => Just(SimpleSql::DiscardAll), + 1 => Just(SimpleSql::CreateTemporaryTable), + 1 => Just(SimpleSql::DropTemporaryTable), + 1 => (1_u8..3).prop_map(SimpleSql::Deallocate), + ]; + let step = prop_oneof![ + 4 => (0_u8..3, extended_sql).prop_map(|(statement, sql)| Step::Parse { statement, sql }), + 4 => (0_u8..3, 0_u8..2).prop_map(|(statement, portal)| Step::Bind { statement, portal }), + 4 => (0_u8..2, prop::bool::weighted(0.1)) + .prop_map(|(portal, fails)| Step::Execute { portal, fails }), + 1 => (0_u8..3).prop_map(Step::CloseStatement), + 1 => (0_u8..2).prop_map(Step::ClosePortal), + 3 => Just(Step::Sync), + 3 => (simple_sql, prop::bool::weighted(0.1)) + .prop_map(|(sql, fails)| Step::Query { sql, fails }), + ]; + prop::collection::vec(step, 1..60) +} + +/// The settings and temporary table a transaction rollback restores +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct Restorable { + mapping_disabled: bool, + local_mapping_disabled: Option, + temporary_table: bool, +} + +impl Restorable { + /// A plain `SET` inside a transaction also replaces any `SET LOCAL` value + fn set_mapping(&mut self, value: bool) { + self.mapping_disabled = value; + self.local_mapping_disabled = None; + } + + fn effective_mapping_disabled(&self) -> bool { + self.local_mapping_disabled.unwrap_or(self.mapping_disabled) + } +} + +/// The state of the PostgreSQL connection behind Proxy +#[derive(Debug, Default)] +struct Postgres { + statements: HashMap, + /// Portals, and whether each has run to completion + portals: HashMap, + current: Restorable, + /// Restored when an explicit transaction rolls back + transaction_start: Option, + /// Whether the explicit transaction has failed + transaction_failed: bool, + /// Restored when the implicit transaction of an extended-protocol batch fails + batch_start: Option, + /// Whether PostgreSQL ignores extended-protocol messages until Sync + skipping: bool, +} + +/// PostgreSQL's reply to one message +enum Reply { + Message(BackendMessage), + Error, +} + +impl Postgres { + fn in_transaction(&self) -> bool { + self.transaction_start.is_some() + } + + fn status(&self) -> TransactionStatus { + match (self.in_transaction(), self.transaction_failed) { + (false, _) => TransactionStatus::Idle, + (true, false) => TransactionStatus::InTransaction, + (true, true) => TransactionStatus::FailedTransaction, + } + } + + fn begin_batch(&mut self) { + if !self.in_transaction() && self.batch_start.is_none() { + self.batch_start = Some(self.current); + } + } + + /// An error inside a transaction or extended-protocol batch + fn fail(&mut self) { + if self.in_transaction() { + self.transaction_failed = true; + } else if let Some(start) = self.batch_start.take() { + self.current = start; + } + } + + fn end_implicit_transaction(&mut self) { + if !self.in_transaction() { + self.batch_start = None; + self.current.local_mapping_disabled = None; + self.portals.clear(); + } + } + + fn end_explicit_transaction(&mut self, commit: bool) { + if let Some(start) = self.transaction_start.take() { + if !commit { + self.current = start; + } + } + self.transaction_failed = false; + self.current.local_mapping_disabled = None; + self.portals.clear(); + } + + fn discard_all(&mut self) { + self.statements.retain(|name, _| name.is_empty()); + self.portals.clear(); + self.current = Restorable::default(); + } + + fn deallocate(&mut self, name: u8) -> Reply { + if self.statements.remove(&statement_name(name)).is_some() { + command(b"DEALLOCATE") + } else { + Reply::Error + } + } + + fn parse(&mut self, name: String, sql: Sql) -> Option { + if self.skipping { + return None; + } + self.begin_batch(); + // PostgreSQL releases the unnamed statement before it parses the new one + if name.is_empty() { + self.statements.remove(&name); + } + let reply = if self.transaction_failed + || sql == Sql::Refusal + || (!name.is_empty() && self.statements.contains_key(&name)) + { + Reply::Error + } else { + self.statements.insert(name, sql); + Reply::Message(BackendMessage::ParseComplete) + }; + Some(self.extended(reply)) + } + + fn bind(&mut self, statement: &[u8], portal: String) -> Option { + if self.skipping { + return None; + } + self.begin_batch(); + let statement = String::from_utf8_lossy(statement).into_owned(); + let reply = match self.statements.get(&statement) { + Some(sql) + if !self.transaction_failed + && (portal.is_empty() || !self.portals.contains_key(&portal)) => + { + self.portals.insert(portal, (*sql, false)); + Reply::Message(BackendMessage::BindComplete) + } + _ => Reply::Error, + }; + Some(self.extended(reply)) + } + + fn execute(&mut self, portal: &[u8], fails: bool) -> Option { + if self.skipping { + return None; + } + self.begin_batch(); + let portal = String::from_utf8_lossy(portal).into_owned(); + let reply = match self.portals.get_mut(&portal) { + Some((Sql::Extended(ExtendedSql::Select), true)) if !self.transaction_failed => { + // A SELECT portal run to completion stays runnable and returns no more rows + command(b"SELECT 0") + } + Some((sql, done)) if !self.transaction_failed && !*done => { + // Any other portal is done once it has run, and running it again fails + *done = true; + let sql = *sql; + match sql { + Sql::Extended(ExtendedSql::Select | ExtendedSql::SetMapping(_)) if fails => { + Reply::Error + } + Sql::Extended(ExtendedSql::Select) => command(b"SELECT 1"), + Sql::Extended(ExtendedSql::SetMapping(value)) => { + self.current.set_mapping(value); + command(b"SET") + } + Sql::Extended(ExtendedSql::Deallocate(name)) => self.deallocate(name), + Sql::Extended(ExtendedSql::DeallocateAll) => { + self.statements.retain(|name, _| name.is_empty()); + command(b"DEALLOCATE ALL") + } + Sql::Simple(_) | Sql::Refusal => Reply::Error, + } + } + _ => Reply::Error, + }; + Some(self.extended(reply)) + } + + /// An extended-protocol error makes PostgreSQL ignore messages until Sync + fn extended(&mut self, reply: Reply) -> Reply { + if matches!(reply, Reply::Error) { + self.fail(); + self.skipping = true; + } + reply + } + + fn sync(&mut self) -> BackendMessage { + self.skipping = false; + self.end_implicit_transaction(); + BackendMessage::ReadyForQuery(self.status()) + } + + fn query(&mut self, sql: Sql, fails: bool) -> Vec { + // A simple Query runs in the unnamed statement and portal, releasing the client's + self.statements.remove(""); + self.portals.remove(""); + let Sql::Simple(sql) = sql else { + self.fail(); + return self.query_end(Reply::Error); + }; + let exits_transaction = matches!(sql, SimpleSql::Commit | SimpleSql::Rollback); + if self.transaction_failed && !exits_transaction { + return self.query_end(Reply::Error); + } + let reply = match sql { + SimpleSql::Select + | SimpleSql::SetMapping(_) + | SimpleSql::SetLocalMapping(_) + | SimpleSql::CreateTemporaryTable + | SimpleSql::DropTemporaryTable + if fails => + { + self.fail(); + Reply::Error + } + SimpleSql::Select => command(b"SELECT 1"), + SimpleSql::SetMapping(value) => { + self.current.set_mapping(value); + command(b"SET") + } + SimpleSql::SetLocalMapping(value) => { + if self.in_transaction() { + self.current.local_mapping_disabled = Some(value); + } + command(b"SET") + } + SimpleSql::Begin => { + if !self.in_transaction() { + self.transaction_start = Some(self.current); + } + command(b"BEGIN") + } + SimpleSql::Commit => { + let failed = self.transaction_failed; + self.end_explicit_transaction(!failed); + command(if failed { b"ROLLBACK" } else { b"COMMIT" }) + } + SimpleSql::Rollback => { + self.end_explicit_transaction(false); + command(b"ROLLBACK") + } + SimpleSql::DiscardAll if self.in_transaction() => { + self.fail(); + Reply::Error + } + SimpleSql::DiscardAll => { + self.discard_all(); + command(b"DISCARD ALL") + } + SimpleSql::CreateTemporaryTable if !self.current.temporary_table => { + self.current.temporary_table = true; + command(b"CREATE TABLE") + } + SimpleSql::DropTemporaryTable if self.current.temporary_table => { + self.current.temporary_table = false; + command(b"DROP TABLE") + } + SimpleSql::Deallocate(name) => match self.deallocate(name) { + Reply::Error => { + self.fail(); + Reply::Error + } + reply => reply, + }, + SimpleSql::CreateTemporaryTable | SimpleSql::DropTemporaryTable => { + self.fail(); + Reply::Error + } + }; + self.query_end(reply) + } + + fn query_end(&mut self, reply: Reply) -> Vec { + if matches!(reply, Reply::Error) && self.in_transaction() { + self.transaction_failed = true; + } + self.end_implicit_transaction(); + vec![ + match reply { + Reply::Message(message) => message, + Reply::Error => error(), + }, + BackendMessage::ReadyForQuery(self.status()), + ] + } +} + +fn command(tag: &'static [u8]) -> Reply { + Reply::Message(command_complete(tag)) +} + +/// Runs `steps` against a connection and the model, checking the connection after each +/// readiness +struct Run { + connection: Connection, + postgres: Postgres, + /// SQL text Proxy may forward, and what PostgreSQL runs for it + sql: HashMap, + batch_open: bool, +} + +impl Run { + fn new(connection: Connection) -> Self { + let mut sql = HashMap::new(); + let select = connection.select.to_owned(); + sql.insert(select.clone(), Sql::Extended(ExtendedSql::Select)); + for value in [true, false] { + sql.insert( + format!( + "SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = {}", + mapping(value) + ), + Sql::Extended(ExtendedSql::SetMapping(value)), + ); + sql.insert( + format!( + "SET LOCAL CIPHERSTASH.UNSAFE_DISABLE_MAPPING = {}", + mapping(value) + ), + Sql::Simple(SimpleSql::SetLocalMapping(value)), + ); + } + for name in 1..3 { + sql.insert( + format!("DEALLOCATE {}", statement_name(name)), + Sql::Extended(ExtendedSql::Deallocate(name)), + ); + } + sql.insert( + "DEALLOCATE ALL".to_owned(), + Sql::Extended(ExtendedSql::DeallocateAll), + ); + for (text, simple) in [ + ("BEGIN", SimpleSql::Begin), + ("COMMIT", SimpleSql::Commit), + ("ROLLBACK", SimpleSql::Rollback), + ("DISCARD ALL", SimpleSql::DiscardAll), + ( + "CREATE TEMPORARY TABLE scratch (id int)", + SimpleSql::CreateTemporaryTable, + ), + ("DROP TABLE scratch", SimpleSql::DropTemporaryTable), + ] { + sql.insert(text.to_owned(), Sql::Simple(simple)); + } + Self { + connection, + postgres: Postgres::default(), + sql, + batch_open: false, + } + } + + fn extended_text(&self, sql: ExtendedSql) -> String { + match sql { + ExtendedSql::Select => self.connection.select.to_owned(), + ExtendedSql::SetMapping(value) => { + format!( + "SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = {}", + mapping(value) + ) + } + ExtendedSql::Deallocate(name) => format!("DEALLOCATE {}", statement_name(name)), + ExtendedSql::DeallocateAll => "DEALLOCATE ALL".to_owned(), + } + } + + fn simple_text(&self, sql: SimpleSql) -> String { + match sql { + SimpleSql::Select => self.connection.select.to_owned(), + SimpleSql::SetMapping(value) => self.extended_text(ExtendedSql::SetMapping(value)), + SimpleSql::SetLocalMapping(value) => format!( + "SET LOCAL CIPHERSTASH.UNSAFE_DISABLE_MAPPING = {}", + mapping(value) + ), + SimpleSql::Begin => "BEGIN".to_owned(), + SimpleSql::Commit => "COMMIT".to_owned(), + SimpleSql::Rollback => "ROLLBACK".to_owned(), + SimpleSql::DiscardAll => "DISCARD ALL".to_owned(), + SimpleSql::CreateTemporaryTable => "CREATE TEMPORARY TABLE scratch (id int)".to_owned(), + SimpleSql::DropTemporaryTable => "DROP TABLE scratch".to_owned(), + SimpleSql::Deallocate(name) => self.extended_text(ExtendedSql::Deallocate(name)), + } + } + + /// What PostgreSQL runs for SQL Proxy forwarded + fn classify(&self, text: &[u8], simple: bool) -> Sql { + let text = String::from_utf8_lossy(text); + match self.sql.get(text.as_ref()) { + Some(Sql::Extended(ExtendedSql::Select)) if simple => Sql::Simple(SimpleSql::Select), + Some(Sql::Extended(ExtendedSql::SetMapping(value))) if simple => { + Sql::Simple(SimpleSql::SetMapping(*value)) + } + Some(Sql::Extended(ExtendedSql::Deallocate(name))) if simple => { + Sql::Simple(SimpleSql::Deallocate(*name)) + } + Some(sql) => *sql, + None if text.contains("__cipherstash_") || text.starts_with("DO ") => Sql::Refusal, + // Proxy rewrote the statement that reads the encrypted column + None if simple => Sql::Simple(SimpleSql::Select), + None => Sql::Extended(ExtendedSql::Select), + } + } + + /// Sends `message`; returns the message Proxy forwarded, if any + async fn send(&mut self, message: FrontendMessage) -> (OperationId, Option) { + let (operation, output) = self.connection.client_output(message).await; + let forwarded = match output { + FrontendMiddlewareOutput::Forward(message) + | FrontendMiddlewareOutput::ForwardThenFlush(message) => Some(message), + FrontendMiddlewareOutput::Suppress(_) => None, + output => panic!("unexpected frontend output {output:?}"), + }; + (operation, forwarded) + } + + async fn reply(&mut self, operation: OperationId, reply: Option) { + match reply { + Some(Reply::Message(message)) => { + self.connection.server(operation, message).await; + } + Some(Reply::Error) => { + self.connection.server(operation, error()).await; + } + None => {} + } + } + + async fn step(&mut self, step: Step) { + match step { + Step::Query { .. } if self.batch_open => { + Box::pin(self.step(Step::Sync)).await; + Box::pin(self.step(step)).await; + } + Step::Parse { statement, sql } => { + self.batch_open = true; + let message = FrontendMessage::Parse(Parse { + statement: Bytes::from(statement_name(statement)), + query: Bytes::from(self.extended_text(sql)), + parameter_types: vec![], + }); + let (operation, forwarded) = self.send(message).await; + let reply = match forwarded { + Some(FrontendMessage::Parse(parse)) => { + let sql = self.classify(&parse.query, false); + let name = String::from_utf8_lossy(&parse.statement).into_owned(); + self.postgres.parse(name, sql) + } + None => None, + Some(other) => panic!("Parse forwarded as {other:?}"), + }; + self.reply(operation, reply).await; + } + Step::Bind { statement, portal } => { + self.batch_open = true; + let message = bind_message_to(statement_name(statement), portal_name(portal)); + let (operation, forwarded) = self.send(message).await; + let reply = match forwarded { + Some(FrontendMessage::Bind(bind)) => { + let portal = String::from_utf8_lossy(&bind.portal).into_owned(); + self.postgres.bind(&bind.statement, portal) + } + None => None, + Some(other) => panic!("Bind forwarded as {other:?}"), + }; + self.reply(operation, reply).await; + } + Step::Execute { portal, fails } => { + self.batch_open = true; + let (operation, forwarded) = self + .send(FrontendMessage::Execute(Execute { + portal: Bytes::from(portal_name(portal)), + max_rows: 0, + })) + .await; + let reply = match forwarded { + Some(FrontendMessage::Execute(execute)) => { + self.postgres.execute(&execute.portal, fails) + } + None => None, + Some(other) => panic!("Execute forwarded as {other:?}"), + }; + self.reply(operation, reply).await; + } + Step::CloseStatement(name) | Step::ClosePortal(name) => { + self.batch_open = true; + let (target, name) = match step { + Step::CloseStatement(_) => (DescribeTarget::Statement, statement_name(name)), + _ => (DescribeTarget::Portal, portal_name(name)), + }; + let (operation, forwarded) = self + .send(FrontendMessage::Close(Close { + target, + name: Bytes::from(name.clone()), + })) + .await; + let reply = forwarded.and_then(|_| { + if self.postgres.skipping { + return None; + } + match target { + DescribeTarget::Statement => self.postgres.statements.remove(&name), + DescribeTarget::Portal => { + self.postgres.portals.remove(&name).map(|(sql, _)| sql) + } + }; + Some(Reply::Message(BackendMessage::CloseComplete)) + }); + self.reply(operation, reply).await; + } + Step::Sync => { + self.batch_open = false; + let (operation, forwarded) = self.send(FrontendMessage::Sync).await; + assert!(matches!(forwarded, Some(FrontendMessage::Sync))); + let ready = self.postgres.sync(); + self.connection.server(operation, ready).await; + self.check(); + } + Step::Query { sql, fails } => { + let text = self.simple_text(sql); + let (operation, forwarded) = + self.send(FrontendMessage::Query(Bytes::from(text))).await; + let Some(FrontendMessage::Query(query)) = forwarded else { + panic!("Query forwarded as {forwarded:?}"); + }; + let sql = self.classify(&query, true); + for message in self.postgres.query(sql, fails) { + self.connection.server(operation, message).await; + } + self.check(); + } + } + } + + /// Proxy's mirror of the connection agrees with PostgreSQL + fn check(&self) { + let postgres = &self.postgres; + let held: HashSet = statement_names(&self.connection).into_iter().collect(); + for name in &held { + assert!( + postgres.statements.contains_key(name), + "Proxy holds statement {name:?}, which PostgreSQL released" + ); + } + if postgres.status() == TransactionStatus::Idle { + let sizes = self.connection.sizes(); + assert_eq!(sizes.operations, 0, "operations held at idle: {sizes:?}"); + assert_eq!(sizes.portals, 0, "portals held at idle: {sizes:?}"); + assert_eq!( + sizes.portal_operations, 0, + "portal operations held at idle: {sizes:?}" + ); + assert_eq!( + sizes.suspended_executions, 0, + "suspended executions held at idle: {sizes:?}" + ); + } + assert_eq!( + self.connection.context.unsafe_disable_mapping().unwrap(), + postgres.current.effective_mapping_disabled(), + "UNSAFE_DISABLE_MAPPING" + ); + let scratch = ObjectName(vec![ObjectNamePart::Identifier(Ident::new("scratch"))]); + assert_eq!( + self.connection + .context + .get_table_resolver() + .resolve_table(&scratch) + .is_ok(), + postgres.current.temporary_table, + "temporary table" + ); + } +} + +fn bind_message_to(statement: String, portal: String) -> FrontendMessage { + FrontendMessage::Bind(Bind { + portal: Bytes::from(portal), + statement: Bytes::from(statement), + parameter_formats: vec![], + parameters: vec![], + result_formats: vec![], + }) +} + +fn run(connect: fn() -> Connection, steps: Vec) { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let mut run = Run::new(connect()); + for step in steps { + run.step(step).await; + } + run.step(Step::Sync).await; + }); +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(256))] + + #[test] + fn plaintext_connection_mirrors_postgresql(steps in steps()) { + run(Connection::new, steps); + } + + #[test] + fn encrypted_connection_mirrors_postgresql(steps in steps()) { + run(Connection::encrypted, steps); + } +}