From 600b7621cdea8e03b7e81811ae13e1ffd78e63fb Mon Sep 17 00:00:00 2001 From: Logan Johnson Date: Tue, 8 Sep 2026 00:09:13 -0400 Subject: [PATCH] test(db): trace one safe tuple from refused activity aggregate Signed-off-by: Logan Johnson --- crates/buzz-db/src/runtime/replica_fence.rs | 159 +++++++++++++++++--- 1 file changed, 136 insertions(+), 23 deletions(-) diff --git a/crates/buzz-db/src/runtime/replica_fence.rs b/crates/buzz-db/src/runtime/replica_fence.rs index 044dc3a58c..0294d4f186 100644 --- a/crates/buzz-db/src/runtime/replica_fence.rs +++ b/crates/buzz-db/src/runtime/replica_fence.rs @@ -544,12 +544,69 @@ pub enum ProbeError { MaskedActivity { /// Number of other client backends with masked/unknown state. masked: i64, + /// Test-only bounded facts from the refused aggregate, never raw labels. + #[cfg(test)] + diagnostic: serde_json::Value, }, /// The single heartbeat row (migration 0026) is missing on the writer. #[error("replica_heartbeat row missing on the writer — migration 0026 not applied?")] HeartbeatRowMissing, } +const ACTIVITY_SQL: &str = r#" + SELECT + least( + (SELECT min(xact_start) + FROM pg_stat_activity + WHERE pid <> pg_backend_pid()), + (SELECT min(prepared) FROM pg_prepared_xacts) + ) AS oldest_xact_start, + (SELECT count(*) + FROM pg_stat_activity + WHERE pid <> pg_backend_pid() + AND (backend_type IS NULL + OR (backend_type = 'client backend' + AND (state IS NULL + OR (state <> 'idle' AND xact_start IS NULL)))) + ) AS masked + "#; + +// Only the aggregate projection changes in test builds: count and one bounded, +// label-free representative consume the SAME filtered rows, not a second scan. +// A single statement alone would not establish that for arbitrary PG statistics +// functions. Here both aggregates share one input from pg_stat_activity, so a +// backend disappearing later cannot replace the refused sample. min(text) keeps +// one fixed-width tuple, rather than collecting every backend into an array. +#[cfg(test)] +const MASKED_DIAGNOSTIC_AGGREGATE: &str = r#"jsonb_build_object( + 'count', count(*), + 'probe_pid', pg_backend_pid(), + 'read_all_stats', pg_has_role(current_user, 'pg_read_all_stats', 'USAGE'), + 'representative', min(jsonb_build_object( + 'pid', pid, + 'same_database', datname = current_database(), + 'same_role', usesysid = current_user::regrole::oid, + 'backend', CASE WHEN backend_type IS NULL THEN 'null' + WHEN backend_type = 'client backend' THEN 'client' + ELSE 'other' END, + 'state', CASE WHEN state IS NULL THEN 'null' + WHEN state = 'idle' THEN 'idle' + WHEN state = 'active' THEN 'active' + WHEN state = 'idle in transaction' THEN 'idle in transaction' + WHEN state = 'idle in transaction (aborted)' THEN 'aborted' + WHEN state = 'disabled' THEN 'disabled' + ELSE 'other' END, + 'xact_start', xact_start, + 'backend_start', backend_start, + 'state_change', state_change + )::text)::jsonb +)"#; + +#[cfg(test)] +fn diagnostic_activity_sql() -> String { + ACTIVITY_SQL.replacen("count(*)", MASKED_DIAGNOSTIC_AGGREGATE, 1) +} + /// Take one ordered writer sample: S, then activity scan, then commit the /// heartbeat token **last**. /// @@ -593,30 +650,24 @@ async fn sample_writer(writer: &PgPool) -> Result { // after the token. Their deferred floor guard already ran at PREPARE, // so `pg_prepared_xacts.prepared` bounds their rows exactly like // `xact_start`; fold it into the same minimum. - let row = sqlx::query( - r#" - SELECT - least( - (SELECT min(xact_start) - FROM pg_stat_activity - WHERE pid <> pg_backend_pid()), - (SELECT min(prepared) FROM pg_prepared_xacts) - ) AS oldest_xact_start, - (SELECT count(*) - FROM pg_stat_activity - WHERE pid <> pg_backend_pid() - AND (backend_type IS NULL - OR (backend_type = 'client backend' - AND (state IS NULL - OR (state <> 'idle' AND xact_start IS NULL)))) - ) AS masked - "#, - ) - .fetch_one(&mut *conn) - .await?; + #[cfg(not(test))] + let query = sqlx::query(ACTIVITY_SQL); + // Only two compile-time literals enter this test-only query builder. + #[cfg(test)] + let query = sqlx::query(sqlx::AssertSqlSafe(diagnostic_activity_sql())); + let row = query.fetch_one(&mut *conn).await?; + #[cfg(not(test))] let masked: i64 = row.get("masked"); + #[cfg(test)] + let diagnostic: serde_json::Value = row.get("masked"); + #[cfg(test)] + let masked = diagnostic["count"].as_i64().expect("count(*) is an i64"); if masked > 0 { - return Err(ProbeError::MaskedActivity { masked }); + return Err(ProbeError::MaskedActivity { + masked, + #[cfg(test)] + diagnostic, + }); } let oldest_xact_start: Option> = row.get("oldest_xact_start"); @@ -802,6 +853,61 @@ pub async fn run_probe(writer: PgPool, fence: Arc) { } } +#[cfg(test)] +mod diagnostic_tests { + use super::*; + use sha2::{Digest, Sha256}; + + #[test] + fn diagnostic_changes_only_the_aggregate_not_the_classified_set() { + // Frozen pre-diagnostic SQL, including whitespace: catches changes to + // the production predicate, other-PID scope, and prepared-xact minimum. + assert_eq!( + hex::encode(Sha256::digest(ACTIVITY_SQL.as_bytes())), + "a11e60f4316a81b8d00809140ed2e52613c3f96c22020eb0fa5c2abb4c25b380" + ); + assert_eq!(ACTIVITY_SQL.matches("count(*)").count(), 1); + let instrumented = diagnostic_activity_sql(); + // Full equality, not a substring test: undoing the single aggregate + // projection must restore every byte of the actual production query. + assert_eq!( + instrumented.replace(MASKED_DIAGNOSTIC_AGGREGATE, "count(*)"), + ACTIVITY_SQL + ); + assert_eq!(instrumented.matches("count(*)").count(), 1); + } + + #[test] + fn diagnostic_projection_is_a_bounded_data_allowlist() { + // Review the complete projection, not a blacklist of possible secrets. + // In particular, database/role identities are comparisons only, state + // and backend type map to literals, and min retains one WHOLE tuple. + let approved = r#"jsonb_build_object( + 'count', count(*), + 'probe_pid', pg_backend_pid(), + 'read_all_stats', pg_has_role(current_user, 'pg_read_all_stats', 'USAGE'), + 'representative', min(jsonb_build_object( + 'pid', pid, + 'same_database', datname = current_database(), + 'same_role', usesysid = current_user::regrole::oid, + 'backend', CASE WHEN backend_type IS NULL THEN 'null' + WHEN backend_type = 'client backend' THEN 'client' ELSE 'other' END, + 'state', CASE WHEN state IS NULL THEN 'null' + WHEN state = 'idle' THEN 'idle' + WHEN state = 'active' THEN 'active' + WHEN state = 'idle in transaction' THEN 'idle in transaction' + WHEN state = 'idle in transaction (aborted)' THEN 'aborted' + WHEN state = 'disabled' THEN 'disabled' ELSE 'other' END, + 'xact_start', xact_start, + 'backend_start', backend_start, + 'state_change', state_change + )::text)::jsonb + )"#; + let tokens = |sql: &str| sql.split_whitespace().collect::>().join(" "); + assert_eq!(tokens(MASKED_DIAGNOSTIC_AGGREGATE), tokens(approved)); + } +} + #[cfg(test)] mod postgres_tests { use super::*; @@ -1072,10 +1178,17 @@ mod postgres_tests { .await .expect_err("masked pg_stat_activity must fail closed"); assert!( - matches!(err, ProbeError::MaskedActivity { masked } if masked >= 1), + matches!(err, ProbeError::MaskedActivity { masked, .. } if masked >= 1), "expected MaskedActivity, got {err:?}" ); + if let ProbeError::MaskedActivity { masked, diagnostic } = &err { + assert_eq!(diagnostic["count"].as_i64(), Some(*masked)); + assert_eq!(diagnostic["read_all_stats"], false); + assert!(diagnostic["representative"]["pid"].is_i64()); + assert_ne!(diagnostic["representative"]["pid"], diagnostic["probe_pid"]); + } + tx.rollback().await.expect("rollback"); unpriv.close().await; sqlx::query(sqlx::AssertSqlSafe(format!("DROP ROLE {role}")))