diff --git a/CHANGELOG.md b/CHANGELOG.md index a201adcd..f3bb7d32 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,10 +14,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Changed +- **PostgreSQL execution lifecycle ownership**: Proxy now tracks prepared statements, portals, Describe operations, execution outcomes, and execution metrics as one connection-local protocol lifecycle. Suspended executions retain their original metrics scope until completion, repeated portal executions receive distinct scopes, and terminal errors replace upstream responses atomically while schema outcomes are reported exactly once. + - **Upstream TLS verification for client traffic**: Connections with `with_tls_verification` enabled now use a cached snapshot of the system root certificates, loaded once when Proxy starts. Unlike Proxy's background database connections, pg-proto client-traffic connections do not apply operating-system revocation checks or enterprise verification policy. Restart Proxy after changing the system trust store. ### Fixed +- **Extended-protocol execution lifecycle regressions**: statement duration and slow-statement metrics now describe each execution rather than the lifetime of its cached prepared statement; distinct portals keep isolated Bind measurements; suspended executions retain their metrics until completion; correlated stale responses and inaccessible connection protocol state close the connection instead of silently omitting metadata transitions; uncorrelated responses retain PostgreSQL passthrough behavior; decryption failures no longer report pending schema changes as successful; and disabling mapping no longer creates empty statement metrics. + - **Query cancellation through Proxy**: Cancellation requests now reach the matching PostgreSQL connection, and their routing entries are removed when the client connection exits. Previously cancellation requests arrived on a separate connection that could not find the original route; retaining those routes globally without cleanup could also leak memory and eventually reject a new connection if PostgreSQL reused a cancellation key. - **Configured default keyset selection**: when a connection has not selected a keyset explicitly, Proxy now scopes encryption and decryption to `CS_DEFAULT_KEYSET_ID`. Previously it passed no keyset to ZeroKMS and could silently use the account default instead, deriving different searchable-encryption terms when the two defaults differed. Before upgrading, verify that the configured and account defaults are intentional; values written by an affected version under the unintended account default must be decrypted with that old keyset and re-encrypted under the configured default. diff --git a/packages/cipherstash-proxy-integration/src/schema_change.rs b/packages/cipherstash-proxy-integration/src/schema_change.rs index 3f443418..0a975905 100644 --- a/packages/cipherstash-proxy-integration/src/schema_change.rs +++ b/packages/cipherstash-proxy-integration/src/schema_change.rs @@ -455,4 +455,21 @@ mod tests { assert!(exists); assert_ciphertext_at_rest(&encrypted, 1, "preserved batch secret").await; } + + #[tokio::test] + async fn multi_statement_simple_query_keeps_the_connection_usable() { + let client = connect_for_test(*PROXY).await; + let first = table("execution_lifecycle_first"); + let second = table("execution_lifecycle_second"); + + client + .batch_execute(&format!( + "CREATE TABLE {first} (id bigint); CREATE TABLE {second} (id bigint)" + )) + .await + .unwrap(); + + let value: i32 = client.query_one("SELECT 1", &[]).await.unwrap().get(0); + assert_eq!(value, 1); + } } diff --git a/packages/cipherstash-proxy/CONTEXT.md b/packages/cipherstash-proxy/CONTEXT.md index c74c9dc5..4e20b975 100644 --- a/packages/cipherstash-proxy/CONTEXT.md +++ b/packages/cipherstash-proxy/CONTEXT.md @@ -34,8 +34,10 @@ A prepared statement bound to parameter values, ready to execute. Either `Encryp carrying the analysed statement, or `Passthrough` when nothing in it is encrypted. **Statement metrics scope**: -The measurement window around a single statement — opened at Parse or Query, closed when -the statement completes. Many of these occur per connection. +The measurement window around one statement execution occurrence. A simple Query owns one; +an extended-protocol Execute owns one from its first execution through any suspension and +resumption until completion or failure, while Parse timing belongs to the prepared Statement +and may be attributed to each execution occurrence. _Avoid_: session (`start_session`, `SessionId` and the `..._statements_session_duration_seconds` metric all use this sense and are misnamed). diff --git a/packages/cipherstash-proxy/docs/adr/0002-context-owns-protocol-metadata-transitions.md b/packages/cipherstash-proxy/docs/adr/0002-context-owns-protocol-metadata-transitions.md new file mode 100644 index 00000000..ace45cbe --- /dev/null +++ b/packages/cipherstash-proxy/docs/adr/0002-context-owns-protocol-metadata-transitions.md @@ -0,0 +1,23 @@ +--- +status: accepted +--- + +# Context owns CipherStash protocol metadata transitions + +`Context` is the connection seam for CipherStash metadata associated with PostgreSQL Statements, +Portals, operations, and statement metrics scopes. Frontend and Backend remain protocol adapters: +they interpret wire messages, while Context applies each correlated metadata transition atomically +and returns the knowledge or effects the adapter needs. pg-proto continues to own protocol ordering +and backpressure, and Schema middleware continues to own transactional schema state under ADR-0001. + +The correlated protocol state is kept in one internal state model rather than independently locked +maps. Context never holds that state lock across asynchronous work, metrics emission, or Schema +middleware calls; a transition first changes protocol state and returns explicit effects, then +Context applies those effects. Passthrough and encrypted traffic use the same lifecycle seam, and +inaccessible, stale, or inconsistent state fails the connection closed rather than silently omitting +a transition. + +Statement metrics scopes belong to execution occurrences, not prepared Statements. A suspended and +resumed execution retains one scope until completion or failure; distinct Portals and repeated +executions receive distinct scopes, while Parse timing remains knowledge of the prepared Statement +that may be attributed to each execution observation. Existing Prometheus metric names remain stable. diff --git a/packages/cipherstash-proxy/src/config/tandem.rs b/packages/cipherstash-proxy/src/config/tandem.rs index 9a331abc..913b8a9c 100644 --- a/packages/cipherstash-proxy/src/config/tandem.rs +++ b/packages/cipherstash-proxy/src/config/tandem.rs @@ -343,6 +343,15 @@ impl TandemConfig { development: None, } } + + #[cfg(test)] + pub fn disable_mapping_for_testing(&mut self) { + self.development = Some(DevelopmentConfig { + disable_mapping: true, + disable_database_tls: false, + enable_mapping_errors: false, + }); + } } impl EncryptConfig { diff --git a/packages/cipherstash-proxy/src/error.rs b/packages/cipherstash-proxy/src/error.rs index cecb1ef0..58b12944 100644 --- a/packages/cipherstash-proxy/src/error.rs +++ b/packages/cipherstash-proxy/src/error.rs @@ -85,6 +85,21 @@ impl Error { #[derive(Error, Debug)] pub enum ContextError { + #[error("Operation has no in-flight Execute state")] + OperationWithoutExecute, + + #[error("Connection protocol state is unavailable")] + ProtocolStateUnavailable, + + #[error("Operation could not be found in connection protocol state")] + UnknownOperation, + + #[error("Operation has no in-flight Describe state")] + UnknownDescribe, + + #[error("Prepared statement has no metrics template")] + StatementMetricsUnavailable, + #[error("Portal could not be found in context")] UnknownPortal, } diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index 0bf57f99..00ad1efb 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -4,7 +4,7 @@ pub mod portal; pub mod statement; pub mod statement_metadata; pub use self::{phase_timing::PhaseTiming, portal::Portal, statement::Statement}; -use super::{column_mapper::ColumnMapper, rewrite::Name, Column}; +use super::{column_mapper::ColumnMapper, rewrite::Name, Column, OperationId}; use crate::{ config::TandemConfig, error::{ConfigError, EncryptError, Error}, @@ -21,7 +21,7 @@ use crate::{ use cipherstash_client::IdentifiedBy; use eql_mapper::{Schema, TableResolver}; use metrics::{counter, histogram}; -use pg_proto::{Describe, DescribeTarget, DiagnosticResponse, OperationId, TransactionStatus}; +use pg_proto::{Describe, DescribeTarget, DiagnosticResponse, TransactionStatus}; use serde_json::json; use sqltk::parser::ast::{Expr, Ident, ObjectName, ObjectNamePart, Set, Value, ValueWithSpan}; pub use statement_metadata::StatementMetadata; @@ -59,11 +59,7 @@ where encryption: T, reload_sender: ReloadSender, schema_middleware: SchemaMiddleware, - statements: Arc>>>, - statement_sessions: Arc>>, - portals: Arc>>>, - operations: Arc>>, - session_metrics: Arc>>, + protocol_state: Arc>, upstream_tls_roots: Arc, unsafe_disable_mapping: bool, keyset_id: Arc>>, @@ -81,6 +77,7 @@ pub struct ExecuteContext { portal: Option>, start: Instant, session_id: Option, + completes_at_readiness: bool, } impl ExecuteContext { @@ -94,9 +91,15 @@ impl ExecuteContext { portal, start: Instant::now(), session_id, + completes_at_readiness: false, } } + fn until_readiness(mut self) -> Self { + self.completes_at_readiness = true; + self + } + fn duration(&self) -> Duration { Instant::now().duration_since(self.start) } @@ -108,15 +111,159 @@ impl ExecuteContext { #[derive(Clone, Debug, Default)] struct OperationContext { - describe_statement: Option>, + describe: Option, execute: Option, error_response: Option, } +#[derive(Clone, Debug)] +struct DescribeContext { + statement: Option>, +} + +#[derive(Debug)] +struct ConnectionProtocolState { + operations: HashMap, + statement_metrics: HashMap, + suspended_executions: HashMap, + statements: HashMap>, + statement_metrics_scopes: HashMap, + portals: HashMap>, + portal_operations: HashMap, +} + +impl Default for ConnectionProtocolState { + fn default() -> Self { + Self { + operations: HashMap::new(), + statement_metrics: HashMap::new(), + suspended_executions: HashMap::new(), + statements: HashMap::new(), + statement_metrics_scopes: HashMap::new(), + portals: HashMap::new(), + portal_operations: HashMap::new(), + } + } +} + +impl ConnectionProtocolState +where + K: Copy + Eq + std::hash::Hash, +{ + fn metrics_referenced(&self, session_id: SessionId) -> bool { + self.statement_metrics_scopes + .values() + .any(|id| *id == session_id) + || self + .portals + .values() + .any(|portal| portal.session_id() == Some(session_id)) + || self.operations.values().any(|operation| { + operation + .execute + .as_ref() + .is_some_and(|execute| execute.session_id() == Some(session_id)) + }) + || self + .suspended_executions + .values() + .any(|(_, execute)| execute.session_id() == Some(session_id)) + } + + fn take_unreferenced_metrics( + &mut self, + candidates: impl IntoIterator, + ) -> Vec { + let mut metrics = Vec::new(); + for session_id in candidates { + if !self.metrics_referenced(session_id) { + metrics.extend(self.statement_metrics.remove(&session_id)); + } + } + metrics + } + + fn finish_execution( + &mut self, + operation: &K, + outcome: ExecutionOutcome, + ) -> Result { + let Some(operation_context) = self.operations.get_mut(operation) else { + return Err(crate::error::ContextError::UnknownOperation); + }; + if outcome == ExecutionOutcome::Completed + && operation_context + .execute + .as_ref() + .is_some_and(|execute| execute.completes_at_readiness) + { + return Ok(ExecutionTransition { + execute: operation_context.execute.clone(), + metrics: None, + finished_metrics: None, + replacement_error: None, + execution_finished: false, + }); + } + let execute = operation_context.execute.take(); + if execute.is_none() && outcome != ExecutionOutcome::Failed { + return Err(crate::error::ContextError::OperationWithoutExecute); + } + let replacement_error = if outcome == ExecutionOutcome::Failed { + operation_context.error_response.take() + } else { + None + }; + let remove_operation = + outcome == ExecutionOutcome::Failed || operation_context.describe.is_none(); + if remove_operation { + self.operations.remove(operation); + } + let metrics = execute + .as_ref() + .and_then(ExecuteContext::session_id) + .and_then(|id| self.statement_metrics.get(&id).cloned()); + let finished_metrics = if outcome == ExecutionOutcome::Suspended { + let execute = execute.as_ref().unwrap(); + self.suspended_executions + .insert(execute.name.clone(), (*operation, execute.clone())); + None + } else { + execute + .as_ref() + .and_then(ExecuteContext::session_id) + .and_then(|id| self.statement_metrics.remove(&id)) + }; + Ok(ExecutionTransition { + execute, + metrics, + finished_metrics, + replacement_error, + execution_finished: true, + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExecutionOutcome { + Completed, + Suspended, + Failed, +} + +struct ExecutionTransition { + execute: Option, + metrics: Option, + finished_metrics: Option, + replacement_error: Option, + execution_finished: bool, +} + #[derive(Clone, Debug)] pub struct SessionMetricsContext { id: SessionId, start: Instant, + records_statement_metrics: bool, pub phase_timing: PhaseTiming, pub metadata: StatementMetadata, } @@ -126,6 +273,7 @@ impl SessionMetricsContext { SessionMetricsContext { id, start: Instant::now(), + records_statement_metrics: true, phase_timing: PhaseTiming::new(), metadata: StatementMetadata::new(), } @@ -177,11 +325,7 @@ where let schema_middleware = SchemaMiddleware::from_store(schema_store); Context { - statements: Arc::new(RwLock::new(HashMap::new())), - statement_sessions: Arc::new(RwLock::new(HashMap::new())), - portals: Arc::new(RwLock::new(HashMap::new())), - operations: Arc::new(RwLock::new(HashMap::new())), - session_metrics: Arc::new(RwLock::new(HashMap::new())), + protocol_state: Arc::new(RwLock::new(ConnectionProtocolState::default())), upstream_tls_roots, client_id, config, @@ -195,55 +339,108 @@ where } } - pub fn set_describe(&mut self, operation: OperationId, describe: Describe) { + pub fn set_describe(&self, operation: OperationId, describe: Describe) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, describe = ?describe); + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; let statement = match &describe { Describe { name, target: DescribeTarget::Portal, - } => self.get_portal_statement(name), + } => match state.portals.get(name).map(Arc::as_ref) { + Some(Portal::Encrypted { statement, .. }) => Some(statement.clone()), + Some(Portal::Passthrough { .. }) | None => None, + }, Describe { name, target: DescribeTarget::Statement, - } => self.get_statement(name), + } => state.statements.get(name).cloned(), }; - let _ = self.operations.write().map(|mut operations| { - operations.entry(operation).or_default().describe_statement = statement; - }); + state.operations.entry(operation).or_default().describe = + Some(DescribeContext { statement }); + Ok(()) + } + + pub fn set_non_execution(&self, operation: OperationId) -> Result<(), Error> { + self.protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)? + .operations + .entry(operation) + .or_default(); + Ok(()) + } + + pub fn complete_non_execution(&self, operation: OperationId) -> Result<(), Error> { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + state + .operations + .remove(&operation) + .map(|_| ()) + .ok_or(crate::error::ContextError::UnknownOperation.into()) } /// /// Marks the current Describe as complete /// Removes the Describe from the Queue /// - pub fn complete_describe(&mut self, operation: OperationId) { + pub fn complete_describe(&self, operation: OperationId) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, msg = "Describe complete"); - let _ = self.operations.write().map(|mut operations| { - if let Some(context) = operations.get_mut(&operation) { - context.describe_statement = None; - if context.execute.is_none() { - operations.remove(&operation); - } + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let Some(context) = state.operations.get_mut(&operation) else { + return Err(crate::error::ContextError::UnknownOperation.into()); + }; + if context.describe.take().is_none() { + if context.execute.is_some() { + return Ok(()); } - }); + state.operations.remove(&operation); + return Err(crate::error::ContextError::UnknownDescribe.into()); + } + if context.execute.is_none() { + state.operations.remove(&operation); + } + Ok(()) } - pub fn start_session(&mut self) -> SessionId { + pub fn start_metrics_scope(&mut self) -> Result { let id = SessionId(self.session_id_counter.fetch_add(1, Ordering::Relaxed)); let ctx = SessionMetricsContext::new(id); - let _ = self - .session_metrics + self.protocol_state .write() - .map(|mut sessions| sessions.insert(id, ctx)); - id + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)? + .statement_metrics + .insert(id, ctx); + Ok(id) } - pub fn finish_session(&mut self, session_id: Option) { + pub fn finish_metrics_scope(&mut self, session_id: Option) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, msg = "Session Metrics finished"); - let session = session_id.and_then(|id| self.session_metrics.write().ok()?.remove(&id)); - if let Some(session) = session { - let duration = session.duration(); - let metadata = &session.metadata; + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let metrics_scope = session_id.and_then(|id| state.statement_metrics.remove(&id)); + drop(state); + self.record_finished_metrics_scope(metrics_scope); + Ok(()) + } + + fn record_finished_metrics_scope(&self, metrics_scope: Option) { + if let Some(metrics_scope) = metrics_scope { + if !metrics_scope.records_statement_metrics { + return; + } + let duration = metrics_scope.duration(); + let metadata = &metrics_scope.metadata; // Get labels for metrics let statement_type = metadata @@ -272,7 +469,7 @@ where if self.config.slow_statements_enabled() && duration > self.config.slow_statement_min_duration() { - let timing = &session.phase_timing; + let timing = &metrics_scope.phase_timing; // Increment slow statements counter counter!(SLOW_STATEMENTS_TOTAL).increment(1); @@ -308,57 +505,113 @@ where operation: OperationId, name: Name, session_id: Option, - ) { + ) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, execute = ?name); + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let portal = state.portals.get(&name).cloned(); + let execute = ExecuteContext::new(name, portal, session_id); + state.operations.entry(operation).or_default().execute = Some(execute); + Ok(()) + } - let portal = self.get_portal(&name); - let ctx = ExecuteContext::new(name, portal, session_id); - let _ = self.operations.write().map(|mut operations| { - operations.entry(operation).or_default().execute = Some(ctx); - }); + pub fn set_simple_query_execute_until_ready( + &mut self, + operation: OperationId, + name: Name, + session_id: Option, + ) -> Result<(), Error> { + debug!(target: CONTEXT, client_id = self.client_id, execute = ?name); + + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let portal = state.portals.get(&name).cloned(); + let execute = ExecuteContext::new(name, portal, session_id).until_readiness(); + state.operations.entry(operation).or_default().execute = Some(execute); + Ok(()) } - /// Set execute state for portal, looking up session ID internally. - pub fn set_execute_for_portal(&mut self, operation: OperationId, name: Name) { - let session_id = self.get_portal_session_id(&name); - self.set_execute(operation, name, session_id); + /// Set execute state for a portal, looking up its metrics scope ID internally. + pub fn set_execute_for_portal( + &mut self, + operation: OperationId, + name: Name, + ) -> Result<(), Error> { + let execution_id = SessionId(self.session_id_counter.fetch_add(1, Ordering::Relaxed)); + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let portal = state.portals.get(&name).cloned(); + let template_id = portal.as_ref().and_then(|portal| portal.session_id()); + let execute = state + .suspended_executions + .remove(&name) + .map(|(_, execute)| execute) + .unwrap_or_else(|| { + let mut metrics = template_id + .and_then(|id| state.statement_metrics.get(&id).cloned()) + .unwrap_or_else(|| SessionMetricsContext::new(execution_id)); + metrics.id = execution_id; + metrics.start = Instant::now(); + metrics.records_statement_metrics = true; + state.statement_metrics.insert(execution_id, metrics); + ExecuteContext::new(name, portal, Some(execution_id)) + }); + state.operations.entry(operation).or_default().execute = Some(execute); + Ok(()) } - /// Marks the current Execution as Complete. - /// - /// If the associated portal is Unnamed, it is closed. - /// - /// From the PostgreSQL Extended Query docs: - /// If successfully created, a named portal object lasts till the end of the current transaction, unless explicitly destroyed. - /// An unnamed portal is destroyed at the end of the transaction, or as soon as the next Bind statement specifying the unnamed portal as destination is issued - /// - /// https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY - pub fn complete_execution(&mut self, operation: OperationId) -> Option { - debug!(target: CONTEXT, client_id = self.client_id, msg = "Execute complete"); - - let execute = self.operations.write().ok().and_then(|mut operations| { - let execute = operations.get_mut(&operation)?.execute.take(); - if operations - .get(&operation) - .is_some_and(|context| context.describe_statement.is_none()) - { - operations.remove(&operation); + /// Applies one terminal or suspended Execute outcome atomically. + pub fn finish_execution( + &self, + operation: OperationId, + outcome: ExecutionOutcome, + ) -> Result, Error> { + debug!(target: CONTEXT, client_id = self.client_id, ?outcome, msg = "Execute outcome"); + + let transition = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + state.finish_execution(&operation, outcome)? + }; + + if transition.execution_finished && outcome != ExecutionOutcome::Suspended { + if let Some(execute) = transition.execute.as_ref() { + self.record_execution_duration(execute, transition.metrics.as_ref()); + + if execute.name.is_empty() { + self.close_portal_if_current(&execute.name, execute.portal.as_ref())?; + } } - execute - }); - if let Some(execute) = execute { - // Get labels from current session metadata - let (statement_type, protocol, mapped, multi_statement) = if let Some(session) = execute - .session_id() - .and_then(|id| self.get_session_metrics(id)) - { - let metadata = &session.metadata; + } + self.record_finished_metrics_scope(transition.finished_metrics); + Ok(transition.replacement_error) + } + + fn record_execution_duration( + &self, + execute: &ExecuteContext, + metrics: Option<&SessionMetricsContext>, + ) { + let (statement_type, protocol, mapped, multi_statement) = metrics + .map(|metrics_scope| { + let metadata = &metrics_scope.metadata; ( metadata .statement_type - .map(|t| t.as_label()) + .map(|kind| kind.as_label()) + .unwrap_or("unknown"), + metadata + .protocol + .map(|kind| kind.as_label()) .unwrap_or("unknown"), - metadata.protocol.map(|p| p.as_label()).unwrap_or("unknown"), if metadata.encrypted { "true" } else { "false" }, if metadata.multi_statement { "true" @@ -366,65 +619,47 @@ where "false" }, ) - } else { - ("unknown", "unknown", "false", "false") - }; - - histogram!( - STATEMENTS_EXECUTION_DURATION_SECONDS, - "statement_type" => statement_type, - "protocol" => protocol, - "mapped" => mapped, - "multi_statement" => multi_statement - ) - .record(execute.duration()); - - if execute.name.is_empty() { - self.close_portal_if_current(&execute.name, execute.portal.as_ref()); - } - return execute.session_id(); - } - None + }) + .unwrap_or(("unknown", "unknown", "false", "false")); + histogram!( + STATEMENTS_EXECUTION_DURATION_SECONDS, + "statement_type" => statement_type, + "protocol" => protocol, + "mapped" => mapped, + "multi_statement" => multi_statement + ) + .record(execute.duration()); } - pub fn add_statement(&mut self, name: Name, statement: Statement) { + pub fn add_statement(&self, name: Name, statement: Statement) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, statement = ?name); - let _ = self - .statements + self.protocol_state .write() - .map(|mut guarded| guarded.insert(name, Arc::new(statement))); + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)? + .statements + .insert(name, Arc::new(statement)); + Ok(()) } - pub fn close_statement(&mut self, name: &Name) { + pub fn close_statement(&self, name: &Name) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, statement = ?name); - let session_id = self.get_statement_session(name); - let _ = self - .statements - .write() - .map(|mut guarded| guarded.remove(name)); - let _ = self - .statement_sessions - .write() - .map(|mut guarded| guarded.remove(name)); - - let session_in_use = session_id.is_some_and(|session_id| { - self.portals.read().is_ok_and(|portals| { - portals - .values() - .any(|portal| portal.session_id() == Some(session_id)) - }) || self.operations.read().is_ok_and(|operations| { - operations.values().any(|operation| { - operation - .execute - .as_ref() - .is_some_and(|execute| execute.session_id() == Some(session_id)) - }) - }) - }); - if !session_in_use { - self.finish_session(session_id); - } + let finished_metrics_scope = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let session_id = state.statement_metrics_scopes.remove(name); + state.statements.remove(name); + let metrics_scope_in_use = session_id.is_some_and(|id| state.metrics_referenced(id)); + if !metrics_scope_in_use { + session_id.and_then(|session_id| state.statement_metrics.remove(&session_id)) + } else { + None + } + }; + self.record_finished_metrics_scope(finished_metrics_scope); + Ok(()) } pub fn transaction_status(&self) -> TransactionStatus { @@ -434,176 +669,422 @@ where .unwrap_or(TransactionStatus::Idle) } - pub fn set_transaction_status(&mut self, status: TransactionStatus) { - if let Ok(mut current) = self.transaction_status.write() { - *current = status; + pub fn ready_for_query( + &self, + status: TransactionStatus, + boundary: Option, + ) -> Result<(), Error> { + *self + .transaction_status + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)? = status; + let Some(boundary) = boundary else { + return Ok(()); + }; + + let (finished_executions, finished_metrics_scopes) = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let mut candidates = Vec::new(); + let mut executions = Vec::new(); + if status == TransactionStatus::Idle { + let portal_names = state + .portal_operations + .iter() + .filter_map(|(name, operation)| { + (*operation <= boundary).then_some(name.clone()) + }) + .collect::>(); + for name in portal_names { + state.portal_operations.remove(&name); + candidates.extend( + state + .portals + .remove(&name) + .and_then(|portal| portal.session_id()), + ); + } + + let suspended_names = state + .suspended_executions + .iter() + .filter_map(|(name, (operation, _))| { + (*operation <= boundary).then_some(name.clone()) + }) + .collect::>(); + for name in suspended_names { + candidates.extend( + state + .suspended_executions + .remove(&name) + .and_then(|(_, execute)| execute.session_id()), + ); + } + } + + let operation_ids = state + .operations + .keys() + .filter(|operation| **operation <= boundary) + .copied() + .collect::>(); + for operation in operation_ids { + if let Some(execute) = state + .operations + .remove(&operation) + .and_then(|operation| operation.execute) + { + let metrics = execute + .session_id() + .and_then(|id| state.statement_metrics.get(&id).cloned()); + candidates.extend(execute.session_id()); + executions.push((execute, metrics)); + } + } + (executions, state.take_unreferenced_metrics(candidates)) + }; + for (execute, metrics) in finished_executions { + self.record_execution_duration(&execute, metrics.as_ref()); + if execute.completes_at_readiness && execute.name.is_empty() { + self.close_portal_if_current(&execute.name, execute.portal.as_ref())?; + } + } + for metrics_scope in finished_metrics_scopes { + self.record_finished_metrics_scope(Some(metrics_scope)); } + Ok(()) } /// Close a statement explicitly requested by the client. /// /// PostgreSQL portals retain the parsed statement they reference and remain /// valid after the statement name is closed, so they must not be removed. - pub fn close_statement_explicit(&mut self, name: &Name) { - self.close_statement(name); - } - - pub fn discard_operation(&mut self, operation: OperationId) { - let _ = self - .operations - .write() - .map(|mut operations| operations.remove(&operation)); - } - - pub fn set_operation_error(&mut self, operation: OperationId, response: DiagnosticResponse) { - let _ = self.operations.write().map(|mut operations| { - operations.entry(operation).or_default().error_response = Some(response); - }); + pub fn close_statement_explicit(&self, name: &Name) -> Result<(), Error> { + self.close_statement(name) + } + + pub fn discard_operation(&mut self, operation: OperationId) -> Result<(), Error> { + let finished_metrics_scopes = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let metrics_scope_id = state + .operations + .remove(&operation) + .and_then(|operation| operation.execute) + .and_then(|execute| execute.session_id()); + state.take_unreferenced_metrics(metrics_scope_id) + }; + for metrics_scope in finished_metrics_scopes { + self.record_finished_metrics_scope(Some(metrics_scope)); + } + Ok(()) } - pub fn take_operation_error(&mut self, operation: OperationId) -> Option { - self.operations + pub fn set_operation_error( + &mut self, + operation: OperationId, + response: DiagnosticResponse, + ) -> Result<(), Error> { + self.protocol_state .write() - .ok()? - .get_mut(&operation)? - .error_response - .take() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)? + .operations + .entry(operation) + .or_default() + .error_response = Some(response); + Ok(()) } - pub fn add_portal(&mut self, name: Name, portal: Portal) { + pub fn add_portal( + &self, + operation: OperationId, + name: Name, + portal: Portal, + ) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, name = ?name, portal = ?portal); - let _ = self - .portals - .write() - .map(|mut portals| portals.insert(name, Arc::new(portal))); + let finished_metrics = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let mut candidates = Vec::new(); + candidates.extend( + state + .portals + .remove(&name) + .and_then(|portal| portal.session_id()), + ); + candidates.extend( + state + .suspended_executions + .remove(&name) + .and_then(|(_, execute)| execute.session_id()), + ); + state.portal_operations.insert(name.clone(), operation); + state.portals.insert(name, Arc::new(portal)); + state.take_unreferenced_metrics(candidates) + }; + for metrics in finished_metrics { + self.record_finished_metrics_scope(Some(metrics)); + } + Ok(()) } - pub fn get_statement(&self, name: &Name) -> Option> { + pub fn get_statement(&self, name: &Name) -> Result>, Error> { debug!(target: CONTEXT, client_id = self.client_id, statement = ?name); - let statements = self.statements.read().ok()?; - statements.get(name).cloned() + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + Ok(state.statements.get(name).cloned()) } - pub fn set_statement_session(&mut self, name: Name, session_id: SessionId) { - let _ = self - .statement_sessions + pub fn set_statement_metrics_scope( + &mut self, + name: Name, + session_id: SessionId, + ) -> Result<(), Error> { + let mut state = self + .protocol_state .write() - .map(|mut sessions| sessions.insert(name, session_id)); + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + if let Some(metrics) = state.statement_metrics.get_mut(&session_id) { + metrics.records_statement_metrics = false; + } + state.statement_metrics_scopes.insert(name, session_id); + Ok(()) } - pub fn get_statement_session(&self, name: &Name) -> Option { - let sessions = self.statement_sessions.read().ok()?; - sessions.get(name).copied() + pub fn get_statement_metrics_scope(&self, name: &Name) -> Result, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + Ok(state.statement_metrics_scopes.get(name).copied()) } - /// Get session for statement, falling back to latest session with warning log. - pub fn get_statement_session_or_latest(&self, name: &Name) -> Option { - if let Some(id) = self.get_statement_session(name) { - return Some(id); - } - - let fallback = self.latest_session_id(); - if fallback.is_some() { - warn!( - target: CONTEXT, - client_id = self.client_id, - prepared_statement = %String::from_utf8_lossy(name), - msg = "Session lookup failed for prepared statement, using latest session" - ); - } - fallback + pub fn start_portal_metrics_scope( + &mut self, + statement: &Name, + ) -> Result, Error> { + let id = SessionId(self.session_id_counter.fetch_add(1, Ordering::Relaxed)); + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let Some(template_id) = state.statement_metrics_scopes.get(statement).copied() else { + return Ok(None); + }; + let Some(mut metrics) = state.statement_metrics.get(&template_id).cloned() else { + return Err(crate::error::ContextError::StatementMetricsUnavailable.into()); + }; + metrics.id = id; + metrics.start = Instant::now(); + metrics.records_statement_metrics = false; + state.statement_metrics.insert(id, metrics); + Ok(Some(id)) } /// /// Close the portal identified by `name` /// Portal is removed from queue /// - pub fn close_portal(&mut self, name: &Name) { + pub fn close_portal(&self, name: &Name) -> Result<(), Error> { debug!(target: CONTEXT, client_id = self.client_id, msg = "Close Portal", name = ?name); - let _ = self.portals.write().map(|mut portals| portals.remove(name)); + let finished_metrics = { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + state.portal_operations.remove(name); + let mut candidates = Vec::new(); + candidates.extend( + state + .portals + .remove(name) + .and_then(|portal| portal.session_id()), + ); + candidates.extend( + state + .suspended_executions + .remove(name) + .and_then(|(_, execute)| execute.session_id()), + ); + state.take_unreferenced_metrics(candidates) + }; + for metrics in finished_metrics { + self.record_finished_metrics_scope(Some(metrics)); + } + Ok(()) } - fn close_portal_if_current(&mut self, name: &Name, expected: Option<&Arc>) { - let _ = self.portals.write().map(|mut portals| { - if expected.is_some_and(|expected| { - portals - .get(name) - .is_some_and(|current| Arc::ptr_eq(current, expected)) - }) { - portals.remove(name); - } - }); + fn close_portal_if_current( + &self, + name: &Name, + expected: Option<&Arc>, + ) -> Result<(), Error> { + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + if expected.is_some_and(|expected| { + state + .portals + .get(name) + .is_some_and(|current| Arc::ptr_eq(current, expected)) + }) { + state.portals.remove(name); + state.portal_operations.remove(name); + } + Ok(()) } - pub fn get_portal(&self, name: &Name) -> Option> { + pub fn get_portal(&self, name: &Name) -> Result>, Error> { debug!(target: CONTEXT, client_id = self.client_id, src = "Get Portal", portal = ?name); - let portals = self.portals.read().ok()?; - - portals.get(name).cloned() + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + Ok(state.portals.get(name).cloned()) } - pub fn get_portal_statement(&self, name: &Name) -> Option> { - let portals = self.portals.read().ok()?; - let portal = portals.get(name)?; + pub fn get_portal_statement(&self, name: &Name) -> Result>, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let Some(portal) = state.portals.get(name) else { + return Ok(None); + }; debug!(target: CONTEXT, client_id = self.client_id, portal = ?portal); - match portal.as_ref() { + Ok(match portal.as_ref() { Portal::Encrypted { statement, .. } => Some(statement.clone()), Portal::Passthrough { .. } => None, - } + }) } - pub fn get_portal_session_id(&self, name: &Name) -> Option { - let portals = self.portals.read().ok()?; - let portal = portals.get(name)?; - portal.session_id() + pub fn get_portal_metrics_scope_id(&self, name: &Name) -> Result, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + Ok(state + .portals + .get(name) + .and_then(|portal| portal.session_id())) } - pub fn get_statement_for_operation(&self, operation: OperationId) -> Option> { - let operations = self.operations.read().ok()?; - let context = operations.get(&operation)?; - if let Some(statement) = &context.describe_statement { - return Some(statement.clone()); - } - match context.execute.as_ref()?.portal.as_deref()? { - Portal::Encrypted { statement, .. } => Some(statement.clone()), - Portal::Passthrough { .. } => None, + pub fn get_statement_for_operation( + &self, + operation: OperationId, + ) -> Result>, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let Some(context) = state.operations.get(&operation) else { + return Err(crate::error::ContextError::UnknownOperation.into()); + }; + if let Some(statement) = context + .describe + .as_ref() + .and_then(|describe| describe.statement.as_ref()) + { + return Ok(Some(statement.clone())); } + Ok( + match context + .execute + .as_ref() + .and_then(|execute| execute.portal.as_deref()) + { + Some(Portal::Encrypted { statement, .. }) => Some(statement.clone()), + Some(Portal::Passthrough { .. }) | None => None, + }, + ) } - pub fn get_statement_from_describe(&self, operation: OperationId) -> Option> { - self.operations + pub fn get_statement_from_describe( + &self, + operation: OperationId, + ) -> Result>, Error> { + let state = self + .protocol_state .read() - .ok()? - .get(&operation)? - .describe_statement - .clone() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let context = state + .operations + .get(&operation) + .ok_or(crate::error::ContextError::UnknownOperation)?; + Ok(context + .describe + .as_ref() + .and_then(|describe| describe.statement.clone())) } - pub fn get_portal_from_execute(&self, operation: OperationId) -> Option> { - self.operations + pub fn get_portal_from_execute( + &self, + operation: OperationId, + ) -> Result>, Error> { + let state = self + .protocol_state .read() - .ok()? - .get(&operation)? + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let context = state + .operations + .get(&operation) + .ok_or(crate::error::ContextError::UnknownOperation)?; + let execute = context .execute - .as_ref()? - .portal - .clone() + .as_ref() + .ok_or(crate::error::ContextError::OperationWithoutExecute)?; + Ok(execute.portal.clone()) } - pub fn get_execute(&self, operation: OperationId) -> Option { - let operations = self.operations.read().ok()?; - let execute_context = operations.get(&operation)?.execute.as_ref()?; + pub fn get_execute(&self, operation: OperationId) -> Result, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let operation_context = state + .operations + .get(&operation) + .ok_or(crate::error::ContextError::UnknownOperation)?; + let execute_context = operation_context + .execute + .as_ref() + .ok_or(crate::error::ContextError::OperationWithoutExecute)?; debug!(target: CONTEXT, client_id = self.client_id, msg = "Get Execute", execute = ?execute_context); - Some(execute_context.to_owned()) + Ok(Some(execute_context.to_owned())) } - pub fn get_session_metrics(&self, session_id: SessionId) -> Option { - let sessions = self.session_metrics.read().ok()?; - let session_context = sessions.get(&session_id)?; - debug!(target: CONTEXT, client_id = self.client_id, msg = "Get Session Metrics", session_metrics = ?session_context); - Some(session_context.to_owned()) + pub fn get_metrics_scope( + &self, + session_id: SessionId, + ) -> Result, Error> { + let state = self + .protocol_state + .read() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + let Some(metrics_scope) = state.statement_metrics.get(&session_id) else { + return Ok(None); + }; + debug!(target: CONTEXT, client_id = self.client_id, msg = "Get statement metrics scope", statement_metrics = ?metrics_scope); + Ok(Some(metrics_scope.to_owned())) + } + + #[cfg(test)] + pub fn active_metrics_scopes(&self) -> Result { + self.protocol_state + .read() + .map(|state| state.statement_metrics.len()) + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable.into()) } /// Returns the resolver for this connection's effective schema snapshot. @@ -645,19 +1126,17 @@ where self.schema_middleware.protocol_boundary(); } - /// Reports successful execution of the next queued statement. - pub fn schema_execution_succeeded(&self) { - self.schema_middleware.execution_succeeded(); + /// Waits for preceding schema-changing executions to resolve. + pub async fn wait_for_schema_execution(&self) { + self.schema_middleware.wait_for_ddl().await; } - /// Reports failed execution of the next queued statement. - pub fn schema_execution_failed(&self) { - self.schema_middleware.execution_failed(); + pub fn report_schema_execution_succeeded(&self) { + self.schema_middleware.execution_succeeded(); } - /// Waits for preceding schema-changing executions to resolve. - pub async fn wait_for_schema_execution(&self) { - self.schema_middleware.wait_for_ddl().await; + pub fn report_schema_execution_failed(&self) { + self.schema_middleware.execution_failed(); } /// Returns whether a schema-changing execution is awaiting a backend outcome. @@ -1050,88 +1529,115 @@ where self.upstream_tls_roots.clone() } - fn with_session_metrics_mut(&mut self, session_id: SessionId, f: F) + fn with_statement_metrics_scope_mut( + &mut self, + session_id: SessionId, + f: F, + ) -> Result<(), Error> where F: FnOnce(&mut SessionMetricsContext), { - if let Ok(mut sessions) = self.session_metrics.write() { - if let Some(session) = sessions.get_mut(&session_id) { - f(session); - } + let mut state = self + .protocol_state + .write() + .map_err(|_| crate::error::ContextError::ProtocolStateUnavailable)?; + if let Some(metrics_scope) = state.statement_metrics.get_mut(&session_id) { + f(metrics_scope); } + Ok(()) } - pub fn latest_session_id(&self) -> Option { - let sessions = self.session_metrics.read().ok()?; - sessions.keys().max().copied() - } - - /// Record parse phase duration for the session (first write wins) - pub fn record_parse_duration(&mut self, session_id: SessionId, duration: Duration) { - self.with_session_metrics_mut(session_id, |session| { - session.phase_timing.record_parse(duration); - }); + /// Record parse phase duration for the statement metrics scope (first write wins) + pub fn record_parse_duration( + &mut self, + session_id: SessionId, + duration: Duration, + ) -> Result<(), Error> { + self.with_statement_metrics_scope_mut(session_id, |metrics_scope| { + metrics_scope.phase_timing.record_parse(duration); + }) } - /// Add encrypt phase duration for the session (accumulate) - pub fn add_encrypt_duration(&mut self, session_id: SessionId, duration: Duration) { - self.with_session_metrics_mut(session_id, |session| { - session.phase_timing.add_encrypt(duration); - }); + /// Add encrypt phase duration for the statement metrics scope (accumulate) + pub fn add_encrypt_duration( + &mut self, + session_id: SessionId, + duration: Duration, + ) -> Result<(), Error> { + self.with_statement_metrics_scope_mut(session_id, |metrics_scope| { + metrics_scope.phase_timing.add_encrypt(duration); + }) } /// Add decrypt phase duration (accumulate) - pub fn add_decrypt_duration(&mut self, session_id: SessionId, duration: Duration) { - self.with_session_metrics_mut(session_id, |session| { - session.phase_timing.add_decrypt(duration); - }); + pub fn add_decrypt_duration( + &mut self, + session_id: SessionId, + duration: Duration, + ) -> Result<(), Error> { + self.with_statement_metrics_scope_mut(session_id, |metrics_scope| { + metrics_scope.phase_timing.add_decrypt(duration); + }) } - /// Update statement metadata for a session - pub fn update_statement_metadata(&mut self, session_id: SessionId, f: F) + /// Update metadata for a statement metrics scope. + pub fn update_statement_metadata(&mut self, session_id: SessionId, f: F) -> Result<(), Error> where F: FnOnce(&mut StatementMetadata), { - self.with_session_metrics_mut(session_id, |session| { - f(&mut session.metadata); - }); + self.with_statement_metrics_scope_mut(session_id, |metrics_scope| { + f(&mut metrics_scope.metadata); + }) } - /// Update statement metadata if session ID is present, no-op otherwise. - pub fn with_session(&mut self, session_id: Option, f: F) + /// Update statement metadata if a metrics scope ID is present, no-op otherwise. + pub fn with_metrics_scope( + &mut self, + session_id: Option, + f: F, + ) -> Result<(), Error> where F: FnOnce(&mut SessionMetricsContext), { if let Some(sid) = session_id { - self.with_session_metrics_mut(sid, f); + self.with_statement_metrics_scope_mut(sid, f)?; } + Ok(()) } - /// Add decrypt phase duration for the current execute session (if any) - pub fn add_decrypt_duration_for_execute(&mut self, operation: OperationId, duration: Duration) { + /// Add decrypt phase duration for the current execution metrics scope (if any) + pub fn add_decrypt_duration_for_execute( + &mut self, + operation: OperationId, + duration: Duration, + ) -> Result<(), Error> { let session_id = self - .get_execute(operation) + .get_execute(operation)? .and_then(|execute| execute.session_id()); if let Some(session_id) = session_id { - self.add_decrypt_duration(session_id, duration); + self.add_decrypt_duration(session_id, duration)?; } + Ok(()) } } #[cfg(test)] mod tests { - use super::{Context, KeysetIdentifier, Portal, Statement}; + use super::{ + ConnectionProtocolState, Context, ExecuteContext, ExecutionOutcome, KeysetIdentifier, + OperationContext, Portal, SessionId, SessionMetricsContext, Statement, + }; use crate::{ config::LogConfig, error::Error, log, - postgresql::{rewrite::Name, Column}, + postgresql::{rewrite::Name, test_operation_id as operation_id, Column}, proxy::{EncryptConfig, EncryptionService}, TandemConfig, }; use cipherstash_client::IdentifiedBy; use eql_mapper::Schema; - use pg_proto::TransactionStatus; + use pg_proto::{Describe, DescribeTarget, TransactionStatus}; use sqltk::parser::{dialect::PostgreSqlDialect, parser::Parser}; use std::sync::Arc; use tokio::sync::mpsc; @@ -1168,6 +1674,280 @@ mod tests { } } + #[test] + fn suspended_execution_retains_its_metrics_scope_until_resumed_completion() { + let session_id = SessionId(1); + let mut state = ConnectionProtocolState::::default(); + state + .statement_metrics + .insert(session_id, SessionMetricsContext::new(session_id)); + state.operations.insert( + 1, + OperationContext { + execute: Some(ExecuteContext::new(Name::new(), None, Some(session_id))), + ..OperationContext::default() + }, + ); + + let suspended = state + .finish_execution(&1, ExecutionOutcome::Suspended) + .unwrap(); + + assert!(suspended.finished_metrics.is_none()); + assert!(state.statement_metrics.contains_key(&session_id)); + assert!(!state.operations.contains_key(&1)); + + state.operations.insert( + 2, + OperationContext { + execute: Some(ExecuteContext::new(Name::new(), None, Some(session_id))), + ..OperationContext::default() + }, + ); + let completed = state + .finish_execution(&2, ExecutionOutcome::Completed) + .unwrap(); + + assert_eq!( + completed.finished_metrics.map(|metrics| metrics.id()), + Some(session_id) + ); + assert!(!state.statement_metrics.contains_key(&session_id)); + assert!(!state.operations.contains_key(&2)); + } + + #[test] + fn each_extended_execution_records_its_own_statement_metrics() { + let mut context = create_context(); + let statement = Name::from("statement"); + let portal = Name::from("portal"); + let template_scope = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(statement, template_scope) + .unwrap(); + context + .add_portal( + operation_id(), + portal.clone(), + Portal::passthrough(Some(template_scope)), + ) + .unwrap(); + let operation = operation_id(); + + context.set_execute_for_portal(operation, portal).unwrap(); + + let execution_scope = context + .get_execute(operation) + .unwrap() + .unwrap() + .session_id() + .unwrap(); + assert!( + context + .get_metrics_scope(execution_scope) + .unwrap() + .unwrap() + .records_statement_metrics + ); + } + + #[test] + fn execution_transition_rejects_stale_and_non_execute_operations() { + let mut state = ConnectionProtocolState::::default(); + + assert!(matches!( + state.finish_execution(&1, ExecutionOutcome::Completed), + Err(crate::error::ContextError::UnknownOperation) + )); + + state.operations.insert(1, OperationContext::default()); + assert!(matches!( + state.finish_execution(&1, ExecutionOutcome::Completed), + Err(crate::error::ContextError::OperationWithoutExecute) + )); + } + + #[test] + fn failed_execution_rejects_an_unknown_operation() { + let context = create_context(); + + assert!(matches!( + context.finish_execution(operation_id(), ExecutionOutcome::Failed), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[test] + fn adding_a_statement_fails_when_protocol_state_is_unavailable() { + let context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.add_statement( + Name::new(), + Statement::new(vec![], vec![], vec![], vec![], vec![]), + ), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn adding_a_portal_fails_when_protocol_state_is_unavailable() { + let context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.add_portal(operation_id(), Name::new(), Portal::passthrough(None)), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn associating_statement_metrics_fails_when_protocol_state_is_unavailable() { + let mut context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.set_statement_metrics_scope(Name::new(), SessionId(1)), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn starting_metrics_fails_when_protocol_state_is_unavailable() { + let mut context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.start_metrics_scope(), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn finishing_metrics_fails_when_protocol_state_is_unavailable() { + let mut context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.finish_metrics_scope(None), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn reading_protocol_metadata_fails_when_protocol_state_is_unavailable() { + let context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.get_statement(&Name::from("statement")), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn closing_a_statement_fails_when_protocol_state_is_unavailable() { + let context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.close_statement(&Name::from("statement")), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn completing_an_unknown_describe_rejects_the_stale_operation() { + let context = create_context(); + + assert!(matches!( + context.complete_describe(operation_id()), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[test] + fn completing_an_operation_without_describe_rejects_and_removes_it() { + let mut context = create_context(); + let operation = operation_id(); + context + .set_operation_error( + operation, + crate::postgresql::diagnostics::invalid_sql_statement("proxy error".to_owned()), + ) + .unwrap(); + + assert!(matches!( + context.complete_describe(operation), + Err(Error::Context(crate::error::ContextError::UnknownDescribe)) + )); + assert!(matches!( + context.complete_describe(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[test] + fn row_description_for_an_execution_does_not_require_describe_state() { + let mut context = create_context(); + let operation = operation_id(); + context.set_execute(operation, Name::new(), None).unwrap(); + + context.complete_describe(operation).unwrap(); + + assert!(context.get_execute(operation).unwrap().is_some()); + } + fn create_context() -> Context { let client_id = 1; let config = Arc::new(TandemConfig::for_testing()); @@ -1215,64 +1995,163 @@ mod tests { } #[test] - fn replacing_a_statement_does_not_finish_an_overlapping_execution_session() { + fn replacing_a_statement_does_not_finish_an_overlapping_execution_metrics_scope() { let mut context = create_context(); let name = Name::default(); - let session_id = context.start_session(); - context.set_statement_session(name.clone(), session_id); - context.add_portal( - Name::from("active_portal"), - Portal::passthrough(Some(session_id)), - ); + let session_id = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(name.clone(), session_id) + .unwrap(); + context + .add_portal( + operation_id(), + Name::from("active_portal"), + Portal::passthrough(Some(session_id)), + ) + .unwrap(); - context.close_statement(&name); + context.close_statement(&name).unwrap(); - assert!(context.get_session_metrics(session_id).is_some()); - assert!(context.get_statement_session(&name).is_none()); + assert!(context.get_metrics_scope(session_id).unwrap().is_some()); + assert!(context + .get_statement_metrics_scope(&name) + .unwrap() + .is_none()); } #[test] - fn closing_an_unreferenced_statement_finishes_its_metrics_session() { + fn replacing_a_statement_does_not_finish_a_suspended_execution_scope() { + let mut context = create_context(); + let statement = Name::from("statement"); + let scope = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(statement.clone(), scope) + .unwrap(); + let operation = operation_id(); + context + .set_execute(operation, Name::from("portal"), Some(scope)) + .unwrap(); + context + .finish_execution(operation, ExecutionOutcome::Suspended) + .unwrap(); + + context.close_statement(&statement).unwrap(); + + assert!(context.get_metrics_scope(scope).unwrap().is_some()); + } + + #[test] + fn readiness_releases_reparsed_statement_metrics_after_their_portal_is_destroyed() { + let mut context = create_context(); + let statement = Name::from("statement"); + let scope = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(statement.clone(), scope) + .unwrap(); + context + .add_portal( + operation_id(), + Name::from("named_portal"), + Portal::passthrough(Some(scope)), + ) + .unwrap(); + + context.close_statement(&statement).unwrap(); + assert!(context.get_metrics_scope(scope).unwrap().is_some()); + + context + .ready_for_query(TransactionStatus::Idle, Some(operation_id())) + .unwrap(); + + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + } + + #[test] + fn closing_an_unexecuted_statement_does_not_record_statement_metrics() { let mut context = create_context(); let name = Name::default(); - let session_id = context.start_session(); - context.set_statement_session(name.clone(), session_id); + let session_id = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(name.clone(), session_id) + .unwrap(); + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); + let handle = recorder.handle(); - context.close_statement_explicit(&name); + metrics::with_local_recorder(&recorder, || { + context.close_statement_explicit(&name).unwrap(); + }); - assert!(context.get_session_metrics(session_id).is_none()); + assert!(context.get_metrics_scope(session_id).unwrap().is_none()); + let rendered = handle.render(); + assert!( + !rendered.contains("cipherstash_proxy_statements_session_duration_seconds_count"), + "{rendered}" + ); } #[test] fn replacing_a_statement_retains_existing_portals() { - let mut context = create_context(); + let context = create_context(); let closed_statement_name = Name::from("closed_statement"); let retained_statement_name = Name::from("retained_statement"); - context.add_statement(closed_statement_name.clone(), statement()); - context.add_statement(retained_statement_name.clone(), statement()); + context + .add_statement(closed_statement_name.clone(), statement()) + .unwrap(); + context + .add_statement(retained_statement_name.clone(), statement()) + .unwrap(); - let closed_statement = context.get_statement(&closed_statement_name).unwrap(); - let retained_statement = context.get_statement(&retained_statement_name).unwrap(); + let closed_statement = context + .get_statement(&closed_statement_name) + .unwrap() + .unwrap(); + let retained_statement = context + .get_statement(&retained_statement_name) + .unwrap() + .unwrap(); let closed_portal_name = Name::from("differently_named_portal"); let retained_portal_name = Name::from("retained_portal"); let passthrough_portal_name = Name::from("passthrough_portal"); - context.add_portal(closed_portal_name.clone(), portal(&closed_statement)); - context.add_portal(retained_portal_name.clone(), portal(&retained_statement)); - context.add_portal(passthrough_portal_name.clone(), Portal::passthrough(None)); + context + .add_portal( + operation_id(), + closed_portal_name.clone(), + portal(&closed_statement), + ) + .unwrap(); + context + .add_portal( + operation_id(), + retained_portal_name.clone(), + portal(&retained_statement), + ) + .unwrap(); + context + .add_portal( + operation_id(), + passthrough_portal_name.clone(), + Portal::passthrough(None), + ) + .unwrap(); - context.close_statement(&closed_statement_name); + context.close_statement(&closed_statement_name).unwrap(); - assert!(context.get_portal(&closed_portal_name).is_some()); - assert!(context.get_portal(&retained_portal_name).is_some()); - assert!(context.get_portal(&passthrough_portal_name).is_some()); + assert!(context.get_portal(&closed_portal_name).unwrap().is_some()); + assert!(context.get_portal(&retained_portal_name).unwrap().is_some()); + assert!(context + .get_portal(&passthrough_portal_name) + .unwrap() + .is_some()); } #[test] fn transaction_status_tracks_backend_ready_state() { - let mut context = create_context(); + let context = create_context(); assert_eq!(context.transaction_status(), TransactionStatus::Idle); - context.set_transaction_status(TransactionStatus::InTransaction); + context + .ready_for_query(TransactionStatus::InTransaction, None) + .unwrap(); assert_eq!( context.transaction_status(), @@ -1280,6 +2159,336 @@ mod tests { ); } + #[test] + fn idle_readiness_finishes_abandoned_suspended_execution_metrics() { + let mut context = create_context(); + let portal = Name::from("limited_portal"); + let template_scope = context.start_metrics_scope().unwrap(); + context + .add_portal( + operation_id(), + portal.clone(), + Portal::passthrough(Some(template_scope)), + ) + .unwrap(); + let operation = operation_id(); + context.set_execute_for_portal(operation, portal).unwrap(); + let execution_scope = context + .get_execute(operation) + .unwrap() + .and_then(|execute| execute.session_id()) + .unwrap(); + context + .finish_execution(operation, ExecutionOutcome::Suspended) + .unwrap(); + + context + .ready_for_query(TransactionStatus::Idle, Some(operation_id())) + .unwrap(); + + assert!(context + .get_metrics_scope(execution_scope) + .unwrap() + .is_none()); + } + + #[test] + fn idle_readiness_finishes_simple_query_execution() { + let mut context = create_context(); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_simple_query_execute_until_ready(operation, Name::new(), Some(scope)) + .unwrap(); + context + .finish_execution(operation, ExecutionOutcome::Completed) + .unwrap(); + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); + let handle = recorder.handle(); + + metrics::with_local_recorder(&recorder, || { + context + .ready_for_query(TransactionStatus::Idle, Some(operation)) + .unwrap(); + }); + + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + assert!(handle + .render() + .contains("cipherstash_proxy_statements_execution_duration_seconds_count")); + } + + #[test] + fn readiness_for_one_query_preserves_later_pipelined_work() { + let mut context = create_context(); + let query_a = operation_id(); + let query_a_scope = context.start_metrics_scope().unwrap(); + context + .set_simple_query_execute_until_ready(query_a, Name::new(), Some(query_a_scope)) + .unwrap(); + + let portal_b = Name::from("pipeline_b"); + let bind_b = operation_id(); + context + .add_portal(bind_b, portal_b.clone(), Portal::passthrough(None)) + .unwrap(); + let describe_b = operation_id(); + context + .set_describe( + describe_b, + Describe { + target: DescribeTarget::Portal, + name: portal_b.clone(), + }, + ) + .unwrap(); + let execute_b = operation_id(); + context + .set_execute_for_portal(execute_b, portal_b.clone()) + .unwrap(); + + context + .ready_for_query(TransactionStatus::Idle, Some(query_a)) + .unwrap(); + + assert!(context.get_metrics_scope(query_a_scope).unwrap().is_none()); + assert!(context.get_portal(&portal_b).unwrap().is_some()); + assert!(context.complete_describe(describe_b).is_ok()); + assert!(context.get_execute(execute_b).unwrap().is_some()); + } + + #[test] + fn transaction_readiness_finishes_one_query_and_preserves_later_pipelined_work() { + let mut context = create_context(); + let query_a = operation_id(); + let query_a_scope = context.start_metrics_scope().unwrap(); + context + .set_simple_query_execute_until_ready(query_a, Name::new(), Some(query_a_scope)) + .unwrap(); + + let query_b = operation_id(); + let query_b_scope = context.start_metrics_scope().unwrap(); + context + .set_simple_query_execute_until_ready(query_b, Name::new(), Some(query_b_scope)) + .unwrap(); + + context + .ready_for_query(TransactionStatus::InTransaction, Some(query_a)) + .unwrap(); + + assert!(matches!( + context.get_execute(query_a), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context.get_metrics_scope(query_a_scope).unwrap().is_none()); + assert!(context.get_execute(query_b).unwrap().is_some()); + assert!(context.get_metrics_scope(query_b_scope).unwrap().is_some()); + } + + #[test] + fn closing_a_suspended_portal_finishes_its_execution_metrics() { + let mut context = create_context(); + let portal = Name::from("limited_portal"); + let template_scope = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(Name::from("statement"), template_scope) + .unwrap(); + context + .add_portal( + operation_id(), + portal.clone(), + Portal::passthrough(Some(template_scope)), + ) + .unwrap(); + let operation = operation_id(); + context + .set_execute_for_portal(operation, portal.clone()) + .unwrap(); + let execution_scope = context + .get_execute(operation) + .unwrap() + .and_then(|execute| execute.session_id()) + .unwrap(); + context + .finish_execution(operation, ExecutionOutcome::Suspended) + .unwrap(); + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); + let handle = recorder.handle(); + + metrics::with_local_recorder(&recorder, || { + context.close_portal(&portal).unwrap(); + }); + + assert!(context + .get_metrics_scope(execution_scope) + .unwrap() + .is_none()); + let rendered = handle.render(); + let count = rendered + .lines() + .find(|line| { + line.starts_with("cipherstash_proxy_statements_session_duration_seconds_count") + }) + .unwrap(); + assert!(count.ends_with(" 1"), "{rendered}"); + } + + #[test] + fn rebinding_a_suspended_portal_starts_a_new_execution_occurrence() { + let mut context = create_context(); + let portal_name = Name::new(); + let first_template = context.start_metrics_scope().unwrap(); + context + .add_portal( + operation_id(), + portal_name.clone(), + Portal::passthrough(Some(first_template)), + ) + .unwrap(); + let operation = operation_id(); + context + .set_execute_for_portal(operation, portal_name.clone()) + .unwrap(); + let first_execution = context + .get_execute(operation) + .unwrap() + .unwrap() + .session_id() + .unwrap(); + context + .finish_execution(operation, ExecutionOutcome::Suspended) + .unwrap(); + + let rebound_template = context.start_metrics_scope().unwrap(); + context + .add_portal( + operation_id(), + portal_name.clone(), + Portal::passthrough(Some(rebound_template)), + ) + .unwrap(); + context + .set_execute_for_portal(operation, portal_name) + .unwrap(); + + assert_eq!( + context + .get_portal_from_execute(operation) + .unwrap() + .and_then(|portal| portal.session_id()), + Some(rebound_template) + ); + assert!(context + .get_metrics_scope(first_execution) + .unwrap() + .is_none()); + } + + #[test] + fn closing_a_portal_fails_when_protocol_state_is_unavailable() { + let context = create_context(); + let protocol_state = context.protocol_state.clone(); + let _ = std::thread::spawn(move || { + let _guard = protocol_state.write().unwrap(); + panic!("poison protocol state"); + }) + .join(); + + assert!(matches!( + context.close_portal(&Name::new()), + Err(Error::Context( + crate::error::ContextError::ProtocolStateUnavailable + )) + )); + } + + #[test] + fn discarding_an_unfinished_operation_releases_its_execution_metrics() { + let mut context = create_context(); + let portal = Name::from("skipped_portal"); + let template_scope = context.start_metrics_scope().unwrap(); + context + .add_portal( + operation_id(), + portal.clone(), + Portal::passthrough(Some(template_scope)), + ) + .unwrap(); + let operation = operation_id(); + context.set_execute_for_portal(operation, portal).unwrap(); + let execution_scope = context + .get_execute(operation) + .unwrap() + .and_then(|execute| execute.session_id()) + .unwrap(); + + context.discard_operation(operation).unwrap(); + + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context + .get_metrics_scope(execution_scope) + .unwrap() + .is_none()); + } + + #[test] + fn discarding_an_operation_preserves_a_metrics_scope_referenced_by_another_operation() { + let mut context = create_context(); + let scope = context.start_metrics_scope().unwrap(); + let first_operation = operation_id(); + let second_operation = operation_id(); + context + .set_execute(first_operation, Name::from("first"), Some(scope)) + .unwrap(); + context + .set_execute(second_operation, Name::from("second"), Some(scope)) + .unwrap(); + + context.discard_operation(first_operation).unwrap(); + + assert!(context.get_metrics_scope(scope).unwrap().is_some()); + + context.discard_operation(second_operation).unwrap(); + + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + } + + #[test] + fn simple_query_execution_finishes_only_at_readiness() { + let mut context = create_context(); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_simple_query_execute_until_ready(operation, Name::new(), Some(scope)) + .unwrap(); + + context + .finish_execution(operation, ExecutionOutcome::Completed) + .unwrap(); + assert!(context.get_execute(operation).unwrap().is_some()); + + context + .finish_execution(operation, ExecutionOutcome::Completed) + .unwrap(); + assert!(context.get_execute(operation).unwrap().is_some()); + + context + .ready_for_query(TransactionStatus::Idle, Some(operation)) + .unwrap(); + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + } + fn get_statement(portal: Arc) -> Arc { match portal.as_ref() { Portal::Encrypted { statement, .. } => statement.clone(), @@ -1293,28 +2502,36 @@ mod tests { pub fn add_and_close_portals() { log::init(LogConfig::default()); - let mut context = create_context(); + let context = create_context(); // Create multiple statements let statement_name_1 = Name::from("statement_1"); let statement_name_2 = Name::from("statement_2"); // Add statements to context - context.add_statement(statement_name_1.clone(), statement()); - context.add_statement(statement_name_2.clone(), statement()); + context + .add_statement(statement_name_1.clone(), statement()) + .unwrap(); + context + .add_statement(statement_name_2.clone(), statement()) + .unwrap(); let portal_name = Name::from("portal"); - let statement_1 = context.get_statement(&statement_name_1).unwrap(); - context.add_portal(portal_name.clone(), portal(&statement_1)); + let statement_1 = context.get_statement(&statement_name_1).unwrap().unwrap(); + context + .add_portal(operation_id(), portal_name.clone(), portal(&statement_1)) + .unwrap(); - let statement_2 = context.get_statement(&statement_name_2).unwrap(); - context.add_portal(portal_name.clone(), portal(&statement_2)); + let statement_2 = context.get_statement(&statement_name_2).unwrap().unwrap(); + context + .add_portal(operation_id(), portal_name.clone(), portal(&statement_2)) + .unwrap(); - let portal = context.get_portal(&portal_name).unwrap(); + let portal = context.get_portal(&portal_name).unwrap().unwrap(); assert_eq!(statement_2, get_statement(portal)); - context.close_portal(&portal_name); - assert!(context.get_portal(&portal_name).is_none()); + context.close_portal(&portal_name).unwrap(); + assert!(context.get_portal(&portal_name).unwrap().is_none()); } fn parse_statement(sql: &str) -> sqltk::parser::ast::Statement { diff --git a/packages/cipherstash-proxy/src/postgresql/driver.rs b/packages/cipherstash-proxy/src/postgresql/driver.rs index 3ccbe305..cc8c43c1 100644 --- a/packages/cipherstash-proxy/src/postgresql/driver.rs +++ b/packages/cipherstash-proxy/src/postgresql/driver.rs @@ -6,11 +6,12 @@ use crate::{ tls, }; use pg_proto::{ - BackendForwarding, BoundedPipeline, CancellationPolicy, Client, ClientTlsConfig, + BackendForwarding, BoundedPipeline, CancelKey, CancellationPolicy, Client, ClientTlsConfig, ClientTlsPolicy, ClientTlsProvider, ConnectTarget, ForwardedMessage, FrontendForwarding, - FrontendMessage, InMemoryCancellationRegistry, InitialServerContext, Intermediary, Server, - ServerIdentity, ServerIdentityProvider, ServerTlsPolicy, SslMode, StartupParameters, - StartupRouteResolver, StaticClientCredentials, StaticMd5ServerCredentials, + FrontendMessage, InMemoryCancellationRegistry, InitialServerContext, Intermediary, + IntermediaryCancellationRegistry, Server, ServerIdentity, ServerIdentityProvider, + ServerTlsPolicy, SslMode, StartupParameters, StartupRouteResolver, StaticClientCredentials, + StaticMd5ServerCredentials, }; use std::{ convert::Infallible, @@ -24,6 +25,25 @@ use tracing::info; static CANCELLATION_REGISTRY: LazyLock = LazyLock::new(InMemoryCancellationRegistry::default); +struct CancellationRegistrationGuard { + registry: InMemoryCancellationRegistry, + key: Option, +} + +impl CancellationRegistrationGuard { + fn new(registry: InMemoryCancellationRegistry, key: Option) -> Self { + Self { registry, key } + } +} + +impl Drop for CancellationRegistrationGuard { + fn drop(&mut self) { + if let Some(key) = self.key.take() { + let _ = self.registry.detach(&key); + } + } +} + #[derive(Clone)] struct Route(String); @@ -102,6 +122,10 @@ where pg_proto::IntermediaryAccept::Session(session) => session, pg_proto::IntermediaryAccept::CancellationForwarded => return Ok(()), }; + let _cancellation_guard = CancellationRegistrationGuard::new( + CANCELLATION_REGISTRY.clone(), + session.cancellation_key().cloned(), + ); let result = async { 'session: loop { let forward = session.forward_next(); @@ -216,9 +240,7 @@ where Ok(()) } .await; - finish_connection(result, || { - let _ = session.detach_cancellation(); - }) + result }}; } @@ -295,14 +317,6 @@ fn invalid_data(error: impl std::fmt::Display) -> Error { std::io::Error::new(std::io::ErrorKind::InvalidData, error.to_string()).into() } -fn finish_connection( - result: Result, - detach_cancellation: impl FnOnce(), -) -> Result { - detach_cancellation(); - result -} - fn upstream_ssl_mode(config: &crate::TandemConfig) -> SslMode { if config.database.with_tls_verification { SslMode::VerifyFull @@ -317,13 +331,10 @@ fn upstream_ssl_mode(config: &crate::TandemConfig) -> SslMode { mod tests { use super::*; use bytes::Bytes; - use pg_proto::{ - CancelKey, CancellationRoute, InMemoryCancellationRegistryError, - IntermediaryCancellationRegistry, - }; + use pg_proto::{CancellationRoute, InMemoryCancellationRegistryError}; #[test] - fn connection_exit_releases_its_cancellation_key() { + fn dropping_a_connection_guard_releases_its_cancellation_key() { for result in [Ok(()), Err("connection failed")] { let registry = InMemoryCancellationRegistry::default(); let key = CancelKey { @@ -337,11 +348,11 @@ mod tests { Err(InMemoryCancellationRegistryError::DuplicateKey) ); - let returned = finish_connection(result, || { - registry.detach(&key).unwrap(); - }); - - assert_eq!(returned, result); + { + let _guard = + CancellationRegistrationGuard::new(registry.clone(), Some(key.clone())); + let _ = result; + } assert_eq!(registry.register(route), Ok(key)); } } diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs index 3b3a2466..58286321 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs @@ -1,4 +1,4 @@ -use super::super::context::Context; +use super::super::context::{Context, ExecutionOutcome}; use super::super::data::to_sql; use super::super::error_handler::PostgreSqlErrorHandler; use super::super::rewrite::UNSPECIFIED_TYPE_OID; @@ -7,6 +7,7 @@ use crate::error::{EncryptError, Error, ProtocolError}; use crate::log::{CONTEXT, DEVELOPMENT, MAPPER, PROTOCOL}; use crate::postgresql::context::Portal; use crate::postgresql::rewrite::data_row; +use crate::postgresql::OperationId; use crate::prometheus::{ DECRYPTED_VALUES_TOTAL, DECRYPTION_DURATION_SECONDS, DECRYPTION_ERROR_TOTAL, DECRYPTION_REQUESTS_TOTAL, ROWS_ENCRYPTED_TOTAL, ROWS_PASSTHROUGH_TOTAL, ROWS_TOTAL, @@ -14,7 +15,7 @@ use crate::prometheus::{ use crate::proxy::EncryptionService; use crate::EqlCiphertext; use metrics::{counter, histogram}; -use pg_proto::{BackendMessage, BackendMiddlewareOutput, DataRow, OperationId}; +use pg_proto::{BackendMessage, BackendMiddlewareOutput, DataRow, DiagnosticResponse}; use std::time::Instant; use tracing::{debug, error, info}; @@ -51,7 +52,7 @@ use tracing::{debug, error, info}; /// /// DataRow messages containing encrypted data are buffered to enable batch decryption: /// - Buffer fills up to a configurable capacity -/// - Flush occurs on buffer full, session end, or non-DataRow message +/// - Flush occurs on buffer full, connection end, or non-DataRow message /// - Batching reduces encryption API round-trips and improves performance /// /// # Message Types Handled @@ -74,11 +75,36 @@ pub struct Backend { const MAX_ENCRYPTED_ROWS: usize = 4096; const MAX_ENCRYPTED_ROW_BYTES: usize = 64 * 1024 * 1024; +fn execution_outcome(message: &BackendMessage) -> Option { + match message { + BackendMessage::CommandComplete(_) | BackendMessage::EmptyQueryResponse => { + Some(ExecutionOutcome::Completed) + } + BackendMessage::PortalSuspended => Some(ExecutionOutcome::Suspended), + _ => None, + } +} + impl Backend { + fn finish_execution( + &self, + operation: Option, + outcome: ExecutionOutcome, + ) -> Result, Error> { + let Some(operation) = operation else { + return Ok(None); + }; + self.context.finish_execution(operation, outcome) + } + fn error_response(&self, err: Error) -> BackendMessage { BackendMessage::ErrorResponse(self.error_to_response(err)) } + fn complete_describe(&self, operation: OperationId) -> Result<(), Error> { + self.context.complete_describe(operation) + } + fn decryption_failure(&mut self, err: Error) -> BackendMiddlewareOutput { self.encrypted_rows.clear(); self.encrypted_rows_operation = None; @@ -120,13 +146,27 @@ impl Backend { ) -> Result { let mut outbound_message = protocol_message.clone(); - if matches!(protocol_message, BackendMessage::ErrorResponse(_)) { - if let Some(response) = operation.and_then(|id| self.context.take_operation_error(id)) { - outbound_message = BackendMessage::ErrorResponse(response); + if matches!( + protocol_message, + BackendMessage::ParseComplete | BackendMessage::BindComplete + ) { + if let Some(operation) = operation { + self.context.complete_non_execution(operation)?; } } if self.context.is_passthrough() { + match &protocol_message { + BackendMessage::CommandComplete(_) + | BackendMessage::EmptyQueryResponse + | BackendMessage::PortalSuspended => { + self.context.report_schema_execution_succeeded(); + } + BackendMessage::ErrorResponse(_) => { + self.context.report_schema_execution_failed(); + } + _ => {} + } debug!(target: DEVELOPMENT, client_id = self.context.client_id, msg = "Passthrough enabled" @@ -138,37 +178,33 @@ impl Backend { // client opening its next connection after ReadyForQuery observes // the newly loaded schema and encrypt configuration. if let BackendMessage::ReadyForQuery(status) = &protocol_message { - self.handle_ready_for_query(*status).await?; + self.handle_ready_for_query(operation, *status).await?; } // CipherStash metadata is operation-keyed even in passthrough mode, // and must be released when pg-proto identifies its terminal response. - match protocol_message { - BackendMessage::CommandComplete(_) - | BackendMessage::EmptyQueryResponse - | BackendMessage::PortalSuspended => { - self.context.schema_execution_succeeded(); - if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); - } + if let Some(outcome) = execution_outcome(&protocol_message) { + if let Some(response) = self.finish_execution(operation, outcome)? { + outbound_message = BackendMessage::ErrorResponse(response); } - BackendMessage::ErrorResponse(_) => { - self.context.schema_execution_failed(); - if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); - } + } else if matches!(protocol_message, BackendMessage::ErrorResponse(_)) + && operation.is_some() + { + if let Some(response) = + self.finish_execution(operation, ExecutionOutcome::Failed)? + { + outbound_message = BackendMessage::ErrorResponse(response); } - BackendMessage::RowDescription(_) | BackendMessage::NoData => { - if let Some(operation) = operation { - self.context.complete_describe(operation); + } else { + match protocol_message { + BackendMessage::RowDescription(_) | BackendMessage::NoData => { + if let Some(operation) = operation { + self.complete_describe(operation)?; + } } + BackendMessage::ReadyForQuery(_) => {} + _ => {} } - BackendMessage::ReadyForQuery(_) => {} - _ => {} } return Ok(BackendMiddlewareOutput::Forward(outbound_message)); @@ -179,24 +215,20 @@ impl Backend { BackendMessage::CommandComplete(_) | BackendMessage::EmptyQueryResponse | BackendMessage::PortalSuspended => { - self.context.schema_execution_succeeded(); + self.context.report_schema_execution_succeeded(); if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); + self.context.discard_operation(operation)?; } } BackendMessage::ErrorResponse(_) => { - self.context.schema_execution_failed(); + self.context.report_schema_execution_failed(); if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); + self.context.discard_operation(operation)?; } } BackendMessage::ReadyForQuery(status) => { self.discard_execution = false; - self.handle_ready_for_query(status).await?; + self.handle_ready_for_query(operation, status).await?; return Ok(BackendMiddlewareOutput::Forward( BackendMessage::ReadyForQuery(status), )); @@ -209,15 +241,12 @@ impl Backend { let mut prefix = if !matches!(protocol_message, BackendMessage::DataRow(_)) && !self.encrypted_rows.is_empty() { + let encrypted_rows_operation = self.encrypted_rows_operation; match self.flush_encrypted_rows().await { Ok(messages) => messages, Err(err) => { self.discard_execution = true; - if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); - } + self.finish_execution(encrypted_rows_operation, ExecutionOutcome::Failed)?; return Ok(self.decryption_failure(err)); } } @@ -225,6 +254,18 @@ impl Backend { Vec::new() }; + match &protocol_message { + BackendMessage::CommandComplete(_) + | BackendMessage::EmptyQueryResponse + | BackendMessage::PortalSuspended => { + self.context.report_schema_execution_succeeded(); + } + BackendMessage::ErrorResponse(_) => { + self.context.report_schema_execution_failed(); + } + _ => {} + } + let keyset_id = self.context.keyset_identifier(); debug!(target: CONTEXT, client_id = ?self.context.client_id, ?keyset_id); @@ -263,32 +304,33 @@ impl Backend { } return match self.flush_encrypted_rows().await { Ok(messages) => Ok(BackendMiddlewareOutput::Expand(messages)), - Err(err) => Ok(self.decryption_failure(err)), + Err(err) => { + self.finish_execution(Some(operation), ExecutionOutcome::Failed)?; + Ok(self.decryption_failure(err)) + } }; } } // Execute phase is always terminated by the appearance of exactly one of these messages: - // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), ErrorResponse, or PortalSuspended. - BackendMessage::CommandComplete(_) + // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), or PortalSuspended. + terminal @ (BackendMessage::CommandComplete(_) | BackendMessage::EmptyQueryResponse - | BackendMessage::PortalSuspended => { - debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "CommandComplete | EmptyQueryResponse | PortalSuspended"); - - if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.discard_operation(operation); - self.context.finish_session(session); + | BackendMessage::PortalSuspended) => { + let outcome = execution_outcome(&terminal).unwrap(); + debug!(target: PROTOCOL, client_id = self.context.client_id, ?outcome, msg = "Execute outcome"); + if let Some(response) = self.finish_execution(operation, outcome)? { + outbound_message = BackendMessage::ErrorResponse(response); } - self.context.schema_execution_succeeded(); } - BackendMessage::ErrorResponse(ref response) => { - self.context.schema_execution_failed(); - self.error_response_handler(response); - - if let Some(operation) = operation { - let session = self.context.complete_execution(operation); - self.context.finish_session(session); + BackendMessage::ErrorResponse(response) => { + self.error_response_handler(&response); + if operation.is_some() { + if let Some(response) = + self.finish_execution(operation, ExecutionOutcome::Failed)? + { + outbound_message = BackendMessage::ErrorResponse(response); + } } } // Describe with Target:Statement @@ -309,14 +351,14 @@ impl Backend { outbound_message = message; } if let Some(operation) = operation { - self.context.complete_describe(operation); + self.complete_describe(operation)?; } } // Describe with Target:Statement or Target::Portal // If the statement returns no rows, NoData is returned instead of a RowDescription BackendMessage::NoData => { if let Some(operation) = operation { - self.context.complete_describe(operation); + self.complete_describe(operation)?; } } // Reload for SompleQuery flow @@ -327,7 +369,7 @@ impl Backend { client_id = self.context.client_id, msg = "ReadyForQuery" ); - self.handle_ready_for_query(status).await?; + self.handle_ready_for_query(operation, status).await?; } _ => { @@ -351,9 +393,10 @@ impl Backend { /// connection-local transaction state from the same authoritative boundary. async fn handle_ready_for_query( &mut self, + operation: Option, status: pg_proto::TransactionStatus, ) -> Result<(), Error> { - self.context.set_transaction_status(status); + self.context.ready_for_query(status, operation)?; if status == pg_proto::TransactionStatus::Idle { self.context.publish_schema_if_changed().await?; } @@ -462,8 +505,10 @@ impl Backend { let mut rows = std::mem::take(&mut self.encrypted_rows); self.encrypted_rows_bytes = 0; - let portal = - operation.and_then(|operation| self.context.get_portal_from_execute(operation)); + let portal = operation + .map(|operation| self.context.get_portal_from_execute(operation)) + .transpose()? + .flatten(); let portal = match portal.as_deref() { Some(Portal::Encrypted { .. }) => portal.unwrap(), _ => { @@ -513,7 +558,7 @@ impl Backend { // Always record for slow-statement diagnostics if let Some(operation) = operation { self.context - .add_decrypt_duration_for_execute(operation, duration); + .add_decrypt_duration_for_execute(operation, duration)?; } // Prometheus metrics remain gated @@ -591,9 +636,11 @@ impl Backend { ) -> Result, Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ParamDescription = ?description); - if let Some(statement) = - operation.and_then(|operation| self.context.get_statement_from_describe(operation)) - { + let statement = operation + .map(|operation| self.context.get_statement_from_describe(operation)) + .transpose()? + .flatten(); + if let Some(statement) = statement { // Describe the params the CLIENT wrote, not the ones PostgreSQL was // sent. A rewrite may have fused or dropped params, in which case // the server's description is both shorter than and shifted from @@ -649,9 +696,11 @@ impl Backend { ) -> Result, Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, RowDescription = ?description); - if let Some(statement) = - operation.and_then(|operation| self.context.get_statement_for_operation(operation)) - { + let statement = operation + .map(|operation| self.context.get_statement_for_operation(operation)) + .transpose()? + .flatten(); + if let Some(statement) = statement { let projection_types = statement .projection_columns .iter() @@ -713,10 +762,11 @@ impl Backend { /// track proxy performance and encryption usage patterns. async fn data_row_handler(&mut self, operation: Option) -> Result { counter!(ROWS_TOTAL).increment(1); - match operation - .and_then(|operation| self.context.get_portal_from_execute(operation)) - .as_deref() - { + let portal = operation + .map(|operation| self.context.get_portal_from_execute(operation)) + .transpose()? + .flatten(); + match portal.as_deref() { Some(Portal::Encrypted { .. }) => { debug!(target: MAPPER, client_id = self.context.client_id, msg = "Encrypted"); @@ -745,11 +795,15 @@ mod tests { use crate::config::TandemConfig; use crate::postgresql::context::KeysetIdentifier; use crate::postgresql::parser::SqlParser; + use crate::postgresql::rewrite::Name; + use crate::postgresql::test_operation_id as operation_id; use crate::proxy::{EncryptConfig, EncryptionService}; + use cipherstash_client::schema::{ColumnConfig, ColumnType}; use eql_mapper::Schema; use std::sync::Arc; use tokio::sync::mpsc; + #[derive(Clone)] struct TestService {} #[async_trait::async_trait] @@ -781,8 +835,18 @@ mod tests { } fn create_backend() -> Backend { + create_backend_with_context().0 + } + + fn create_backend_with_context() -> (Backend, Context) { + create_backend_with_encrypt_config(EncryptConfig::default()) + } + + fn create_backend_with_encrypt_config( + encrypt_config: EncryptConfig, + ) -> (Backend, Context) { let config = Arc::new(TandemConfig::for_testing()); - let encrypt_config = Arc::new(EncryptConfig::default()); + let encrypt_config = Arc::new(encrypt_config); let schema = Arc::new(Schema::new("public")); let (reload_sender, _) = mpsc::unbounded_channel(); let context = Context::new( @@ -794,7 +858,7 @@ mod tests { TestService {}, reload_sender, ); - Backend::new(context) + (Backend::new(context.clone()), context) } #[test] @@ -828,7 +892,7 @@ mod tests { ); let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); context.execute_simple_schema_statements(&[ddl]); - context.schema_execution_succeeded(); + context.report_schema_execution_succeeded(); let mut backend = Backend::new(context); let expected = @@ -839,6 +903,364 @@ mod tests { assert!(reload_receiver.try_recv().is_err()); } + #[tokio::test] + async fn database_error_without_an_execution_is_forwarded() { + let mut backend = create_backend(); + let response = crate::postgresql::diagnostics::invalid_sql_statement( + "syntax error at or near SELECT".to_owned(), + ); + let message = BackendMessage::ErrorResponse(response.clone()); + + let output = backend.intercept(None, message.clone()).await.unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Forward(message)); + } + + #[tokio::test] + async fn terminal_message_without_an_operation_is_forwarded() { + let (mut backend, context) = create_backend_with_context(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + assert!(context.schema_ddl_in_flight()); + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"SELECT 1")); + + let output = backend.intercept(None, message.clone()).await.unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Forward(message)); + assert!(!context.schema_ddl_in_flight()); + } + + #[tokio::test] + async fn correlated_terminal_message_for_an_unknown_operation_fails_closed() { + let mut backend = create_backend(); + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"SELECT 1")); + + let result = backend + .intercept(Some(operation_id()), message.clone()) + .await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn terminal_message_for_a_non_execution_operation_fails_closed() { + let (mut backend, context) = create_backend_with_context(); + let operation = operation_id(); + context.set_non_execution(operation).unwrap(); + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"SELECT 1")); + + let result = backend.intercept(Some(operation), message.clone()).await; + + assert!(matches!( + result, + Err(Error::Context( + crate::error::ContextError::OperationWithoutExecute + )) + )); + } + + #[tokio::test] + async fn no_data_for_an_untracked_describe_fails_closed() { + let mut backend = create_backend(); + + let result = backend + .intercept(Some(operation_id()), BackendMessage::NoData) + .await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn row_description_for_an_untracked_mapped_operation_fails_closed() { + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert( + crate::Identifier::new("records", "secret"), + ColumnConfig::build("secret".to_owned()).casts_as(ColumnType::Text), + ); + let (mut backend, _) = create_backend_with_encrypt_config(encrypt_config); + let message = BackendMessage::RowDescription(pg_proto::RowDescription { fields: vec![] }); + + let result = backend + .intercept(Some(operation_id()), message.clone()) + .await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn data_row_for_an_unknown_correlated_operation_fails_closed() { + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert( + crate::Identifier::new("records", "secret"), + ColumnConfig::build("secret".to_owned()).casts_as(ColumnType::Text), + ); + let (mut backend, _) = create_backend_with_encrypt_config(encrypt_config); + + let message = BackendMessage::DataRow(DataRow { + columns: vec![Some(bytes::Bytes::from_static(b"ciphertext"))], + }); + let result = backend + .intercept(Some(operation_id()), message.clone()) + .await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(!backend.discard_execution); + } + + #[tokio::test] + async fn database_error_for_a_non_execution_operation_is_forwarded() { + let (mut backend, context) = create_backend_with_context(); + let operation = operation_id(); + context.set_non_execution(operation).unwrap(); + let response = crate::postgresql::diagnostics::invalid_sql_statement( + "syntax error at or near SELECT".to_owned(), + ); + let message = BackendMessage::ErrorResponse(response); + + let output = backend + .intercept(Some(operation), message.clone()) + .await + .unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Forward(message)); + } + + #[tokio::test] + async fn database_error_for_an_unregistered_operation_fails_closed() { + let (mut backend, context) = create_backend_with_context(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + assert!(context.schema_ddl_in_flight()); + let message = BackendMessage::ErrorResponse( + crate::postgresql::diagnostics::invalid_sql_statement("database error".to_owned()), + ); + + let result = backend + .intercept(Some(operation_id()), message.clone()) + .await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(!context.schema_ddl_in_flight()); + } + + #[tokio::test] + async fn parse_completion_releases_non_execution_error_correlation() { + let (mut backend, context) = create_backend_with_context(); + let operation = operation_id(); + context.set_non_execution(operation).unwrap(); + + let output = backend + .intercept(Some(operation), BackendMessage::ParseComplete) + .await + .unwrap(); + + assert_eq!( + output, + BackendMiddlewareOutput::Forward(BackendMessage::ParseComplete) + ); + assert!(matches!( + context.finish_execution(operation, ExecutionOutcome::Failed), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn correlated_non_execution_completion_for_an_unknown_operation_fails_closed() { + for message in [BackendMessage::ParseComplete, BackendMessage::BindComplete] { + let mut backend = create_backend(); + + let result = backend.intercept(Some(operation_id()), message).await; + + assert!(matches!( + result, + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + } + + #[tokio::test] + async fn stored_proxy_error_replaces_a_non_execution_database_error() { + let (mut backend, mut context) = create_backend_with_context(); + let operation = operation_id(); + let replacement = + crate::postgresql::diagnostics::invalid_sql_statement("proxy parse error".to_owned()); + context + .set_operation_error(operation, replacement.clone()) + .unwrap(); + let database_error = + BackendMessage::ErrorResponse(crate::postgresql::diagnostics::invalid_sql_statement( + "database parse error".to_owned(), + )); + + let output = backend + .intercept(Some(operation), database_error) + .await + .unwrap(); + + assert_eq!( + output, + BackendMiddlewareOutput::Forward(BackendMessage::ErrorResponse(replacement)) + ); + } + + #[tokio::test] + async fn successful_execution_is_not_replaced_by_a_stored_error() { + let (mut backend, mut context) = create_backend_with_context(); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_execute(operation, Name::new(), Some(scope)) + .unwrap(); + context + .set_operation_error( + operation, + crate::postgresql::diagnostics::invalid_sql_statement("proxy error".to_owned()), + ) + .unwrap(); + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"SELECT 1")); + + let output = backend + .intercept(Some(operation), message.clone()) + .await + .unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Forward(message)); + } + + #[tokio::test] + async fn decryption_recovery_discards_suppressed_execution_state() { + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert( + crate::Identifier::new("records", "secret"), + ColumnConfig::build("secret".to_owned()).casts_as(ColumnType::Text), + ); + let (mut backend, mut context) = create_backend_with_encrypt_config(encrypt_config); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_execute(operation, Name::new(), Some(scope)) + .unwrap(); + backend.discard_execution = true; + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"SELECT 1")); + + let output = backend + .intercept(Some(operation), message.clone()) + .await + .unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Suppress(message)); + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + } + + #[tokio::test] + async fn discarded_execution_response_resolves_its_schema_execution() { + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert( + crate::Identifier::new("records", "secret"), + ColumnConfig::build("secret".to_owned()).casts_as(ColumnType::Text), + ); + let (mut backend, mut context) = create_backend_with_encrypt_config(encrypt_config); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_execute(operation, Name::new(), Some(scope)) + .unwrap(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + assert!(context.schema_ddl_in_flight()); + backend.discard_execution = true; + + let message = BackendMessage::CommandComplete(bytes::Bytes::from_static(b"CREATE TABLE")); + let output = backend + .intercept(Some(operation), message.clone()) + .await + .unwrap(); + + assert_eq!(output, BackendMiddlewareOutput::Suppress(message)); + assert!(!context.schema_ddl_in_flight()); + } + + #[tokio::test] + async fn uncorrelated_flush_failure_returns_a_postgresql_error() { + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert( + crate::Identifier::new("records", "secret"), + ColumnConfig::build("secret".to_owned()).casts_as(ColumnType::Text), + ); + let (mut backend, mut context) = create_backend_with_encrypt_config(encrypt_config); + let operation = operation_id(); + let scope = context.start_metrics_scope().unwrap(); + context + .set_execute(operation, Name::new(), Some(scope)) + .unwrap(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + assert!(context.schema_ddl_in_flight()); + backend.encrypted_rows_operation = Some(operation); + backend.encrypted_rows.push(DataRow { + columns: vec![Some(bytes::Bytes::from_static(b"invalid ciphertext"))], + }); + + let output = backend + .intercept( + None, + BackendMessage::CommandComplete(bytes::Bytes::from_static(b"CREATE TABLE")), + ) + .await + .unwrap(); + + assert!(matches!( + output, + BackendMiddlewareOutput::Expand(messages) + if matches!(messages.as_slice(), [BackendMessage::ErrorResponse(_)]) + )); + assert!(backend.discard_execution); + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + assert!(context.get_metrics_scope(scope).unwrap().is_none()); + assert!(context.schema_ddl_in_flight()); + } + + #[tokio::test] + async fn suspended_portal_resolves_one_schema_execution() { + let (mut backend, context) = create_backend_with_context(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + assert!(context.schema_ddl_in_flight()); + + let output = backend + .intercept(None, BackendMessage::PortalSuspended) + .await + .unwrap(); + + assert_eq!( + output, + BackendMiddlewareOutput::Forward(BackendMessage::PortalSuspended) + ); + assert!(!context.schema_ddl_in_flight()); + } + #[tokio::test] async fn publication_failure_closes_connection_before_idle_readiness() { let config = Arc::new(TandemConfig::for_testing()); @@ -856,7 +1278,7 @@ mod tests { ); let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); context.execute_simple_schema_statements(&[ddl]); - context.schema_execution_succeeded(); + context.report_schema_execution_succeeded(); let reload_task = tokio::spawn(async move { let Some(crate::proxy::ReloadCommand::DatabaseSchema(responder)) = diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs index 23d72aa7..fa03097e 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs @@ -17,6 +17,7 @@ use crate::postgresql::data::{ use crate::postgresql::inbound_eql; use crate::postgresql::rewrite::Name; use crate::postgresql::rewrite::UNSPECIFIED_TYPE_OID; +use crate::postgresql::OperationId; use crate::prometheus::{ ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS, ENCRYPTION_ERROR_TOTAL, ENCRYPTION_REQUESTS_TOTAL, STATEMENTS_ENCRYPTED_TOTAL, @@ -66,7 +67,7 @@ use tracing::{debug, info, warn}; /// - **Query Transformation**: Rewrite SQL to use EQL functions for encrypted operations /// - **Protocol Handling**: Manage PostgreSQL extended query protocol (Parse/Bind/Execute) /// - **Error Management**: Convert encryption errors to PostgreSQL-compatible error responses -/// - **Context Management**: Track statements, portals, and session state +/// - **Context Management**: Track statements, portals, and connection state /// /// # Supported PostgreSQL Messages /// @@ -106,10 +107,55 @@ impl Frontend { pub async fn intercept( &mut self, - operation: pg_proto::OperationId, + operation: OperationId, protocol_message: FrontendMessage, ) -> Result { + if matches!( + protocol_message, + FrontendMessage::Parse(_) | FrontendMessage::Bind(_) + ) { + self.context.set_non_execution(operation)?; + } if self.context.mapping_disabled() { + match &protocol_message { + FrontendMessage::Query(_) => { + let metrics_scope = self.context.start_metrics_scope()?; + self.context.add_portal( + operation, + Name::new(), + Portal::passthrough(Some(metrics_scope)), + )?; + self.context.set_simple_query_execute_until_ready( + operation, + Name::new(), + Some(metrics_scope), + )?; + } + FrontendMessage::Bind(bind) => { + let session_id = self.context.get_statement_metrics_scope(&bind.statement)?; + self.context.add_portal( + operation, + bind.portal.clone(), + Portal::passthrough(session_id), + )?; + } + FrontendMessage::Execute(execute) => self + .context + .set_execute_for_portal(operation, execute.portal.clone())?, + FrontendMessage::Describe(describe) => { + self.context.set_describe(operation, describe.clone())?; + } + FrontendMessage::Close(close) => match close.target { + DescribeTarget::Portal => self.context.close_portal(&close.name)?, + DescribeTarget::Statement => { + self.context.close_statement_explicit(&close.name)?; + } + }, + FrontendMessage::Parse(parse) => { + self.context.close_statement(&parse.statement)?; + } + _ => {} + } return Ok(FrontendMiddlewareOutput::Forward(protocol_message)); } @@ -119,7 +165,7 @@ impl Frontend { self.context.mark_schema_protocol_boundary(); return Ok(FrontendMiddlewareOutput::Forward(protocol_message)); } - self.context.discard_operation(operation); + self.context.discard_operation(operation)?; return Ok(FrontendMiddlewareOutput::Suppress(protocol_message)); } @@ -128,12 +174,13 @@ impl Frontend { match protocol_message { FrontendMessage::Query(query) => { - match self.query_handler(operation, query).await { + let session_id = self.context.start_metrics_scope()?; + match self.query_handler(operation, query, session_id).await { Ok(Some(mapped)) => outbound_message = mapped, // No mapping needed, don't change the bytes Ok(None) => (), Err(err) => { - let session_id = self.context.latest_session_id(); + self.context.finish_metrics_scope(Some(session_id))?; warn!( client_id = self.context.client_id, msg = "Query Handler Error", @@ -141,8 +188,8 @@ impl Frontend { ); let response = self.error_to_response(err); self.context - .set_operation_error(operation, response.clone()); - self.context.set_execute(operation, Name::new(), session_id); + .set_operation_error(operation, response.clone())?; + self.context.set_execute(operation, Name::new(), None)?; self.context.mark_schema_protocol_boundary(); outbound_message = self.simple_error_query(&response); } @@ -162,7 +209,7 @@ impl Frontend { // No mapping needed, don't change the bytes Ok(None) => (), Err(err) => { - self.context.close_statement(&statement); + self.context.close_statement(&statement)?; warn!( client_id = self.context.client_id, msg = "Parse Handler Error", @@ -170,12 +217,7 @@ impl Frontend { ); let response = self.error_to_response(err); self.context - .set_operation_error(operation, response.clone()); - self.context.set_execute( - operation, - Name::new(), - self.context.latest_session_id(), - ); + .set_operation_error(operation, response.clone())?; self.failed_extended_batch = true; outbound_message = self.extended_parse_error(failed_parse, &response); } @@ -183,7 +225,7 @@ impl Frontend { } FrontendMessage::Bind(bind) => { let failed_bind = bind.clone(); - match self.bind_handler(Bind::try_from(bind)?).await { + match self.bind_handler(operation, Bind::try_from(bind)?).await { Ok(Some(mapped)) => outbound_message = mapped, // No mapping needed, don't change the bytes Ok(None) => (), @@ -195,12 +237,7 @@ impl Frontend { ); let response = self.error_to_response(err); self.context - .set_operation_error(operation, response.clone()); - self.context.set_execute( - operation, - Name::new(), - self.context.latest_session_id(), - ); + .set_operation_error(operation, response.clone())?; self.failed_extended_batch = true; outbound_message = self.extended_bind_error(failed_bind, &response); } @@ -211,12 +248,7 @@ impl Frontend { ); let response = self.error_to_response(err); self.context - .set_operation_error(operation, response.clone()); - self.context.set_execute( - operation, - Name::new(), - self.context.latest_session_id(), - ); + .set_operation_error(operation, response.clone())?; self.failed_extended_batch = true; outbound_message = self.extended_bind_error(failed_bind, &response); } @@ -228,12 +260,7 @@ impl Frontend { ); let response = self.error_to_response(err); self.context - .set_operation_error(operation, response.clone()); - self.context.set_execute( - operation, - Name::new(), - self.context.latest_session_id(), - ); + .set_operation_error(operation, response.clone())?; self.failed_extended_batch = true; outbound_message = self.extended_bind_error(failed_bind, &response); } @@ -268,32 +295,34 @@ impl Frontend { async fn describe_handler( &mut self, - operation: pg_proto::OperationId, + operation: OperationId, describe: Describe, ) -> Result<(), Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ?describe); - self.context.set_describe(operation, describe); + self.context.set_describe(operation, describe)?; Ok(()) } async fn close_handler(&mut self, close: Close) -> Result<(), Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ?close); match close.target { - DescribeTarget::Portal => self.context.close_portal(&close.name), - DescribeTarget::Statement => self.context.close_statement_explicit(&close.name), + DescribeTarget::Portal => self.context.close_portal(&close.name)?, + DescribeTarget::Statement => { + self.context.close_statement_explicit(&close.name)?; + } } Ok(()) } async fn execute_handler( &mut self, - operation: pg_proto::OperationId, + operation: OperationId, execute: Execute, ) -> Result { debug!(target: PROTOCOL, client_id = self.context.client_id, ?execute); let executes_ddl = self.context.execute_schema_portal(&execute.portal); self.context - .set_execute_for_portal(operation, execute.portal.to_owned()); + .set_execute_for_portal(operation, execute.portal.to_owned())?; Ok(executes_ddl) } @@ -331,16 +360,16 @@ impl Frontend { /// - `Err(error)` - Processing failed, error should be sent to client async fn query_handler( &mut self, - operation: pg_proto::OperationId, + operation: OperationId, query: bytes::Bytes, + session_id: SessionId, ) -> Result, Error> { let handler_start = Instant::now(); - let session_id = self.context.start_session(); // Set protocol type for diagnostics self.context.update_statement_metadata(session_id, |m| { m.protocol = Some(ProtocolType::Simple); - }); + })?; let parse_timer = PhaseTimer::start(); @@ -413,7 +442,7 @@ impl Frontend { session_id, Portal::passthrough(Some(session_id)), &parsed_statements, - ); + )?; return Ok(None); }; } @@ -431,7 +460,7 @@ impl Frontend { // Record parse duration before encryption work starts if !parse_duration_recorded { self.context - .record_parse_duration(session_id, parse_timer.elapsed()); + .record_parse_duration(session_id, parse_timer.elapsed())?; parse_duration_recorded = true; } @@ -469,7 +498,7 @@ impl Frontend { portal = Portal::encrypted(Arc::new(statement), Some(session_id)); self.context.update_statement_metadata(session_id, |m| { m.encrypted = true; - }); + })?; } None => { debug!(target: MAPPER, @@ -485,7 +514,7 @@ impl Frontend { // Record parse/typecheck duration (if not already recorded before encryption) if !parse_duration_recorded { self.context - .record_parse_duration(session_id, parse_timer.elapsed()); + .record_parse_duration(session_id, parse_timer.elapsed())?; } // Set statement type based on parsed statements @@ -500,14 +529,14 @@ impl Frontend { self.context.update_statement_metadata(session_id, |m| { m.statement_type = Some(statement_type); m.set_multi_statement(parsed_statements.len() > 1); - }); + })?; // Set query fingerprint self.context.update_statement_metadata(session_id, |m| { m.set_query_fingerprint(&query_text); - }); + })?; - self.record_simple_schema_execution(operation, session_id, portal, &forwarded_statements); + self.record_simple_schema_execution(operation, session_id, portal, &forwarded_statements)?; if encrypted { let transformed_statement = forwarded_statements @@ -550,15 +579,19 @@ impl Frontend { /// Records the portal and schema intents for the statements actually sent to PostgreSQL. fn record_simple_schema_execution( &mut self, - operation: pg_proto::OperationId, + operation: OperationId, session_id: SessionId, portal: Portal, statements: &[ast::Statement], - ) { - self.context.add_portal(Name::new(), portal); - self.context - .set_execute(operation, Name::new(), Some(session_id)); + ) -> Result<(), Error> { + self.context.add_portal(operation, Name::new(), portal)?; + self.context.set_simple_query_execute_until_ready( + operation, + Name::new(), + Some(session_id), + )?; self.context.execute_simple_schema_statements(statements); + Ok(()) } /// Encrypts literal values found in SQL statements. @@ -649,14 +682,14 @@ impl Frontend { let duration = Instant::now().duration_since(start); // Add to phase timing diagnostics (accumulate) - self.context.add_encrypt_duration(session_id, duration); + self.context.add_encrypt_duration(session_id, duration)?; // Update metadata with encrypted values count let encrypted_count = encrypted.iter().filter(|e| e.is_some()).count(); self.context.update_statement_metadata(session_id, |m| { m.encrypted = true; m.set_encrypted_values_count(encrypted_count); - }); + })?; counter!(ENCRYPTION_REQUESTS_TOTAL).increment(1); counter!(ENCRYPTED_VALUES_TOTAL).increment(encrypted_count as u64); @@ -761,12 +794,12 @@ impl Frontend { mut message: Parse, ) -> Result, Error> { let original_query = message.query.clone(); - let session_id = self.context.start_session(); + let session_id = self.context.start_metrics_scope()?; // Set protocol type self.context.update_statement_metadata(session_id, |m| { m.protocol = Some(ProtocolType::Extended); - }); + })?; let parse_timer = PhaseTimer::start(); @@ -792,13 +825,13 @@ impl Frontend { // statement for every command, so the `END` at the close of a // transaction binds against the `SELECT` that preceded it. // - // Closing also drops the name's session mapping, so it must happen - // BEFORE the new session is recorded — the other way around wipes the + // Closing also drops the name's metrics scope mapping, so it must happen + // BEFORE the new scope is recorded — the other way around wipes the // mapping that was just written and every Bind falls back to the - // latest-session guess. - self.context.close_statement(&message.statement); + // latest-scope guess. + self.context.close_statement(&message.statement)?; self.context - .set_statement_session(message.statement.to_owned(), session_id); + .set_statement_metrics_scope(message.statement.to_owned(), session_id)?; let mut statement_text = String::from_utf8_lossy(&message.query).into_owned(); let statement = SqlParser::parse_statement(&statement_text)?; @@ -818,7 +851,7 @@ impl Frontend { self.context.update_statement_metadata(session_id, |m| { m.statement_type = Some(StatementType::from_statement(&statement)); m.set_query_fingerprint(&statement_text); - }); + })?; if let Some(mapping_disabled) = self.context.maybe_set_unsafe_disable_mapping(&statement) { warn!( @@ -866,7 +899,7 @@ impl Frontend { if typed_statement.requires_transform() { // Record parse duration before encryption work starts self.context - .record_parse_duration(session_id, parse_timer.elapsed()); + .record_parse_duration(session_id, parse_timer.elapsed())?; parse_duration_recorded = true; let encrypted_literals = self @@ -901,7 +934,7 @@ impl Frontend { message.parameter_types = rewrite_parse_param_types(&client_param_types, &statement.output_params); self.context - .add_statement(message.statement.to_owned(), statement); + .add_statement(message.statement.to_owned(), statement)?; } _ => { debug!(target: MAPPER, @@ -915,7 +948,7 @@ impl Frontend { // Record parse duration (if not already recorded before encryption) if !parse_duration_recorded { self.context - .record_parse_duration(session_id, parse_timer.elapsed()); + .record_parse_duration(session_id, parse_timer.elapsed())?; } if message.query != original_query || message.parameter_types != client_param_types { @@ -1058,63 +1091,82 @@ impl Frontend { /// - `Ok(Some(bytes))` - Modified Bind message with encrypted parameter values /// - `Ok(None)` - No parameter encryption needed, forward original message /// - `Err(error)` - Processing failed, error should be sent to client - async fn bind_handler(&mut self, mut bind: Bind) -> Result, Error> { + async fn bind_handler( + &mut self, + operation: OperationId, + mut bind: Bind, + ) -> Result, Error> { self.context .bind_schema_statement(bind.portal.to_owned(), &bind.prepared_statement); if self.context.unsafe_disable_mapping() { warn!(msg = "Encrypted statement mapping is not enabled"); counter!(STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL).increment(1); counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1); + self.context.add_portal( + operation, + bind.portal.to_owned(), + Portal::passthrough(None), + )?; return Ok(None); } let session_id = self .context - .get_statement_session_or_latest(&bind.prepared_statement); + .start_portal_metrics_scope(&bind.prepared_statement)?; - // Track param bytes for diagnostics - let param_bytes: usize = bind.param_values.iter().map(|p| p.bytes.len()).sum(); - self.context - .with_session(session_id, |m| m.metadata.set_param_bytes(param_bytes)); + let result = async { + // Track param bytes for diagnostics + let param_bytes: usize = bind.param_values.iter().map(|p| p.bytes.len()).sum(); + self.context + .with_metrics_scope(session_id, |m| m.metadata.set_param_bytes(param_bytes))?; - debug!(target: PROTOCOL, client_id = self.context.client_id, bind = ?bind); + debug!(target: PROTOCOL, client_id = self.context.client_id, bind = ?bind); - let mut portal = Portal::passthrough(session_id); + let mut portal = Portal::passthrough(session_id); - if let Some(statement) = self.context.get_statement(&bind.prepared_statement) { - debug!(target:MAPPER, client_id = self.context.client_id, ?statement); + if let Some(statement) = self.context.get_statement(&bind.prepared_statement)? { + debug!(target:MAPPER, client_id = self.context.client_id, ?statement); - if statement.has_params() { - let encrypted = self.encrypt_params(session_id, &bind, &statement).await?; - bind.rewrite(&statement.output_params, encrypted)?; - } - if statement.has_projection() { - portal = Portal::encrypted_with_format_codes( - statement, - bind.result_columns_format_codes.to_owned(), - session_id, - ); - self.context - .with_session(session_id, |m| m.metadata.encrypted = true); - } - }; + if statement.has_params() { + let encrypted = self.encrypt_params(session_id, &bind, &statement).await?; + bind.rewrite(&statement.output_params, encrypted)?; + } + if statement.has_projection() { + portal = Portal::encrypted_with_format_codes( + statement, + bind.result_columns_format_codes.to_owned(), + session_id, + ); + self.context + .with_metrics_scope(session_id, |m| m.metadata.encrypted = true)?; + } + }; - debug!(target: MAPPER, client_id = self.context.client_id, portal = ?portal); - self.context.add_portal(bind.portal.to_owned(), portal); + debug!(target: MAPPER, client_id = self.context.client_id, portal = ?portal); + self.context + .add_portal(operation, bind.portal.to_owned(), portal)?; - if bind.requires_rewrite() { - let message = FrontendMessage::from(bind); - debug!( - target: MAPPER, - client_id = self.context.client_id, - msg = "Rewrite Bind", - ?message - ); + if bind.requires_rewrite() { + let message = FrontendMessage::from(bind); + debug!( + target: MAPPER, + client_id = self.context.client_id, + msg = "Rewrite Bind", + ?message + ); - Ok(Some(message)) - } else { - Ok(None) + Ok(Some(message)) + } else { + Ok(None) + } + } + .await; + + if result.is_err() { + self.context.finish_metrics_scope(session_id)?; } + + result } /// @@ -1169,13 +1221,13 @@ impl Frontend { // Record timing and metadata for this encryption operation let encrypted_count = encrypted.iter().filter(|e| e.is_some()).count(); - self.context.with_session(session_id, |m| { + self.context.with_metrics_scope(session_id, |m| { // Add to phase timing diagnostics (accumulate) m.phase_timing.add_encrypt(duration); // Always update metadata for slow-statement logging m.metadata.encrypted = true; m.metadata.set_encrypted_values_count(encrypted_count); - }); + })?; // Prometheus metrics remain gated if self.context.prometheus_enabled() { @@ -1634,19 +1686,28 @@ mod tests { use super::{quote_literal, Frontend}; use crate::config::TandemConfig; use crate::error::{EncryptError, Error, MappingError}; + use crate::postgresql::context::statement::{OutputParam, OutputParamSource}; + use crate::postgresql::context::Statement; use crate::postgresql::context::{Context, KeysetIdentifier}; use crate::postgresql::error_handler::PostgreSqlErrorHandler; use crate::postgresql::inbound_eql::InboundEql; + use crate::postgresql::test_operation_id as operation_id; use crate::postgresql::Column; use crate::proxy::{EncryptConfig, EncryptionService}; + use crate::Identifier; use cipherstash_client::eql::{EncryptedPayloadV3, EQL_SCHEMA_VERSION_V3}; + use cipherstash_client::schema::{ColumnConfig, ColumnMode, ColumnType}; use cipherstash_client::zerokms::EncryptedRecord; + use eql_mapper::EqlTermVariant; use eql_mapper::Schema; - use pg_proto::{Bind, FrontendMessage, Parse}; + use pg_proto::{Bind, Describe, DescribeTarget, FrontendMessage, Parse}; use std::sync::Arc; use tokio::sync::mpsc; - struct TestService; + #[derive(Clone)] + struct TestService { + fail_encrypt: bool, + } #[async_trait::async_trait] impl EncryptionService for TestService { @@ -1656,6 +1717,9 @@ mod tests { _plaintexts: Vec>, _columns: &[Option], ) -> Result>, Error> { + if self.fail_encrypt { + return Err(EncryptError::InvalidInboundEqlPayload.into()); + } Ok(Vec::new()) } @@ -1677,17 +1741,41 @@ mod tests { } fn frontend() -> Frontend { + frontend_with_config(TandemConfig::for_testing()).0 + } + + fn frontend_with_config(config: TandemConfig) -> (Frontend, Context) { + frontend_with_service( + config, + TestService { + fail_encrypt: false, + }, + ) + } + + fn frontend_with_service( + config: TandemConfig, + service: TestService, + ) -> (Frontend, Context) { + frontend_with_encrypt_config_and_service(config, EncryptConfig::default(), service) + } + + fn frontend_with_encrypt_config_and_service( + config: TandemConfig, + encrypt_config: EncryptConfig, + service: TestService, + ) -> (Frontend, Context) { let (reload_sender, _) = mpsc::unbounded_channel(); let context = Context::new( 1, - Arc::new(TandemConfig::for_testing()), - Arc::new(EncryptConfig::default()), + Arc::new(config), + Arc::new(encrypt_config), Arc::new(Schema::new("public")), Arc::new(rustls::RootCertStore::empty()), - TestService, + service, reload_sender, ); - Frontend::new(context) + (Frontend::new(context.clone()), context) } fn inbound_storage_payload() -> crate::EqlCiphertext { @@ -1728,6 +1816,21 @@ mod tests { )); } + #[tokio::test] + async fn failed_simple_query_releases_its_metrics_scope() { + let (mut frontend, context) = frontend_with_config(TandemConfig::for_testing()); + + frontend + .intercept( + operation_id(), + FrontendMessage::Query(bytes::Bytes::from_static(b"select 'unterminated")), + ) + .await + .unwrap(); + + assert_eq!(context.active_metrics_scopes().unwrap(), 0); + } + #[test] fn extended_errors_preserve_the_original_protocol_message_kind() { let frontend = frontend(); @@ -1771,4 +1874,365 @@ mod tests { Err(Error::Encrypt(EncryptError::InvalidInboundEqlPayload)) )); } + + #[tokio::test] + async fn mapping_disabled_extended_protocol_does_not_create_statement_metrics() { + let mut config = TandemConfig::for_testing(); + config.disable_mapping_for_testing(); + let (mut frontend, context) = frontend_with_config(config); + let statement = bytes::Bytes::from_static(b"statement"); + let portal = bytes::Bytes::from_static(b"portal"); + + frontend + .intercept( + operation_id(), + FrontendMessage::Parse(Parse { + statement: statement.clone(), + query: bytes::Bytes::from_static(b"select 1"), + parameter_types: Vec::new(), + }), + ) + .await + .unwrap(); + frontend + .intercept( + operation_id(), + FrontendMessage::Bind(Bind { + portal: portal.clone(), + statement: statement.clone(), + parameter_formats: Vec::new(), + parameters: Vec::new(), + result_formats: Vec::new(), + }), + ) + .await + .unwrap(); + + assert!(context + .get_statement_metrics_scope(&statement) + .unwrap() + .is_none()); + assert!(context + .get_portal_metrics_scope_id(&portal) + .unwrap() + .is_none()); + assert_eq!(context.active_metrics_scopes().unwrap(), 0); + } + + #[tokio::test] + async fn portals_bound_from_one_statement_have_isolated_parameter_metrics() { + let (mut frontend, context) = frontend_with_config(TandemConfig::for_testing()); + let statement = bytes::Bytes::from_static(b"statement"); + let first_portal = bytes::Bytes::from_static(b"first_portal"); + let second_portal = bytes::Bytes::from_static(b"second_portal"); + frontend + .intercept( + operation_id(), + FrontendMessage::Parse(Parse { + statement: statement.clone(), + query: bytes::Bytes::from_static(b"select $1::text"), + parameter_types: Vec::new(), + }), + ) + .await + .unwrap(); + + for (portal, parameter) in [ + (first_portal.clone(), bytes::Bytes::from_static(b"a")), + ( + second_portal.clone(), + bytes::Bytes::from_static(b"different"), + ), + ] { + frontend + .intercept( + operation_id(), + FrontendMessage::Bind(Bind { + portal, + statement: statement.clone(), + parameter_formats: Vec::new(), + parameters: vec![Some(parameter)], + result_formats: Vec::new(), + }), + ) + .await + .unwrap(); + } + + let first_scope = context + .get_portal_metrics_scope_id(&first_portal) + .unwrap() + .unwrap(); + let second_scope = context + .get_portal_metrics_scope_id(&second_portal) + .unwrap() + .unwrap(); + assert_ne!(first_scope, second_scope); + assert_eq!( + context + .get_metrics_scope(first_scope) + .unwrap() + .unwrap() + .metadata + .param_bytes, + 1 + ); + assert_eq!( + context + .get_metrics_scope(second_scope) + .unwrap() + .unwrap() + .metadata + .param_bytes, + 9 + ); + } + + #[tokio::test] + async fn failed_bind_releases_its_portal_metrics_scope() { + let column_config = ColumnConfig { + name: "secret".to_owned(), + in_place: false, + cast_type: ColumnType::Json, + indexes: vec![], + mode: ColumnMode::PlaintextDuplicate, + }; + let column = Column { + identifier: Identifier::new("records", "secret"), + config: column_config.clone(), + postgres_type: postgres_types::Type::JSONB, + eql_term: EqlTermVariant::JsonValueSelector, + }; + let mut encrypt_config = EncryptConfig::default(); + encrypt_config.insert(Identifier::new("records", "secret"), column_config); + let (mut frontend, mut context) = frontend_with_encrypt_config_and_service( + TandemConfig::for_testing(), + encrypt_config, + TestService { fail_encrypt: true }, + ); + let statement_name = bytes::Bytes::from_static(b"statement"); + let template_scope = context.start_metrics_scope().unwrap(); + context + .set_statement_metrics_scope(statement_name.clone(), template_scope) + .unwrap(); + context + .add_statement( + statement_name.clone(), + Statement::new( + vec![Some(column.clone())], + vec![OutputParam { + column: Some(column), + source: OutputParamSource::Input(0), + query_operand: false, + }], + vec![], + vec![], + vec![], + ), + ) + .unwrap(); + + frontend + .intercept( + operation_id(), + FrontendMessage::Bind(Bind { + portal: bytes::Bytes::from_static(b"portal"), + statement: statement_name, + parameter_formats: Vec::new(), + parameters: vec![Some(bytes::Bytes::from_static(b"\"value\""))], + result_formats: Vec::new(), + }), + ) + .await + .unwrap(); + + assert_eq!(context.active_metrics_scopes().unwrap(), 1); + } + + #[tokio::test] + async fn parse_registers_non_execution_error_correlation() { + let mut config = TandemConfig::for_testing(); + config.disable_mapping_for_testing(); + let (mut frontend, context) = frontend_with_config(config); + let operation = operation_id(); + + frontend + .intercept( + operation, + FrontendMessage::Parse(Parse { + statement: bytes::Bytes::from_static(b"statement"), + query: bytes::Bytes::from_static(b"select 1"), + parameter_types: Vec::new(), + }), + ) + .await + .unwrap(); + + assert!(context + .finish_execution( + operation, + crate::postgresql::context::ExecutionOutcome::Failed, + ) + .is_ok()); + } + + #[tokio::test] + async fn bind_registers_non_execution_error_correlation() { + let mut config = TandemConfig::for_testing(); + config.disable_mapping_for_testing(); + let (mut frontend, context) = frontend_with_config(config); + let operation = operation_id(); + + frontend + .intercept( + operation, + FrontendMessage::Bind(Bind { + portal: bytes::Bytes::from_static(b"portal"), + statement: bytes::Bytes::from_static(b"statement"), + parameter_formats: Vec::new(), + parameters: Vec::new(), + result_formats: Vec::new(), + }), + ) + .await + .unwrap(); + + assert!(context + .finish_execution( + operation, + crate::postgresql::context::ExecutionOutcome::Failed, + ) + .is_ok()); + } + + #[tokio::test] + async fn connection_disabled_mapping_allows_bound_portals_to_be_described() { + let (mut frontend, _) = frontend_with_config(TandemConfig::for_testing()); + let portal = bytes::Bytes::from_static(b"portal"); + + frontend + .intercept( + operation_id(), + FrontendMessage::Query(bytes::Bytes::from_static( + b"SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true", + )), + ) + .await + .unwrap(); + frontend + .intercept( + operation_id(), + FrontendMessage::Bind(Bind { + portal: portal.clone(), + statement: bytes::Bytes::from_static(b"unknown_statement"), + parameter_formats: Vec::new(), + parameters: Vec::new(), + result_formats: Vec::new(), + }), + ) + .await + .unwrap(); + + let result = frontend + .intercept( + operation_id(), + FrontendMessage::Describe(Describe { + target: DescribeTarget::Portal, + name: portal, + }), + ) + .await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn mapping_disabled_simple_query_waits_for_every_statement_completion() { + let mut config = TandemConfig::for_testing(); + config.disable_mapping_for_testing(); + let (mut frontend, context) = frontend_with_config(config); + let operation = operation_id(); + + frontend + .intercept( + operation, + FrontendMessage::Query(bytes::Bytes::from_static(b"SELECT 1; SELECT 2")), + ) + .await + .unwrap(); + + context + .finish_execution( + operation, + crate::postgresql::context::ExecutionOutcome::Completed, + ) + .unwrap(); + assert!(context.get_execute(operation).unwrap().is_some()); + + context + .finish_execution( + operation, + crate::postgresql::context::ExecutionOutcome::Completed, + ) + .unwrap(); + assert!(context.get_execute(operation).unwrap().is_some()); + context + .ready_for_query(pg_proto::TransactionStatus::Idle, Some(operation)) + .unwrap(); + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn mapping_disabled_unparsed_simple_query_waits_for_readiness() { + let mut config = TandemConfig::for_testing(); + config.disable_mapping_for_testing(); + let (mut frontend, context) = frontend_with_config(config); + let operation = operation_id(); + + frontend + .intercept( + operation, + FrontendMessage::Query(bytes::Bytes::from_static( + b"DO $$ BEGIN RAISE NOTICE 'x'; END $$; SELECT 1", + )), + ) + .await + .unwrap(); + + context + .finish_execution( + operation, + crate::postgresql::context::ExecutionOutcome::Completed, + ) + .unwrap(); + assert!(context.get_execute(operation).unwrap().is_some()); + + context + .ready_for_query(pg_proto::TransactionStatus::Idle, Some(operation)) + .unwrap(); + assert!(matches!( + context.get_execute(operation), + Err(Error::Context(crate::error::ContextError::UnknownOperation)) + )); + } + + #[tokio::test] + async fn unknown_describe_target_is_forwarded_to_postgresql() { + let mut frontend = frontend(); + + let result = frontend + .intercept( + operation_id(), + FrontendMessage::Describe(Describe { + target: DescribeTarget::Portal, + name: bytes::Bytes::from_static(b"unknown_portal"), + }), + ) + .await; + + assert!(result.is_ok()); + } } diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs index f76f4da6..60b1bb43 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs @@ -57,7 +57,7 @@ where operation: OperationId, message: FrontendMessage, ) -> Result { - self.frontend.intercept(operation, message).await + self.frontend.intercept(operation.into(), message).await } async fn backend_operation( @@ -68,6 +68,8 @@ where operation: Option, message: BackendMessage, ) -> Result { - self.backend.intercept(operation, message).await + self.backend + .intercept(operation.map(Into::into), message) + .await } } diff --git a/packages/cipherstash-proxy/src/postgresql/mod.rs b/packages/cipherstash-proxy/src/postgresql/mod.rs index a7922b1a..6f029795 100644 --- a/packages/cipherstash-proxy/src/postgresql/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/mod.rs @@ -15,3 +15,38 @@ pub use context::Context; pub use context::KeysetIdentifier; pub use driver::handler; pub(crate) use rewrite::Name; + +/// Proxy-owned identity for correlating PostgreSQL protocol operations. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct OperationId(OperationIdInner); + +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +enum OperationIdInner { + Protocol(pg_proto::OperationId), + #[cfg(test)] + Test(u64), +} + +impl From for OperationId { + fn from(id: pg_proto::OperationId) -> Self { + Self(OperationIdInner::Protocol(id)) + } +} + +#[cfg(test)] +pub(crate) fn test_operation_id() -> OperationId { + use std::sync::atomic::{AtomicU64, Ordering}; + + static NEXT_ID: AtomicU64 = AtomicU64::new(1); + OperationId(OperationIdInner::Test( + NEXT_ID.fetch_add(1, Ordering::Relaxed), + )) +} + +#[cfg(test)] +mod tests { + #[test] + fn test_operation_ids_are_distinct() { + assert_ne!(super::test_operation_id(), super::test_operation_id()); + } +} diff --git a/proxy-architecture-walkthrough.html b/proxy-architecture-walkthrough.html new file mode 100644 index 00000000..6f4b515f --- /dev/null +++ b/proxy-architecture-walkthrough.html @@ -0,0 +1,287 @@ + + + + + + CIP-3933 — Execution Lifecycle Walkthrough + + + +
+ +
+
+
CIP-3933 · Standalone PR walkthrough · fixed snapshot
+

