//! Audit fix #6: a data write and the trigger fan-out it produces must commit //! ATOMICALLY. //! //! Before this, a service wrote the row (commit #1), then asked the emitter to //! resolve matching triggers and insert outbox rows (commit #2). If commit #2 //! failed, the row was permanently in the store with its trigger having never //! fired — unobservable to the caller, unrecoverable by any retry, and only //! logged. `PostgresKvWriter` runs both on ONE connection in ONE transaction, so //! an emit failure rolls the write back and surfaces as an error. //! //! Fault injection: a Postgres BEFORE-INSERT trigger on `outbox` that raises, //! scoped to THIS test's freshly-minted `app_id`, so it cannot affect any test //! running in parallel against the same database. Skips when `DATABASE_URL` is //! unset. use std::sync::Arc; use picloud_manager_core::atomic_write::{ GroupKvTarget, GroupKvWriter, KvWriter, PostgresGroupKvWriter, PostgresKvWriter, }; use picloud_manager_core::group_quota::GroupWriteQuota; use picloud_shared::{AppId, ExecutionId, GroupId, GroupKvError, RequestId, ScriptId, SdkCallCx}; use sqlx::postgres::PgPoolOptions; use sqlx::PgPool; use uuid::Uuid; async fn pool_or_skip() -> Option { let Ok(url) = std::env::var("DATABASE_URL") else { eprintln!("atomic_write: DATABASE_URL unset — skipping"); return None; }; let pool = PgPoolOptions::new() .max_connections(24) .connect(&url) .await .expect("connect"); sqlx::migrate!("./migrations") .run(&pool) .await .expect("migrate"); Some(pool) } fn cx(app_id: Uuid) -> SdkCallCx { SdkCallCx { app_id: AppId::from(app_id), script_id: ScriptId::new(), principal: None, execution_id: ExecutionId::new(), request_id: RequestId::new(), trigger_depth: 0, root_execution_id: ExecutionId::new(), is_dead_letter_handler: false, event: None, } } /// An app with one script and one enabled KV trigger watching `collection`. struct Fixture { app: Uuid, } async fn setup(pool: &PgPool, collection: &str) -> Fixture { let uniq = Uuid::new_v4().simple().to_string(); let (group,): (Uuid,) = sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id") .bind(format!("aw-grp-{uniq}")) .fetch_one(pool) .await .expect("group"); let (app,): (Uuid,) = sqlx::query_as("INSERT INTO apps (slug, name, group_id) VALUES ($1, $1, $2) RETURNING id") .bind(format!("aw-app-{uniq}")) .bind(group) .fetch_one(pool) .await .expect("app"); let (admin,): (Uuid,) = sqlx::query_as( "INSERT INTO admin_users (username, password_hash) VALUES ($1, 'x') RETURNING id", ) .bind(format!("aw-admin-{uniq}")) .fetch_one(pool) .await .expect("admin"); let (script,): (Uuid,) = sqlx::query_as( "INSERT INTO scripts (app_id, name, source) VALUES ($1, 'handler', '') RETURNING id", ) .bind(app) .fetch_one(pool) .await .expect("script"); let (trigger,): (Uuid,) = sqlx::query_as( "INSERT INTO triggers ( \ app_id, script_id, kind, retry_max_attempts, retry_backoff, \ retry_base_ms, registered_by_principal \ ) VALUES ($1, $2, 'kv', 3, 'exponential', 1000, $3) RETURNING id", ) .bind(app) .bind(script) .bind(admin) .fetch_one(pool) .await .expect("trigger"); sqlx::query( "INSERT INTO kv_trigger_details (trigger_id, collection_glob, ops) \ VALUES ($1, $2, '{}')", ) .bind(trigger) .bind(collection) .execute(pool) .await .expect("kv_trigger_details"); Fixture { app } } /// Make every outbox insert for `app` fail. Scoped by app_id so parallel tests /// (which use their own fresh apps) are untouched. async fn break_outbox_for(pool: &PgPool, app: Uuid, tag: &str) { sqlx::query(&format!( "CREATE FUNCTION break_outbox_{tag}() RETURNS TRIGGER AS $$ \ BEGIN \ IF NEW.app_id = '{app}'::uuid THEN \ RAISE EXCEPTION 'injected outbox failure'; \ END IF; \ RETURN NEW; \ END; $$ LANGUAGE plpgsql" )) .execute(pool) .await .expect("create fn"); sqlx::query(&format!( "CREATE TRIGGER break_outbox_{tag} BEFORE INSERT ON outbox \ FOR EACH ROW EXECUTE FUNCTION break_outbox_{tag}()" )) .execute(pool) .await .expect("create trigger"); } async fn unbreak_outbox(pool: &PgPool, tag: &str) { sqlx::query(&format!( "DROP TRIGGER IF EXISTS break_outbox_{tag} ON outbox" )) .execute(pool) .await .expect("drop trigger"); sqlx::query(&format!("DROP FUNCTION IF EXISTS break_outbox_{tag}()")) .execute(pool) .await .expect("drop fn"); } async fn kv_count(pool: &PgPool, app: Uuid) -> i64 { let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM kv_entries WHERE app_id = $1") .bind(app) .fetch_one(pool) .await .expect("count kv"); n } async fn outbox_count(pool: &PgPool, app: Uuid) -> i64 { let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM outbox WHERE app_id = $1") .bind(app) .fetch_one(pool) .await .expect("count outbox"); n } /// `scripts` is RESTRICT on app delete (code is not data), so drop it first. async fn cleanup(pool: &PgPool, app: Uuid) { sqlx::query("DELETE FROM scripts WHERE app_id = $1") .bind(app) .execute(pool) .await .expect("cleanup scripts"); sqlx::query("DELETE FROM apps WHERE id = $1") .bind(app) .execute(pool) .await .expect("cleanup app"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn a_failed_fan_out_rolls_the_kv_write_back() { let Some(pool) = pool_or_skip().await else { return; }; let f = setup(&pool, "watched").await; let writer = PostgresKvWriter::new(pool.clone()); let cx = cx(f.app); // A distinct, valid SQL identifier for this run's trigger/function names. let tag = Uuid::new_v4().simple().to_string(); // --- Control: with the outbox healthy, the write lands AND fans out. ----- writer .set(&cx, "watched", "k", serde_json::json!(1)) .await .expect("healthy set succeeds"); assert_eq!(kv_count(&pool, f.app).await, 1, "the row is written"); assert_eq!( outbox_count(&pool, f.app).await, 1, "the matching trigger fanned out" ); // --- The fix: a fan-out failure must roll the write back. --------------- break_outbox_for(&pool, f.app, &tag).await; let err = writer .set(&cx, "watched", "k2", serde_json::json!(2)) .await .expect_err("an outbox failure must surface as an error, not be swallowed"); unbreak_outbox(&pool, &tag).await; assert!( format!("{err}").contains("event emit"), "the error should name the emit failure, got: {err}" ); assert_eq!( kv_count(&pool, f.app).await, 1, "k2 must NOT be in the store — the write is rolled back with its fan-out. \ A committed row whose trigger never fired is exactly the durability hole \ the transactional outbox closes." ); assert_eq!( outbox_count(&pool, f.app).await, 1, "and no outbox row for the rolled-back write" ); // --- An unwatched collection still writes (no trigger → no fan-out). ---- writer .set(&cx, "unwatched", "k", serde_json::json!(3)) .await .expect("a write with no matching trigger needs no outbox row"); assert_eq!(kv_count(&pool, f.app).await, 2); assert_eq!(outbox_count(&pool, f.app).await, 1, "still no new fan-out"); cleanup(&pool, f.app).await; } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn a_failed_fan_out_rolls_a_delete_back() { let Some(pool) = pool_or_skip().await else { return; }; let f = setup(&pool, "watched").await; let writer = PostgresKvWriter::new(pool.clone()); let cx = cx(f.app); let tag = Uuid::new_v4().simple().to_string(); writer .set(&cx, "watched", "doomed", serde_json::json!(1)) .await .expect("seed"); assert_eq!(kv_count(&pool, f.app).await, 1); // A delete that can't fan out must leave the row in place — otherwise the // key is gone and nothing downstream ever learns it was removed. break_outbox_for(&pool, f.app, &tag).await; let err = writer .delete(&cx, "watched", "doomed") .await .expect_err("outbox failure surfaces"); unbreak_outbox(&pool, &tag).await; assert!(format!("{err}").contains("event emit")); assert_eq!( kv_count(&pool, f.app).await, 1, "the delete rolled back — the key is still there" ); cleanup(&pool, f.app).await; } // ---------------------------------------------------------------------------- // Audit fix #8: the per-group quota must be a BOUND, not a hint. // ---------------------------------------------------------------------------- async fn mk_group(pool: &PgPool) -> Uuid { let (g,): (Uuid,) = sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id") .bind(format!("q-grp-{}", Uuid::new_v4().simple())) .fetch_one(pool) .await .expect("group"); g } async fn group_kv_count(pool: &PgPool, group: Uuid) -> i64 { let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_kv_entries WHERE group_id = $1") .bind(group) .fetch_one(pool) .await .expect("count"); n } /// The service read the group's usage on one pooled connection and wrote on /// another, so N concurrent writers each saw the same pre-write count, each /// decided it had room, and together blew past the ceiling. The writer now holds /// a per-group advisory lock across the check AND the write, so they serialize /// and each sees its predecessor's row. /// /// Note this is NOT fixed by merely putting the check in the same transaction: /// under READ COMMITTED each transaction's `COUNT(*)` still sees a snapshot /// without the others' uncommitted rows. The lock is what does the work. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_writers_cannot_push_a_group_past_its_row_quota() { const MAX_ROWS: u64 = 5; let Some(pool) = pool_or_skip().await else { return; }; let group = mk_group(&pool).await; let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone())); let quota = GroupWriteQuota { max_rows: MAX_ROWS, max_total_bytes: u64::MAX, max_value_bytes: 256 * 1024, }; // 20 writers race to insert 20 DISTINCT new keys into an empty group whose // ceiling is 5. Exactly 5 may win. let app = Uuid::new_v4(); // cx.app_id is not used by the group writer's SQL let mut set = tokio::task::JoinSet::new(); for i in 0..20 { let w = Arc::clone(&writer); set.spawn(async move { let cx = cx(app); w.set( &cx, GroupKvTarget { group_id: GroupId::from(group), collection: "shared", key: &format!("k{i}"), }, serde_json::json!(i), quota, ) .await }); } let mut ok = 0; let mut denied = 0; while let Some(r) = set.join_next().await { match r.expect("task") { Ok(()) => ok += 1, Err(GroupKvError::QuotaExceeded { .. }) => denied += 1, Err(e) => panic!("unexpected error: {e}"), } } assert_eq!( group_kv_count(&pool, group).await, i64::try_from(MAX_ROWS).unwrap(), "the group must hold EXACTLY its ceiling — a check-then-write race would overshoot" ); let max_rows = usize::try_from(MAX_ROWS).unwrap(); assert_eq!(ok, max_rows, "exactly {MAX_ROWS} writers may win"); assert_eq!(denied, 20 - max_rows, "the rest are refused"); sqlx::query("DELETE FROM groups WHERE id = $1") .bind(group) .execute(&pool) .await .expect("cleanup"); } /// The byte ceiling has the same race, and the same fix. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_writers_cannot_push_a_group_past_its_byte_quota() { let Some(pool) = pool_or_skip().await else { return; }; let group = mk_group(&pool).await; let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone())); // A ~100-byte value, and a ceiling that fits about 4 of them. `max_value_bytes` // is set low enough that the upper-bound fast path can't short-circuit the // scan (rows_after * max_value_bytes must exceed the ceiling immediately). let value = serde_json::json!("x".repeat(100)); let quota = GroupWriteQuota { max_total_bytes: 450, max_rows: u64::MAX, max_value_bytes: 200, }; let app = Uuid::new_v4(); let mut set = tokio::task::JoinSet::new(); for i in 0..16 { let w = Arc::clone(&writer); let v = value.clone(); set.spawn(async move { let cx = cx(app); w.set( &cx, GroupKvTarget { group_id: GroupId::from(group), collection: "shared", key: &format!("k{i}"), }, v, quota, ) .await }); } while let Some(r) = set.join_next().await { match r.expect("task") { Ok(()) | Err(GroupKvError::TotalBytesQuotaExceeded { .. }) => {} Err(e) => panic!("unexpected error: {e}"), } } let (bytes,): (i64,) = sqlx::query_as( "SELECT COALESCE(SUM(octet_length(value::text)), 0)::BIGINT \ FROM group_kv_entries WHERE group_id = $1", ) .bind(group) .fetch_one(&pool) .await .expect("bytes"); assert!( bytes <= 450, "stored bytes ({bytes}) must not exceed the 450-byte ceiling — \ concurrent writers must not each see the same pre-write total" ); assert!(bytes > 0, "some writers should have succeeded"); sqlx::query("DELETE FROM groups WHERE id = $1") .bind(group) .execute(&pool) .await .expect("cleanup"); }