Files
codex/prs/bolinfest/PR-2340.md
2025-09-02 15:17:45 -07:00

19 KiB
Raw Blame History

PR #2340: fix: introduce MutexExt::lock_unchecked() so we stop ignoring unwrap() throughout codex.rs

Description

This way we are sure a dangerous unwrap() does not sneak in!


Stack created with Sapling. Best reviewed with ReviewStack.

  • #2345
  • #2329
  • #2343
  • -> #2340
  • #2338

Full Diff

diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs
index 21e80406e5..fcccb40f8c 100644
--- a/codex-rs/core/src/apply_patch.rs
+++ b/codex-rs/core/src/apply_patch.rs
@@ -8,7 +8,6 @@ use crate::safety::assess_patch_safety;
 use codex_apply_patch::ApplyPatchAction;
 use codex_apply_patch::ApplyPatchFileChange;
 use std::collections::HashMap;
-use std::path::Path;
 use std::path::PathBuf;
 
 pub const CODEX_APPLY_PATCH_ARG1: &str = "--codex-run-as-apply-patch";
@@ -45,12 +44,10 @@ pub(crate) async fn apply_patch(
     call_id: &str,
     action: ApplyPatchAction,
 ) -> InternalApplyPatchInvocation {
-    let writable_roots_snapshot = sess.get_writable_roots().to_vec();
-
     match assess_patch_safety(
         &action,
         sess.get_approval_policy(),
-        &writable_roots_snapshot,
+        sess.get_sandbox_policy(),
         sess.get_cwd(),
     ) {
         SafetyCheck::AutoApprove { .. } => {
@@ -124,30 +121,3 @@ pub(crate) fn convert_apply_patch_to_protocol(
     }
     result
 }
-
-pub(crate) fn get_writable_roots(cwd: &Path) -> Vec<PathBuf> {
-    let mut writable_roots = Vec::new();
-    if cfg!(target_os = "macos") {
-        // On macOS, $TMPDIR is private to the user.
-        writable_roots.push(std::env::temp_dir());
-
-        // Allow pyenv to update its shims directory. Without this, any tool
-        // that happens to be managed by `pyenv` will fail with an error like:
-        //
-        //   pyenv: cannot rehash: $HOME/.pyenv/shims isn't writable
-        //
-        // which is emitted every time `pyenv` tries to run `rehash` (for
-        // example, after installing a new Python package that drops an entry
-        // point). Although the sandbox is intentionally readonly by default,
-        // writing to the user's local `pyenv` directory is safe because it
-        // is already userwritable and scoped to the current user account.
-        if let Ok(home_dir) = std::env::var("HOME") {
-            let pyenv_dir = PathBuf::from(home_dir).join(".pyenv");
-            writable_roots.push(pyenv_dir);
-        }
-    }
-
-    writable_roots.push(cwd.to_path_buf());
-
-    writable_roots
-}
diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs
index 482ac2f1ab..6b6ba9ad7a 100644
--- a/codex-rs/core/src/codex.rs
+++ b/codex-rs/core/src/codex.rs
@@ -1,6 +1,3 @@
-// Poisoned mutex should fail the program
-#![expect(clippy::unwrap_used)]
-
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -8,6 +5,7 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::Mutex;
+use std::sync::MutexGuard;
 use std::sync::atomic::AtomicU64;
 use std::time::Duration;
 
@@ -31,12 +29,11 @@ use tracing::warn;
 use uuid::Uuid;
 
 use crate::ModelProviderInfo;
+use crate::apply_patch;
 use crate::apply_patch::ApplyPatchExec;
 use crate::apply_patch::CODEX_APPLY_PATCH_ARG1;
 use crate::apply_patch::InternalApplyPatchInvocation;
 use crate::apply_patch::convert_apply_patch_to_protocol;
-use crate::apply_patch::get_writable_roots;
-use crate::apply_patch::{self};
 use crate::client::ModelClient;
 use crate::client_common::Prompt;
 use crate::client_common::ResponseEvent;
@@ -108,6 +105,21 @@ use crate::turn_diff_tracker::TurnDiffTracker;
 use crate::user_notification::UserNotification;
 use crate::util::backoff;
 
+// A convenience extension trait for acquiring mutex locks where poisoning is
+// unrecoverable and should abort the program. This avoids scattered `.unwrap()`
+// calls on `lock()` while still surfacing a clear panic message when a lock is
+// poisoned.
+trait MutexExt<T> {
+    fn lock_unchecked(&self) -> MutexGuard<'_, T>;
+}
+
+impl<T> MutexExt<T> for Mutex<T> {
+    fn lock_unchecked(&self) -> MutexGuard<'_, T> {
+        #[expect(clippy::expect_used)]
+        self.lock().expect("poisoned lock")
+    }
+}
+
 /// The high-level interface to the Codex system.
 /// It operates as a queue pair where you send submissions and receive events.
 pub struct Codex {
@@ -230,7 +242,6 @@ pub(crate) struct Session {
     approval_policy: AskForApproval,
     sandbox_policy: SandboxPolicy,
     shell_environment_policy: ShellEnvironmentPolicy,
-    writable_roots: Vec<PathBuf>,
     disable_response_storage: bool,
     tools_config: ToolsConfig,
 
@@ -409,8 +420,6 @@ impl Session {
             state.history.record_items(&restored_items);
         }
 
-        let writable_roots = get_writable_roots(&cwd);
-
         // Handle MCP manager result and record any startup failures.
         let (mcp_connection_manager, failed_clients) = match mcp_res {
             Ok((mgr, failures)) => (mgr, failures),
@@ -463,7 +472,6 @@ impl Session {
             sandbox_policy,
             shell_environment_policy: config.shell_environment_policy.clone(),
             cwd,
-            writable_roots,
             mcp_connection_manager,
             notify,
             state: Mutex::new(state),
@@ -507,14 +515,14 @@ impl Session {
         Ok(sess)
     }
 
-    pub(crate) fn get_writable_roots(&self) -> &[PathBuf] {
-        &self.writable_roots
-    }
-
     pub(crate) fn get_approval_policy(&self) -> AskForApproval {
         self.approval_policy
     }
 
+    pub(crate) fn get_sandbox_policy(&self) -> &SandboxPolicy {
+        &self.sandbox_policy
+    }
+
     pub(crate) fn get_cwd(&self) -> &Path {
         &self.cwd
     }
@@ -526,7 +534,7 @@ impl Session {
     }
 
     pub fn set_task(&self, task: AgentTask) {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         if let Some(current_task) = state.current_task.take() {
             current_task.abort();
         }
@@ -534,7 +542,7 @@ impl Session {
     }
 
     pub fn remove_task(&self, sub_id: &str) {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         if let Some(task) = &state.current_task {
             if task.sub_id == sub_id {
                 state.current_task.take();
@@ -570,7 +578,7 @@ impl Session {
         };
         let _ = self.tx_event.send(event).await;
         {
-            let mut state = self.state.lock().unwrap();
+            let mut state = self.state.lock_unchecked();
             state.pending_approvals.insert(sub_id, tx_approve);
         }
         rx_approve
@@ -596,21 +604,21 @@ impl Session {
         };
         let _ = self.tx_event.send(event).await;
         {
-            let mut state = self.state.lock().unwrap();
+            let mut state = self.state.lock_unchecked();
             state.pending_approvals.insert(sub_id, tx_approve);
         }
         rx_approve
     }
 
     pub fn notify_approval(&self, sub_id: &str, decision: ReviewDecision) {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         if let Some(tx_approve) = state.pending_approvals.remove(sub_id) {
             tx_approve.send(decision).ok();
         }
     }
 
     pub fn add_approved_command(&self, cmd: Vec<String>) {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         state.approved_commands.insert(cmd);
     }
 
@@ -620,14 +628,14 @@ impl Session {
         debug!("Recording items for conversation: {items:?}");
         self.record_state_snapshot(items).await;
 
-        self.state.lock().unwrap().history.record_items(items);
+        self.state.lock_unchecked().history.record_items(items);
     }
 
     async fn record_state_snapshot(&self, items: &[ResponseItem]) {
         let snapshot = { crate::rollout::SessionStateSnapshot {} };
 
         let recorder = {
-            let guard = self.rollout.lock().unwrap();
+            let guard = self.rollout.lock_unchecked();
             guard.as_ref().cloned()
         };
 
@@ -805,12 +813,12 @@ impl Session {
     /// Build the full turn input by concatenating the current conversation
     /// history with additional items for this turn.
     pub fn turn_input_with_history(&self, extra: Vec<ResponseItem>) -> Vec<ResponseItem> {
-        [self.state.lock().unwrap().history.contents(), extra].concat()
+        [self.state.lock_unchecked().history.contents(), extra].concat()
     }
 
     /// Returns the input if there was no task running to inject into
     pub fn inject_input(&self, input: Vec<InputItem>) -> Result<(), Vec<InputItem>> {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         if state.current_task.is_some() {
             state.pending_input.push(input.into());
             Ok(())
@@ -820,7 +828,7 @@ impl Session {
     }
 
     pub fn get_pending_input(&self) -> Vec<ResponseInputItem> {
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         if state.pending_input.is_empty() {
             Vec::with_capacity(0)
         } else {
@@ -844,7 +852,7 @@ impl Session {
 
     fn abort(&self) {
         info!("Aborting existing session");
-        let mut state = self.state.lock().unwrap();
+        let mut state = self.state.lock_unchecked();
         state.pending_approvals.clear();
         state.pending_input.clear();
         if let Some(task) = state.current_task.take() {
@@ -1048,7 +1056,7 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
 
                 // Gracefully flush and shutdown rollout recorder on session end so tests
                 // that inspect the rollout file do not race with the background writer.
-                let recorder_opt = sess.rollout.lock().unwrap().take();
+                let recorder_opt = sess.rollout.lock_unchecked().take();
                 if let Some(rec) = recorder_opt {
                     if let Err(e) = rec.shutdown().await {
                         warn!("failed to shutdown rollout recorder: {e}");
@@ -1464,7 +1472,7 @@ async fn try_run_turn(
             }
             ResponseEvent::OutputTextDelta(delta) => {
                 {
-                    let mut st = sess.state.lock().unwrap();
+                    let mut st = sess.state.lock_unchecked();
                     st.history.append_assistant_text(&delta);
                 }
 
@@ -1580,7 +1588,7 @@ async fn run_compact_task(
     };
     sess.send_event(event).await;
 
-    let mut state = sess.state.lock().unwrap();
+    let mut state = sess.state.lock_unchecked();
     state.history.keep_last_messages(1);
 }
 
@@ -1620,8 +1628,9 @@ async fn handle_response_item(
                 };
                 sess.tx_event.send(event).await.ok();
             }
-            if sess.show_raw_agent_reasoning && content.is_some() {
-                let content = content.unwrap();
+            if sess.show_raw_agent_reasoning
+                && let Some(content) = content
+            {
                 for item in content {
                     let text = match item {
                         ReasoningItemContent::ReasoningText { text } => text,
@@ -1891,7 +1900,7 @@ async fn handle_container_exec_with_params(
         }
         None => {
             let safety = {
-                let state = sess.state.lock().unwrap();
+                let state = sess.state.lock_unchecked();
                 assess_command_safety(
                     &params.command,
                     sess.approval_policy,
@@ -2231,7 +2240,7 @@ async fn drain_to_completed(sess: &Session, sub_id: &str, prompt: &Prompt) -> Co
         match event {
             Ok(ResponseEvent::OutputItemDone(item)) => {
                 // Record only to in-memory conversation history; avoid state snapshot.
-                let mut state = sess.state.lock().unwrap();
+                let mut state = sess.state.lock_unchecked();
                 state.history.record_items(std::slice::from_ref(&item));
             }
             Ok(ResponseEvent::Completed {
diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs
index ac95b6a20a..1d264d3ed1 100644
--- a/codex-rs/core/src/protocol.rs
+++ b/codex-rs/core/src/protocol.rs
@@ -156,10 +156,31 @@ pub enum SandboxPolicy {
 /// not modified by the agent.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct WritableRoot {
+    /// Absolute path, by construction.
     pub root: PathBuf,
+
+    /// Also absolute paths, by construction.
     pub read_only_subpaths: Vec<PathBuf>,
 }
 
+impl WritableRoot {
+    pub(crate) fn is_path_writable(&self, path: &Path) -> bool {
+        // Check if the path is under the root.
+        if !path.starts_with(&self.root) {
+            return false;
+        }
+
+        // Check if the path is under any of the read-only subpaths.
+        for subpath in &self.read_only_subpaths {
+            if path.starts_with(subpath) {
+                return false;
+            }
+        }
+
+        true
+    }
+}
+
 impl FromStr for SandboxPolicy {
     type Err = serde_json::Error;
 
diff --git a/codex-rs/core/src/safety.rs b/codex-rs/core/src/safety.rs
index 74872ddc4f..c878a71110 100644
--- a/codex-rs/core/src/safety.rs
+++ b/codex-rs/core/src/safety.rs
@@ -21,7 +21,7 @@ pub enum SafetyCheck {
 pub fn assess_patch_safety(
     action: &ApplyPatchAction,
     policy: AskForApproval,
-    writable_roots: &[PathBuf],
+    sandbox_policy: &SandboxPolicy,
     cwd: &Path,
 ) -> SafetyCheck {
     if action.is_empty() {
@@ -45,7 +45,7 @@ pub fn assess_patch_safety(
     // is possible that paths in the patch are hard links to files outside the
     // writable roots, so we should still run `apply_patch` in a sandbox in that
     // case.
-    if is_write_patch_constrained_to_writable_paths(action, writable_roots, cwd)
+    if is_write_patch_constrained_to_writable_paths(action, sandbox_policy, cwd)
         || policy == AskForApproval::OnFailure
     {
         // Only autoapprove when we can actually enforce a sandbox. Otherwise
@@ -171,13 +171,19 @@ pub fn get_platform_sandbox() -> Option<SandboxType> {
 
 fn is_write_patch_constrained_to_writable_paths(
     action: &ApplyPatchAction,
-    writable_roots: &[PathBuf],
+    sandbox_policy: &SandboxPolicy,
     cwd: &Path,
 ) -> bool {
     // Earlyexit if there are no declared writable roots.
-    if writable_roots.is_empty() {
-        return false;
-    }
+    let writable_roots = match sandbox_policy {
+        SandboxPolicy::ReadOnly => {
+            return false;
+        }
+        SandboxPolicy::DangerFullAccess => {
+            return true;
+        }
+        SandboxPolicy::WorkspaceWrite { .. } => sandbox_policy.get_writable_roots_with_cwd(cwd),
+    };
 
     // Normalize a path by removing `.` and resolving `..` without touching the
     // filesystem (works even if the file does not exist).
@@ -209,15 +215,9 @@ fn is_write_patch_constrained_to_writable_paths(
             None => return false,
         };
 
-        writable_roots.iter().any(|root| {
-            let root_abs = if root.is_absolute() {
-                root.clone()
-            } else {
-                normalize(&cwd.join(root)).unwrap_or_else(|| cwd.join(root))
-            };
-
-            abs.starts_with(&root_abs)
-        })
+        writable_roots
+            .iter()
+            .any(|writable_root| writable_root.is_path_writable(&abs))
     };
 
     for (path, change) in action.changes() {
@@ -246,38 +246,56 @@ fn is_write_patch_constrained_to_writable_paths(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tempfile::TempDir;
 
     #[test]
     fn test_writable_roots_constraint() {
-        let cwd = std::env::current_dir().unwrap();
+        // Use a temporary directory as our workspace to avoid touching
+        // the real current working directory.
+        let tmp = TempDir::new().unwrap();
+        let cwd = tmp.path().to_path_buf();
         let parent = cwd.parent().unwrap().to_path_buf();
 
-        // Helper to build a singleentry map representing a patch that adds a
-        // file at `p`.
+        // Helper to build a singleentry patch that adds a file at `p`.
         let make_add_change = |p: PathBuf| ApplyPatchAction::new_add_for_test(&p, "".to_string());
 
         let add_inside = make_add_change(cwd.join("inner.txt"));
         let add_outside = make_add_change(parent.join("outside.txt"));
 
+        // Policy limited to the workspace only; exclude system temp roots so
+        // only `cwd` is writable by default.
+        let policy_workspace_only = SandboxPolicy::WorkspaceWrite {
+            writable_roots: vec![],
+            network_access: false,
+            exclude_tmpdir_env_var: true,
+            exclude_slash_tmp: true,
+        };
+
         assert!(is_write_patch_constrained_to_writable_paths(
             &add_inside,
-            &[PathBuf::from(".")],
+            &policy_workspace_only,
             &cwd,
         ));
 
-        let add_outside_2 = make_add_change(parent.join("outside.txt"));
         assert!(!is_write_patch_constrained_to_writable_paths(
-            &add_outside_2,
-            &[PathBuf::from(".")],
+            &add_outside,
+            &policy_workspace_only,
             &cwd,
         ));
 
-        // With parent dir added as writable root, it should pass.
+        // With the parent dir explicitly added as a writable root, the
+        // outside write should be permitted.
+        let policy_with_parent = SandboxPolicy::WorkspaceWrite {
+            writable_roots: vec![parent.clone()],
+            network_access: false,
+            exclude_tmpdir_env_var: true,
+            exclude_slash_tmp: true,
+        };
         assert!(is_write_patch_constrained_to_writable_paths(
             &add_outside,
-            &[PathBuf::from("..")],
+            &policy_with_parent,
             &cwd,
-        ))
+        ));
     }
 
     #[test]

Review Comments

codex-rs/core/src/codex.rs

@@ -107,6 +105,21 @@ use crate::turn_diff_tracker::TurnDiffTracker;
 use crate::user_notification::UserNotification;
 use crate::util::backoff;
 
+// A convenience extension trait for acquiring mutex locks where poisoning is
+// unrecoverable and should abort the program. This avoids scattered `.unwrap()`
+// calls on `lock()` while still surfacing a clear panic message when a lock is
+// poisoned.
+trait MutexExt<T> {
+    fn lock_unchecked(&self) -> MutexGuard<'_, T>;
+}
+
+impl<T> MutexExt<T> for Mutex<T> {
+    fn lock_unchecked(&self) -> MutexGuard<'_, T> {
+        #[expect(clippy::expect_used)]
+        self.lock().expect("poisoned lock")
+    }
+}

Let's do that as needed. For now, this is the only file in core that needs this.

In most other cases, I expect RwLock is more appropriate than Mutex, but I audited our access and it isn't that read-heavy.