From 52d935d744ffb80ae440d1d8193b62619af3e94c Mon Sep 17 00:00:00 2001 From: Zaki Date: Sat, 21 Mar 2026 10:05:09 -0700 Subject: [PATCH] fix(agent): eliminate TOCTOU race in approval thread resolution (#1486) Restructure process_approval() to hold the session lock for the entire check-take-verify-mutate sequence. Previously, the lock was dropped after taking the pending approval and re-acquired to verify request_id and set state, creating a window where concurrent operations could modify thread state or lose the pending approval. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/thread_ops.rs | 223 +++++++++++++++++++++++++++++++--------- 1 file changed, 176 insertions(+), 47 deletions(-) diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 0fb968f1..3087ee73 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -868,7 +868,9 @@ impl Agent { approved: bool, always: bool, ) -> Result { - // Get pending approval for this thread + // Get pending approval for this thread. + // Hold the session lock for the entire check-take-verify-mutate sequence + // to eliminate the TOCTOU race window (fixes #1486). let pending = { let mut sess = session.lock().await; let thread = sess @@ -886,54 +888,52 @@ impl Agent { return Ok(SubmissionResult::ok_with_message("")); } - thread.take_pending_approval() - }; + let pending = match thread.take_pending_approval() { + Some(p) => p, + None => { + tracing::debug!( + %thread_id, + "Ignoring stale approval: no pending approval found" + ); + return Ok(SubmissionResult::ok_with_message("")); + } + }; - let pending = match pending { - Some(p) => p, - None => { - tracing::debug!( - %thread_id, - "Ignoring stale approval: no pending approval found" - ); - return Ok(SubmissionResult::ok_with_message("")); - } - }; - - // Verify request ID if provided - if let Some(req_id) = request_id - && req_id != pending.request_id - { - // Put it back and return error - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { + // Verify request ID while still holding the lock + if let Some(req_id) = request_id + && req_id != pending.request_id + { + // Restore the pending approval atomically without releasing the lock thread.await_approval(pending); + return Ok(SubmissionResult::error( + "Request ID mismatch. Use the correct request ID.", + )); } - return Ok(SubmissionResult::error( - "Request ID mismatch. Use the correct request ID.", - )); - } - if approved { - // If always, add to auto-approved set - if always { - let mut sess = session.lock().await; - sess.auto_approve_tool(&pending.tool_name); + // If approved with "always", set auto-approve while we still have the lock. + // Drop the thread borrow before calling sess methods to satisfy the borrow checker. + if approved && always { + let tool_name = pending.tool_name.clone(); + sess.auto_approve_tool(&tool_name); tracing::info!( "Auto-approved tool '{}' for session {}", - pending.tool_name, + tool_name, sess.id ); } - // Reset thread state to processing + // Set thread state to Processing for approved requests. + // Re-borrow the thread since the previous borrow was dropped. + if approved + && let Some(thread) = sess.threads.get_mut(&thread_id) { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.state = ThreadState::Processing; - } + thread.state = ThreadState::Processing; } + pending + }; + + if approved { // Execute the approved tool and continue the loop let mut job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session") @@ -1460,17 +1460,25 @@ impl Agent { ); { let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.clear_pending_approval(); - thread.complete_turn(&rejection); - // User message already persisted at turn start; save rejection response - self.persist_assistant_response( - thread_id, - &message.channel, - &message.user_id, - &rejection, - ) - .await; + match sess.threads.get_mut(&thread_id) { + Some(thread) => { + thread.clear_pending_approval(); + thread.complete_turn(&rejection); + // User message already persisted at turn start; save rejection response + self.persist_assistant_response( + thread_id, + &message.channel, + &message.user_id, + &rejection, + ) + .await; + } + None => { + tracing::warn!( + %thread_id, + "Thread disappeared during approval rejection" + ); + } } } @@ -2012,6 +2020,127 @@ mod tests { } } + /// Regression test for #1486: verify that request_id mismatch atomically + /// restores the pending approval (no TOCTOU window where it could be lost). + #[test] + fn test_request_id_mismatch_preserves_pending_state() { + use crate::agent::session::{PendingApproval, Session, Thread, ThreadState}; + use uuid::Uuid; + + let session_id = Uuid::new_v4(); + let thread_id = Uuid::new_v4(); + let correct_request_id = Uuid::new_v4(); + let wrong_request_id = Uuid::new_v4(); + + let mut thread = Thread::with_id(thread_id, session_id); + let pending = PendingApproval { + request_id: correct_request_id, + tool_name: "shell".to_string(), + parameters: serde_json::json!({"command": "rm -rf /"}), + display_parameters: serde_json::json!({"command": "[REDACTED]"}), + description: "Execute dangerous command".to_string(), + tool_call_id: "call_42".to_string(), + context_messages: vec![], + deferred_tool_calls: vec![], + user_timezone: None, + allow_always: false, + }; + thread.await_approval(pending); + + let mut session = Session::new("test-user"); + session.threads.insert(thread_id, thread); + + // Verify initial state + assert_eq!( + session.threads[&thread_id].state, + ThreadState::AwaitingApproval + ); + assert!(session.threads[&thread_id].pending_approval.is_some()); + + // Simulate the atomic check-take-verify sequence from the fixed code: + // take the approval, then verify request_id, and restore on mismatch. + let thread = session.threads.get_mut(&thread_id).unwrap(); + let taken = thread.take_pending_approval().unwrap(); + assert!(thread.pending_approval.is_none(), "take should clear it"); + + // Wrong request_id -- restore atomically (within same "lock scope" in production) + assert_ne!(wrong_request_id, taken.request_id); + thread.await_approval(taken); + + // Verify the pending approval is fully restored + assert_eq!( + session.threads[&thread_id].state, + ThreadState::AwaitingApproval + ); + let restored = session.threads[&thread_id] + .pending_approval + .as_ref() + .expect("pending approval must be restored after request_id mismatch"); + assert_eq!(restored.request_id, correct_request_id); + assert_eq!(restored.tool_name, "shell"); + assert_eq!(restored.tool_call_id, "call_42"); + } + + /// Regression test for #1486: verify that approval with "always" flag + /// correctly transitions thread state and registers auto-approve in a + /// single lock scope. + #[test] + fn test_approval_with_always_sets_auto_approve_and_processing() { + use crate::agent::session::{PendingApproval, Session, Thread, ThreadState}; + use uuid::Uuid; + + let session_id = Uuid::new_v4(); + let thread_id = Uuid::new_v4(); + let request_id = Uuid::new_v4(); + + let mut thread = Thread::with_id(thread_id, session_id); + let pending = PendingApproval { + request_id, + tool_name: "web_fetch".to_string(), + parameters: serde_json::json!({"url": "https://example.com"}), + display_parameters: serde_json::json!({"url": "https://example.com"}), + description: "Fetch a web page".to_string(), + tool_call_id: "call_99".to_string(), + context_messages: vec![], + deferred_tool_calls: vec![], + user_timezone: None, + allow_always: true, + }; + thread.await_approval(pending); + + let mut session = Session::new("test-user"); + session.threads.insert(thread_id, thread); + + // Simulate the combined lock scope: take + verify + auto_approve + state transition + // Simulate the combined lock scope: take + verify + auto_approve + state transition. + // Use separate scopes to mirror the borrow pattern in production code. + let taken = { + let thread = session.threads.get_mut(&thread_id).unwrap(); + thread.take_pending_approval().unwrap() + }; + + // Request ID matches + assert_eq!(taken.request_id, request_id); + + // Auto-approve the tool (thread borrow dropped, so session is accessible) + session.auto_approve_tool(&taken.tool_name); + + // Re-borrow thread and set state to Processing + { + let thread = session.threads.get_mut(&thread_id).unwrap(); + thread.state = ThreadState::Processing; + } + + // Verify both mutations happened atomically + assert!(session.is_tool_auto_approved("web_fetch")); + assert_eq!( + session.threads[&thread_id].state, + ThreadState::Processing + ); + // Pending approval should be consumed (not restored) + assert!(session.threads[&thread_id].pending_approval.is_none()); + } + // Helper function to extract the approval message without needing a full Agent instance fn extract_approval_message( session: &crate::agent::session::Session,