Execution lifecycle
belongs to Context

+

A guided review of cip-3933-execution-lifecycle: why Proxy moves correlated PostgreSQL metadata behind one connection-local seam, how execution occurrences now survive suspension and resume, and where to concentrate review effort.

+
+
1focused refactor
+
12code and contract files
+
+2,410additions
+
−495deletions
+
+

Code snapshot: merge base 195344ec12da6ab6ddf8158264cf3f66002a0074 through implementation head fddd09ddf56ee66ba5d813457df01f068d5c6a3e. The walkthrough itself was added in the following documentation commit, so its figures intentionally exclude its own file. See the pinned implementation comparison and CIP-3933.

+
Publishing status: this is the standalone source walkthrough in the Proxy branch. It follows the information architecture of the Suite PR-walkthrough project, but it is not yet an Astro/MDX content entry and is not linked from that site’s index.
+
+ +
+

Start here

+
The design move: Frontend and Backend still interpret PostgreSQL messages, but they no longer coordinate related state by reaching into separate maps. Context now applies each connection-local lifecycle transition atomically and returns the effects the adapters need.
+

This is not a new wire protocol and not a schema-lifecycle rewrite. It is an ownership refactor that makes existing behavior explicit, then fixes edge cases that could not be represented cleanly before: suspended execution, repeated portal execution, terminal error replacement, and statement metrics scoped to an execution occurrence. Stale terminal and Describe responses remain deliberately tolerated by Backend and are called out below.

