From 3c1b1bbc2c74d0defd32d567af40629737269638 Mon Sep 17 00:00:00 2001 From: Adam J Esslinger Date: Wed, 16 Sep 2026 14:46:49 -0400 Subject: [PATCH] fix: bound startup-script preflight with a 30s timeout (#85) preflight_startup_script ran the connect + transaction + batch_execute + rollback sequence with no timeout, so a startup script that blocks (an advisory lock, a slow function, pg_sleep) or a stalled host wedged pool creation -- and therefore every RPC that needs a pool -- indefinitely. The builtin driver wraps the same preflight in a 30-second tokio::time::timeout with a clearly attributed error message; this ports that guard exactly (same STARTUP_SCRIPT_TIMEOUT_MS value, same error text). TDD: added hung_startup_script_times_out_instead_of_hanging_pool_creation (a pg_sleep(35) startup script) to tests/live_db.rs. Confirmed it hangs past 40s against the pre-fix code (killed by an outer `timeout 40`, exit 124) and passes in ~30.3s against the fix, with the exact expected "Timed out running PostgreSQL startup script after 30000 ms" message. Full live_db suite (25 passed, 1 pre-existing pgvector-only test ignored) and the 309-test unit suite both pass; clippy and fmt clean. --- src/client.rs | 37 +++++++++++++++++++++++++++++++++++++ tests/live_db.rs | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/client.rs b/src/client.rs index 3ad117b..7a1dd0b 100644 --- a/src/client.rs +++ b/src/client.rs @@ -21,6 +21,7 @@ use std::collections::HashMap; use std::str::FromStr; use std::sync::{LazyLock, Mutex}; +use std::time::Duration; use deadpool_postgres::{Config, ManagerConfig, Pool, RecyclingMethod, Runtime, SslMode}; use tokio_postgres::types::{ToSql, Type}; @@ -426,6 +427,12 @@ fn startup_script_error(err: tokio_postgres::Error) -> String { format!("Startup script failed: {}", format_pg_error(&err)) } +/// The startup-script preflight opens a real network connection, so bound +/// it: a broken script or a stalled host must never wedge pool creation +/// indefinitely. Matches the builtin driver's +/// `POSTGRES_STARTUP_SCRIPT_TIMEOUT_MS` (`src-tauri/src/pool_manager.rs`). +const STARTUP_SCRIPT_TIMEOUT_MS: u64 = 30_000; + /// Build the `post_create` hook that runs the startup script on every new /// pooled connection (matches the builtin driver's `post_create` hook — see /// `src-tauri/src/pool_manager.rs`). @@ -449,6 +456,11 @@ fn startup_script_hook(script: &str) -> deadpool_postgres::Hook { /// preflight exists only for early, well-labelled failures — the per-pool /// `post_create` hook is the single place the script actually takes effect. /// Matches the builtin driver's `run_postgres_startup_script` preflight. +/// +/// Bounded by [`STARTUP_SCRIPT_TIMEOUT_MS`]: this opens a real network +/// connection and runs caller-supplied SQL, so a broken script or a stalled +/// host must never wedge pool creation (and therefore every RPC that needs +/// a pool) indefinitely. async fn preflight_startup_script(cfg: &Config, tls: T, script: &str) -> Result<(), String> where T: tokio_postgres::tls::MakeTlsConnect + Clone + Sync + Send + 'static, @@ -459,6 +471,31 @@ where let pg_config = cfg .get_pg_config() .map_err(|e| format!("Pool creation failed: {e}"))?; + let timeout = Duration::from_millis(STARTUP_SCRIPT_TIMEOUT_MS); + tokio::time::timeout( + timeout, + run_startup_script_preflight(pg_config, tls, script), + ) + .await + .map_err(|_| { + format!( + "Timed out running PostgreSQL startup script after {} ms", + timeout.as_millis() + ) + })? +} + +async fn run_startup_script_preflight( + pg_config: tokio_postgres::Config, + tls: T, + script: &str, +) -> Result<(), String> +where + T: tokio_postgres::tls::MakeTlsConnect + Clone + Sync + Send + 'static, + T::Stream: Sync + Send, + T::TlsConnect: Sync + Send, + >::Future: Send, +{ let (mut client, connection) = pg_config .connect(tls) .await diff --git a/tests/live_db.rs b/tests/live_db.rs index 55f9c06..281dad9 100644 --- a/tests/live_db.rs +++ b/tests/live_db.rs @@ -803,6 +803,41 @@ fn broken_startup_script_fails_fast_with_clear_attribution() { ); } +// Coverage for #85: preflight_startup_script ran the caller-supplied script +// with no timeout, so a script that blocks (or a stalled host) wedged pool +// creation — and therefore every RPC that needs a pool — indefinitely. +// pg_sleep(35) outlasts the 30-second timeout, so this must fail with a +// clear "Timed out" message well before 35s, not hang until the script +// itself finishes. +#[test] +fn hung_startup_script_times_out_instead_of_hanging_pool_creation() { + let mut plugin = Plugin::spawn(); + let mut params = conn_params(); + // Use a startup_script unique to this test so the pool cache key (which + // folds in startup_script) can't reuse a pool already validated by + // another test — a fresh identity guarantees the preflight actually runs. + params["startup_script"] = json!("SELECT pg_sleep(35)"); + + let started = std::time::Instant::now(); + let response = plugin.call("test_connection", json!({ "params": params })); + let elapsed = started.elapsed(); + + let error = response + .get("error") + .and_then(|e| e.get("message")) + .and_then(Value::as_str) + .expect("a hung startup script must produce a JSON-RPC error, not hang forever"); + assert_eq!( + error, "Timed out running PostgreSQL startup script after 30000 ms", + "error should clearly attribute the failure to the startup-script timeout" + ); + assert!( + elapsed < std::time::Duration::from_secs(33), + "preflight should fail at the 30s timeout, not wait for the 35s pg_sleep to finish \ + (took {elapsed:?})" + ); +} + // Coverage for #43: build_pool never called cfg.ssl_mode(...), so // tokio_postgres's own default (SslMode::Prefer) applied regardless of the // plugin's ssl_mode value, letting ssl_mode=require silently connect over