From 3bf280a66830ee7c36137cab6f5394c7cabdb6c9 Mon Sep 17 00:00:00 2001 From: Brad Seiler Date: Wed, 4 Nov 2026 16:31:26 -0500 Subject: [PATCH] fix(db): bound live authorization statements and cancellation Apply a transaction-local server deadline below the caller timeout, preserve tighter operator limits, and close cancelled or failed checkouts. Verify blocked-query cancellation, pool recovery, and writer timeout restoration on the same backend. Co-authored-by: Honor Signed-off-by: Brad Seiler --- crates/buzz-db/src/lib.rs | 4 +- .../buzz-db/src/store/live_authorization.rs | 234 +++++++++++++++++- 2 files changed, 231 insertions(+), 7 deletions(-) diff --git a/crates/buzz-db/src/lib.rs b/crates/buzz-db/src/lib.rs index f29f81a9cf..9cfd930b72 100644 --- a/crates/buzz-db/src/lib.rs +++ b/crates/buzz-db/src/lib.rs @@ -105,7 +105,9 @@ pub use community::{ pub use error::{DbError, Result}; pub use event::{ChannelHeadPrecondition, ChannelHeadWriteStatus, DeletionTarget}; pub use event::{EventQuery, DEFAULT_MAX_PAGE_LIMIT}; -pub use live_authorization::{LiveAuthorizationState, LiveAuthorizationTarget}; +pub use live_authorization::{ + LiveAuthorizationState, LiveAuthorizationTarget, LIVE_AUTHORIZATION_QUERY_TIMEOUT, +}; pub use nip_fi_binding::{NipFiBindingRecord, NipFiUnbind, RowCapExceeded}; pub use reaction::ReactionEventInsertOutcome; pub use reminder::DueReminder; diff --git a/crates/buzz-db/src/store/live_authorization.rs b/crates/buzz-db/src/store/live_authorization.rs index c2963a93d9..bbe0734b6d 100644 --- a/crates/buzz-db/src/store/live_authorization.rs +++ b/crates/buzz-db/src/store/live_authorization.rs @@ -2,11 +2,67 @@ use buzz_core::CommunityId; use buzz_datastore_tracing::datastore_span; -use sqlx::Row as _; +use sqlx::{Connection as _, Row as _}; use uuid::Uuid; use crate::{observability, Db, DbError, Result}; +/// End-to-end deadline shared with the live-session revalidator. +pub const LIVE_AUTHORIZATION_QUERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1); +const STATEMENT_TIMEOUT: std::time::Duration = + LIVE_AUTHORIZATION_QUERY_TIMEOUT.saturating_sub(std::time::Duration::from_millis(100)); + +// A cancelled SQLx query must not wait for its unfinished statement in the +// pool's return-to-pool ping. SQLx bounds close_on_drop; PostgreSQL's own +// deadline below also terminates work if it has not noticed the disconnect. +struct AuthorizationConnection { + connection: sqlx::pool::PoolConnection, + reusable: bool, +} + +impl Drop for AuthorizationConnection { + fn drop(&mut self) { + if !self.reusable { + self.connection.close_on_drop(); + } + } +} + +async fn fetch_authorization_rows( + pool: &sqlx::PgPool, + query: sqlx::query::Query<'_, sqlx::Postgres, sqlx::postgres::PgArguments>, +) -> Result> { + let deadline = tokio::time::Instant::now() + LIVE_AUTHORIZATION_QUERY_TIMEOUT; + let mut guard = AuthorizationConnection { + connection: observability::acquire_writer_until( + pool, + observability::WriterOperation::Authorization, + deadline, + ) + .await?, + reusable: false, + }; + // BEGIN and tighten travel in one message. Keep a stricter operator cap; + // rollback restores the writer's session setting on successful reuse. + let begin = format!( + "BEGIN; SELECT set_config('statement_timeout', CASE \ + WHEN current_setting('statement_timeout')::interval = interval '0' \ + OR current_setting('statement_timeout')::interval > interval '{0}ms' \ + THEN '{0}ms' ELSE current_setting('statement_timeout') END, true)", + STATEMENT_TIMEOUT.as_millis(), + ); + // The only interpolation is the integer millisecond value of the + // compile-time deadline above; no request data enters this SQL. + let mut tx = guard + .connection + .begin_with(sqlx::AssertSqlSafe(begin)) + .await?; + let rows = query.fetch_all(&mut *tx).await?; + tx.rollback().await?; + guard.reusable = true; + Ok(rows) +} + /// One admitted principal and captured owner whose live authorization should be revalidated. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct LiveAuthorizationTarget { @@ -67,9 +123,7 @@ async fn live_authorization_state_batch( .iter() .map(|target| target.session_owner_pubkey.map(|owner| owner.to_vec())) .collect::>(); - let mut connection = - observability::acquire_writer(pool, observability::WriterOperation::Authorization).await?; - let rows = sqlx::query( + let rows = fetch_authorization_rows(pool, sqlx::query( r#" WITH requested AS ( SELECT DISTINCT community_id, pubkey, session_owner_pubkey @@ -117,8 +171,7 @@ async fn live_authorization_state_batch( .bind(community_ids) .bind(pubkeys) .bind(session_owner_pubkeys) - .fetch_all(&mut *connection) - .await?; + ).await?; rows.into_iter() .map(|row| { @@ -198,6 +251,175 @@ mod postgres_tests { Keys::generate().public_key().to_bytes() } + async fn wait_for_batch_activity(monitor: &PgPool, application: &str, active: bool) { + tokio::time::timeout(std::time::Duration::from_millis(1500), async { + loop { + let running: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity \ + WHERE application_name = $1 AND state = 'active' AND wait_event_type = 'Lock')", + ) + .bind(application) + .fetch_one(monitor) + .await + .expect("observe batch backend"); + if running == active { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("batch backend reaches expected activity within deadline"); + } + + async fn blocked_batch_fixture() -> ( + Db, + PgPool, + sqlx::Transaction<'static, sqlx::Postgres>, + String, + ) { + use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + let monitor = setup_db().await; + let application = format!("live-auth-deadline-{}", Uuid::new_v4()); + let options: PgConnectOptions = crate::test_support::database_url().parse().unwrap(); + let pool = PgPoolOptions::new() + .max_connections(1) + .after_connect(|conn, _| { + Box::pin(async move { + sqlx::raw_sql("SET statement_timeout = '120s'; SET lock_timeout = '0'") + .execute(conn) + .await?; + Ok(()) + }) + }) + .connect_with(options.application_name(&application)) + .await + .unwrap(); + let mut lock = monitor.pool.begin().await.unwrap(); + sqlx::query("LOCK TABLE users IN ACCESS EXCLUSIVE MODE") + .execute(&mut *lock) + .await + .unwrap(); + (Db::from_pool(pool), monitor.pool, lock, application) + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn batch_deadline_preserves_tighter_cap_and_restores_writer_session() { + let (db, _monitor, lock, _application) = blocked_batch_fixture().await; + lock.rollback().await.unwrap(); + for (configured, effective, restored) in + [("120s", "900ms", "2min"), ("200ms", "200ms", "200ms")] + { + let mut conn = db.pool.acquire().await.unwrap(); + sqlx::query("SELECT set_config('statement_timeout', $1, false)") + .bind(configured) + .execute(&mut *conn) + .await + .unwrap(); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *conn) + .await + .unwrap(); + drop(conn); + let rows = fetch_authorization_rows( + &db.pool, + sqlx::query("SELECT current_setting('statement_timeout') AS cap"), + ) + .await + .unwrap(); + assert_eq!( + rows[0].get::("cap"), + effective, + "preserve tighter cap inside read" + ); + db.live_authorization_state_batch(&[LiveAuthorizationTarget { + community_id: CommunityId::from_uuid(Uuid::new_v4()), + pubkey: pubkey(), + session_owner_pubkey: None, + }]) + .await + .unwrap(); + let mut conn = db.pool.acquire().await.unwrap(); + let after: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .unwrap(); + let after_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!( + after_pid, pid, + "successful batch reuses the same writer backend" + ); + assert_eq!( + after, restored, + "writer writes retain their session timeout" + ); + } + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn batch_statement_deadline_releases_writer_and_cancels_backend() { + let (db, monitor, lock, application) = blocked_batch_fixture().await; + let target = LiveAuthorizationTarget { + community_id: CommunityId::from_uuid(Uuid::new_v4()), + pubkey: pubkey(), + session_owner_pubkey: None, + }; + let task_db = db.clone(); + let task = tokio::spawn(async move { + tokio::time::timeout( + std::time::Duration::from_millis(1500), + task_db.live_authorization_state_batch(&[target]), + ) + .await + }); + wait_for_batch_activity(&monitor, &application, true).await; + let error = task + .await + .unwrap() + .expect("server cancels before client deadline") + .expect_err("locked batch must fail"); + let DbError::Sqlx(sqlx::Error::Database(error)) = error else { + panic!("expected PostgreSQL statement cancellation, got {error:?}"); + }; + assert_eq!(error.code().as_deref(), Some("57014")); + let _available = + tokio::time::timeout(std::time::Duration::from_millis(1500), db.pool.acquire()) + .await + .expect("writer slot released promptly") + .unwrap(); + wait_for_batch_activity(&monitor, &application, false).await; + lock.rollback().await.unwrap(); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn cancelled_batch_closes_connection_and_stops_backend_within_deadline() { + let (db, monitor, lock, application) = blocked_batch_fixture().await; + let target = LiveAuthorizationTarget { + community_id: CommunityId::from_uuid(Uuid::new_v4()), + pubkey: pubkey(), + session_owner_pubkey: None, + }; + let task_db = db.clone(); + let task = + tokio::spawn(async move { task_db.live_authorization_state_batch(&[target]).await }); + wait_for_batch_activity(&monitor, &application, true).await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + let _available = + tokio::time::timeout(std::time::Duration::from_millis(1500), db.pool.acquire()) + .await + .expect("cancelled checkout does not strand writer slot") + .unwrap(); + wait_for_batch_activity(&monitor, &application, false).await; + lock.rollback().await.unwrap(); + } + #[tokio::test] #[ignore = "requires Postgres"] async fn batch_is_tenant_scoped_reads_owner_bans_ignores_timeouts_and_scales_to_10k() {