+
pg-proto orders messages→adapters interpret→Context transitions metadata→effects run outside lock
+
+ +
+

Suggested reading order

+
+

Read the contract first

Start with ADR 0002 and the Proxy glossary. They define the intended seam and settle “statement metrics scope” as one execution occurrence. Treat tolerated untracked responses as an adapter exception to inspect.

+

Understand the state model

In context/mod.rs, review ConnectionProtocolState, ExecuteContext, and the transition/effect split before reading individual methods.

+

Trace one extended execution

Follow Parse → Bind → Execute → PortalSuspended → Execute → CommandComplete through Frontend, Context, and Backend.

+

Trace terminal and simple-query paths

Check ErrorResponse replacement and the “complete at ReadyForQuery” state used by simple Query.

+

Finish at the boundaries

Confirm pg-proto preserves backpressure, Schema middleware still owns transaction schema, and the one added integration regression keeps a connection usable after a multi-statement simple Query.

+
+
+ +
+

Change map

+ + + + + + + + + + + + + +
Files and concerns in the fixed snapshot
AreaRole in this changeWhat to review
context/mod.rsNew connection-local protocol state model and atomic transitions.ownershiplockingmetricstests
middleware/frontend.rsTurns client messages into Context transitions.ParseBindExecutesimple Query
middleware/backend.rsTurns database responses into completion, suspension, or failure transitions.correlationreplacementreadiness
driver.rsAdjacent cancellation cleanup: an RAII guard detaches the connection’s cancellation key on every exit path. Review separately from lifecycle ownership.adjacent changecancellationcleanup
postgresql/mod.rsAdds a Proxy-owned OperationId wrapper and test-only identity generator used to exercise Context transitions without pg-proto internals.test seamidentity
error.rs · config/tandem.rsAdds explicit state-unavailable/config behavior needed by the seam.customer errorcompatibility
schema_change.rsAdds one integration regression: a multi-statement simple Query leaves the connection usable. The detailed lifecycle matrix remains unit coverage.integrationsimple Query
CONTEXT.md · ADR · CHANGELOGRecords language, architectural ownership, and user-visible consequences.contractscope
+
+ +
+

