diff --git a/CHANGELOG.md b/CHANGELOG.md index e7d00852..0c47cf18 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - **PostgreSQL protocol error handling after the pg-proto migration**: Proxy now rejects `require_tls` configurations that omit a certificate, preserves PostgreSQL transaction state when statement mapping fails, returns decryption failures as PostgreSQL errors without dropping the connection, and reloads changed schemas only after PostgreSQL confirms the transaction boundary. Prepared-statement replacement also preserves existing portals and overlapping statement metrics. +### Security + +- **DDL now updates encryption metadata transactionally**: Proxy applies schema changes only after PostgreSQL confirms execution, keeps successful changes connection-local until commit, and atomically publishes schema and EQL domain metadata before reporting idle readiness. Extended-protocol DDL, explicit transactions, savepoints, rollbacks, one-`Sync` pipelining, and already-open connections now observe the correct schema generation. Unmodelled DDL, simple-query batches whose DDL may change encryption metadata before a dependent statement, and failed catalog publication fail closed instead of risking plaintext writes through stale metadata; encryption-neutral DDL and native temporary-table batches remain compatible. + ## [3.0.1] - 2026-08-05 ### Added diff --git a/Cargo.lock b/Cargo.lock index 6d3f842e..1c10c355 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3157,9 +3157,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pg-proto" -version = "0.11.1" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72f5b88d40736045b71bb72d114a6f212c98095e1f43cfc998c948e429b55012" +checksum = "d3e010cd297396de7654638926d75d89a0a313cc9e526fa42a480f1c0e0f0197" dependencies = [ "base64 0.23.1", "bytes", @@ -3180,9 +3180,9 @@ dependencies = [ [[package]] name = "pg-proto-fsm" -version = "0.11.1" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc37a28331c337145f6500402d7857565ff807891712d79833e0fbc00567c48d" +checksum = "bb7664de912f117f8ce3e2abf9a032cd69e577b5a8edd0efe72b48559337d6fb" dependencies = [ "proc-macro2", "quote", diff --git a/docs/errors.md b/docs/errors.md index d7267cce..7f085306 100644 --- a/docs/errors.md +++ b/docs/errors.md @@ -12,6 +12,8 @@ - [Invalid SQL statement](#mapping-invalid-sql-statement) - [Unsupported parameter type](#mapping-unsupported-parameter-type) - [Statement could not be type checked](#mapping-statement-could-not-be-type-checked) + - [Dependent statement after DDL](#mapping-dependent-statement-after-ddl) + - [Unmodelled DDL](#mapping-unmodelled-ddl) - [Unmappable encrypted column](#mapping-unmappable-encrypted-column) - [Internal Error](#mapping-internal-error) @@ -264,6 +266,38 @@ If the error persists, please contact CipherStash [support](https://cipherstash. + + + +## Dependent statement after DDL + +A simple-query batch contains a schema-dependent statement after DDL that may change encryption +metadata. Proxy cannot observe the DDL execution result between statements in one simple-query +message, so it refuses the complete batch before PostgreSQL executes any part of it. Native DDL +and native temporary-table batches continue to pass through because they introduce no encryption +obligation. + +### How to fix + +Send the DDL and the dependent statement as separate queries. Extended-protocol clients may +pipeline them; Proxy defers dependent mapping until PostgreSQL reports the DDL outcome. + + + + + +## Unmodelled DDL + +PostgreSQL successfully executed a schema change whose connection-local effect Proxy cannot model +safely, such as conditional or cascading DDL. Schema-dependent statements are refused for the +rest of that transaction. + +### How to fix + +Roll back the transaction, or commit it and wait for Proxy to publish an authoritative catalog +snapshot before issuing schema-dependent statements. + + diff --git a/packages/cipherstash-proxy-integration/src/lib.rs b/packages/cipherstash-proxy-integration/src/lib.rs index d2a713d5..63b1e022 100644 --- a/packages/cipherstash-proxy-integration/src/lib.rs +++ b/packages/cipherstash-proxy-integration/src/lib.rs @@ -25,6 +25,7 @@ mod multitenant; mod ore_order_helpers; mod passthrough; mod pipeline; +/// Database-backed transaction-aware schema middleware regressions. mod schema_change; mod select; mod set_keyset_error; diff --git a/packages/cipherstash-proxy-integration/src/schema_change.rs b/packages/cipherstash-proxy-integration/src/schema_change.rs index ff1fa042..e610219b 100644 --- a/packages/cipherstash-proxy-integration/src/schema_change.rs +++ b/packages/cipherstash-proxy-integration/src/schema_change.rs @@ -1,25 +1,243 @@ #[cfg(test)] +/// End-to-end schema-change tests through Proxy and directly against PostgreSQL. mod tests { - use crate::common::{connect_with_tls, random_id, PROXY}; + use crate::common::{connect, connect_with_tls, get_database_port, random_id, PROXY}; + use tokio_postgres::Client; + + async fn connect_for_test(port: u16) -> Client { + if std::env::var("CS_TEST_USE_TLS").as_deref() == Ok("false") { + connect(port).await + } else { + connect_with_tls(port).await + } + } + + fn table(prefix: &str) -> String { + format!("{prefix}_{}", random_id()) + } + + fn create_encrypted_table(table: &str) -> String { + format!("CREATE TABLE {table} (id bigint PRIMARY KEY, secret eql_v3_text_search NOT NULL)") + } + + async fn assert_ciphertext_at_rest(table: &str, id: i64, plaintext: &str) { + let postgres = connect_for_test(get_database_port()).await; + let sql = format!("SELECT secret::text FROM {table} WHERE id = $1"); + let stored: String = postgres.query_one(&sql, &[&id]).await.unwrap().get(0); + + assert!( + !stored.contains(plaintext), + "plaintext reached PostgreSQL: {stored}" + ); + let payload: serde_json::Value = serde_json::from_str(&stored).unwrap(); + assert!( + payload.get("c").is_some(), + "missing record ciphertext: {payload}" + ); + } + + async fn insert_secret(client: &Client, table: &str, id: i64, plaintext: &str) { + let sql = format!("INSERT INTO {table} (id, secret) VALUES ($1, $2)"); + assert_eq!(client.execute(&sql, &[&id, &plaintext]).await.unwrap(), 1); + } #[tokio::test] - async fn schema_change_reloads_schema() { - let client = connect_with_tls(*PROXY).await; + async fn later_connection_encrypts_immediately_after_extended_protocol_ddl() { + let ddl_connection = connect_for_test(*PROXY).await; + let already_open_connection = connect_for_test(*PROXY).await; + let table = table("bug_308_extended"); - let id = random_id(); + ddl_connection + .execute(&create_encrypted_table(&table), &[]) + .await + .unwrap(); - let sql = format!( - "CREATE TABLE table_{id} ( - id bigint, - PRIMARY KEY(id) - );" + insert_secret(&already_open_connection, &table, 1, "classified").await; + assert_ciphertext_at_rest(&table, 1, "classified").await; + } + + #[tokio::test] + async fn explicit_transaction_uses_successful_ddl_overlay_before_commit() { + let client = connect_for_test(*PROXY).await; + let table = table("bug_308_transaction"); + + client.batch_execute("BEGIN").await.unwrap(); + client + .execute(&create_encrypted_table(&table), &[]) + .await + .unwrap(); + insert_secret(&client, &table, 1, "inside transaction").await; + client.batch_execute("COMMIT").await.unwrap(); + + assert_ciphertext_at_rest(&table, 1, "inside transaction").await; + } + + #[tokio::test] + async fn encryption_neutral_alter_table_keeps_transaction_mappable() { + let client = connect_for_test(*PROXY).await; + let table = table("bug_308_safe_alter"); + + client + .execute(&create_encrypted_table(&table), &[]) + .await + .unwrap(); + client.batch_execute("BEGIN").await.unwrap(); + client + .batch_execute(&format!( + "ALTER TABLE {table} ALTER COLUMN secret SET NOT NULL" + )) + .await + .unwrap(); + insert_secret(&client, &table, 1, "after safe alter").await; + client.batch_execute("COMMIT").await.unwrap(); + + assert_ciphertext_at_rest(&table, 1, "after safe alter").await; + } + + #[tokio::test] + async fn pipelined_statement_waits_for_extended_ddl_activation() { + let client = connect_for_test(*PROXY).await; + let table = table("bug_308_pipeline"); + let create = create_encrypted_table(&table); + let insert = format!("INSERT INTO {table} (id, secret) VALUES ($1, $2)"); + let create = client.prepare(&create).await.unwrap(); + + let (created, inserted) = tokio::join!( + client.execute(&create, &[]), + client.execute(&insert, &[&1_i64, &"pipelined"]), ); + created.unwrap(); + assert_eq!(inserted.unwrap(), 1); + + assert_ciphertext_at_rest(&table, 1, "pipelined").await; + } + + #[tokio::test] + async fn rollback_discards_successful_ddl_overlay() { + let client = connect_for_test(*PROXY).await; + let postgres = connect_for_test(get_database_port()).await; + let table = table("bug_308_rollback"); + + client.batch_execute("BEGIN").await.unwrap(); + client + .execute(&create_encrypted_table(&table), &[]) + .await + .unwrap(); + client.batch_execute("ROLLBACK").await.unwrap(); + + let exists: bool = postgres + .query_one("SELECT to_regclass($1) IS NOT NULL", &[&table]) + .await + .unwrap() + .get(0); + assert!(!exists); + } + + #[tokio::test] + async fn rollback_to_savepoint_restores_schema_and_encryption_overlay() { + let client = connect_for_test(*PROXY).await; + let postgres = connect_for_test(get_database_port()).await; + let retained = table("bug_308_retained"); + let reverted = table("bug_308_reverted"); + + client.batch_execute("BEGIN").await.unwrap(); + client + .execute(&create_encrypted_table(&retained), &[]) + .await + .unwrap(); + client + .batch_execute("SAVEPOINT Before_Reverted") + .await + .unwrap(); + client + .execute(&create_encrypted_table(&reverted), &[]) + .await + .unwrap(); + client + .batch_execute("ROLLBACK TO SAVEPOINT before_reverted") + .await + .unwrap(); + insert_secret(&client, &retained, 1, "savepoint secret").await; + client.batch_execute("COMMIT").await.unwrap(); + + assert_ciphertext_at_rest(&retained, 1, "savepoint secret").await; + let exists: bool = postgres + .query_one("SELECT to_regclass($1) IS NOT NULL", &[&reverted]) + .await + .unwrap() + .get(0); + assert!(!exists); + } - let _ = client.execute(&sql, &[]).await.unwrap(); + #[tokio::test] + async fn simple_query_batch_with_dependent_post_ddl_statement_fails_closed() { + let client = connect_for_test(*PROXY).await; + let postgres = connect_for_test(get_database_port()).await; + let table = table("bug_308_simple_batch"); + let batch = format!( + "{}; INSERT INTO {table} (id, secret) VALUES (1, 'plaintext')", + create_encrypted_table(&table) + ); + + assert!(client.simple_query(&batch).await.is_err()); + + let exists: bool = postgres + .query_one("SELECT to_regclass($1) IS NOT NULL", &[&table]) + .await + .unwrap() + .get(0); + assert!(!exists); + } + + #[tokio::test] + async fn compatibility_fallback_tracks_native_ddl_for_a_later_encrypted_alter() { + let client = connect_for_test(*PROXY).await; + let table = table("bug_308_fallback"); + + client + .batch_execute(&format!( + "CREATE TABLE {table} (id bigint PRIMARY KEY); \ + INSERT INTO {table} (id) VALUES (1)" + )) + .await + .unwrap(); + client + .execute( + &format!("ALTER TABLE {table} ADD COLUMN secret eql_v3_text_search"), + &[], + ) + .await + .unwrap(); + + insert_secret(&client, &table, 2, "after fallback").await; + assert_ciphertext_at_rest(&table, 2, "after fallback").await; + } + + #[tokio::test] + async fn rewritten_simple_query_preserves_native_ddl_and_its_intent() { + let client = connect_for_test(*PROXY).await; + let encrypted = table("bug_308_rewritten"); + let staging = table("bug_308_staging"); - let sql = format!("SELECT id FROM table_{id}"); - let rows = client.query(&sql, &[]).await.unwrap(); + client + .execute(&create_encrypted_table(&encrypted), &[]) + .await + .unwrap(); + client + .batch_execute(&format!( + "CREATE TABLE {staging} (id bigint PRIMARY KEY); \ + INSERT INTO {encrypted} (id, secret) VALUES (1, 'preserved batch secret')" + )) + .await + .unwrap(); - assert!(rows.is_empty()); + let postgres = connect_for_test(get_database_port()).await; + let exists: bool = postgres + .query_one("SELECT to_regclass($1) IS NOT NULL", &[&staging]) + .await + .unwrap() + .get(0); + assert!(exists); + assert_ciphertext_at_rest(&encrypted, 1, "preserved batch secret").await; } } diff --git a/packages/cipherstash-proxy/CONTEXT.md b/packages/cipherstash-proxy/CONTEXT.md index 9289b0a0..c74c9dc5 100644 --- a/packages/cipherstash-proxy/CONTEXT.md +++ b/packages/cipherstash-proxy/CONTEXT.md @@ -103,8 +103,40 @@ Proxy's in-band control API, intercepted rather than forwarded — `KEYSET_ID`, that print `CIPHERSTASH.DISABLE_MAPPING` are wrong. **Reload**: -Re-reading state from the database after observed DDL. Two independent things reload: the -database schema, and the column encrypt config. +Re-reading authoritative schema state from PostgreSQL after observed DDL. A reload produces +one **committed schema snapshot**; it does not merge Proxy's inferred DDL effects into shared +state. + +**Committed schema snapshot**: +An immutable, monotonically versioned pair of database structure and column encryption +metadata loaded from PostgreSQL. The pair is published atomically because a table without its +encryption policy (or an encryption policy without its table) is not a valid observable state. + +**Transaction schema overlay**: +The confirmed effects of successful DDL executions in one connection's current transaction. +It is checkpointed by savepoints, restored by `ROLLBACK TO SAVEPOINT`, and discarded by a full +rollback. Parsed or prepared DDL is only intent; it enters the overlay after PostgreSQL reports +successful execution. + +**Effective schema**: +The committed schema snapshot pinned when a transaction starts, with that transaction's schema +overlay applied. EQL Mapper type-checks and transforms against this view. An idle connection +adopts the latest committed snapshot before its next transaction. + +**Schema publication**: +Atomically replacing the shared committed schema snapshot after the outermost transaction +containing DDL commits and an authoritative catalog reload succeeds. Proxy completes publication +before forwarding `ReadyForQuery(I)`, so a connection opened after readiness observes the new +schema and encryption metadata. Failed publication is fail-closed: the affected connection is +closed without forwarding readiness, and the dirty publication remains eligible for retry. + +**Schema middleware**: +The owner of transactional schema state. Frontend and Backend report protocol lifecycle events; +they do not directly change overlays or dirty flags. The middleware owns DDL detection, prepared +DDL effects, successful-execution activation, savepoint and transaction transitions, +effective-schema resolution, and the decision that publication is required. `Context` performs the +authoritative reload round trip, while `SchemaManager` coalesces reloads and orders their +generations. See `docs/adr/0001-transaction-aware-schema-middleware.md`. ## Note on `session` diff --git a/packages/cipherstash-proxy/Cargo.toml b/packages/cipherstash-proxy/Cargo.toml index f694b0f0..a33baf77 100644 --- a/packages/cipherstash-proxy/Cargo.toml +++ b/packages/cipherstash-proxy/Cargo.toml @@ -28,7 +28,7 @@ md-5 = "0.10.6" metrics = "0.24.3" metrics-exporter-prometheus = "0.17" moka = { version = "0.12", features = ["future"] } -pg-proto = "0.11.1" +pg-proto = "0.12.0" postgres-protocol = "0.6.7" postgres-types = { version = "0.2.8", features = ["with-serde_json-1"] } rand = "0.9" diff --git a/packages/cipherstash-proxy/docs/adr/0001-transaction-aware-schema-middleware.md b/packages/cipherstash-proxy/docs/adr/0001-transaction-aware-schema-middleware.md new file mode 100644 index 00000000..614e9f13 --- /dev/null +++ b/packages/cipherstash-proxy/docs/adr/0001-transaction-aware-schema-middleware.md @@ -0,0 +1,102 @@ +--- +status: accepted +issue: BUG-308 +--- + +# Transaction-aware schema middleware + +## Context + +Proxy must map a statement using the schema PostgreSQL makes visible to that statement. DDL is +transactional: its effects become visible inside the transaction after successful execution, +but other connections cannot observe them until the outermost transaction commits. + +The existing design marks schema change while parsing SQL and reloads through a separate database +connection. In the extended protocol, a client's preparation `Sync` can consume that marker before +the DDL is executed. Even moving the reload to every `ReadyForQuery` is insufficient: a +`ReadyForQuery(T)` occurs inside an explicit transaction, where the loader connection still cannot +see uncommitted DDL. Schema and column encryption configuration are also loaded and published +separately, while existing connections retain snapshots taken when their contexts were created. + +These behaviours can make later connections use stale mapping or encryption metadata. They also +make connection-local behaviour depend on speculative DDL inferred at `Parse`, even if execution +later fails. + +## Decision + +Introduce standalone schema middleware as the sole owner of the transactional schema lifecycle. +Frontend and Backend report protocol events to it; neither manipulates schema-change flags or +reload managers directly. + +### State model + +- A **committed schema snapshot** is immutable and monotonically versioned. It contains database + structure and encryption metadata derived from EQL domain types as one atomic value. +- A transaction pins the current committed snapshot. Existing idle connections adopt the latest + snapshot before starting their next transaction. +- A **transaction schema overlay** contains only effects confirmed by successful DDL execution. + The connection's **effective schema** is its pinned snapshot plus this overlay. +- Savepoints checkpoint the overlay. `ROLLBACK TO SAVEPOINT` restores its checkpoint, full rollback + discards it, and release preserves its effects in the enclosing transaction. +- Parsed or prepared DDL records intent with the prepared statement. Each successful execution + applies its effect; `Parse` alone never changes schema state. +- If a successful DDL cannot be modelled accurately, later schema-dependent statements in that + transaction fail closed. + +### Protocol ordering + +After a DDL `Execute` is forwarded, protocol-control messages required to complete that execution +continue to flow, but later schema-dependent operations wait until its success or failure is known. +Proxy injects `Flush` after the DDL `Execute`, allowing PostgreSQL to return `CommandComplete` +without waiting for a client `Sync`. The client's `Sync` remains the sole synchronization boundary, +so it still receives exactly one `ReadyForQuery`. This supports clients that pipeline DDL and a +dependent `Parse` in one batch without speculative mapping or a protocol deadlock. + +A simple-query message fails closed when DDL may change encryption metadata and a later statement +requires that metadata for mapping. Native DDL and native temporary-table batches continue to pass +through because they introduce no encryption obligation. Proxy derives execution intents from the +statements it actually forwards, including when a mapping error uses the compatibility passthrough +fallback, so backend outcomes cannot become misaligned with phantom intents. + +### Publication + +When the outermost transaction containing successful DDL commits, one reload coordinator reads the +authoritative PostgreSQL catalog. Concurrent publication requests are coalesced, and generation +ordering prevents an older reload from replacing newer state. The coordinator atomically publishes +the combined schema and encryption snapshot before Proxy forwards `ReadyForQuery(I)`. + +Proxy does not merge its inferred overlay into shared state. PostgreSQL remains authoritative for +cascades, conditional DDL, server-side effects, and the final outcome of the transaction. + +If publication fails after PostgreSQL has committed, Proxy retains the dirty publication for retry, +does not forward successful readiness, and closes the affected client connection. The database +commit cannot be undone, but Proxy must not imply that stale encryption metadata is safe to use. + +## Consequences + +- DDL becomes visible to later statements on the same connection immediately after successful + execution, including within an explicit transaction. +- Other connections observe DDL only after commit and successful publication. +- Every transaction maps against a stable schema and encryption-policy generation. +- Frontend and Backend become protocol adapters around a testable schema state machine. +- Extended-protocol pipelining requires bounded deferral after DDL execution. +- Native temporary tables are connection-local and absent from authoritative reloads. Proxy ignores + them only when they cannot shadow an encrypted table or introduce EQL columns; unsafe cases fail + closed for the rest of the connection (unless rollback to an earlier savepoint removes the + object), because a global catalog reload cannot prove connection-local state disappeared. +- The local overlay deliberately models only deterministic schema changes. Encryption-neutral + constraints, defaults, nullability, ownership, trigger/rule state, and row-level-security state + are accepted, while conditional, cascading, table-rewriting, view, and type-changing operations + remain unmodelled within a transaction. +- Availability is intentionally sacrificed when committed schema state cannot be published safely. +- Schema and encryption managers can no longer publish independent observable states. + +## Verification + +State-machine tests cover successful execution, execution failure, explicit commit, full rollback, +savepoint rollback, generation ordering, deferral, unmodelled DDL, reload failure, and safe native +batches. Database-backed tests cover extended-protocol autocommit, explicit transactions, an +already-open second connection, pipelining with a single client `Sync`, direct ciphertext +verification, safe `ALTER TABLE`, native temporary tables, and both compatibility-fallback and +rewritten-batch intent bookkeeping. Backend-adapter tests prove that publication failure closes +the connection before idle readiness. diff --git a/packages/cipherstash-proxy/src/config/tandem.rs b/packages/cipherstash-proxy/src/config/tandem.rs index 4329a163..9a331abc 100644 --- a/packages/cipherstash-proxy/src/config/tandem.rs +++ b/packages/cipherstash-proxy/src/config/tandem.rs @@ -278,7 +278,7 @@ impl TandemConfig { /// In order of precedence /// config if explicitly set /// RUST_MIN_STACK env var if set - /// DEBUG_THREAD_STACK_SIZE if log level is Debug or Trace + /// DEBUG_THREAD_STACK_SIZE for debug builds or Debug/Trace logging /// otherwise set to DEFAULT_THREAD_STACK_SIZE (2MiB) /// pub fn thread_stack_size(&self) -> usize { @@ -289,7 +289,7 @@ impl TandemConfig { // If the environment variable is set, use that value if let Ok(stack_size) = env::var("RUST_MIN_STACK") { - stack_size + return stack_size .parse() .inspect_err(|err| { println!("Could not parse env var RUST_MIN_STACK: {err}"); @@ -298,7 +298,10 @@ impl TandemConfig { .unwrap_or(DEFAULT_THREAD_STACK_SIZE); } - if self.log.level == LogLevel::Debug || self.log.level == LogLevel::Trace { + if cfg!(debug_assertions) + || self.log.level == LogLevel::Debug + || self.log.level == LogLevel::Trace + { return DEBUG_THREAD_STACK_SIZE; } @@ -954,4 +957,28 @@ mod tests { ); }); } + + #[test] + fn thread_stack_size_honors_rust_min_stack() { + temp_env::with_var("RUST_MIN_STACK", Some("8388608"), || { + assert_eq!( + TandemConfig::for_testing().thread_stack_size(), + 8 * 1024 * 1024 + ); + }); + } + + #[test] + fn debug_builds_use_the_larger_thread_stack() { + temp_env::with_var("RUST_MIN_STACK", None::<&str>, || { + assert_eq!( + TandemConfig::for_testing().thread_stack_size(), + if cfg!(debug_assertions) { + crate::config::DEBUG_THREAD_STACK_SIZE + } else { + crate::config::DEFAULT_THREAD_STACK_SIZE + } + ); + }); + } } diff --git a/packages/cipherstash-proxy/src/error.rs b/packages/cipherstash-proxy/src/error.rs index c1b6ec5f..fa47dfef 100644 --- a/packages/cipherstash-proxy/src/error.rs +++ b/packages/cipherstash-proxy/src/error.rs @@ -74,8 +74,11 @@ impl Error { // Forwarding a statement that references a legacy EQL v2 column // stores plaintext in a column its operator believes is encrypted // (CIP-3688). No configuration may turn that back on. - Error::Mapping(MappingError::UnmappableEncryptedColumn { .. }) - | Error::Encrypt(EncryptError::InvalidInboundEqlPayload) + Error::Mapping( + MappingError::UnmappableEncryptedColumn { .. } + | MappingError::DependentStatementAfterDdl + | MappingError::UnmodelledDdl + ) | Error::Encrypt(EncryptError::InvalidInboundEqlPayload) ) } } @@ -100,6 +103,14 @@ pub enum ZeroKMSError { #[derive(Error, Debug)] pub enum MappingError { + /// A simple-query batch would map against encryption metadata changed earlier in the batch. + #[error("A simple-query batch cannot contain a schema-dependent statement after DDL that may change encryption metadata. Send the DDL and dependent statement as separate queries. For help visit {}#mapping-dependent-statement-after-ddl", ERROR_DOC_BASE_URL)] + DependentStatementAfterDdl, + + /// Confirmed DDL cannot be represented safely by the transaction overlay. + #[error("A successful schema change in this transaction cannot be modelled safely. Roll back the transaction before issuing schema-dependent statements. For help visit {}#mapping-unmodelled-ddl", ERROR_DOC_BASE_URL)] + UnmodelledDdl, + #[error("Invalid parameter for column '{}' of type '{}' in table '{}' (OID {}). For help visit {}#mapping-invalid-parameter", _0.column_name(), _0.cast_type(), _0.table_name(), _0.oid(), ERROR_DOC_BASE_URL)] InvalidParameter(Box), diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index 6d2b3436..0bf57f99 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -7,13 +7,16 @@ pub use self::{phase_timing::PhaseTiming, portal::Portal, statement::Statement}; use super::{column_mapper::ColumnMapper, rewrite::Name, Column}; use crate::{ config::TandemConfig, - error::{EncryptError, Error}, + error::{ConfigError, EncryptError, Error}, log::{CONTEXT, SLOW_STATEMENTS}, prometheus::{ SLOW_STATEMENTS_TOTAL, STATEMENTS_EXECUTION_DURATION_SECONDS, STATEMENTS_SESSION_DURATION_SECONDS, }, - proxy::{EncryptConfig, EncryptionService, ReloadCommand, ReloadSender}, + proxy::{ + schema::{CommittedSchemaStore, SchemaMiddleware}, + EncryptConfig, EncryptionService, ReloadCommand, ReloadSender, + }, }; use cipherstash_client::IdentifiedBy; use eql_mapper::{Schema, TableResolver}; @@ -25,7 +28,7 @@ pub use statement_metadata::StatementMetadata; use std::{ collections::HashMap, sync::{ - atomic::{AtomicBool, AtomicU64, Ordering}, + atomic::{AtomicU64, Ordering}, Arc, LazyLock, RwLock, }, time::{Duration, Instant}, @@ -53,17 +56,14 @@ where { pub client_id: i32, config: Arc, - encrypt_config: Arc, encryption: T, reload_sender: ReloadSender, - column_mapper: ColumnMapper, + schema_middleware: SchemaMiddleware, statements: Arc>>>, statement_sessions: Arc>>, portals: Arc>>>, operations: Arc>>, - schema_changed: Arc, session_metrics: Arc>>, - table_resolver: Arc, upstream_tls_roots: Arc, unsafe_disable_mapping: bool, keyset_id: Arc>>, @@ -153,21 +153,39 @@ where encryption: T, reload_sender: ReloadSender, ) -> Context { - let column_mapper = ColumnMapper::new(encrypt_config.clone()); + let schema_store = + CommittedSchemaStore::from_parts((*schema).clone(), (*encrypt_config).clone()); + Self::new_with_schema_store( + client_id, + config, + schema_store, + upstream_tls_roots, + encryption, + reload_sender, + ) + } + + /// Constructs a connection context over the shared committed schema store. + pub fn new_with_schema_store( + client_id: i32, + config: Arc, + schema_store: CommittedSchemaStore, + upstream_tls_roots: Arc, + encryption: T, + reload_sender: ReloadSender, + ) -> Context { + 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())), - schema_changed: Arc::new(AtomicBool::new(false)), session_metrics: Arc::new(RwLock::new(HashMap::new())), - table_resolver: Arc::new(TableResolver::new_editable(schema)), upstream_tls_roots, client_id, config, - encrypt_config, - column_mapper, + schema_middleware, encryption, reload_sender, unsafe_disable_mapping: false, @@ -588,20 +606,96 @@ where Some(session_context.to_owned()) } - pub fn set_schema_changed(&self) { - debug!(target: CONTEXT, - client_id = self.client_id, - msg = "Schema changed" - ); - self.schema_changed.store(true, Ordering::Release); + /// Returns the resolver for this connection's effective schema snapshot. + pub fn get_table_resolver(&self) -> Arc { + self.schema_middleware.resolver() } - pub fn take_schema_changed(&self) -> bool { - self.schema_changed.swap(false, Ordering::AcqRel) + /// Records schema intent for a parsed prepared statement. + pub fn prepare_schema_statement(&self, name: Name, statement: sqltk::parser::ast::Statement) { + self.schema_middleware.prepare(name, statement); } - pub fn get_table_resolver(&self) -> Arc { - self.table_resolver.clone() + /// Associates a bound portal with its prepared statement's schema intent. + pub fn bind_schema_statement(&self, portal: Name, prepared_statement: &Name) { + self.schema_middleware.bind(portal, prepared_statement); + } + + /// Records a portal execution and returns whether its DDL needs an injected flush. + pub fn execute_schema_portal(&self, portal: &Name) -> bool { + self.schema_middleware.execute(portal) + } + + /// Records statements in one simple-query protocol message. + pub fn execute_simple_schema_statements(&self, statements: &[sqltk::parser::ast::Statement]) { + self.schema_middleware.simple_query(statements); + } + + /// Returns whether a simple-query batch must be rejected to protect encryption. + pub fn simple_query_requires_fail_closed( + &self, + statements: &[sqltk::parser::ast::Statement], + ) -> bool { + self.schema_middleware + .simple_query_requires_fail_closed(statements) + } + + /// Records an extended-protocol synchronization boundary. + pub fn mark_schema_protocol_boundary(&self) { + self.schema_middleware.protocol_boundary(); + } + + /// Reports successful execution of the next queued statement. + pub fn schema_execution_succeeded(&self) { + self.schema_middleware.execution_succeeded(); + } + + /// Reports failed execution of the next queued statement. + pub fn schema_execution_failed(&self) { + self.schema_middleware.execution_failed(); + } + + /// Waits for preceding schema-changing executions to resolve. + pub async fn wait_for_schema_execution(&self) { + self.schema_middleware.wait_for_ddl().await; + } + + /// Returns whether a schema-changing execution is awaiting a backend outcome. + pub fn schema_ddl_in_flight(&self) -> bool { + self.schema_middleware.ddl_in_flight() + } + + /// Refuses schema-dependent work after confirmed unmodelled DDL. + pub fn ensure_schema_modelled(&self) -> Result<(), Error> { + if self.schema_middleware.has_unmodelled_ddl() { + return Err(crate::error::MappingError::UnmodelledDdl.into()); + } + Ok(()) + } + + /// Replaces idle connection-local state with the latest committed snapshot. + pub fn adopt_latest_schema(&self) { + self.schema_middleware.adopt_latest(); + } + + /// Publishes pending shared state, then prepares the effective schema for mapping. + pub async fn prepare_schema_for_statement(&self) -> Result<(), Error> { + if self + .schema_middleware + .requires_publication_before_statement() + { + if !self.reload_schema().await { + return Err(ConfigError::SchemaCouldNotBeLoaded.into()); + } + self.schema_middleware.publication_succeeded(); + } + self.schema_middleware.before_statement(); + Ok(()) + } + + /// Reports a readiness boundary and PostgreSQL transaction status. + pub fn schema_ready_for_query(&self, status: crate::proxy::schema::TransactionStatus) { + self.schema_middleware.ready_for_query(status); } /// Examines a [`sqltk::parser::ast::Statement`] and if it is precisely equal to `SET UNSAFE_DISABLE_MAPPING = {boolean};` @@ -833,16 +927,26 @@ where } /// Reload schema if it has changed since last check. - pub async fn reload_schema_if_changed(&self) { - if self.take_schema_changed() && !self.reload_schema().await { - // Preserve the dirty state when the reload task is unavailable so - // a later statement can retry instead of silently losing the DDL. - self.set_schema_changed(); + pub async fn publish_schema_if_changed(&self) -> Result<(), Error> { + if !self.schema_middleware.needs_publication() { + self.adopt_latest_schema(); + return Ok(()); } + + if self.schema_middleware.has_local_changes() { + self.schema_middleware.mark_publication_pending(); + } + + if !self.reload_schema().await { + return Err(ConfigError::SchemaCouldNotBeLoaded.into()); + } + + self.schema_middleware.publication_succeeded(); + Ok(()) } pub fn is_passthrough(&self) -> bool { - self.encrypt_config.is_empty() || self.config.mapping_disabled() + self.schema_middleware.encrypt_config().is_empty() || self.config.mapping_disabled() } // Column processing delegation methods @@ -850,28 +954,31 @@ where &self, typed_statement: &eql_mapper::TypeCheckedStatement<'_>, ) -> Result>, Error> { - self.column_mapper.get_projection_columns(typed_statement) + ColumnMapper::new(self.schema_middleware.encrypt_config()) + .get_projection_columns(typed_statement) } pub fn get_param_columns( &self, typed_statement: &eql_mapper::TypeCheckedStatement<'_>, ) -> Result>, Error> { - self.column_mapper.get_param_columns(typed_statement) + ColumnMapper::new(self.schema_middleware.encrypt_config()) + .get_param_columns(typed_statement) } pub fn get_output_param_columns( &self, plan: &eql_mapper::ParamPlan, ) -> Result>, Error> { - self.column_mapper.get_output_param_columns(plan) + ColumnMapper::new(self.schema_middleware.encrypt_config()).get_output_param_columns(plan) } pub fn get_literal_columns( &self, typed_statement: &eql_mapper::TypeCheckedStatement<'_>, ) -> Result>, Error> { - self.column_mapper.get_literal_columns(typed_statement) + ColumnMapper::new(self.schema_middleware.encrypt_config()) + .get_literal_columns(typed_statement) } // Direct config access methods @@ -1019,7 +1126,7 @@ mod tests { error::Error, log, postgresql::{rewrite::Name, Column}, - proxy::{EncryptConfig, EncryptionService, ReloadCommand}, + proxy::{EncryptConfig, EncryptionService}, TandemConfig, }; use cipherstash_client::IdentifiedBy; @@ -1082,70 +1189,6 @@ mod tests { ) } - #[tokio::test] - async fn successful_schema_reload_consumes_change_flag_once() { - let config = Arc::new(TandemConfig::for_testing()); - let encrypt_config = Arc::new(EncryptConfig::default()); - let schema = Arc::new(Schema::new("public")); - let (reload_sender, mut reload_receiver) = mpsc::unbounded_channel(); - let context = Context::new( - 1, - config, - encrypt_config, - schema, - Arc::new(rustls::RootCertStore::empty()), - TestService {}, - reload_sender, - ); - let reload_task = tokio::spawn(async move { - let Some(ReloadCommand::DatabaseSchema(responder)) = reload_receiver.recv().await - else { - panic!("expected database schema reload"); - }; - responder.send(true).expect("reload receiver is alive"); - tokio::time::timeout(std::time::Duration::from_millis(20), reload_receiver.recv()) - .await - .is_err() - }); - - context.set_schema_changed(); - context.reload_schema_if_changed().await; - context.reload_schema_if_changed().await; - - assert!(!context.take_schema_changed()); - assert!(reload_task.await.expect("reload task did not panic")); - } - - #[tokio::test] - async fn failed_schema_reload_keeps_change_flag_for_retry() { - let config = Arc::new(TandemConfig::for_testing()); - let encrypt_config = Arc::new(EncryptConfig::default()); - let schema = Arc::new(Schema::new("public")); - let (reload_sender, mut reload_receiver) = mpsc::unbounded_channel(); - let context = Context::new( - 1, - config, - encrypt_config, - schema, - Arc::new(rustls::RootCertStore::empty()), - TestService {}, - reload_sender, - ); - let reload_task = tokio::spawn(async move { - let Some(ReloadCommand::DatabaseSchema(responder)) = reload_receiver.recv().await - else { - panic!("expected database schema reload"); - }; - responder.send(false).expect("reload receiver is alive"); - }); - - context.set_schema_changed(); - context.reload_schema_if_changed().await; - - reload_task.await.expect("reload task did not panic"); - assert!(context.take_schema_changed()); - } - #[tokio::test] async fn empty_plaintext_batch_does_not_call_encryption_service() { let context = create_context(); @@ -1157,7 +1200,6 @@ mod tests { assert_eq!(output.len(), 2); assert!(output.iter().all(Option::is_none)); } - fn statement() -> Statement { Statement { param_columns: vec![], diff --git a/packages/cipherstash-proxy/src/postgresql/driver.rs b/packages/cipherstash-proxy/src/postgresql/driver.rs index 58648016..da157542 100644 --- a/packages/cipherstash-proxy/src/postgresql/driver.rs +++ b/packages/cipherstash-proxy/src/postgresql/driver.rs @@ -160,6 +160,40 @@ where Err(pg_proto::ForwardError::Middleware(error)) => return Err(error), Err(error) => return Err(invalid_data(error)), }; + if matches!(&forwarded, ForwardedMessage::FrontendExpanded { .. }) { + // A DDL Execute is expanded to Execute + Flush so PostgreSQL can + // report its outcome before later schema-dependent frontend work. + // Keep the single-task intermediary progressing on the backend + // leg until that DDL resolves; otherwise a later frontend mapping + // call can wait for an outcome this task has not read yet. + // Box the whole secondary forwarding state machine so it + // does not enlarge the normal frontend driver's stack frame. + Box::pin(async { + while context.schema_ddl_in_flight() { + let backend = session.forward_backend(); + let drained = match connection_timeout { + Some(duration) => tokio::time::timeout(duration, backend) + .await + .map_err(|_| Error::ConnectionTimeout { duration })?, + None => backend.await, + }; + match drained { + Ok( + BackendForwarding::Forwarded(_) + | BackendForwarding::Expanded { .. } + | BackendForwarding::Suppressed(_) + | BackendForwarding::Held, + ) => {} + Err(pg_proto::ForwardError::Middleware(error)) => { + return Err(error) + } + Err(error) => return Err(invalid_data(error)), + } + } + Ok(()) + }) + .await?; + } if matches!( forwarded, ForwardedMessage::Frontend(FrontendMessage::Terminate) diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs index a8c9aa36..bff9e078 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs @@ -137,8 +137,8 @@ impl Backend { // forwarding ReadyForQuery. This ordering also guarantees that a // client opening its next connection after ReadyForQuery observes // the newly loaded schema and encrypt configuration. - if matches!(&protocol_message, BackendMessage::ReadyForQuery(_)) { - self.context.reload_schema_if_changed().await; + if let BackendMessage::ReadyForQuery(status) = &protocol_message { + self.handle_ready_for_query(*status).await?; } // CipherStash metadata is operation-keyed even in passthrough mode, @@ -147,6 +147,7 @@ impl Backend { 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); @@ -154,6 +155,7 @@ impl Backend { } } BackendMessage::ErrorResponse(_) => { + self.context.schema_execution_failed(); if let Some(operation) = operation { let session = self.context.complete_execution(operation); self.context.discard_operation(operation); @@ -165,9 +167,7 @@ impl Backend { self.context.complete_describe(operation); } } - BackendMessage::ReadyForQuery(status) => { - self.context.set_transaction_status(status); - } + BackendMessage::ReadyForQuery(_) => {} _ => {} } @@ -179,6 +179,7 @@ impl Backend { 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); @@ -186,6 +187,7 @@ impl Backend { } } BackendMessage::ErrorResponse(_) => { + self.context.schema_execution_failed(); if let Some(operation) = operation { let session = self.context.complete_execution(operation); self.context.discard_operation(operation); @@ -194,8 +196,7 @@ impl Backend { } BackendMessage::ReadyForQuery(status) => { self.discard_execution = false; - self.context.set_transaction_status(status); - self.context.reload_schema_if_changed().await; + self.handle_ready_for_query(status).await?; return Ok(BackendMiddlewareOutput::Forward( BackendMessage::ReadyForQuery(status), )); @@ -279,8 +280,10 @@ impl Backend { self.context.discard_operation(operation); self.context.finish_session(session); } + self.context.schema_execution_succeeded(); } BackendMessage::ErrorResponse(ref response) => { + self.context.schema_execution_failed(); self.error_response_handler(response); if let Some(operation) = operation { @@ -320,12 +323,11 @@ impl Backend { // Reload is potentially triggered by a FrontEnd Sync message. // However, the SimpleQuery flow does not use Sync so we check here as well BackendMessage::ReadyForQuery(status) => { - self.context.set_transaction_status(status); debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "ReadyForQuery" ); - self.context.reload_schema_if_changed().await; + self.handle_ready_for_query(status).await?; } _ => { @@ -345,6 +347,29 @@ impl Backend { } } + /// Publishes committed DDL before exposing idle readiness, then updates + /// connection-local transaction state from the same authoritative boundary. + async fn handle_ready_for_query( + &mut self, + status: pg_proto::TransactionStatus, + ) -> Result<(), Error> { + self.context.set_transaction_status(status); + if status == pg_proto::TransactionStatus::Idle { + self.context.publish_schema_if_changed().await?; + } + let schema_status = match status { + pg_proto::TransactionStatus::Idle => crate::proxy::schema::TransactionStatus::Idle, + pg_proto::TransactionStatus::InTransaction => { + crate::proxy::schema::TransactionStatus::InTransaction + } + pg_proto::TransactionStatus::FailedTransaction => { + crate::proxy::schema::TransactionStatus::FailedTransaction + } + }; + self.context.schema_ready_for_query(schema_status); + Ok(()) + } + /// Handles PostgreSQL ErrorResponse messages from the server. /// /// ErrorResponse messages indicate that an error occurred during SQL execution. @@ -719,6 +744,7 @@ mod tests { use super::*; use crate::config::TandemConfig; use crate::postgresql::context::KeysetIdentifier; + use crate::postgresql::parser::SqlParser; use crate::proxy::{EncryptConfig, EncryptionService}; use eql_mapper::Schema; use std::sync::Arc; @@ -786,7 +812,7 @@ mod tests { } #[tokio::test] - async fn passthrough_reloads_changed_schema_before_ready_for_query() { + async fn publication_failure_closes_connection_before_idle_readiness() { let config = Arc::new(TandemConfig::for_testing()); let encrypt_config = Arc::new(EncryptConfig::default()); let schema = Arc::new(Schema::new("public")); @@ -800,7 +826,9 @@ mod tests { TestService {}, reload_sender, ); - context.set_schema_changed(); + let ddl = SqlParser::parse_statement("create table reports (id bigint)").unwrap(); + context.execute_simple_schema_statements(&[ddl]); + context.schema_execution_succeeded(); let reload_task = tokio::spawn(async move { let Some(crate::proxy::ReloadCommand::DatabaseSchema(responder)) = @@ -808,17 +836,13 @@ mod tests { else { panic!("expected a database schema reload command"); }; - responder.send(true).expect("reload receiver must be open"); + responder.send(false).unwrap(); }); let mut backend = Backend::new(context); let ready = BackendMessage::ReadyForQuery(pg_proto::TransactionStatus::Idle); - let output = backend.intercept(None, ready.clone()).await.unwrap(); + let result = backend.intercept(None, ready).await; reload_task.await.unwrap(); - assert!(matches!( - output, - BackendMiddlewareOutput::Forward(message) if message == ready - )); - assert!(!backend.context.take_schema_changed()); + assert!(result.is_err()); } } diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs index cb6db2d5..23d72aa7 100644 --- a/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs @@ -116,6 +116,7 @@ impl Frontend { if self.failed_extended_batch { if matches!(protocol_message, FrontendMessage::Sync) { self.failed_extended_batch = false; + self.context.mark_schema_protocol_boundary(); return Ok(FrontendMiddlewareOutput::Forward(protocol_message)); } self.context.discard_operation(operation); @@ -123,6 +124,7 @@ impl Frontend { } let mut outbound_message = protocol_message.clone(); + let mut flush_after_forward = false; match protocol_message { FrontendMessage::Query(query) => { @@ -141,6 +143,7 @@ impl Frontend { self.context .set_operation_error(operation, response.clone()); self.context.set_execute(operation, Name::new(), session_id); + self.context.mark_schema_protocol_boundary(); outbound_message = self.simple_error_query(&response); } } @@ -149,7 +152,7 @@ impl Frontend { self.describe_handler(operation, describe).await?; } FrontendMessage::Execute(execute) => { - self.execute_handler(operation, execute).await?; + flush_after_forward = self.execute_handler(operation, execute).await?; } FrontendMessage::Parse(parse) => { let statement = parse.statement.clone(); @@ -242,6 +245,7 @@ impl Frontend { client_id = self.context.client_id, message = ?protocol_message, ); + self.context.mark_schema_protocol_boundary(); } FrontendMessage::Close(close) => { self.close_handler(close).await?; @@ -255,7 +259,11 @@ impl Frontend { } } - Ok(FrontendMiddlewareOutput::Forward(outbound_message)) + if flush_after_forward { + Ok(FrontendMiddlewareOutput::ForwardThenFlush(outbound_message)) + } else { + Ok(FrontendMiddlewareOutput::Forward(outbound_message)) + } } async fn describe_handler( @@ -281,11 +289,12 @@ impl Frontend { &mut self, operation: pg_proto::OperationId, execute: Execute, - ) -> Result<(), Error> { + ) -> 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()); - Ok(()) + Ok(executes_ddl) } /// Handles PostgreSQL Query messages (simple query protocol). @@ -338,7 +347,21 @@ impl Frontend { // Simple Query may contain many statements let query_text = String::from_utf8_lossy(&query).into_owned(); let parsed_statements = SqlParser::parse_statements(&query_text)?; - let mut transformed_statements = vec![]; + self.context.prepare_schema_for_statement().await?; + if self + .context + .simple_query_requires_fail_closed(&parsed_statements) + { + return Err(MappingError::DependentStatementAfterDdl.into()); + } + if parsed_statements + .iter() + .any(eql_mapper::requires_type_check) + { + self.context.wait_for_schema_execution().await; + self.context.ensure_schema_modelled()?; + } + let mut forwarded_statements = vec![]; debug!(target: MAPPER, client_id = self.context.client_id, @@ -362,15 +385,15 @@ impl Frontend { warn!(msg = "Encrypted statement mapping is not enabled"); counter!(STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL).increment(1); counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1); + forwarded_statements.push(statement.clone()); continue; } self.handle_set_keyset(statement)?; - self.check_for_schema_change(statement); - if !eql_mapper::requires_type_check(statement) { counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1); + forwarded_statements.push(statement.clone()); continue; } @@ -385,6 +408,12 @@ impl Frontend { if self.context.mapping_errors_enabled() || err.must_fail_closed() { return Err(err); } else { + self.record_simple_schema_execution( + operation, + session_id, + Portal::passthrough(Some(session_id)), + &parsed_statements, + ); return Ok(None); }; } @@ -397,6 +426,7 @@ impl Frontend { msg = "Encryptable Statement", ); + let mut transformed = false; if typed_statement.requires_transform() { // Record parse duration before encryption work starts if !parse_duration_recorded { @@ -413,9 +443,8 @@ impl Frontend { ) .await?; - if let Some(transformed_statement) = self - .transform_statement(&typed_statement, &encrypted_literals) - .await? + if let Some(transformed_statement) = + self.transform_statement(&typed_statement, &encrypted_literals)? { debug!(target: MAPPER, client_id = self.context.client_id, @@ -424,11 +453,16 @@ impl Frontend { // The simple protocol has no params, so the plan is // always empty here — only the SQL is needed. - transformed_statements.push(transformed_statement.statement); + forwarded_statements.push(transformed_statement.statement); encrypted = true; + transformed = true; } } + if !transformed { + forwarded_statements.push(typed_statement.statement.clone()); + } + counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1); // Set Encrypted portal and mark as mapped @@ -443,7 +477,7 @@ impl Frontend { msg = "Passthrough Statement" ); counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1); - transformed_statements.push(statement.clone()); + forwarded_statements.push(statement.clone()); } }; } @@ -473,12 +507,10 @@ impl Frontend { m.set_query_fingerprint(&query_text); }); - self.context.add_portal(Name::new(), portal); - self.context - .set_execute(operation, Name::new(), Some(session_id)); + self.record_simple_schema_execution(operation, session_id, portal, &forwarded_statements); if encrypted { - let transformed_statement = transformed_statements + let transformed_statement = forwarded_statements .iter() .map(|s| s.to_string()) .collect::>() @@ -515,6 +547,20 @@ 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, + 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)); + self.context.execute_simple_schema_statements(statements); + } + /// Encrypts literal values found in SQL statements. /// /// Takes literal values extracted from SQL statements and encrypts those that @@ -624,7 +670,7 @@ impl Frontend { /// - rewrites any encrypted literal values /// - wraps any nodes in appropriate EQL function /// - async fn transform_statement( + fn transform_statement( &mut self, typed_statement: &TypeCheckedStatement<'_>, encrypted_literals: &Vec>, @@ -757,6 +803,14 @@ impl Frontend { let mut statement_text = String::from_utf8_lossy(&message.query).into_owned(); let statement = SqlParser::parse_statement(&statement_text)?; + self.context.prepare_schema_for_statement().await?; + if eql_mapper::requires_type_check(&statement) { + self.context.wait_for_schema_execution().await; + self.context.ensure_schema_modelled()?; + } + self.context + .prepare_schema_statement(message.statement.to_owned(), statement.clone()); + // Record diagnostics before any passthrough path can return early. // Statements that do not require EQL type checking (for example, // plaintext INSERTs and SELECT pg_sleep(...)) still need their real @@ -782,8 +836,6 @@ impl Frontend { self.handle_set_keyset(&statement)?; - self.check_for_schema_change(&statement); - if !eql_mapper::requires_type_check(&statement) { counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1); return Ok(None); @@ -821,9 +873,8 @@ impl Frontend { .encrypt_literals(session_id, &typed_statement, &statement.literal_columns) .await?; - if let Some(transformed_statement) = self - .transform_statement(&typed_statement, &encrypted_literals) - .await? + if let Some(transformed_statement) = + self.transform_statement(&typed_statement, &encrypted_literals)? { debug!(target: MAPPER, client_id = self.context.client_id, @@ -881,19 +932,6 @@ impl Frontend { } } - /// - /// Check the Statement AST for DDL - /// Sets a schema changed flag in the Context - /// - /// - fn check_for_schema_change(&self, statement: &ast::Statement) { - let schema_changed = eql_mapper::collect_ddl(self.context.get_table_resolver(), statement); - - if schema_changed { - self.context.set_schema_changed(); - } - } - /// /// Handles `SET CIPHERSTASH KEYSET_*` statements /// @@ -1021,6 +1059,8 @@ impl Frontend { /// - `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> { + 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); diff --git a/packages/cipherstash-proxy/src/postgresql/mod.rs b/packages/cipherstash-proxy/src/postgresql/mod.rs index c3751ba1..a7922b1a 100644 --- a/packages/cipherstash-proxy/src/postgresql/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/mod.rs @@ -14,3 +14,4 @@ pub use context::column::Column; pub use context::Context; pub use context::KeysetIdentifier; pub use driver::handler; +pub(crate) use rewrite::Name; diff --git a/packages/cipherstash-proxy/src/proxy/encrypt_config/config.rs b/packages/cipherstash-proxy/src/proxy/encrypt_config/config.rs new file mode 100644 index 00000000..5f3dbe6d --- /dev/null +++ b/packages/cipherstash-proxy/src/proxy/encrypt_config/config.rs @@ -0,0 +1,95 @@ +use cipherstash_client::eql; +use cipherstash_client::schema::ColumnConfig; +use std::collections::HashMap; + +/// +/// Column configuration keyed by table name and column name +/// - key: `{table_name}.{column_name}` +/// +type EncryptConfigMap = HashMap; + +#[derive(Clone, Debug, PartialEq)] +/// Encryption policies indexed by their resolved table and column names. +pub struct EncryptConfig { + config: EncryptConfigMap, +} + +impl EncryptConfig { + /// Constructs encryption metadata from an already indexed configuration map. + pub fn new_from_config(config: EncryptConfigMap) -> Self { + Self { config } + } + + /// Constructs an empty encryption configuration. + pub fn new() -> Self { + Self { + config: HashMap::new(), + } + } + + /// Returns whether the snapshot contains no encrypted columns. + pub fn is_empty(&self) -> bool { + self.config.is_empty() + } + + /// Returns the encryption policy for one resolved column. + pub fn get_column_config(&self, identifier: &eql::Identifier) -> Option { + self.config.get(identifier).cloned() + } + + /// Returns whether any encrypted column belongs to `table`. + pub(crate) fn contains_table(&self, table: &str) -> bool { + self.config + .keys() + .any(|identifier| identifier.table == table) + } + + /// Inserts or replaces encryption metadata for one column. + pub(crate) fn insert(&mut self, identifier: eql::Identifier, config: ColumnConfig) { + self.config.insert(identifier, config); + } + + /// Removes encryption metadata for one column. + pub(crate) fn remove_column(&mut self, table: &str, column: &str) { + self.config + .remove(&eql::Identifier::new(table.to_owned(), column.to_owned())); + } + + /// Removes all encryption metadata for a table. + pub(crate) fn remove_table(&mut self, table: &str) { + self.config + .retain(|identifier, _| identifier.table != table); + } + + /// Moves encryption metadata to a renamed column identifier. + pub(crate) fn rename_column(&mut self, table: &str, from: &str, to: &str) { + let from = eql::Identifier::new(table.to_owned(), from.to_owned()); + if let Some(config) = self.config.remove(&from) { + self.config.insert( + eql::Identifier::new(table.to_owned(), to.to_owned()), + config, + ); + } + } + + /// Moves all encryption metadata to a renamed table identifier. + pub(crate) fn rename_table(&mut self, from: &str, to: &str) { + let renamed = self + .config + .iter() + .filter(|(identifier, _)| identifier.table == from) + .map(|(identifier, config)| (identifier.column.clone(), config.clone())) + .collect::>(); + self.remove_table(from); + for (column, config) in renamed { + self.config + .insert(eql::Identifier::new(to.to_owned(), column), config); + } + } +} + +impl Default for EncryptConfig { + fn default() -> Self { + Self::new() + } +} diff --git a/packages/cipherstash-proxy/src/proxy/encrypt_config/from_domain.rs b/packages/cipherstash-proxy/src/proxy/encrypt_config/from_domain.rs index 1fc0958f..d8713033 100644 --- a/packages/cipherstash-proxy/src/proxy/encrypt_config/from_domain.rs +++ b/packages/cipherstash-proxy/src/proxy/encrypt_config/from_domain.rs @@ -3,7 +3,8 @@ //! EQL v3 domain types are self-configuring: a column's Postgres domain (e.g. //! `eql_v3_text_search`) encodes both the plaintext token type and the //! searchable-encryption terms it stores. That is enough to build the -//! [`ColumnConfig`] the encrypt pipeline needs, so `eql_v2.add_search_config` +//! [`cipherstash_client::schema::ColumnConfig`] the encrypt pipeline needs, so +//! `eql_v2.add_search_config` //! and the `eql_v2_configuration` table are redundant. //! //! The SEM term → index mapping mirrors the client's indexers (verified against diff --git a/packages/cipherstash-proxy/src/proxy/encrypt_config/manager.rs b/packages/cipherstash-proxy/src/proxy/encrypt_config/manager.rs deleted file mode 100644 index 6b923fc5..00000000 --- a/packages/cipherstash-proxy/src/proxy/encrypt_config/manager.rs +++ /dev/null @@ -1,230 +0,0 @@ -use super::from_domain::column_config_from_domain; -use crate::{ - config::DatabaseConfig, connect, error::Error, log::ENCRYPT_CONFIG, proxy::SCHEMA_QUERY, -}; -use arc_swap::ArcSwap; -use cipherstash_client::eql; -use cipherstash_client::schema::ColumnConfig; -use std::{collections::HashMap, sync::Arc, time::Duration}; -use tokio::{task::JoinHandle, time}; -use tracing::{debug, error, info, warn}; - -/// -/// Column configuration keyed by table name and column name -/// - key: `{table_name}.{column_name}` -/// -type EncryptConfigMap = HashMap; - -#[derive(Clone, Debug)] -pub struct EncryptConfig { - config: EncryptConfigMap, -} - -impl EncryptConfig { - pub fn new_from_config(config: EncryptConfigMap) -> Self { - Self { config } - } - - pub fn new() -> Self { - Self { - config: HashMap::new(), - } - } - - pub fn is_empty(&self) -> bool { - self.config.is_empty() - } - - pub fn get_column_config(&self, identifier: &eql::Identifier) -> Option { - self.config.get(identifier).cloned() - } -} - -impl Default for EncryptConfig { - fn default() -> Self { - Self::new() - } -} - -#[derive(Clone, Debug)] -pub struct EncryptConfigManager { - config: DatabaseConfig, - encrypt_config: Arc>, - _reload_handle: Arc>, -} - -impl EncryptConfigManager { - pub async fn init(config: &DatabaseConfig) -> Result { - let config = config.clone(); - init_reloader(config).await - } - - pub fn load(&self) -> Arc { - self.encrypt_config.load().clone() - } - - pub fn is_empty(&self) -> bool { - self.encrypt_config.load().is_empty() - } - - pub async fn reload(&self) -> bool { - match load_encrypt_config_with_retry(&self.config).await { - Ok(reloaded) => { - debug!(target: ENCRYPT_CONFIG, msg = "Reloaded encrypt configuration"); - self.encrypt_config.swap(Arc::new(reloaded)); - true - } - Err(err) => { - warn!( - msg = "Error reloading encrypt configuration", - error = err.to_string() - ); - false - } - } - } -} - -async fn init_reloader(config: DatabaseConfig) -> Result { - // Skip retries on startup as the likely failure mode is configuration - // Only warn on startup, otherwise warning on every reload - let encrypt_config = match load_encrypt_config(&config).await { - Ok(encrypt_config) => encrypt_config, - Err(err) => { - // Encrypt config is inferred from the schema (EQL v3 self-configuring - // domains), so a load error here is a database/connection failure, not - // a missing config table. A schema with no encrypted columns is a - // successful (empty) load, warned about below. - error!( - msg = "Error loading Encrypt configuration", - error = err.to_string() - ); - return Err(err); - } - }; - - debug!(target: ENCRYPT_CONFIG, ?encrypt_config); - - if encrypt_config.is_empty() { - warn!(msg = "ENCRYPT CONFIGURATION NOT LOADED"); - warn!(msg = "No active Encrypt configuration found in database."); - warn!(msg = "Data is not protected with encryption"); - } else { - info!(msg = "Loaded Encrypt configuration"); - } - - let encrypt_config = Arc::new(ArcSwap::new(Arc::new(encrypt_config))); - - let config_ref = config.clone(); - - let dataset_ref = encrypt_config.clone(); - let reload_handle = tokio::spawn(async move { - let reload_interval = tokio::time::Duration::from_secs(config_ref.config_reload_interval); - - let mut interval = tokio::time::interval_at( - tokio::time::Instant::now() + reload_interval, - reload_interval, - ); - - loop { - interval.tick().await; - - match load_encrypt_config_with_retry(&config_ref).await { - Ok(reloaded) => { - debug!(target: ENCRYPT_CONFIG, msg = "Reloaded Encrypt configuration"); - dataset_ref.swap(Arc::new(reloaded)); - } - Err(err) => { - warn!( - msg = "Error reloading Encrypt configuration", - error = err.to_string() - ); - } - } - } - }); - - Ok(EncryptConfigManager { - config, - encrypt_config, - _reload_handle: Arc::new(reload_handle), - }) -} - -/// Fetch the dataset and retry on any error -/// -/// When databases and the proxy start up at the same time they might not be ready to accept connections before the -/// proxy tries to query the schema. To give the proxy the best chance of initialising correctly this method will -/// retry the query a few times before passing on the error. -async fn load_encrypt_config_with_retry(config: &DatabaseConfig) -> Result { - let mut retry_count = 0; - let max_retry_count = 10; - let max_backoff = Duration::from_secs(2); - - loop { - match load_encrypt_config(config).await { - Ok(encrypt_config) => { - return Ok(encrypt_config); - } - - Err(err) => { - if retry_count >= max_retry_count { - debug!( - ENCRYPT_CONFIG, - msg = "Encrypt configuration could not beloaded", - retries = retry_count, - error = err.to_string() - ); - return Err(err); - } - } - } - - let sleep_duration_ms = (100 * 2_u64.pow(retry_count)).min(max_backoff.as_millis() as _); - - time::sleep(Duration::from_millis(sleep_duration_ms)).await; - - retry_count += 1; - } -} - -/// Loads the encrypt configuration by inferring it from the database schema. -/// -/// EQL v3 columns are self-configuring domain types, so each encrypted column's -/// `ColumnConfig` is derived from its Postgres domain typname (see -/// [`column_config_from_domain`]). There is no `eql_v2_configuration` table or -/// `add_search_config` in v3 — the schema is the single source of truth. -pub async fn load_encrypt_config(config: &DatabaseConfig) -> Result { - let client = connect::database(config).await?; - - let tables = client.query(SCHEMA_QUERY, &[]).await?; - - let mut map = EncryptConfigMap::new(); - - for table in tables { - let table_name: String = table.get("table_name"); - let columns: Vec = table.get("columns"); - let column_domain_names: Vec> = table.get("column_domain_names"); - - for (column, domain) in columns.iter().zip(column_domain_names) { - let Some(domain) = domain else { continue }; - if let Some(column_config) = column_config_from_domain(&table_name, column, &domain) { - debug!( - target: ENCRYPT_CONFIG, - msg = "Encrypted column", - table = table_name, - column = column, - domain = domain - ); - // First wins. The query returns the search path in precedence - // order, so if a table name does appear in more than one - // searched schema this keeps the one PostgreSQL itself would - // resolve to, rather than whichever happened to come last. - map.entry(eql::Identifier::new(table_name.clone(), column.clone())) - .or_insert(column_config); - } - } - } - - Ok(EncryptConfig::new_from_config(map)) -} diff --git a/packages/cipherstash-proxy/src/proxy/encrypt_config/mod.rs b/packages/cipherstash-proxy/src/proxy/encrypt_config/mod.rs index 4edf5db8..be7203b1 100644 --- a/packages/cipherstash-proxy/src/proxy/encrypt_config/mod.rs +++ b/packages/cipherstash-proxy/src/proxy/encrypt_config/mod.rs @@ -1,4 +1,6 @@ -mod from_domain; -mod manager; +/// Immutable encryption-policy snapshots. +mod config; +/// Derives encryption metadata from EQL domain declarations. +pub(crate) mod from_domain; -pub use manager::{EncryptConfig, EncryptConfigManager}; +pub use config::EncryptConfig; diff --git a/packages/cipherstash-proxy/src/proxy/mod.rs b/packages/cipherstash-proxy/src/proxy/mod.rs index 8bf051e0..0ca35886 100644 --- a/packages/cipherstash-proxy/src/proxy/mod.rs +++ b/packages/cipherstash-proxy/src/proxy/mod.rs @@ -5,7 +5,7 @@ use crate::{ connect, error::{ConfigError, Error, TlsConfigError}, postgresql::{Column, Context, KeysetIdentifier}, - proxy::{encrypt_config::EncryptConfigManager, schema::SchemaManager}, + proxy::schema::SchemaManager, tls, }; use cipherstash_client::encryption::Plaintext; @@ -14,7 +14,8 @@ use tokio::sync::oneshot::Sender; use tracing::{debug, warn}; mod encrypt_config; -mod schema; +/// Transaction-aware schema snapshots, overlays, and publication. +pub(crate) mod schema; mod zerokms; pub use encrypt_config::EncryptConfig; @@ -38,7 +39,6 @@ const AGGREGATE_QUERY: &str = include_str!("./sql/select_aggregates.sql"); #[derive(Debug)] pub enum ReloadCommand { DatabaseSchema(ReloadResponder), - EncryptSchema(ReloadResponder), } /// @@ -46,7 +46,6 @@ pub enum ReloadCommand { /// pub struct Proxy { pub config: Arc, - pub encrypt_config_manager: EncryptConfigManager, pub schema_manager: SchemaManager, /// The EQL version installed in the database or `None` if it was not present pub eql_version: Option, @@ -75,24 +74,17 @@ impl Proxy { // Ensures error on start if credential or network issue zerokms.init_cipher(None).await?; - let encrypt_config_manager = EncryptConfigManager::init(&config.database).await?; - let schema_manager = SchemaManager::init(&config.database).await?; let eql_version = Proxy::eql_version(&config).await?; let (reload_sender, reload_receiver) = mpsc::unbounded_channel(); - Proxy::receive( - reload_receiver, - schema_manager.clone(), - encrypt_config_manager.clone(), - ); + Proxy::receive(reload_receiver, schema_manager.clone()); Ok(Proxy { config: Arc::new(config), zerokms, - encrypt_config_manager, schema_manager, eql_version, upstream_tls_roots, @@ -119,23 +111,17 @@ impl Proxy { Ok(version) } - pub fn receive( - mut reload_receiver: ReloadReceiver, - schema_manager: SchemaManager, - encrypt_config_manager: EncryptConfigManager, - ) { + /// Starts the asynchronous coordinator for schema reload requests. + pub fn receive(mut reload_receiver: ReloadReceiver, schema_manager: SchemaManager) { tokio::task::spawn(async move { while let Some(command) = reload_receiver.recv().await { debug!(msg = "ReloadCommand received", ?command); match command { ReloadCommand::DatabaseSchema(responder) => { - let schema_reloaded = schema_manager.reload().await; - let encrypt_config_reloaded = encrypt_config_manager.reload().await; - let _ = responder.send(schema_reloaded && encrypt_config_reloaded); - } - ReloadCommand::EncryptSchema(responder) => { - let reloaded = encrypt_config_manager.reload().await; - let _ = responder.send(reloaded); + let schema_manager = schema_manager.clone(); + tokio::task::spawn(async move { + let _ = responder.send(schema_manager.reload().await); + }); } } } @@ -147,16 +133,13 @@ impl Proxy { /// pub fn context(&self, client_id: i32) -> Context { let config = self.config.clone(); - let encrypt_config = self.encrypt_config_manager.load(); - let schema = self.schema_manager.load(); let reload_sender = self.reload_sender.clone(); let encryption = self.zerokms.clone(); - Context::new( + Context::new_with_schema_store( client_id, config, - encrypt_config, - schema, + self.schema_manager.store(), self.upstream_tls_roots.clone(), encryption, reload_sender, diff --git a/packages/cipherstash-proxy/src/proxy/schema/manager.rs b/packages/cipherstash-proxy/src/proxy/schema/manager.rs index 94eea3c5..e6bd39f8 100644 --- a/packages/cipherstash-proxy/src/proxy/schema/manager.rs +++ b/packages/cipherstash-proxy/src/proxy/schema/manager.rs @@ -1,60 +1,258 @@ +#![deny(missing_docs)] + use super::eql_domains; use crate::config::DatabaseConfig; use crate::error::Error; +use crate::proxy::encrypt_config::from_domain::column_config_from_domain; +use crate::proxy::EncryptConfig; use crate::proxy::{AGGREGATE_QUERY, SCHEMA_QUERY}; use crate::{connect, log::SCHEMA}; use arc_swap::ArcSwap; +use cipherstash_client::eql::Identifier; use eql_mapper::{Column, Schema, Table}; use sqltk::parser::ast::Ident; -use std::sync::Arc; +use std::future::Future; +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, +}; use std::time::Duration; -use tokio::{task::JoinHandle, time}; +use tokio::{sync::Mutex, task::JoinHandle, time}; use tracing::{debug, info, warn}; #[derive(Clone, Debug)] +/// An immutable, atomically published schema and encryption-metadata generation. +pub struct CommittedSchemaSnapshot { + version: u64, + schema: Arc, + encrypt_config: Arc, +} + +#[derive(Clone, Debug)] +/// Shared access to committed snapshots and their publication generations. +pub struct CommittedSchemaStore { + snapshot: Arc>, + requested_publication: Arc, + published_publication: Arc, +} + +impl CommittedSchemaStore { + /// Creates the initial committed generation from aligned schema metadata. + pub(crate) fn from_parts(schema: Schema, encrypt_config: EncryptConfig) -> Self { + Self { + snapshot: Arc::new(ArcSwap::new(Arc::new(CommittedSchemaSnapshot::new( + 1, + schema, + encrypt_config, + )))), + requested_publication: Arc::new(AtomicU64::new(0)), + published_publication: Arc::new(AtomicU64::new(0)), + } + } + + /// Loads one internally consistent committed snapshot. + pub fn load(&self) -> Arc { + self.snapshot.load().clone() + } + + /// Returns whether a requested publication has not yet completed. + pub fn publication_pending(&self) -> bool { + self.requested_publication.load(Ordering::Acquire) + > self.published_publication.load(Ordering::Acquire) + } + + /// Advances the requested publication generation. + pub fn mark_publication_pending(&self) { + self.requested_publication.fetch_add(1, Ordering::AcqRel); + } + + /// Marks all currently requested publication generations as satisfied. + pub(crate) fn publication_succeeded(&self) { + self.published_publication.store( + self.requested_publication.load(Ordering::Acquire), + Ordering::Release, + ); + } + + #[cfg(test)] + /// Creates a committed store without connecting to PostgreSQL. + pub fn for_testing(schema: Schema, encrypt_config: EncryptConfig) -> Self { + Self::from_parts(schema, encrypt_config) + } + + #[cfg(test)] + /// Publishes an aligned test snapshot as the next version. + pub fn publish_for_testing(&self, schema: Schema, encrypt_config: EncryptConfig) { + let version = self.load().version() + 1; + self.snapshot.store(Arc::new(CommittedSchemaSnapshot::new( + version, + schema, + encrypt_config, + ))); + self.publication_succeeded(); + } +} + +impl CommittedSchemaSnapshot { + /// Creates an immutable snapshot at an explicit monotonic version. + fn new(version: u64, schema: Schema, encrypt_config: EncryptConfig) -> Self { + Self { + version, + schema: Arc::new(schema), + encrypt_config: Arc::new(encrypt_config), + } + } + + /// Returns the monotonic snapshot version. + pub fn version(&self) -> u64 { + self.version + } + + /// Returns the structural schema from this generation. + pub fn schema(&self) -> Arc { + self.schema.clone() + } + + /// Returns encryption metadata from the same generation as the schema. + pub fn encrypt_config(&self) -> Arc { + self.encrypt_config.clone() + } +} + +#[derive(Clone, Debug)] +/// Loads, versions, and atomically publishes authoritative committed snapshots. pub struct SchemaManager { config: DatabaseConfig, - schema: Arc>, + snapshot: Arc>, + requested_generation: Arc, + requested_publication: Arc, + published_publication: Arc, + reload_lock: Arc>, _reload_handle: Arc>, } impl SchemaManager { + /// Loads the initial snapshot and starts periodic authoritative refreshes. pub async fn init(config: &DatabaseConfig) -> Result { let config = config.clone(); init_reloader(config).await } - pub fn load(&self) -> Arc { - self.schema.load().clone() + /// Loads the current atomic committed snapshot. + pub fn load(&self) -> Arc { + self.snapshot.load().clone() } + /// Returns a cloneable store for per-connection schema middleware. + pub fn store(&self) -> CommittedSchemaStore { + CommittedSchemaStore { + snapshot: self.snapshot.clone(), + requested_publication: self.requested_publication.clone(), + published_publication: self.published_publication.clone(), + } + } + + /// Requests a coalesced authoritative reload and reports whether it published safely. pub async fn reload(&self) -> bool { - match load_schema_with_retry(&self.config).await { - Ok(reloaded) => { - debug!(target: SCHEMA, msg = "Reloaded database schema"); - self.schema.swap(Arc::new(reloaded)); - true - } - Err(err) => { - warn!( - msg = "Error reloading database schema", - error = err.to_string() - ); - false - } + coalesced_reload( + self.snapshot.clone(), + self.requested_generation.clone(), + self.requested_publication.clone(), + self.published_publication.clone(), + self.reload_lock.clone(), + || load_snapshot_with_retry(&self.config), + ) + .await + } +} + +/// Coalesces concurrent reload requests while preserving generation ordering. +async fn coalesced_reload( + snapshot: Arc>, + requested_generation: Arc, + requested_publication: Arc, + published_publication: Arc, + reload_lock: Arc>, + load: F, +) -> bool +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + let requested = requested_generation.fetch_add(1, Ordering::AcqRel) + 1; + let _guard = reload_lock.lock().await; + let publication_generation = requested_publication.load(Ordering::Acquire); + + if snapshot.load().version() >= requested { + published_publication.fetch_max(publication_generation, Ordering::AcqRel); + return true; + } + + // A catalog read can only satisfy requests already visible when the read + // begins. A request arriving while it is in progress must trigger a later + // read, because PostgreSQL may have committed that DDL after this read's + // transaction snapshot was established. + let loaded_generation = requested_generation.load(Ordering::Acquire); + + match load().await { + Ok((schema, encrypt_config)) => { + debug!(target: SCHEMA, msg = "Reloaded committed schema snapshot", version = loaded_generation); + publish_if_newer( + &snapshot, + CommittedSchemaSnapshot::new(loaded_generation, schema, encrypt_config), + ); + published_publication.fetch_max(publication_generation, Ordering::AcqRel); + true + } + Err(err) => { + warn!( + msg = "Error reloading committed schema snapshot", + error = err.to_string() + ); + false } } } +/// Publishes a candidate only when it advances the committed version. +fn publish_if_newer( + store: &ArcSwap, + candidate: CommittedSchemaSnapshot, +) -> bool { + if candidate.version() <= store.load().version() { + return false; + } + store.store(Arc::new(candidate)); + true +} + +/// Loads the initial snapshot and starts periodic authoritative refreshes. async fn init_reloader(config: DatabaseConfig) -> Result { // Skip retries on startup as the likely failure mode is configuration - let schema = load_schema(&config).await?; - info!(msg = "Loaded database schema"); + let (schema, encrypt_config) = load_snapshot(&config).await?; + if encrypt_config.is_empty() { + warn!(msg = "ENCRYPT CONFIGURATION NOT LOADED"); + warn!(msg = "No active Encrypt configuration found in database."); + warn!(msg = "Data is not protected with encryption"); + } + info!(msg = "Loaded committed schema snapshot"); - let schema = Arc::new(ArcSwap::new(Arc::new(schema))); + let snapshot = Arc::new(ArcSwap::new(Arc::new(CommittedSchemaSnapshot::new( + 1, + schema, + encrypt_config, + )))); + let requested_generation = Arc::new(AtomicU64::new(1)); + let requested_publication = Arc::new(AtomicU64::new(0)); + let published_publication = Arc::new(AtomicU64::new(0)); + let reload_lock = Arc::new(Mutex::new(())); let config_ref = config.clone(); - let schema_ref = schema.clone(); + let snapshot_ref = snapshot.clone(); + let generation_ref = requested_generation.clone(); + let requested_publication_ref = requested_publication.clone(); + let published_publication_ref = published_publication.clone(); + let reload_lock_ref = reload_lock.clone(); let reload_handle = tokio::spawn(async move { let reload_interval = tokio::time::Duration::from_secs(config_ref.config_reload_interval); @@ -67,41 +265,45 @@ async fn init_reloader(config: DatabaseConfig) -> Result { loop { interval.tick().await; - match load_schema_with_retry(&config_ref).await { - Ok(reloaded) => { - schema_ref.swap(Arc::new(reloaded)); - } - Err(err) => { - warn!( - msg = "Error loading database schema", - error = err.to_string() - ); - } - } + coalesced_reload( + snapshot_ref.clone(), + generation_ref.clone(), + requested_publication_ref.clone(), + published_publication_ref.clone(), + reload_lock_ref.clone(), + || load_snapshot_with_retry(&config_ref), + ) + .await; } }); Ok(SchemaManager { config, - schema, + snapshot, + requested_generation, + requested_publication, + published_publication, + reload_lock, _reload_handle: Arc::new(reload_handle), }) } -/// Fetch the dataset and retry on any error +/// Fetches an atomic snapshot and retries transient catalog-read failures. /// /// When databases and the proxy start up at the same time they might not be ready to accept connections before the /// proxy tries to query the schema. To give the proxy the best chance of initialising correctly this method will /// retry the query a few times before passing on the error. -async fn load_schema_with_retry(config: &DatabaseConfig) -> Result { +async fn load_snapshot_with_retry( + config: &DatabaseConfig, +) -> Result<(Schema, EncryptConfig), Error> { let mut retry_count = 0; let max_retry_count = 10; let max_backoff = Duration::from_secs(2); loop { - match load_schema(config).await { - Ok(schema) => { - return Ok(schema); + match load_snapshot(config).await { + Ok(snapshot) => { + return Ok(snapshot); } Err(e) => { @@ -177,17 +379,26 @@ fn classify_column( Column::native(ident) } +/// Loads the current structural schema without exposing encryption metadata. pub async fn load_schema(config: &DatabaseConfig) -> Result { + load_snapshot(config).await.map(|(schema, _)| schema) +} + +/// Reads schema and encryption metadata in one repeatable-read transaction. +async fn load_snapshot(config: &DatabaseConfig) -> Result<(Schema, EncryptConfig), Error> { let client = connect::database(config).await?; + client + .batch_execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") + .await?; let tables = client.query(SCHEMA_QUERY, &[]).await?; let mut schema = Schema::new("public"); + let mut encrypt_config = EncryptConfig::new(); if tables.is_empty() { warn!(msg = "Database schema contains no tables"); - return Ok(schema); - }; + } for table in tables { let table_name: String = table.get("table_name"); @@ -210,6 +421,13 @@ pub async fn load_schema(config: &DatabaseConfig) -> Result { ); table.add_column(Arc::new(column)); + + if let Some(domain) = column_domain_name.as_deref() { + if let Some(config) = column_config_from_domain(&table_name, col, domain) { + encrypt_config + .insert(Identifier::new(table_name.clone(), col.clone()), config); + } + } }); schema.add_table(table); @@ -224,13 +442,16 @@ pub async fn load_schema(config: &DatabaseConfig) -> Result { }) .collect(); - Ok(schema) + client.batch_execute("COMMIT").await?; + Ok((schema, encrypt_config)) } #[cfg(test)] mod test { use super::*; use eql_mapper::ColumnKind; + use std::sync::atomic::AtomicUsize; + use tokio::sync::Notify; /// The shape `information_schema.columns` reports for a column declared with /// the EQL v2 composite type, verified against PostgreSQL 17: `udt_name` is @@ -292,4 +513,98 @@ mod test { ColumnKind::Native ); } + + #[test] + fn an_older_generation_cannot_replace_a_newer_snapshot() { + let snapshot = ArcSwap::new(Arc::new(CommittedSchemaSnapshot::new( + 3, + Schema::new("public"), + EncryptConfig::new(), + ))); + + assert!(!publish_if_newer( + &snapshot, + CommittedSchemaSnapshot::new(2, Schema::new("stale"), EncryptConfig::new(),), + )); + assert_eq!(snapshot.load().version(), 3); + } + + #[tokio::test] + async fn requests_arriving_during_a_reload_are_coalesced_into_one_follow_up() { + let snapshot = Arc::new(ArcSwap::new(Arc::new(CommittedSchemaSnapshot::new( + 1, + Schema::new("public"), + EncryptConfig::new(), + )))); + let requested_generation = Arc::new(AtomicU64::new(1)); + let requested_publication = Arc::new(AtomicU64::new(1)); + let published_publication = Arc::new(AtomicU64::new(0)); + let reload_lock = Arc::new(Mutex::new(())); + let loads = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + + let first = tokio::spawn(coalesced_reload( + snapshot.clone(), + requested_generation.clone(), + requested_publication.clone(), + published_publication.clone(), + reload_lock.clone(), + { + let loads = loads.clone(); + let started = started.clone(); + let release = release.clone(); + move || async move { + loads.fetch_add(1, Ordering::AcqRel); + started.notify_one(); + release.notified().await; + Ok((Schema::new("public"), EncryptConfig::new())) + } + }, + )); + + started.notified().await; + requested_publication.fetch_add(1, Ordering::AcqRel); + let second = tokio::spawn(coalesced_reload( + snapshot.clone(), + requested_generation.clone(), + requested_publication.clone(), + published_publication.clone(), + reload_lock.clone(), + { + let loads = loads.clone(); + move || async move { + loads.fetch_add(1, Ordering::AcqRel); + Ok((Schema::new("public"), EncryptConfig::new())) + } + }, + )); + + let third = tokio::spawn(coalesced_reload( + snapshot.clone(), + requested_generation.clone(), + requested_publication.clone(), + published_publication.clone(), + reload_lock.clone(), + { + let loads = loads.clone(); + move || async move { + loads.fetch_add(1, Ordering::AcqRel); + Ok((Schema::new("public"), EncryptConfig::new())) + } + }, + )); + + while requested_generation.load(Ordering::Acquire) < 4 { + tokio::task::yield_now().await; + } + release.notify_one(); + + assert!(first.await.unwrap()); + assert!(second.await.unwrap()); + assert!(third.await.unwrap()); + assert_eq!(loads.load(Ordering::Acquire), 2); + assert_eq!(snapshot.load().version(), 4); + assert_eq!(published_publication.load(Ordering::Acquire), 2); + } } diff --git a/packages/cipherstash-proxy/src/proxy/schema/middleware.rs b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs new file mode 100644 index 00000000..1ea381bb --- /dev/null +++ b/packages/cipherstash-proxy/src/proxy/schema/middleware.rs @@ -0,0 +1,1375 @@ +//! Transaction-aware coordination between PostgreSQL protocol events and schema mapping. +//! +//! # Principles of operation +//! +//! PostgreSQL is the authority. Proxy never publishes a schema inferred from client SQL to +//! other connections. Instead, [`CommittedSchemaStore`] exposes immutable, monotonically +//! versioned snapshots loaded from PostgreSQL. Each snapshot contains both structural schema +//! and encryption metadata, so a mapper can never observe a table definition from one catalog +//! generation and encryption rules from another. +//! +//! A connection adopts a committed snapshot while idle and pins it when work begins. Confirmed +//! DDL is applied to a connection-local overlay, producing the connection's effective schema: +//! the pinned committed snapshot plus successful changes in the current transaction. Savepoints +//! checkpoint both the structural overlay and encryption metadata. Full rollback discards the +//! overlay; rollback to a savepoint restores its checkpoint; release keeps the changes in the +//! enclosing transaction. +//! +//! Protocol events, rather than parsing alone, drive state changes: +//! +//! 1. `Parse` records DDL intent under the prepared-statement name. +//! 2. `Bind` associates that intent with a portal. +//! 3. `Execute` queues the intent and marks DDL as in flight. +//! 4. Backend success activates the change; backend failure discards it. +//! 5. `Sync` and simple-query boundaries segment the execution queue because PostgreSQL skips +//! the rest of an extended-protocol batch after an error. +//! +//! Schema-dependent work waits behind in-flight DDL, which makes pipelining correct without +//! guessing whether PostgreSQL will accept the change. At outermost commit, the connection asks +//! the schema manager to reload PostgreSQL before successful idle readiness reaches the client. +//! Publication requests are shared, coalesced, and generation ordered. If publication fails, +//! the committed database transaction cannot be undone; Proxy therefore retains the pending +//! publication, withholds successful readiness, and closes that client connection. +//! +//! # Trade-offs and limitations +//! +//! * The overlay intentionally models a small, deterministic DDL subset. Conditional DDL, +//! cascading changes, `CREATE TABLE AS`/`LIKE`/`CLONE`/`INHERITS`, views, and unsupported +//! `ALTER TABLE` operations are treated as unmodelled. After such DDL succeeds, later +//! schema-dependent statements fail closed until rollback or authoritative publication. +//! * A simple-query message has one PostgreSQL response boundary, so Proxy cannot safely remap a +//! later statement after observing an earlier statement's result. Batches that may change +//! encryption metadata and then perform schema-dependent work are rejected. Explicitly native +//! table DDL remains compatible because it cannot create an encryption obligation. +//! * Temporary and other connection-local catalog objects are invisible to the separate +//! publication connection. Native temporary-table batches may pass through, but encrypted +//! temporary objects cannot be represented or authoritatively published and are unsupported. +//! A temporary table that could shadow encrypted metadata leaves the connection fail-closed +//! until an earlier savepoint removes it or the connection ends; an authoritative global reload +//! cannot prove that a connection-local object disappeared. +//! * SQL executed indirectly by procedures, extensions, or dynamic SQL cannot be inferred from +//! the wire statement. Periodic authoritative reloads eventually discover committed global +//! changes; a transaction still keeps its pinned view for consistency. +//! * Prepared-statement and portal intents are retained for the connection lifetime. Reusing a +//! protocol name replaces its entry, but closing many unique names can retain bounded metadata +//! until the client disconnects. +//! +//! For supported schema changes, these conservative refusals trade some PostgreSQL surface-area +//! compatibility for the core invariant: Proxy must never forward plaintext because it +//! speculated about uncommitted or incompletely modelled schema state. + +#![deny(missing_docs)] + +use super::eql_domains; +use super::manager::CommittedSchemaStore; +use crate::postgresql::Name; +use crate::proxy::encrypt_config::from_domain::column_config_from_domain; +use crate::proxy::EncryptConfig; +use cipherstash_client::eql::Identifier; +use eql_mapper::{ColumnKind, Schema, SchemaWithEdits, TableResolver}; +use sqltk::parser::ast::{ + AlterColumnOperation, AlterTableOperation, ColumnDef, DataType, DropBehavior, Ident, + ObjectName, ObjectNamePart, ObjectType, Statement, +}; +use std::collections::{HashMap, VecDeque}; +use std::sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, RwLock, +}; +use tokio::sync::Notify; + +#[derive(Clone, Debug)] +/// The schema-relevant meaning recorded for a parsed statement. +struct Intent { + statement: Statement, + ddl: bool, + modelled: bool, +} + +#[derive(Clone, Debug)] +/// One execution or synchronization boundary awaiting a backend outcome. +enum PendingExecution { + /// An execution, optionally carrying schema intent. + Execute(Option>), + /// The boundary terminated by the next `ReadyForQuery` message. + ReadyBoundary, +} + +#[derive(Clone, Debug)] +/// A transaction savepoint and its corresponding schema checkpoint. +struct Savepoint { + name: Ident, + schema: SchemaWithEdits, + encrypt_config: EncryptConfig, + unmodelled: bool, + local_unmodelled: bool, +} + +/// PostgreSQL's transaction state carried by a `ReadyForQuery` message. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum TransactionStatus { + /// The connection is outside a transaction block. + Idle, + /// The connection is inside a transaction block that can accept commands. + InTransaction, + /// The connection is inside a failed transaction block awaiting rollback. + FailedTransaction, + /// An unrecognized protocol status, treated conservatively as non-idle. + Unknown(u8), +} + +impl TransactionStatus { + /// Returns whether PostgreSQL is outside a transaction block. + pub fn is_idle(self) -> bool { + self == Self::Idle + } +} + +impl From for TransactionStatus { + fn from(status: u8) -> Self { + match status { + b'I' => Self::Idle, + b'T' => Self::InTransaction, + b'E' => Self::FailedTransaction, + status => Self::Unknown(status), + } + } +} + +/// Connection-local owner of the effective database schema. +/// +/// Protocol adapters report DDL execution outcomes through this interface; +/// parsing alone never changes the resolver visible to later statements. +#[derive(Clone, Debug)] +pub struct SchemaMiddleware { + store: CommittedSchemaStore, + base: Arc>>, + base_encrypt_config: Arc>>, + encrypt_config: Arc>>, + resolver: Arc>>, + prepared: Arc>>, + portals: Arc>>, + executions: Arc>>, + in_flight_ddl: Arc, + execution_finished: Arc, + dirty: Arc, + unmodelled: Arc, + local_unmodelled: Arc, + transaction_active: Arc, + savepoints: Arc>>, +} + +impl SchemaMiddleware { + #[cfg(test)] + /// Constructs middleware around a schema-only snapshot for unit tests. + pub fn new(schema: Arc) -> Self { + Self::from_store(CommittedSchemaStore::for_testing( + (*schema).clone(), + EncryptConfig::new(), + )) + } + + /// Constructs connection-local middleware backed by the shared committed store. + pub fn from_store(store: CommittedSchemaStore) -> Self { + let snapshot = store.load(); + let schema = snapshot.schema(); + Self { + store, + base: Arc::new(RwLock::new(schema.clone())), + base_encrypt_config: Arc::new(RwLock::new(snapshot.encrypt_config())), + encrypt_config: Arc::new(RwLock::new(snapshot.encrypt_config())), + resolver: Arc::new(RwLock::new(Arc::new(TableResolver::new_editable(schema)))), + prepared: Arc::new(RwLock::new(HashMap::new())), + portals: Arc::new(RwLock::new(HashMap::new())), + executions: Arc::new(RwLock::new(VecDeque::new())), + in_flight_ddl: Arc::new(AtomicUsize::new(0)), + execution_finished: Arc::new(Notify::new()), + dirty: Arc::new(AtomicBool::new(false)), + unmodelled: Arc::new(AtomicBool::new(false)), + local_unmodelled: Arc::new(AtomicBool::new(false)), + transaction_active: Arc::new(AtomicBool::new(false)), + savepoints: Arc::new(RwLock::new(Vec::new())), + } + } + + /// Returns the resolver for the connection's current effective schema. + pub fn resolver(&self) -> Arc { + self.resolver.read().unwrap().clone() + } + + /// Returns encryption metadata aligned with the current effective schema. + pub fn encrypt_config(&self) -> Arc { + self.encrypt_config.read().unwrap().clone() + } + + /// Replaces connection-local state with the latest committed snapshot. + pub fn adopt_latest(&self) { + let snapshot = self.store.load(); + let schema = snapshot.schema(); + *self.base.write().unwrap() = schema.clone(); + let encrypt_config = snapshot.encrypt_config(); + *self.base_encrypt_config.write().unwrap() = encrypt_config.clone(); + *self.encrypt_config.write().unwrap() = encrypt_config; + *self.resolver.write().unwrap() = Arc::new(TableResolver::new_editable(schema)); + self.savepoints.write().unwrap().clear(); + self.unmodelled.store(false, Ordering::Release); + } + + /// Returns whether authoritative catalog publication is still required. + pub fn needs_publication(&self) -> bool { + self.dirty.load(Ordering::Acquire) || self.store.publication_pending() + } + + /// Returns whether this connection has confirmed, unpublished DDL changes. + pub fn has_local_changes(&self) -> bool { + self.dirty.load(Ordering::Acquire) + } + + /// Records a shared publication request for the connection's committed DDL. + pub fn mark_publication_pending(&self) { + self.store.mark_publication_pending(); + } + + /// Returns whether an idle connection must publish pending catalog state before use. + pub fn requires_publication_before_statement(&self) -> bool { + !self.transaction_active.load(Ordering::Acquire) + && !self.has_local_changes() + && self.store.publication_pending() + } + + /// Clears local dirty state and adopts the newly published snapshot. + pub fn publication_succeeded(&self) { + self.dirty.store(false, Ordering::Release); + self.adopt_latest(); + } + + /// Returns whether confirmed DDL cannot be represented by the local overlay. + pub fn has_unmodelled_ddl(&self) -> bool { + self.unmodelled.load(Ordering::Acquire) || self.local_unmodelled.load(Ordering::Acquire) + } + + /// Adopts a newer committed snapshot before an idle connection starts work. + pub fn before_statement(&self) { + if !self.transaction_active.load(Ordering::Acquire) + && !self.has_local_changes() + && !self.store.publication_pending() + && self.in_flight_ddl.load(Ordering::Acquire) == 0 + { + self.adopt_latest(); + } + } + + /// Applies a PostgreSQL readiness boundary and its transaction status. + pub fn ready_for_query(&self, status: TransactionStatus) { + self.discard_skipped_executions(); + self.transaction_active + .store(!status.is_idle(), Ordering::Release); + if status.is_idle() && !self.needs_publication() { + self.adopt_latest(); + } + } + + /// Records schema intent for a named prepared statement. + pub fn prepare(&self, name: Name, statement: Statement) { + self.prepared.write().unwrap().insert( + name, + Intent { + ddl: self.is_schema_ddl(&statement), + modelled: is_modelled_ddl(&statement), + statement, + }, + ); + } + + /// Associates a portal with the schema intent of its prepared statement. + pub fn bind(&self, portal: Name, prepared_statement: &Name) { + let intent = self + .prepared + .read() + .unwrap() + .get(prepared_statement) + .cloned(); + let mut portals = self.portals.write().unwrap(); + match intent { + Some(intent) => { + portals.insert(portal, intent); + } + None => { + portals.remove(&portal); + } + } + } + + /// Records a portal execution awaiting a backend result and returns whether + /// the protocol adapter must flush DDL immediately to unblock dependent work. + pub fn execute(&self, portal: &Name) -> bool { + let intent = self.portals.read().unwrap().get(portal).cloned(); + let executes_ddl = intent.as_ref().is_some_and(|intent| intent.ddl); + self.transaction_active.store(true, Ordering::Release); + if executes_ddl { + self.in_flight_ddl.fetch_add(1, Ordering::AcqRel); + } + self.executions + .write() + .unwrap() + .push_back(PendingExecution::Execute(intent.map(Box::new))); + executes_ddl + } + + /// Records every statement and readiness boundary in a simple-query message. + pub fn simple_query(&self, statements: &[Statement]) { + self.transaction_active.store(true, Ordering::Release); + let mut executions = self.executions.write().unwrap(); + for statement in statements { + let intent = Intent { + ddl: self.is_schema_ddl(statement), + modelled: is_modelled_ddl(statement), + statement: statement.clone(), + }; + if intent.ddl { + self.in_flight_ddl.fetch_add(1, Ordering::AcqRel); + } + executions.push_back(PendingExecution::Execute(Some(Box::new(intent)))); + } + executions.push_back(PendingExecution::ReadyBoundary); + } + + /// Returns whether mapping a later statement in this simple-query batch + /// could observe encryption metadata changed by an earlier DDL statement. + pub fn simple_query_requires_fail_closed(&self, statements: &[Statement]) -> bool { + let mut encryption_changing_ddl_seen = false; + + for statement in statements { + if encryption_changing_ddl_seen && eql_mapper::requires_type_check(statement) { + return true; + } + encryption_changing_ddl_seen |= self.ddl_may_change_encryption(statement); + } + + false + } + + /// Conservatively classifies DDL that can change encryption metadata. + fn ddl_may_change_encryption(&self, statement: &Statement) -> bool { + match statement { + Statement::CreateTable(create) => { + self.encrypt_config() + .contains_table(&postgres_object_name(&create.name)) + || create.query.is_some() + || create.like.is_some() + || create.clone.is_some() + || create.inherits.is_some() + || create + .columns + .iter() + .any(|column| matches!(column_kind(column), ColumnKind::Eql(_, _))) + } + Statement::AlterTable { + name, operations, .. + } => { + let table = postgres_object_name(name); + let encrypt_config = self.encrypt_config(); + operations.iter().any(|operation| match operation { + AlterTableOperation::AddColumn { column_def, .. } => { + matches!(column_kind(column_def), ColumnKind::Eql(_, _)) + } + operation if is_encryption_neutral_alter_operation(operation) => false, + AlterTableOperation::DropColumn { column_name, .. } + | AlterTableOperation::RenameColumn { + old_column_name: column_name, + .. + } => encrypt_config + .get_column_config(&Identifier::new( + table.clone(), + postgres_identifier(column_name), + )) + .is_some(), + AlterTableOperation::RenameTable { .. } => { + encrypt_config.contains_table(&table) + } + _ => true, + }) + } + Statement::Drop { + object_type: ObjectType::Table, + names, + .. + } => names.iter().any(|name| { + self.encrypt_config() + .contains_table(&postgres_object_name(name)) + }), + Statement::CreateView { .. } + | Statement::Drop { + object_type: ObjectType::View, + .. + } => true, + _ => false, + } + } + + /// Returns whether a statement changes tracked schema state for this connection. + /// + /// Native temporary tables are connection-local and cannot be published from the + /// manager's catalog connection, so they are ignored unless they could shadow an + /// encrypted table or introduce an encryption domain. Those unsafe cases remain DDL + /// and therefore fail closed as unmodelled changes after successful execution. + fn is_schema_ddl(&self, statement: &Statement) -> bool { + match statement { + Statement::CreateTable(create) if create.temporary => { + self.ddl_may_change_encryption(statement) + } + _ => is_catalog_schema_ddl(statement), + } + } + + /// Marks the `Sync` boundary whose `ReadyForQuery` terminates an extended + /// protocol batch. PostgreSQL skips the remaining executions in that batch + /// after an error, but may already have a later batch queued behind it. + pub fn protocol_boundary(&self) { + self.executions + .write() + .unwrap() + .push_back(PendingExecution::ReadyBoundary); + } + + /// Waits until all earlier schema-changing executions have resolved. + pub async fn wait_for_ddl(&self) { + loop { + let notified = self.execution_finished.notified(); + if self.in_flight_ddl.load(Ordering::Acquire) == 0 { + return; + } + notified.await; + } + } + + /// Returns whether a schema-changing execution is awaiting a backend outcome. + pub fn ddl_in_flight(&self) -> bool { + self.in_flight_ddl.load(Ordering::Acquire) != 0 + } + + #[cfg(test)] + /// Records a direct statement execution for state-machine tests. + pub fn execution_started(&self, statement: Statement) { + let ddl = self.is_schema_ddl(&statement); + if ddl { + self.in_flight_ddl.fetch_add(1, Ordering::AcqRel); + } + self.executions + .write() + .unwrap() + .push_back(PendingExecution::Execute(Some(Box::new(Intent { + modelled: is_modelled_ddl(&statement), + statement, + ddl, + })))); + } + + /// Applies the next queued execution after PostgreSQL confirms success. + pub fn execution_succeeded(&self) { + if let Some(Some(intent)) = self.pop_execution() { + let statement = intent.statement; + match statement { + Statement::Savepoint { name } => { + let resolver = self.resolver(); + let overlay = resolver.as_schema_with_edits().unwrap(); + let checkpoint = overlay.read().unwrap().clone(); + self.savepoints.write().unwrap().push(Savepoint { + name, + schema: checkpoint, + encrypt_config: (*self.encrypt_config()).clone(), + unmodelled: self.unmodelled.load(Ordering::Acquire), + local_unmodelled: self.local_unmodelled.load(Ordering::Acquire), + }); + } + Statement::ReleaseSavepoint { name } => { + let mut savepoints = self.savepoints.write().unwrap(); + if let Some(index) = savepoints + .iter() + .rposition(|savepoint| identifiers_equal(&savepoint.name, &name)) + { + savepoints.truncate(index); + } else { + self.mark_unmodelled(false); + } + } + Statement::Rollback { + savepoint: Some(name), + .. + } => { + let mut savepoints = self.savepoints.write().unwrap(); + if let Some(index) = savepoints + .iter() + .rposition(|savepoint| identifiers_equal(&savepoint.name, &name)) + { + let checkpoint = savepoints[index].schema.clone(); + let encrypt_config = savepoints[index].encrypt_config.clone(); + let unmodelled = savepoints[index].unmodelled; + let local_unmodelled = savepoints[index].local_unmodelled; + savepoints.truncate(index + 1); + let resolver = self.resolver(); + let overlay = resolver.as_schema_with_edits().unwrap(); + *overlay.write().unwrap() = checkpoint; + *self.encrypt_config.write().unwrap() = Arc::new(encrypt_config); + self.unmodelled.store(unmodelled, Ordering::Release); + self.local_unmodelled + .store(local_unmodelled, Ordering::Release); + self.dirty.store( + unmodelled || resolver.has_schema_changed(), + Ordering::Release, + ); + } else { + self.mark_unmodelled(false); + } + } + Statement::Rollback { + savepoint: None, .. + } => { + *self.resolver.write().unwrap() = Arc::new(TableResolver::new_editable( + self.base.read().unwrap().clone(), + )); + *self.encrypt_config.write().unwrap() = + self.base_encrypt_config.read().unwrap().clone(); + self.savepoints.write().unwrap().clear(); + self.dirty.store(false, Ordering::Release); + self.unmodelled.store(false, Ordering::Release); + } + statement if intent.ddl => { + if !intent.modelled { + let local_only = matches!( + statement, + Statement::CreateTable(ref create) if create.temporary + ); + self.mark_unmodelled(local_only); + } else { + self.apply_ddl(&statement); + } + } + _ => {} + } + if intent.ddl { + self.dirty.store( + self.unmodelled.load(Ordering::Acquire) || self.resolver().has_schema_changed(), + Ordering::Release, + ); + self.in_flight_ddl.fetch_sub(1, Ordering::AcqRel); + self.execution_finished.notify_waiters(); + } + } + } + + /// Applies modelled DDL to schema and encryption overlays as one operation. + fn apply_ddl(&self, statement: &Statement) { + eql_mapper::collect_ddl_with_column_kind(self.resolver(), statement, &|column| { + column_kind(column) + }); + + let mut config = (*self.encrypt_config()).clone(); + apply_encrypt_config(&mut config, statement); + *self.encrypt_config.write().unwrap() = Arc::new(config); + } + + /// Discards the next queued execution after PostgreSQL reports failure. + pub fn execution_failed(&self) { + if let Some(Some(intent)) = self.pop_execution() { + if intent.ddl { + self.in_flight_ddl.fetch_sub(1, Ordering::AcqRel); + self.execution_finished.notify_waiters(); + } + } + } + + /// Removes the next execution without crossing a readiness boundary. + fn pop_execution(&self) -> Option> { + let mut executions = self.executions.write().unwrap(); + if matches!(executions.front(), Some(PendingExecution::Execute(_))) { + match executions.pop_front().unwrap() { + PendingExecution::Execute(intent) => Some(intent.map(|intent| *intent)), + PendingExecution::ReadyBoundary => unreachable!(), + } + } else { + None + } + } + + /// Discards executions PostgreSQL skipped after an error up to readiness. + fn discard_skipped_executions(&self) { + let mut discarded_ddl = 0; + let mut executions = self.executions.write().unwrap(); + while let Some(execution) = executions.pop_front() { + match execution { + PendingExecution::Execute(Some(intent)) if intent.ddl => discarded_ddl += 1, + PendingExecution::Execute(_) => {} + PendingExecution::ReadyBoundary => break, + } + } + drop(executions); + + if discarded_ddl > 0 { + self.in_flight_ddl + .fetch_sub(discarded_ddl, Ordering::AcqRel); + self.execution_finished.notify_waiters(); + } + } + + /// Prevents subsequent mapping when local state cannot be reconstructed exactly. + fn mark_unmodelled(&self, connection_local: bool) { + if connection_local { + self.local_unmodelled.store(true, Ordering::Release); + } else { + self.unmodelled.store(true, Ordering::Release); + self.dirty.store(true, Ordering::Release); + } + } +} + +/// Returns the unqualified PostgreSQL domain name declared for a column. +fn column_domain(column: &ColumnDef) -> String { + match &column.data_type { + DataType::Custom(name, _) => postgres_object_name(name), + data_type => data_type.to_string().to_ascii_lowercase(), + } +} + +/// Classifies a declared column as native or as an EQL domain. +fn column_kind(column: &ColumnDef) -> ColumnKind { + eql_domains::resolve(&column_domain(column)) + .map(|(identity, traits)| ColumnKind::Eql(traits, identity)) + .unwrap_or(ColumnKind::Native) +} + +/// Returns an identifier exactly as PostgreSQL stores it in the catalog. +fn postgres_identifier(identifier: &Ident) -> String { + if identifier.quote_style.is_some() { + identifier.value.clone() + } else { + identifier.value.to_ascii_lowercase() + } +} + +/// Returns the catalog spelling of the final component of an object name. +fn postgres_object_name(name: &ObjectName) -> String { + match name.0.last() { + Some(ObjectNamePart::Identifier(name)) => postgres_identifier(name), + _ => String::new(), + } +} + +/// Compares identifiers according to PostgreSQL's quoted-identifier rules. +fn identifiers_equal(left: &Ident, right: &Ident) -> bool { + postgres_identifier(left) == postgres_identifier(right) +} + +/// Adds inferred encryption metadata for one newly declared column. +fn add_column_config(config: &mut EncryptConfig, table: &str, column: &ColumnDef) { + let domain = column_domain(column); + let column_name = postgres_identifier(&column.name); + if let Some(column_config) = column_config_from_domain(table, &column_name, &domain) { + config.insert( + Identifier::new(table.to_owned(), column_name), + column_config, + ); + } +} + +/// Applies modelled DDL to a mutable encryption-metadata overlay. +fn apply_encrypt_config(config: &mut EncryptConfig, statement: &Statement) { + match statement { + Statement::CreateTable(create) => { + let table = postgres_object_name(&create.name); + config.remove_table(&table); + for column in &create.columns { + add_column_config(config, &table, column); + } + } + Statement::AlterTable { + name, operations, .. + } => { + let table = postgres_object_name(name); + for operation in operations { + match operation { + AlterTableOperation::AddColumn { column_def, .. } => { + add_column_config(config, &table, column_def); + } + AlterTableOperation::DropColumn { column_name, .. } => { + config.remove_column(&table, &postgres_identifier(column_name)); + } + AlterTableOperation::RenameColumn { + old_column_name, + new_column_name, + } => { + config.rename_column( + &table, + &postgres_identifier(old_column_name), + &postgres_identifier(new_column_name), + ); + } + AlterTableOperation::RenameTable { table_name } => { + config.rename_table(&table, &postgres_object_name(table_name)); + } + _ => {} + } + } + } + Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + names, + .. + } => { + for name in names { + config.remove_table(&postgres_object_name(name)); + } + } + _ => {} + } +} + +/// Returns whether a statement changes schema state tracked by the middleware. +fn is_catalog_schema_ddl(statement: &Statement) -> bool { + matches!( + statement, + Statement::CreateTable(_) + | Statement::CreateView { .. } + | Statement::AlterTable { .. } + | Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + .. + } + ) +} + +/// Returns whether an `ALTER TABLE` operation cannot affect encryption metadata. +fn is_encryption_neutral_alter_operation(operation: &AlterTableOperation) -> bool { + matches!( + operation, + AlterTableOperation::AddConstraint(_) + | AlterTableOperation::DropConstraint { + drop_behavior: None | Some(DropBehavior::Restrict), + .. + } + | AlterTableOperation::RenameConstraint { .. } + | AlterTableOperation::DisableRowLevelSecurity + | AlterTableOperation::DisableRule { .. } + | AlterTableOperation::DisableTrigger { .. } + | AlterTableOperation::EnableAlwaysRule { .. } + | AlterTableOperation::EnableAlwaysTrigger { .. } + | AlterTableOperation::EnableReplicaRule { .. } + | AlterTableOperation::EnableReplicaTrigger { .. } + | AlterTableOperation::EnableRowLevelSecurity + | AlterTableOperation::EnableRule { .. } + | AlterTableOperation::EnableTrigger { .. } + | AlterTableOperation::OwnerTo { .. } + | AlterTableOperation::AlterColumn { + op: AlterColumnOperation::SetNotNull + | AlterColumnOperation::DropNotNull + | AlterColumnOperation::SetDefault { .. } + | AlterColumnOperation::DropDefault + | AlterColumnOperation::AddGenerated { .. }, + .. + } + ) +} + +/// Returns whether a DDL statement can be represented exactly by the overlay. +fn is_modelled_ddl(statement: &Statement) -> bool { + match statement { + Statement::CreateTable(create) => { + !create.or_replace + && !create.if_not_exists + && !create.temporary + && create.query.is_none() + && create.like.is_none() + && create.clone.is_none() + && create.inherits.is_none() + && create.on_commit.is_none() + } + Statement::AlterTable { operations, .. } => { + operations.iter().all(|operation| match operation { + AlterTableOperation::AddColumn { if_not_exists, .. } => !if_not_exists, + AlterTableOperation::RenameColumn { .. } + | AlterTableOperation::RenameTable { .. } => true, + AlterTableOperation::DropColumn { drop_behavior, .. } => { + *drop_behavior != Some(DropBehavior::Cascade) + } + operation => is_encryption_neutral_alter_operation(operation), + }) + } + Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + cascade, + .. + } => !cascade, + Statement::CreateView { .. } => false, + _ => true, + } +} + +#[cfg(test)] +/// Unit coverage for transaction and protocol state transitions. +mod tests { + use super::*; + use eql_mapper::Table; + use sqltk::parser::ast::{Ident, ObjectName, ObjectNamePart}; + use sqltk::parser::{dialect::PostgreSqlDialect, parser::Parser}; + + fn parse(sql: &str) -> Statement { + Parser::new(&PostgreSqlDialect {}) + .try_with_sql(sql) + .unwrap() + .parse_statement() + .unwrap() + } + + fn table(name: &str) -> ObjectName { + ObjectName(vec![ObjectNamePart::Identifier(Ident::new(name))]) + } + + #[test] + fn ddl_changes_effective_schema_only_after_successful_execution() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + let ddl = parse("create table reports (id bigint)"); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + + middleware.execution_started(ddl); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + + middleware.execution_succeeded(); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + } + + #[test] + fn full_rollback_discards_successful_transaction_ddl() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse("create table reports (id bigint)")); + middleware.execution_succeeded(); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + + middleware.execution_started(parse("rollback")); + middleware.execution_succeeded(); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + } + + #[test] + fn explicit_commit_keeps_changes_dirty_until_authoritative_publication() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.simple_query(&[parse("begin")]); + middleware.execution_succeeded(); + middleware.execution_started(parse("create table reports (id bigint)")); + middleware.execution_succeeded(); + middleware.ready_for_query(TransactionStatus::InTransaction); + middleware.simple_query(&[parse("commit")]); + middleware.execution_succeeded(); + middleware.ready_for_query(TransactionStatus::Idle); + + assert!(middleware.needs_publication()); + middleware.publication_succeeded(); + assert!(!middleware.needs_publication()); + } + + #[test] + fn rollback_to_savepoint_restores_the_overlay_checkpoint() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse("create table accounts (id bigint)")); + middleware.execution_succeeded(); + middleware.execution_started(parse("savepoint before_reports")); + middleware.execution_succeeded(); + middleware.execution_started(parse("create table reports (id bigint)")); + middleware.execution_succeeded(); + + middleware.execution_started(parse("rollback to savepoint before_reports")); + middleware.execution_succeeded(); + + assert!(middleware + .resolver() + .resolve_table(&table("accounts")) + .is_ok()); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + } + + #[test] + fn release_savepoint_preserves_changes_in_the_enclosing_transaction() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse("create table accounts (id bigint)")); + middleware.execution_succeeded(); + middleware.execution_started(parse("savepoint before_reports")); + middleware.execution_succeeded(); + middleware.execution_started(parse("create table reports (id bigint)")); + middleware.execution_succeeded(); + middleware.execution_started(parse("release savepoint before_reports")); + middleware.execution_succeeded(); + + assert!(middleware + .resolver() + .resolve_table(&table("accounts")) + .is_ok()); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + assert!(middleware.needs_publication()); + } + + #[test] + fn rollback_to_savepoint_preserves_an_earlier_unmodelled_change() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse("alter table secrets alter column value type text")); + middleware.execution_succeeded(); + middleware.execution_started(parse("savepoint after_unmodelled")); + middleware.execution_succeeded(); + middleware.execution_started(parse("create table reports (id bigint)")); + middleware.execution_succeeded(); + middleware.execution_started(parse("rollback to savepoint after_unmodelled")); + middleware.execution_succeeded(); + + assert!(middleware.has_unmodelled_ddl()); + assert!(middleware.needs_publication()); + } + + #[test] + fn idle_connection_adopts_a_newly_published_snapshot() { + let store = CommittedSchemaStore::for_testing(Schema::new("public"), EncryptConfig::new()); + let middleware = SchemaMiddleware::from_store(store.clone()); + let mut published = Schema::new("public"); + published.add_table(Table::new(Ident::new("reports"))); + + store.publish_for_testing(published, EncryptConfig::new()); + middleware.adopt_latest(); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + } + + #[test] + fn active_transaction_keeps_its_pinned_snapshot_until_idle() { + let store = CommittedSchemaStore::for_testing(Schema::new("public"), EncryptConfig::new()); + let middleware = SchemaMiddleware::from_store(store.clone()); + let begin = parse("begin"); + middleware.simple_query(&[begin]); + middleware.execution_succeeded(); + + let mut published = Schema::new("public"); + published.add_table(Table::new(Ident::new("reports"))); + store.publish_for_testing(published, EncryptConfig::new()); + middleware.before_statement(); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + middleware.ready_for_query(TransactionStatus::Idle); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + } + + #[test] + fn ordinary_extended_execution_pins_the_snapshot_until_readiness() { + let store = CommittedSchemaStore::for_testing(Schema::new("public"), EncryptConfig::new()); + let middleware = SchemaMiddleware::from_store(store.clone()); + let statement = Name::from("statement"); + let portal = Name::from("portal"); + middleware.prepare(statement.clone(), parse("select 1")); + middleware.bind(portal.clone(), &statement); + assert!(!middleware.execute(&portal)); + + let mut published = Schema::new("public"); + published.add_table(Table::new(Ident::new("reports"))); + store.publish_for_testing(published, EncryptConfig::new()); + middleware.before_statement(); + + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + middleware.execution_succeeded(); + middleware.protocol_boundary(); + middleware.ready_for_query(TransactionStatus::Idle); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_ok()); + } + + #[test] + fn extended_ddl_execution_requests_an_immediate_flush() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + let statement = Name::from("statement"); + let portal = Name::from("portal"); + middleware.prepare(statement.clone(), parse("create table reports (id bigint)")); + middleware.bind(portal.clone(), &statement); + + assert!(middleware.execute(&portal)); + } + + #[test] + fn publication_failure_is_visible_to_other_connections() { + let store = CommittedSchemaStore::for_testing(Schema::new("public"), EncryptConfig::new()); + let publisher = SchemaMiddleware::from_store(store.clone()); + let idle_connection = SchemaMiddleware::from_store(store); + + publisher.execution_started(parse("create table reports (id bigint)")); + publisher.execution_succeeded(); + publisher.mark_publication_pending(); + + assert!(idle_connection.requires_publication_before_statement()); + } + + #[test] + fn successful_ddl_activates_schema_and_encryption_metadata_together() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse( + "create table secrets (id bigint, value eql_v3_text_search)", + )); + middleware.execution_succeeded(); + + let column = middleware + .resolver() + .resolve_table_column(&table("secrets"), &Ident::new("value")) + .unwrap(); + assert!(matches!(column.kind, ColumnKind::Eql(_, _))); + assert!(middleware + .encrypt_config() + .get_column_config(&Identifier::new("secrets".to_owned(), "value".to_owned())) + .is_some()); + } + + #[test] + fn successful_unmodelled_ddl_refuses_later_schema_use_until_rollback() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + + middleware.execution_started(parse("alter table secrets alter column value type text")); + middleware.execution_succeeded(); + + assert!(middleware.has_unmodelled_ddl()); + assert!(middleware.needs_publication()); + + middleware.execution_started(parse("rollback")); + middleware.execution_succeeded(); + assert!(!middleware.has_unmodelled_ddl()); + } + + #[test] + fn create_table_as_is_unmodelled() { + assert!(!is_modelled_ddl(&parse( + "create table archived_reports as select 1 as id", + ))); + } + + #[test] + fn cascading_column_drop_is_unmodelled() { + assert!(!is_modelled_ddl(&parse( + "alter table reports drop column account_id cascade", + ))); + } + + #[test] + fn conditional_create_and_add_are_unmodelled() { + assert!(!is_modelled_ddl(&parse( + "create table if not exists reports (id bigint)", + ))); + assert!(!is_modelled_ddl(&parse( + "alter table reports add column if not exists account_id bigint", + ))); + } + + #[test] + fn unquoted_domain_names_follow_postgresql_case_folding() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + let statements = vec![ + parse("create table secrets (value EQL_V3_TEXT_SEARCH)"), + parse("insert into secrets (value) values ('classified')"), + ]; + + assert!(matches!( + column_kind(match &statements[0] { + Statement::CreateTable(create) => &create.columns[0], + _ => unreachable!(), + }), + ColumnKind::Eql(_, _) + )); + assert!(middleware.simple_query_requires_fail_closed(&statements)); + } + + #[test] + fn encryption_overlay_uses_catalog_spelling_for_unquoted_identifiers() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("create table Secrets (Secret EQL_V3_TEXT_SEARCH)")); + middleware.execution_succeeded(); + + assert!(middleware + .encrypt_config() + .get_column_config(&Identifier::new("secrets", "secret")) + .is_some()); + } + + #[test] + fn dependent_insert_transforms_against_encryption_overlay() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse( + "create table secrets (id bigint, secret eql_v3_text_search)", + )); + middleware.execution_succeeded(); + + let statement = parse("insert into secrets (id, secret) values ($1, $2)"); + let typed = eql_mapper::type_check(middleware.resolver(), &statement).unwrap(); + let transformed = typed.transform(Default::default()).unwrap(); + + assert_eq!( + transformed.statement.to_string(), + "INSERT INTO secrets (id, secret) VALUES ($1, $2::JSONB::public.eql_v3_text_search)" + ); + } + + #[test] + fn quoted_domain_names_preserve_case() { + let column = match parse("create table secrets (value \"EQL_V3_TEXT_SEARCH\")") { + Statement::CreateTable(create) => create.columns.into_iter().next().unwrap(), + _ => unreachable!(), + }; + + assert!(matches!(column_kind(&column), ColumnKind::Native)); + } + + #[test] + fn savepoint_names_follow_postgresql_case_folding() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse( + "create table users (id bigint, secret eql_v3_text_search)", + )); + middleware.execution_succeeded(); + middleware.execution_started(parse("savepoint Foo")); + middleware.execution_succeeded(); + middleware.execution_started(parse("alter table users drop column secret")); + middleware.execution_succeeded(); + middleware.execution_started(parse("rollback to savepoint foo")); + middleware.execution_succeeded(); + + assert!(middleware + .resolver() + .resolve_table_column(&table("users"), &Ident::new("secret")) + .is_ok()); + } + + #[test] + fn savepoint_state_desynchronization_fails_closed() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("rollback to savepoint missing")); + middleware.execution_succeeded(); + + assert!(middleware.has_unmodelled_ddl()); + assert!(middleware.needs_publication()); + } + + #[test] + fn encryption_neutral_alter_table_operations_are_modelled() { + for sql in [ + "alter table users add constraint users_email_uq unique (email)", + "alter table users alter column email set not null", + "alter table users alter column email drop not null", + "alter table users alter column email set default 'unknown'", + "alter table users alter column email drop default", + "alter table users owner to current_user", + "alter table users enable row level security", + ] { + assert!(is_modelled_ddl(&parse(sql)), "not modelled: {sql}"); + } + } + + #[test] + fn encryption_neutral_alter_on_encrypted_table_allows_a_dependent_batch_statement() { + let mut encrypt_config = EncryptConfig::new(); + let column = match parse("create table users (secret eql_v3_text_search)") { + Statement::CreateTable(create) => create.columns.into_iter().next().unwrap(), + _ => unreachable!(), + }; + add_column_config(&mut encrypt_config, "users", &column); + let store = CommittedSchemaStore::for_testing(Schema::new("public"), encrypt_config); + let middleware = SchemaMiddleware::from_store(store); + let statements = vec![ + parse("alter table users alter column secret set not null"), + parse("insert into users (secret) values ('classified')"), + ]; + + assert!(!middleware.simple_query_requires_fail_closed(&statements)); + } + + #[test] + fn native_temporary_table_execution_does_not_dirty_schema_state() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("create temporary table names (name text)")); + middleware.execution_succeeded(); + + assert!(!middleware.has_unmodelled_ddl()); + assert!(!middleware.needs_publication()); + assert!(middleware + .resolver() + .resolve_table(&table("names")) + .is_err()); + } + + #[test] + fn temporary_table_cannot_shadow_an_encrypted_table_in_a_batch() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse( + "create table users (id bigint, email eql_v3_text_search)", + )); + middleware.execution_succeeded(); + let statements = vec![ + parse("create temporary table users (id bigint, email text)"), + parse("insert into users (id, email) values (1, 'alice@example.com')"), + ]; + + assert!(middleware.simple_query_requires_fail_closed(&statements)); + } + + #[test] + fn temporary_shadowing_remains_fail_closed_after_global_publication() { + let mut encrypt_config = EncryptConfig::new(); + let column = match parse("create table users (email eql_v3_text_search)") { + Statement::CreateTable(create) => create.columns.into_iter().next().unwrap(), + _ => unreachable!(), + }; + add_column_config(&mut encrypt_config, "users", &column); + let store = CommittedSchemaStore::for_testing(Schema::new("public"), encrypt_config); + let middleware = SchemaMiddleware::from_store(store); + + middleware.execution_started(parse("create temporary table users (email text)")); + middleware.execution_succeeded(); + assert!(middleware.has_unmodelled_ddl()); + assert!(!middleware.needs_publication()); + + middleware.protocol_boundary(); + middleware.ready_for_query(TransactionStatus::Idle); + middleware.publication_succeeded(); + + assert!(middleware.has_unmodelled_ddl()); + } + + #[test] + fn native_temporary_table_batch_does_not_fail_closed() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + let statements = vec![ + parse("create temporary table names (name text)"), + parse("insert into names (name) values ('Ada')"), + ]; + + assert!(!middleware.simple_query_requires_fail_closed(&statements)); + } + + #[test] + fn encrypted_table_batch_fails_closed_before_dependent_insert() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + let statements = vec![ + parse("create table secrets (value eql_v3_text_search)"), + parse("insert into secrets (value) values ('classified')"), + ]; + + assert!(middleware.simple_query_requires_fail_closed(&statements)); + } + + #[tokio::test] + async fn dependent_mapping_waits_for_ddl_execution_outcome() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("create table reports (id bigint)")); + + assert!(tokio::time::timeout( + std::time::Duration::from_millis(20), + middleware.wait_for_ddl(), + ) + .await + .is_err()); + + middleware.execution_failed(); + + tokio::time::timeout( + std::time::Duration::from_millis(20), + middleware.wait_for_ddl(), + ) + .await + .unwrap(); + assert!(middleware + .resolver() + .resolve_table(&table("reports")) + .is_err()); + } + + #[tokio::test] + async fn readiness_discards_pipelined_executions_skipped_after_an_error() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("create table first_table (id bigint)")); + middleware.execution_started(parse("create table skipped_table (id bigint)")); + middleware.protocol_boundary(); + + middleware.execution_failed(); + middleware.ready_for_query(TransactionStatus::Idle); + + tokio::time::timeout( + std::time::Duration::from_millis(20), + middleware.wait_for_ddl(), + ) + .await + .unwrap(); + assert!(middleware + .resolver() + .resolve_table(&table("skipped_table")) + .is_err()); + } + + #[tokio::test] + async fn readiness_does_not_discard_a_later_pipelined_batch() { + let middleware = SchemaMiddleware::new(Arc::new(Schema::new("public"))); + middleware.execution_started(parse("create table failed_table (id bigint)")); + middleware.execution_started(parse("create table skipped_table (id bigint)")); + middleware.protocol_boundary(); + middleware.execution_started(parse("create table later_table (id bigint)")); + middleware.protocol_boundary(); + + middleware.execution_failed(); + middleware.ready_for_query(TransactionStatus::Idle); + + assert!(tokio::time::timeout( + std::time::Duration::from_millis(20), + middleware.wait_for_ddl(), + ) + .await + .is_err()); + middleware.execution_succeeded(); + middleware.ready_for_query(TransactionStatus::Idle); + middleware.wait_for_ddl().await; + assert!(middleware + .resolver() + .resolve_table(&table("later_table")) + .is_ok()); + } +} diff --git a/packages/cipherstash-proxy/src/proxy/schema/mod.rs b/packages/cipherstash-proxy/src/proxy/schema/mod.rs index c34d83ce..2696a112 100644 --- a/packages/cipherstash-proxy/src/proxy/schema/mod.rs +++ b/packages/cipherstash-proxy/src/proxy/schema/mod.rs @@ -1,4 +1,11 @@ +//! Schema loading, committed snapshots, and connection-local transactional overlays. + +/// Resolves EQL domain identities and capabilities. mod eql_domains; +/// Loads and atomically publishes authoritative committed snapshots. mod manager; +/// Coordinates schema state with PostgreSQL transaction and protocol events. +mod middleware; -pub use manager::SchemaManager; +pub use manager::{CommittedSchemaStore, SchemaManager}; +pub use middleware::{SchemaMiddleware, TransactionStatus}; diff --git a/packages/eql-mapper/src/model/schema_delta.rs b/packages/eql-mapper/src/model/schema_delta.rs index f78d953b..22198bac 100644 --- a/packages/eql-mapper/src/model/schema_delta.rs +++ b/packages/eql-mapper/src/model/schema_delta.rs @@ -23,7 +23,7 @@ use super::{ /// /// All table and column lookups during EQL mapping will go through via the overlay scheme, falling back to the /// loaded schema. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct SchemaWithEdits { schema: Arc, overlays: HashMap, @@ -114,7 +114,7 @@ impl SchemaWithEdits { } /// Acts like a mask over a table or an existing table that has been dropped in the current transaction. -#[derive(Debug)] +#[derive(Clone, Debug)] enum Overlay { /// Hides the existence of table in the main [`Schema`] causing resolution of that table to fail. Dropped, @@ -190,10 +190,24 @@ impl From<&OverlayTable> for Table { /// /// Returns `true` if `statement` contained relevant DDL (regardless of `TableResolver` variant). pub fn collect_ddl(table_resolver: Arc, statement: &Statement) -> bool { + collect_ddl_with_column_kind(table_resolver, statement, &|_| ColumnKind::Native) +} + +/// Applies DDL using the caller's classification for newly declared columns. +/// +/// The mapper itself defaults columns to native because PostgreSQL domain +/// identity belongs to Proxy. Proxy supplies the catalog-backed classifier so +/// a transaction-local EQL domain has the same type as its eventual catalog row. +pub fn collect_ddl_with_column_kind( + table_resolver: Arc, + statement: &Statement, + classify: &dyn Fn(&ColumnDef) -> ColumnKind, +) -> bool { if let Some(schema_with_edits) = table_resolver.as_schema_with_edits() { let mut visitor = DdlCollector { schema: schema_with_edits, changed: false, + classify, }; let _ = statement.accept(&mut visitor); return visitor.changed; @@ -202,12 +216,13 @@ pub fn collect_ddl(table_resolver: Arc, statement: &Statement) -> table_resolver.has_schema_changed() } -struct DdlCollector { +struct DdlCollector<'a> { schema: Arc>, changed: bool, + classify: &'a dyn Fn(&ColumnDef) -> ColumnKind, } -impl DdlCollector { +impl DdlCollector<'_> { fn capture_create_view(&self, name: &ObjectName, columns: &[ViewColumnDef]) { let name = name.clone(); let mut table = OverlayTable::new(name.clone()); @@ -228,7 +243,7 @@ impl DdlCollector { for def in columns { table.add_column(Column { name: def.name.clone(), - kind: ColumnKind::Native, + kind: (self.classify)(def), }); } @@ -244,7 +259,7 @@ impl DdlCollector { if let Overlay::Table(table) = overlay { table.add_column(Column { name: column_def.name.clone(), - kind: ColumnKind::Native, + kind: (self.classify)(column_def), }); } } @@ -317,7 +332,7 @@ impl DdlCollector { } } -impl<'ast> Visitor<'ast> for DdlCollector { +impl<'ast> Visitor<'ast> for DdlCollector<'_> { type Error = Infallible; fn enter(&mut self, node: &'ast N) -> ControlFlow> { diff --git a/tests/integration/golang/pgx_test.go b/tests/integration/golang/pgx_test.go index 18f0cd8c..6d2348f6 100644 --- a/tests/integration/golang/pgx_test.go +++ b/tests/integration/golang/pgx_test.go @@ -69,6 +69,39 @@ INSERT INTO t (name) VALUES require.Equal("Ada", result) } +func TestPgxBatchPipelinesEncryptedDDLAndDependentInsert(t *testing.T) { + t.Parallel() + require := require.New(t) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + config, err := pgx.ParseConfig(os.Getenv("DATABASE_URL")) + require.NoError(err) + config.DefaultQueryExecMode = pgx.QueryExecModeExec + conn, err := pgx.ConnectConfig(ctx, config) + require.NoError(err) + defer conn.Close(ctx) + + table := fmt.Sprintf("bug_308_pgx_pipeline_%d", rand.Int()) + batch := &pgx.Batch{} + batch.Queue(fmt.Sprintf( + "CREATE TABLE %s (id bigint PRIMARY KEY, secret EQL_V3_TEXT_SEARCH NOT NULL)", + table, + )) + batch.Queue( + fmt.Sprintf("INSERT INTO %s (id, secret) VALUES ($1, $2)", table), + 1, + "batched", + ) + + results := conn.SendBatch(ctx, batch) + _, err = results.Exec() + require.NoError(err) + command, err := results.Exec() + require.NoError(err) + require.Equal(int64(1), command.RowsAffected()) + require.NoError(results.Close()) +} + func TestPgxEncryptedMapText(t *testing.T) { t.Parallel() conn := setupPgxConnection(t) @@ -261,9 +294,9 @@ func TestPgxInsertEncryptedWithStructScan(t *testing.T) { // EncryptedRowWithJsonb represents a row with id, encrypted_text and encrypted_jsonb fields type EncryptedRowWithJsonb struct { - ID int `db:"id"` - EncryptedText string `db:"encrypted_text"` - EncryptedJsonb map[string]interface{} `db:"encrypted_jsonb"` + ID int `db:"id"` + EncryptedText string `db:"encrypted_text"` + EncryptedJsonb map[string]interface{} `db:"encrypted_jsonb"` } // Scan implements the sql.Scanner interface for EncryptedRowWithJsonb @@ -304,4 +337,3 @@ func TestPgxInsertEncryptedWithJsonbStructScan(t *testing.T) { }) } } -