diff --git a/CHANGELOG.md b/CHANGELOG.md index 945ae5fa0..b0687ddaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -210,6 +210,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed +- **Security (`ShadowSentinel`)**: `check_tool_call` awaited its two pre-tool-dispatch DB reads + (`get_trajectory`, `get_tool_history`) with no timeout, so a stalled DB connection (e.g. a + slow/unresponsive Postgres backend) could block dispatch of every `Shell`/`FileWrite`/ + `ExfilCapable`/`McpUnclassified` tool call for the whole session. Both reads are now wrapped + in `tokio::time::timeout`, bounded by the existing `probe_timeout_ms.min(2000)` (no new config + field). A timeout logs a warning and falls back to the same empty/partial trajectory the + pre-existing DB-error branch already produced — fail-open, matching `ShadowSentinel`'s + documented defence-in-depth contract; the primary `PolicyGateExecutor`/`TrajectorySentinel` + gates are unaffected and continue to run regardless (#6269). - **Worktree**: `--bare` silently skipped the entire worktree subsystem bootstrap (`WorktreeManager` construction, `probe_capabilities`) with no warning when `worktree.enabled = true` in the active config — the 6th confirmed instance of the `--bare` diff --git a/crates/zeph-core/src/agent/shadow_sentinel.rs b/crates/zeph-core/src/agent/shadow_sentinel.rs index 90f505817..3bdc7f5cc 100644 --- a/crates/zeph-core/src/agent/shadow_sentinel.rs +++ b/crates/zeph-core/src/agent/shadow_sentinel.rs @@ -695,64 +695,47 @@ impl ShadowSentinel { false } - /// Evaluate a proposed tool call and return a probe verdict. - /// - /// Returns `ProbeVerdict::Skip` when: - /// - The tool is not in a high-risk category. - /// - The feature is disabled. - /// - The per-turn probe budget (`max_probes_per_turn`) is exhausted. - /// - /// `ToolRiskCategory::ExfilCapable` calls (#5749) draw from their own independent, higher - /// budget (`2 * max_probes_per_turn`) instead of the shared counter, so unrelated earlier - /// probes in the same turn can never wave one through — but the budget is still finite, not - /// unconditional. `ToolRiskCategory::McpUnclassified` calls (#5750) get a reduced share of - /// the shared budget — at least one slot is always reserved for keyword-matched - /// (higher-confidence) categories so a burst of low-signal MCP engagement cannot starve them - /// out within the same turn, at any `max_probes_per_turn` value. + /// Load the trajectory + cross-session tool history used as probe context. /// - /// This method takes `&self` so it can be called from parallel tool dispatch. + /// Filters out `probe_result` events — exposing probe verdicts to the LLM would allow + /// prompt injection attacks that craft tool outputs to manipulate perceived safety. /// - /// # Errors + /// Each DB read is independently bounded by `probe_timeout_ms.min(2000)`: a stalled DB + /// connection must never block dispatch of every high-risk tool call for the session. A + /// timeout or DB error falls back to an empty/partial result (fail-open), matching the + /// probe's own fail-open default. /// - /// Does not return errors; probe failures are handled internally (fail-open or - /// fail-closed depending on `deny_on_timeout`). - #[tracing::instrument(name = "security.shadow.check", skip(self, tool_args), fields(tool_id = %qualified_tool_id))] - pub async fn check_tool_call( - &self, - qualified_tool_id: &str, - tool_args: &JsonValue, - turn_number: u64, - current_risk_level: &str, - ) -> ProbeVerdict { - if !self.config.enabled { - return ProbeVerdict::Skip; - } - - let category = self.classify_tool(qualified_tool_id); - if category == ToolRiskCategory::Low { - return ProbeVerdict::Skip; - } - - if self.probe_budget_exhausted(category) { - return ProbeVerdict::Skip; - } - - // Load recent trajectory for probe context. - // Filter out probe_result events — exposing probe verdicts to the LLM would allow - // prompt injection attacks that craft tool outputs to manipulate perceived safety. - let mut trajectory: Vec = match self - .store - .get_trajectory(&self.session_id, self.config.max_context_events) - .await + /// The two reads run sequentially, each with its own independent timeout budget, and the + /// LLM probe call in [`check_tool_call`](Self::check_tool_call) has its own separate, + /// uncapped `probe_timeout_ms` timeout on top — worst-case `check_tool_call` latency is + /// therefore additive across all three: `2 * probe_timeout_ms.min(2000) + probe_timeout_ms` + /// (~6s at the 2000ms default), not a single shared ~2s bound. + async fn load_probe_context(&self, qualified_tool_id: &str) -> Vec { + let db_timeout_ms = self.config.probe_timeout_ms.min(2000); + let db_timeout = std::time::Duration::from_millis(db_timeout_ms); + + let mut trajectory: Vec = match tokio::time::timeout( + db_timeout, + self.store + .get_trajectory(&self.session_id, self.config.max_context_events), + ) + .await { - Ok(t) => t + Ok(Ok(t)) => t .into_iter() .filter(|e| e.event_type != "probe_result") .collect(), - Err(e) => { + Ok(Err(e)) => { tracing::warn!(error = %e, "ShadowSentinel: failed to load trajectory, proceeding without context"); vec![] } + Err(_) => { + tracing::warn!( + timeout_ms = db_timeout_ms, + "ShadowSentinel: trajectory load timed out, proceeding without context" + ); + vec![] + } }; // Reserve half the total budget for cross-session history so recurring risk patterns @@ -772,16 +755,17 @@ impl ShadowSentinel { // other sessions inform the probe, not just the current session (#5449). The // current session is excluded in SQL (not just filtered client-side) so its own // activity can never crowd genuinely cross-session rows out of the LIMIT clip. - match self - .store - .get_tool_history( + match tokio::time::timeout( + db_timeout, + self.store.get_tool_history( qualified_tool_id, self.session_id.as_str(), self.config.max_context_events, - ) - .await + ), + ) + .await { - Ok(history) => { + Ok(Ok(history)) => { // get_tool_history is DESC (newest first); reverse to ASC to match // trajectory ordering, then prepend so trajectory stays oldest-first. let mut cross_session: Vec = history @@ -795,11 +779,64 @@ impl ShadowSentinel { } trajectory.splice(0..0, cross_session); } - Err(e) => { + Ok(Err(e)) => { tracing::warn!(error = %e, "ShadowSentinel: failed to load cross-session tool history, proceeding without it"); } + Err(_) => { + tracing::warn!( + timeout_ms = db_timeout_ms, + "ShadowSentinel: cross-session tool history load timed out, proceeding without it" + ); + } + } + + trajectory + } + + /// Evaluate a proposed tool call and return a probe verdict. + /// + /// Returns `ProbeVerdict::Skip` when: + /// - The tool is not in a high-risk category. + /// - The feature is disabled. + /// - The per-turn probe budget (`max_probes_per_turn`) is exhausted. + /// + /// `ToolRiskCategory::ExfilCapable` calls (#5749) draw from their own independent, higher + /// budget (`2 * max_probes_per_turn`) instead of the shared counter, so unrelated earlier + /// probes in the same turn can never wave one through — but the budget is still finite, not + /// unconditional. `ToolRiskCategory::McpUnclassified` calls (#5750) get a reduced share of + /// the shared budget — at least one slot is always reserved for keyword-matched + /// (higher-confidence) categories so a burst of low-signal MCP engagement cannot starve them + /// out within the same turn, at any `max_probes_per_turn` value. + /// + /// This method takes `&self` so it can be called from parallel tool dispatch. + /// + /// # Errors + /// + /// Does not return errors; probe failures are handled internally (fail-open or + /// fail-closed depending on `deny_on_timeout`). + #[tracing::instrument(name = "security.shadow.check", skip(self, tool_args), fields(tool_id = %qualified_tool_id))] + pub async fn check_tool_call( + &self, + qualified_tool_id: &str, + tool_args: &JsonValue, + turn_number: u64, + current_risk_level: &str, + ) -> ProbeVerdict { + if !self.config.enabled { + return ProbeVerdict::Skip; } + let category = self.classify_tool(qualified_tool_id); + if category == ToolRiskCategory::Low { + return ProbeVerdict::Skip; + } + + if self.probe_budget_exhausted(category) { + return ProbeVerdict::Skip; + } + + let trajectory = self.load_probe_context(qualified_tool_id).await; + let verdict = self .probe .evaluate(qualified_tool_id, tool_args, &trajectory) @@ -1995,6 +2032,77 @@ mod tests { ); } + // ── #6269: DB-read timeout fail-open ───────────────────────────────────── + + /// #6269 regression: both DB reads inside `load_probe_context` (`get_trajectory` and + /// `get_tool_history`, reached via `check_tool_call`) must fail open when the DB pool + /// stalls, exactly like their existing `Err` (DB-error) branches and the LLM probe's own + /// timeout branch. A real stall is forced — not a synthetic sleep race — by exhausting + /// the in-memory `SQLite` pool's sole connection: `test_pool()` connects with `":memory:"`, + /// which `crates/zeph-db/src/pool.rs` hard-caps at `max_connections(1)`, so holding one + /// `BEGIN IMMEDIATE` transaction open blocks both `fetch_all(&pool)` calls on + /// `pool.acquire()` until `probe_timeout_ms` elapses. + #[tokio::test] + async fn check_tool_call_falls_open_when_both_db_reads_stall() { + use tracing_subscriber::layer::SubscriberExt as _; + + let pool = test_pool().await; + let raw_pool = pool.clone(); + let store = ShadowEventStore::new(pool); + + // Seed real rows so an empty captured trajectory can only be explained by the + // timeout fallback below, not by the store genuinely having nothing to return. + let base = unix_now(); + seed_events(&store, "current-session", "builtin:shell", "own", base, 2).await; + seed_events(&store, "other-session", "builtin:shell", "cross", base, 2).await; + + let messages: Arc>> = + Arc::new(std::sync::Mutex::new(Vec::new())); + let layer = MessageCaptureLayer { + messages: messages.clone(), + }; + let subscriber = tracing_subscriber::registry().with(layer); + let _guard = tracing::subscriber::set_default(subscriber); + + let config = zeph_config::ShadowSentinelConfig { + enabled: true, + probe_timeout_ms: 50, + ..zeph_config::ShadowSentinelConfig::default() + }; + + // Hold the sole in-memory SQLite connection so both `fetch_all` calls inside + // `load_probe_context` block on `pool.acquire()` for the full 50ms probe timeout. + let tx = zeph_db::begin_write(&raw_pool) + .await + .expect("hold sole in-memory sqlite connection"); + + let trajectory = + capture_check_tool_call_trajectory(store, config, "current-session", "builtin:shell") + .await; + + drop(tx); + + assert!( + trajectory.is_empty(), + "trajectory passed to the probe must be empty when both get_trajectory and \ + get_tool_history time out, despite real seeded data existing; got: {trajectory:?}" + ); + + let captured_logs = messages.lock().unwrap(); + assert!( + captured_logs + .iter() + .any(|m| m.contains("trajectory load timed out")), + "expected a warn log for the timed-out get_trajectory read, got: {captured_logs:?}" + ); + assert!( + captured_logs + .iter() + .any(|m| m.contains("cross-session tool history load timed out")), + "expected a warn log for the timed-out get_tool_history read, got: {captured_logs:?}" + ); + } + // ── #5766: record_tool_event had zero test coverage ───────────────────────── #[tokio::test]