Architecture step-through

+

The two diagrams below are the original architecture audit, preserved in this walkthrough. Read them as an ownership comparison: the external message path barely changes; the metadata seam does.

+ +

Before: distributed lifecycle orchestration

+
+ + + Client app + Driver + Frontend adapter + Backend adapter + PostgreSQL + Statements + Portals + Operations + Statement scopes + Metrics + Five independent locks + Schema middlewaretransactional schema state + ZeroKMSencryption adapter + + + + + + +
+

Context already owned the maps, but Frontend and Backend jointly orchestrated lifecycle changes across their independent locks. A single logical event—such as completing an execution—required several reads and writes, making atomic correlation difficult. Missing or poisoned state was commonly treated as absent. Schema middleware separately owned transactional schema state, and Context invoked the ZeroKMS encryption adapter.

+ +

After: Context becomes the connection seam

+
+ + + + Client app + Driver + Frontend adapter + Backend adapter + PostgreSQL + + Contextconnection seam + Atomic lifecycletransitions + ConnectionProtocolState · one lock + Operationsportal correlation + Statements · portals + Suspended executes + Metrics scopes + + + + + Schema middlewareowns transaction schema + ZeroKMSencryption adapter + + +
+

Frontend and Backend remain wire-protocol adapters. Context now owns correlated CipherStash metadata and applies transitions under one lock. pg-proto retains protocol ordering and backpressure; Schema middleware retains transactional schema state.

+
+ +
+

Execution lifecycle before the refactor

+
+ + + Preparedstarts one scope + Portalinherits scope + Executingoperation lookup + Complete immediatelyincluding PortalSuspended + Finish shared scopefirst terminal response + ReadyForQuerytransaction status only + + + + ParseBindExecuteCommandCompleteErrorResponsesimple-query command mayfinish before readinesslater protocol event + + Missing lifecycle states + No suspended execution retained for resume + No explicit “complete at readiness” state + +
+

A prepared Statement created a metrics scope that its Portals and Execute operations reused. Backend terminal messages independently completed execution and finished that shared scope. PortalSuspended did not retain an execution occurrence for resumption, and simple-query execution had no explicit state saying that its scope remained open until ReadyForQuery.

+
+ +
+

Execution lifecycle after the refactor

+
+ + + PreparedPortalExecutingSuspendedAwait readinessCompleteFailed + + ParseBindExecutePortalSuspendedsimple queryCommandComplete / EmptyQueryResponseReadyForQueryErrorResponsenew Execute + +
+

An Execute now owns a distinct metrics scope. PortalSuspended retains that occurrence so the next Execute resumes it; CommandComplete or EmptyQueryResponse completes it; an ErrorResponse fails it. Simple Query explicitly waits for ReadyForQuery, keeping transaction-boundary behavior separate from an individual command response.

+
+ +
+

The invariants this refactor establishes

+
+

One correlated transition

Statements, portals, operations, suspended executes, and metric scopes change under one protocol-state lock.

+

One scope per occurrence

Repeated execution does not reuse the prepared Statement’s measurement window. Parse timing remains reusable Statement knowledge.

+

Suspension is not completion

A suspended portal retains the same execution occurrence and metric identity until resume reaches a terminal result.

+

Effects happen after mutation

The lock is released before metrics emission, asynchronous work, or Schema middleware calls.

+

Errors are atomic

A terminal local failure replaces the upstream response as part of the same correlated lifecycle decision.

+

The snapshot has a tolerance boundary

Context protocol-state methods reject unavailable or inconsistent state, while Backend deliberately ignores unknown operations, Executes, and Describes for untracked responses. The unrelated keyset accessor remains optional, but correlated protocol metadata no longer treats a poisoned lock as absence.

+
+
+ +
+

Review hotspots

+
+

Transition completeness

For every terminal backend response, verify the operation, execution occurrence, portal relationship, metric scope, and schema outcome are each consumed exactly once.

+

Suspension and resume identity

Confirm a resumed Execute cannot accidentally allocate a second scope or complete a stale occurrence.

+

Simple Query readiness

A command can complete before ReadyForQuery. Review multi-statement and error paths for premature metric completion or duplicate schema reporting.

+

Lock boundaries

Every transition should return effects before calling metrics, Schema middleware, encryption, or other asynchronous work.

+

Tolerated untracked responses

Backend converts unknown-operation, operation-without-Execute, and unknown-Describe errors into warnings. Verify this exception is necessary and cannot consume or conceal correlated state; other transition failures still propagate.

+

Metrics compatibility

The meaning becomes more precise, but existing Prometheus metric names remain stable. Check dashboards need no rename while values now represent executions.

+
+
+ +
+

Validation map

+ + + + + + + + + + + + +
Coverage present in the fixed snapshot, and remaining review gaps
ScenarioCoverageExpected lifecycle
Suspended portal resumesUnitOne execution occurrence and one metrics scope survive suspension through final completion.
Portal executes repeatedlyUnitEach occurrence receives a distinct metrics scope while reusing prepared Statement knowledge.
Unknown or stale operationUnit + review gapContext rejects stale transition requests; Backend deliberately warns and ignores untracked terminal and Describe responses.
Protocol state lock unavailableUnit matrixCorrelated protocol metadata reads, mutations, and cleanup return ProtocolStateUnavailable rather than silently omitting a transition.
ErrorResponseUnitThe execution fails once, local terminal replacement is atomic, and schema failure is reported once.
Simple multi-statement QueryIntegrationCommand responses may arrive independently; the connection remains usable and lifecycle completion waits for readiness where required.
Statement or portal closeUnitOnly unreferenced metric state is released; PostgreSQL’s surviving portal semantics remain intact.
+
Validation principle: the branch’s unit tests pin most transitions. Its one new integration regression pins connection usability after a multi-statement simple Query; it does not independently exercise the full lifecycle matrix.
+
+ +
+

Where this PR deliberately stops

+

Context owns CipherStash’s correlated metadata, not the PostgreSQL protocol itself. pg-proto still owns ordering and backpressure. Schema middleware still owns transaction-aware schema overlays and publication. ZeroKMS remains the encryption adapter. The refactor also keeps existing Prometheus metric names, despite clarifying their semantics.

+
A reviewer should be able to approve the ownership change without also approving a new protocol engine, schema state machine, encryption boundary, or metrics migration. Those are explicitly outside this branch.
+
+ + + +
+
+ +