From a8f0059ace40a95856b314a9b14be6d5c1cdbad4 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 23 Feb 2026 09:54:37 -0800 Subject: [PATCH] feat: implement zsh shell tool via shell-escalation --- codex-rs/Cargo.lock | 4 +- codex-rs/app-server/src/main.rs | 5 - codex-rs/arg0/Cargo.toml | 3 + codex-rs/arg0/src/lib.rs | 30 + codex-rs/cli/src/main.rs | 5 - codex-rs/core/Cargo.toml | 3 + codex-rs/core/src/codex.rs | 17 +- codex-rs/core/src/config/mod.rs | 8 +- codex-rs/core/src/lib.rs | 23 +- codex-rs/core/src/sandboxing/mod.rs | 8 +- codex-rs/core/src/state/service.rs | 4 +- codex-rs/core/src/tools/runtimes/shell.rs | 329 ++++++++++- codex-rs/core/src/zsh_exec_bridge/mod.rs | 557 ------------------ codex-rs/exec-server/Cargo.toml | 3 + codex-rs/exec-server/src/unix.rs | 4 +- codex-rs/exec-server/src/unix/mcp.rs | 84 ++- .../src/unix/mcp_escalation_policy.rs | 6 +- codex-rs/shell-escalation/Cargo.toml | 5 +- codex-rs/shell-escalation/src/lib.rs | 18 +- .../src/unix/core_shell_escalation.rs | 2 +- .../src/unix/escalate_server.rs | 87 +-- codex-rs/shell-escalation/src/unix/mod.rs | 2 +- codex-rs/shell-escalation/src/unix/socket.rs | 49 +- 23 files changed, 516 insertions(+), 740 deletions(-) delete mode 100644 codex-rs/core/src/zsh_exec_bridge/mod.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b5e5cd5ac3..e333b82e78 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1408,6 +1408,7 @@ dependencies = [ "anyhow", "codex-apply-patch", "codex-linux-sandbox", + "codex-shell-escalation", "codex-utils-home-dir", "dotenvy", "tempfile", @@ -1653,6 +1654,7 @@ dependencies = [ "codex-rmcp-client", "codex-secrets", "codex-shell-command", + "codex-shell-escalation", "codex-skills", "codex-state", "codex-utils-absolute-path", @@ -1799,6 +1801,7 @@ dependencies = [ "shlex", "tempfile", "tokio", + "tokio-util", "tracing", "tracing-subscriber", ] @@ -2205,7 +2208,6 @@ version = "0.0.0" dependencies = [ "anyhow", "async-trait", - "codex-core", "codex-execpolicy", "codex-protocol", "libc", diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index c56c3c8f98..5c4e5eacc7 100644 --- a/codex-rs/app-server/src/main.rs +++ b/codex-rs/app-server/src/main.rs @@ -24,11 +24,6 @@ struct AppServerArgs { fn main() -> anyhow::Result<()> { arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move { - // Run wrapper mode only after arg0 dispatch so `codex-linux-sandbox` - // invocations don't get misclassified as zsh exec-wrapper calls. - if codex_core::maybe_run_zsh_exec_wrapper_mode()? { - return Ok(()); - } let args = AppServerArgs::parse(); let managed_config_path = managed_config_path_from_debug_env(); let loader_overrides = LoaderOverrides { diff --git a/codex-rs/arg0/Cargo.toml b/codex-rs/arg0/Cargo.toml index c5d9681327..31cee35994 100644 --- a/codex-rs/arg0/Cargo.toml +++ b/codex-rs/arg0/Cargo.toml @@ -19,3 +19,6 @@ codex-utils-home-dir = { workspace = true } dotenvy = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"] } + +[target.'cfg(unix)'.dependencies] +codex-shell-escalation = { workspace = true } diff --git a/codex-rs/arg0/src/lib.rs b/codex-rs/arg0/src/lib.rs index 3ac7e017f2..fa352d9067 100644 --- a/codex-rs/arg0/src/lib.rs +++ b/codex-rs/arg0/src/lib.rs @@ -12,6 +12,8 @@ use tempfile::TempDir; const LINUX_SANDBOX_ARG0: &str = "codex-linux-sandbox"; const APPLY_PATCH_ARG0: &str = "apply_patch"; const MISSPELLED_APPLY_PATCH_ARG0: &str = "applypatch"; +#[cfg(unix)] +const EXECVE_WRAPPER_ARG0: &str = "codex-execve-wrapper"; const LOCK_FILENAME: &str = ".lock"; const TOKIO_WORKER_STACK_SIZE_BYTES: usize = 16 * 1024 * 1024; @@ -39,6 +41,32 @@ pub fn arg0_dispatch() -> Option { .and_then(|s| s.to_str()) .unwrap_or(""); + #[cfg(unix)] + if exe_name == EXECVE_WRAPPER_ARG0 { + let mut args = std::env::args(); + let _ = args.next(); + let file = match args.next() { + Some(file) => file, + None => std::process::exit(1), + }; + let argv = args.collect::>(); + + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(_) => std::process::exit(1), + }; + let exit_code = runtime.block_on(codex_shell_escalation::unix::escalate_client::run( + file, argv, + )); + match exit_code { + Ok(exit_code) => std::process::exit(exit_code), + Err(_) => std::process::exit(1), + } + } + if exe_name == LINUX_SANDBOX_ARG0 { // Safety: [`run_main`] never returns. codex_linux_sandbox::run_main(); @@ -227,6 +255,8 @@ pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result &'static str { fn main() -> anyhow::Result<()> { arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move { - // Run wrapper mode only after arg0 dispatch so `codex-linux-sandbox` - // invocations don't get misclassified as zsh exec-wrapper calls. - if codex_core::maybe_run_zsh_exec_wrapper_mode()? { - return Ok(()); - } cli_main(codex_linux_sandbox_exe).await?; Ok(()) }) diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 15f9d7d11c..bd571248c7 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -138,6 +138,9 @@ windows-sys = { version = "0.52", features = [ [target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies] keyring = { workspace = true, features = ["sync-secret-service"] } +[target.'cfg(unix)'.dependencies] +codex-shell-escalation = { workspace = true } + [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 249cb99776..e61b7af320 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -249,7 +249,6 @@ use crate::turn_diff_tracker::TurnDiffTracker; use crate::unified_exec::UnifiedExecProcessManager; use crate::util::backoff; use crate::windows_sandbox::WindowsSandboxLevelExt; -use crate::zsh_exec_bridge::ZshExecBridge; use codex_async_utils::OrCancelExt; use codex_otel::OtelManager; use codex_otel::TelemetryAuthMode; @@ -1208,7 +1207,8 @@ impl Session { "zsh fork feature enabled, but `zsh_path` is not configured; set `zsh_path` in config.toml" ) })?; - shell::get_shell(shell::ShellType::Zsh, Some(zsh_path)).ok_or_else(|| { + let zsh_path = zsh_path.to_path_buf(); + shell::get_shell(shell::ShellType::Zsh, Some(&zsh_path)).ok_or_else(|| { anyhow::anyhow!( "zsh fork feature enabled, but zsh_path `{}` is not usable; set `zsh_path` to a valid zsh executable", zsh_path.display() @@ -1287,12 +1287,6 @@ impl Session { (None, None) }; - let zsh_exec_bridge = - ZshExecBridge::new(config.zsh_path.clone(), config.codex_home.clone()); - zsh_exec_bridge - .initialize_for_session(&conversation_id.to_string()) - .await; - let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -1308,7 +1302,7 @@ impl Session { unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), - zsh_exec_bridge, + shell_zsh_path: config.zsh_path.clone(), analytics_events_client: AnalyticsEventsClient::new( Arc::clone(&config), Arc::clone(&auth_manager), @@ -4227,7 +4221,6 @@ mod handlers { .unified_exec_manager .terminate_all_processes() .await; - sess.services.zsh_exec_bridge.shutdown().await; info!("Shutting down Codex instance"); let history = sess.clone_history().await; let turn_count = history @@ -7895,7 +7888,7 @@ mod tests { unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), - zsh_exec_bridge: ZshExecBridge::default(), + shell_zsh_path: None, analytics_events_client: AnalyticsEventsClient::new( Arc::clone(&config), Arc::clone(&auth_manager), @@ -8048,7 +8041,7 @@ mod tests { unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), - zsh_exec_bridge: ZshExecBridge::default(), + shell_zsh_path: None, analytics_events_client: AnalyticsEventsClient::new( Arc::clone(&config), Arc::clone(&auth_manager), diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ffd35e036a..2d34577223 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -372,7 +372,7 @@ pub struct Config { pub js_repl_node_module_dirs: Vec, /// Optional absolute path to patched zsh used by zsh-exec-bridge-backed shell execution. - pub zsh_path: Option, + pub zsh_path: Option, /// Value to use for `reasoning.effort` when making a request using the /// Responses API. @@ -1484,7 +1484,7 @@ pub struct ConfigOverrides { pub codex_linux_sandbox_exe: Option, pub js_repl_node_path: Option, pub js_repl_node_module_dirs: Option>, - pub zsh_path: Option, + pub zsh_path: Option, pub base_instructions: Option, pub developer_instructions: Option, pub personality: Option, @@ -1905,8 +1905,8 @@ impl Config { }) .unwrap_or_default(); let zsh_path = zsh_path_override - .or(config_profile.zsh_path.map(Into::into)) - .or(cfg.zsh_path.map(Into::into)); + .or(config_profile.zsh_path) + .or(cfg.zsh_path); let review_model = override_review_model.or(cfg.review_model); diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e40905e314..a9ca37e700 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -109,7 +109,6 @@ pub mod terminal; mod tools; pub mod turn_diff_tracker; mod turn_metadata; -mod zsh_exec_bridge; pub use rollout::ARCHIVED_SESSIONS_SUBDIR; pub use rollout::INTERACTIVE_SESSION_SOURCES; pub use rollout::RolloutRecorder; @@ -144,7 +143,17 @@ pub(crate) use codex_shell_command::is_safe_command; pub(crate) use codex_shell_command::parse_command; pub(crate) use codex_shell_command::powershell; +pub use client::ModelClient; +pub use client::ModelClientSession; +pub use client::ResponsesWebsocketVersion; pub use client::X_CODEX_TURN_METADATA_HEADER; +pub use client::ws_version_from_features; +pub use client_common::Prompt; +pub use client_common::REVIEW_PROMPT; +pub use client_common::ResponseEvent; +pub use client_common::ResponseStream; +pub use compact::content_items_to_text; +pub use event_mapping::parse_turn_item; pub use exec_policy::ExecPolicyError; pub use exec_policy::check_execpolicy_for_warnings; pub use exec_policy::format_exec_policy_error_with_source; @@ -153,18 +162,6 @@ pub use file_watcher::FileWatcherEvent; pub use safety::get_platform_sandbox; pub use tools::spec::parse_tool_input_schema; pub use turn_metadata::build_turn_metadata_header; -pub use zsh_exec_bridge::maybe_run_zsh_exec_wrapper_mode; - -pub use client::ModelClient; -pub use client::ModelClientSession; -pub use client::ResponsesWebsocketVersion; -pub use client::ws_version_from_features; -pub use client_common::Prompt; -pub use client_common::REVIEW_PROMPT; -pub use client_common::ResponseEvent; -pub use client_common::ResponseStream; -pub use compact::content_items_to_text; -pub use event_mapping::parse_turn_item; pub mod compact; pub mod memory_trace; pub mod otel_init; diff --git a/codex-rs/core/src/sandboxing/mod.rs b/codex-rs/core/src/sandboxing/mod.rs index f9c3c171e7..87f7298fd0 100644 --- a/codex-rs/core/src/sandboxing/mod.rs +++ b/codex-rs/core/src/sandboxing/mod.rs @@ -163,19 +163,13 @@ impl SandboxManager { SandboxType::MacosSeatbelt => { let mut seatbelt_env = HashMap::new(); seatbelt_env.insert(CODEX_SANDBOX_ENV_VAR.to_string(), "seatbelt".to_string()); - let zsh_exec_bridge_wrapper_socket = env - .get(crate::zsh_exec_bridge::ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR) - .map(PathBuf::from); - let zsh_exec_bridge_allowed_unix_sockets = zsh_exec_bridge_wrapper_socket - .as_ref() - .map_or_else(Vec::new, |path| vec![path.clone()]); let mut args = create_seatbelt_command_args( command.clone(), policy, sandbox_policy_cwd, enforce_managed_network, network, - &zsh_exec_bridge_allowed_unix_sockets, + &[], ); let mut full_command = Vec::with_capacity(1 + args.len()); full_command.push(MACOS_PATH_TO_SEATBELT_EXECUTABLE.to_string()); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index fe82ec84c0..96d420de05 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,9 +15,9 @@ use crate::state_db::StateDbHandle; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; -use crate::zsh_exec_bridge::ZshExecBridge; use codex_hooks::Hooks; use codex_otel::OtelManager; +use codex_utils_absolute_path::AbsolutePathBuf; use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio::sync::watch; @@ -27,7 +27,7 @@ pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, - pub(crate) zsh_exec_bridge: ZshExecBridge, + pub(crate) shell_zsh_path: Option, pub(crate) analytics_events_client: AnalyticsEventsClient, pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 11ad669169..3daf5e2ed2 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -5,7 +5,11 @@ Executes shell requests under the orchestrator: asks for approval when needed, builds a CommandSpec, and runs it under the current SandboxAttempt. */ use crate::command_canonicalization::canonicalize_command_for_approval; +use crate::error::CodexErr; +use crate::error::SandboxErr; use crate::exec::ExecToolCallOutput; +use crate::exec::SandboxType; +use crate::exec::is_likely_sandbox_denied; use crate::features::Feature; use crate::powershell::prefix_powershell_script_with_utf8; use crate::sandboxing::SandboxPermissions; @@ -26,19 +30,56 @@ use crate::tools::sandboxing::ToolCtx; use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::with_cached_approval; -use crate::zsh_exec_bridge::ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR; +use codex_execpolicy::Decision; +use codex_execpolicy::Policy; +use codex_execpolicy::RuleMatch; use codex_network_proxy::NetworkProxy; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::ReviewDecision; +use codex_protocol::protocol::SandboxPolicy; +use codex_shell_command::bash::parse_shell_lc_plain_commands; +use codex_shell_command::bash::parse_shell_lc_single_command_prefix; +#[cfg(unix)] +use codex_shell_escalation::unix::core_shell_escalation::ShellActionProvider; +#[cfg(unix)] +use codex_shell_escalation::unix::core_shell_escalation::ShellPolicyFactory; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_protocol::EscalateAction; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_server::ExecParams; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_server::ExecResult; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_server::SandboxState; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_server::ShellCommandExecutor; +#[cfg(unix)] +use codex_shell_escalation::unix::escalate_server::run_escalate_server; +#[cfg(unix)] +use codex_shell_escalation::unix::stopwatch::Stopwatch; +#[cfg(unix)] +use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; +use shlex::try_join as shlex_try_join; +use std::collections::HashMap; +use std::path::Path; use std::path::PathBuf; +#[cfg(unix)] +use std::sync::Arc; +#[cfg(unix)] +use std::time::Duration; +#[cfg(unix)] +use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] pub struct ShellRequest { pub command: Vec, pub cwd: PathBuf, pub timeout_ms: Option, - pub env: std::collections::HashMap, - pub explicit_env_overrides: std::collections::HashMap, + pub env: HashMap, + pub explicit_env_overrides: HashMap, pub network: Option, pub sandbox_permissions: SandboxPermissions, pub justification: Option, @@ -73,6 +114,7 @@ impl Sandboxable for ShellRuntime { fn sandbox_preference(&self) -> SandboxablePreference { SandboxablePreference::Auto } + fn escalate_on_failure(&self) -> bool { true } @@ -146,6 +188,209 @@ impl Approvable for ShellRuntime { } } +#[cfg(unix)] +struct CoreShellActionProvider { + policy: Arc>, + session: std::sync::Arc, + turn: std::sync::Arc, + call_id: String, + approval_policy: AskForApproval, + sandbox_policy: SandboxPolicy, + sandbox_permissions: SandboxPermissions, +} + +#[cfg(unix)] +impl CoreShellActionProvider { + fn decision_driven_by_policy(matched_rules: &[RuleMatch], decision: Decision) -> bool { + matched_rules.iter().any(|rule_match| { + !matches!(rule_match, RuleMatch::HeuristicsRuleMatch { .. }) + && rule_match.decision() == decision + }) + } + + async fn prompt( + &self, + command: &[String], + workdir: &Path, + stopwatch: &Stopwatch, + ) -> anyhow::Result { + let command = command.to_vec(); + let workdir = workdir.to_path_buf(); + let session = self.session.clone(); + let turn = self.turn.clone(); + let call_id = self.call_id.clone(); + Ok(stopwatch + .pause_for(async move { + session + .request_command_approval( + &turn, call_id, None, command, workdir, None, None, None, + ) + .await + }) + .await) + } +} + +#[cfg(unix)] +#[async_trait::async_trait] +impl ShellActionProvider for CoreShellActionProvider { + async fn determine_action( + &self, + file: &Path, + argv: &[String], + workdir: &Path, + stopwatch: &Stopwatch, + ) -> anyhow::Result { + let command = std::iter::once(file.to_string_lossy().to_string()) + .chain(argv.iter().cloned()) + .collect::>(); + let (commands, used_complex_parsing) = + if let Some(commands) = parse_shell_lc_plain_commands(&command) { + (commands, false) + } else if let Some(single_command) = parse_shell_lc_single_command_prefix(&command) { + (vec![single_command], true) + } else { + (vec![command.clone()], false) + }; + + let policy = self.policy.read().await; + let fallback = |cmd: &[String]| { + crate::exec_policy::render_decision_for_unmatched_command( + self.approval_policy, + &self.sandbox_policy, + cmd, + self.sandbox_permissions, + used_complex_parsing, + ) + }; + let evaluation = policy.check_multiple(commands.iter(), &fallback); + let decision_driven_by_policy = + Self::decision_driven_by_policy(&evaluation.matched_rules, evaluation.decision); + let needs_escalation = + self.sandbox_permissions.requires_escalated_permissions() || decision_driven_by_policy; + + Ok(match evaluation.decision { + Decision::Forbidden => EscalateAction::Deny { + reason: Some("Execution forbidden by policy".to_string()), + }, + Decision::Prompt => { + if self.approval_policy == AskForApproval::Never { + EscalateAction::Deny { + reason: Some("Execution forbidden by policy".to_string()), + } + } else if decision_driven_by_policy { + EscalateAction::Escalate + } else { + match self.prompt(&command, workdir, stopwatch).await? { + ReviewDecision::Approved + | ReviewDecision::ApprovedExecpolicyAmendment { .. } + | ReviewDecision::ApprovedForSession => { + if needs_escalation { + EscalateAction::Escalate + } else { + EscalateAction::Run + } + } + ReviewDecision::Denied => EscalateAction::Deny { + reason: Some("User denied execution".to_string()), + }, + ReviewDecision::Abort => EscalateAction::Deny { + reason: Some("User cancelled execution".to_string()), + }, + } + } + } + Decision::Allow => EscalateAction::Run, + }) + } +} + +#[cfg(unix)] +struct CoreShellCommandExecutor; + +#[cfg(unix)] +#[async_trait::async_trait] +impl ShellCommandExecutor for CoreShellCommandExecutor { + async fn run( + &self, + command: Vec, + cwd: PathBuf, + env: HashMap, + cancel_rx: CancellationToken, + sandbox_state: &SandboxState, + ) -> anyhow::Result { + let result = crate::exec::process_exec_tool_call( + crate::exec::ExecParams { + command, + cwd, + expiration: crate::exec::ExecExpiration::Cancellation(cancel_rx), + env, + network: None, + sandbox_permissions: SandboxPermissions::UseDefault, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + justification: None, + arg0: None, + }, + &sandbox_state.sandbox_policy, + &sandbox_state.sandbox_cwd, + &sandbox_state.codex_linux_sandbox_exe, + sandbox_state.use_linux_sandbox_bwrap, + None, + ) + .await?; + + Ok(ExecResult { + exit_code: result.exit_code, + output: result.aggregated_output.text, + duration: result.duration, + timed_out: result.timed_out, + }) + } +} + +#[cfg(unix)] +fn shell_execve_wrapper() -> anyhow::Result { + let exe = std::env::current_exe()?; + exe.parent() + .map(|parent| parent.join("codex-execve-wrapper")) + .ok_or_else(|| anyhow::anyhow!("failed to determine codex-execve-wrapper path")) +} + +#[cfg(unix)] +fn shell_exec_zsh_path(path: &AbsolutePathBuf) -> PathBuf { + path.to_path_buf() +} + +#[cfg(unix)] +fn map_exec_result( + sandbox: SandboxType, + result: ExecResult, +) -> Result { + let output = ExecToolCallOutput { + exit_code: result.exit_code, + stdout: crate::exec::StreamOutput::new(result.output.clone()), + stderr: crate::exec::StreamOutput::new(String::new()), + aggregated_output: crate::exec::StreamOutput::new(result.output.clone()), + duration: result.duration, + timed_out: result.timed_out, + }; + + if result.timed_out { + return Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Timeout { + output: Box::new(output), + }))); + } + + if is_likely_sandbox_denied(sandbox, &output) { + return Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + }))); + } + + Ok(output) +} + impl ToolRuntime for ShellRuntime { fn network_approval_spec( &self, @@ -165,15 +410,13 @@ impl ToolRuntime for ShellRuntime { attempt: &SandboxAttempt<'_>, ctx: &ToolCtx, ) -> Result { - let base_command = &req.command; - let session_shell = ctx.session.user_shell(); let command = maybe_wrap_shell_lc_with_snapshot( - base_command, - session_shell.as_ref(), + &req.command, + ctx.session.user_shell().as_ref(), &req.cwd, &req.explicit_env_overrides, ); - let command = if matches!(session_shell.shell_type, ShellType::PowerShell) + let command = if matches!(ctx.session.user_shell().shell_type, ShellType::PowerShell) && ctx.session.features().enabled(Feature::PowershellUtf8) { prefix_powershell_script_with_utf8(&command) @@ -181,21 +424,15 @@ impl ToolRuntime for ShellRuntime { command }; - if ctx.session.features().enabled(Feature::ShellZshFork) { - let wrapper_socket_path = ctx - .session - .services - .zsh_exec_bridge - .next_wrapper_socket_path(); - let mut zsh_fork_env = req.env.clone(); - zsh_fork_env.insert( - ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR.to_string(), - wrapper_socket_path.to_string_lossy().to_string(), - ); + #[cfg(unix)] + if let Some(shell_zsh_path) = ctx.session.services.shell_zsh_path.as_ref() + && ctx.session.features().enabled(Feature::ShellZshFork) + && matches!(ctx.session.user_shell().shell_type, ShellType::Zsh) + { let spec = build_command_spec( &command, &req.cwd, - &zsh_fork_env, + &req.env, req.timeout_ms.into(), req.sandbox_permissions, req.justification.clone(), @@ -203,12 +440,52 @@ impl ToolRuntime for ShellRuntime { let env = attempt .env_for(spec, req.network.as_ref()) .map_err(|err| ToolError::Codex(err.into()))?; - return ctx - .session - .services - .zsh_exec_bridge - .execute_shell_request(&env, &ctx.session, &ctx.turn, &ctx.call_id) - .await; + let (_, args) = env + .command + .split_first() + .ok_or_else(|| ToolError::Rejected("command args are empty".to_string()))?; + let script = shlex_try_join(args.iter().map(String::as_str)) + .map_err(|err| ToolError::Rejected(format!("serialize shell script: {err}")))?; + let effective_timeout = Duration::from_millis( + req.timeout_ms + .unwrap_or(crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS), + ); + let exec_policy = Arc::new(RwLock::new( + ctx.session.services.exec_policy.current().as_ref().clone(), + )); + let sandbox_state = SandboxState { + sandbox_policy: ctx.turn.sandbox_policy.get().clone(), + codex_linux_sandbox_exe: attempt.codex_linux_sandbox_exe.cloned(), + sandbox_cwd: req.cwd.clone(), + use_linux_sandbox_bwrap: attempt.use_linux_sandbox_bwrap, + }; + let exec_result = run_escalate_server( + ExecParams { + command: script, + workdir: req.cwd.to_string_lossy().to_string(), + timeout_ms: Some(effective_timeout.as_millis() as u64), + login: Some(false), + }, + &sandbox_state, + shell_exec_zsh_path(shell_zsh_path), + shell_execve_wrapper().map_err(|err| ToolError::Rejected(format!("{err}")))?, + exec_policy.clone(), + ShellPolicyFactory::new(CoreShellActionProvider { + policy: Arc::clone(&exec_policy), + session: Arc::clone(&ctx.session), + turn: Arc::clone(&ctx.turn), + call_id: ctx.call_id.clone(), + approval_policy: ctx.turn.approval_policy.value(), + sandbox_policy: attempt.policy.clone(), + sandbox_permissions: req.sandbox_permissions, + }), + effective_timeout, + &CoreShellCommandExecutor, + ) + .await + .map_err(|err| ToolError::Rejected(err.to_string()))?; + + return map_exec_result(attempt.sandbox, exec_result); } let spec = build_command_spec( diff --git a/codex-rs/core/src/zsh_exec_bridge/mod.rs b/codex-rs/core/src/zsh_exec_bridge/mod.rs deleted file mode 100644 index 8094a4b786..0000000000 --- a/codex-rs/core/src/zsh_exec_bridge/mod.rs +++ /dev/null @@ -1,557 +0,0 @@ -use crate::exec::ExecToolCallOutput; -use crate::tools::sandboxing::ToolError; -use std::path::PathBuf; -use tokio::sync::Mutex; -use uuid::Uuid; - -#[cfg(unix)] -use crate::error::CodexErr; -#[cfg(unix)] -use crate::error::SandboxErr; -#[cfg(unix)] -use crate::protocol::EventMsg; -#[cfg(unix)] -use crate::protocol::ExecCommandOutputDeltaEvent; -#[cfg(unix)] -use crate::protocol::ExecOutputStream; -#[cfg(unix)] -use crate::protocol::ReviewDecision; -#[cfg(unix)] -use anyhow::Context as _; -#[cfg(unix)] -use codex_protocol::approvals::ExecPolicyAmendment; -#[cfg(unix)] -use codex_utils_pty::process_group::kill_child_process_group; -#[cfg(unix)] -use serde::Deserialize; -#[cfg(unix)] -use serde::Serialize; -#[cfg(unix)] -use std::io::Read; -#[cfg(unix)] -use std::io::Write; -#[cfg(unix)] -use std::time::Instant; -#[cfg(unix)] -use tokio::io::AsyncReadExt; -#[cfg(unix)] -use tokio::net::UnixListener; -#[cfg(unix)] -use tokio::net::UnixStream; - -pub(crate) const ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR: &str = - "CODEX_ZSH_EXEC_BRIDGE_WRAPPER_SOCKET"; -pub(crate) const ZSH_EXEC_WRAPPER_MODE_ENV_VAR: &str = "CODEX_ZSH_EXEC_WRAPPER_MODE"; -#[cfg(unix)] -pub(crate) const EXEC_WRAPPER_ENV_VAR: &str = "EXEC_WRAPPER"; - -#[derive(Debug, Clone, PartialEq, Eq, Default)] -pub(crate) struct ZshExecBridgeSessionState { - pub(crate) initialized_session_id: Option, -} - -#[derive(Debug, Default)] -pub(crate) struct ZshExecBridge { - zsh_path: Option, - state: Mutex, -} - -#[cfg(unix)] -#[derive(Debug, Deserialize, Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum WrapperIpcRequest { - ExecRequest { - request_id: String, - file: String, - argv: Vec, - cwd: String, - }, -} - -#[cfg(unix)] -#[derive(Debug, Deserialize, Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum WrapperIpcResponse { - ExecResponse { - request_id: String, - action: WrapperExecAction, - reason: Option, - }, -} - -#[cfg(unix)] -#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -enum WrapperExecAction { - Run, - Deny, -} - -impl ZshExecBridge { - pub(crate) fn new(zsh_path: Option, _codex_home: PathBuf) -> Self { - Self { - zsh_path, - state: Mutex::new(ZshExecBridgeSessionState::default()), - } - } - - pub(crate) async fn initialize_for_session(&self, session_id: &str) { - let mut state = self.state.lock().await; - state.initialized_session_id = Some(session_id.to_string()); - } - - pub(crate) async fn shutdown(&self) { - let mut state = self.state.lock().await; - state.initialized_session_id = None; - } - - pub(crate) fn next_wrapper_socket_path(&self) -> PathBuf { - let socket_id = Uuid::new_v4().as_simple().to_string(); - let temp_dir = std::env::temp_dir(); - let canonical_temp_dir = temp_dir.canonicalize().unwrap_or(temp_dir); - canonical_temp_dir.join(format!("czs-{}.sock", &socket_id[..12])) - } - - #[cfg(not(unix))] - pub(crate) async fn execute_shell_request( - &self, - _req: &crate::sandboxing::ExecRequest, - _session: &crate::codex::Session, - _turn: &crate::codex::TurnContext, - _call_id: &str, - ) -> Result { - let _ = &self.zsh_path; - Err(ToolError::Rejected( - "shell_zsh_fork is only supported on unix".to_string(), - )) - } - - #[cfg(unix)] - pub(crate) async fn execute_shell_request( - &self, - req: &crate::sandboxing::ExecRequest, - session: &crate::codex::Session, - turn: &crate::codex::TurnContext, - call_id: &str, - ) -> Result { - let zsh_path = self.zsh_path.clone().ok_or_else(|| { - ToolError::Rejected( - "shell_zsh_fork enabled, but zsh_path is not configured".to_string(), - ) - })?; - - let command = req.command.clone(); - if command.is_empty() { - return Err(ToolError::Rejected("command args are empty".to_string())); - } - - let wrapper_socket_path = req - .env - .get(ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR) - .map(PathBuf::from) - .unwrap_or_else(|| self.next_wrapper_socket_path()); - - let listener = { - let _ = std::fs::remove_file(&wrapper_socket_path); - UnixListener::bind(&wrapper_socket_path).map_err(|err| { - ToolError::Rejected(format!( - "bind wrapper socket at {}: {err}", - wrapper_socket_path.display() - )) - })? - }; - - let wrapper_path = std::env::current_exe().map_err(|err| { - ToolError::Rejected(format!("resolve current executable path: {err}")) - })?; - - let mut cmd = tokio::process::Command::new(&command[0]); - #[cfg(unix)] - if let Some(arg0) = &req.arg0 { - cmd.arg0(arg0); - } - if command.len() > 1 { - cmd.args(&command[1..]); - } - cmd.current_dir(&req.cwd); - cmd.stdin(std::process::Stdio::null()); - cmd.stdout(std::process::Stdio::piped()); - cmd.stderr(std::process::Stdio::piped()); - cmd.kill_on_drop(true); - cmd.env_clear(); - cmd.envs(&req.env); - cmd.env( - ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR, - wrapper_socket_path.to_string_lossy().to_string(), - ); - cmd.env(EXEC_WRAPPER_ENV_VAR, &wrapper_path); - cmd.env(ZSH_EXEC_WRAPPER_MODE_ENV_VAR, "1"); - - let mut child = cmd.spawn().map_err(|err| { - ToolError::Rejected(format!( - "failed to start zsh fork command {} with zsh_path {}: {err}", - command[0], - zsh_path.display() - )) - })?; - - let (stream_tx, mut stream_rx) = - tokio::sync::mpsc::unbounded_channel::<(ExecOutputStream, Vec)>(); - - if let Some(mut out) = child.stdout.take() { - let tx = stream_tx.clone(); - tokio::spawn(async move { - let mut buf = [0_u8; 8192]; - loop { - let read = match out.read(&mut buf).await { - Ok(0) => break, - Ok(n) => n, - Err(err) => { - tracing::warn!("zsh fork stdout read error: {err}"); - break; - } - }; - let _ = tx.send((ExecOutputStream::Stdout, buf[..read].to_vec())); - } - }); - } - - if let Some(mut err) = child.stderr.take() { - let tx = stream_tx.clone(); - tokio::spawn(async move { - let mut buf = [0_u8; 8192]; - loop { - let read = match err.read(&mut buf).await { - Ok(0) => break, - Ok(n) => n, - Err(err) => { - tracing::warn!("zsh fork stderr read error: {err}"); - break; - } - }; - let _ = tx.send((ExecOutputStream::Stderr, buf[..read].to_vec())); - } - }); - } - drop(stream_tx); - - let mut stdout_bytes = Vec::new(); - let mut stderr_bytes = Vec::new(); - let mut child_exit = None; - let mut timed_out = false; - let mut stream_open = true; - let mut user_rejected = false; - let start = Instant::now(); - - let expiration = req.expiration.clone().wait(); - tokio::pin!(expiration); - - while child_exit.is_none() || stream_open { - tokio::select! { - result = child.wait(), if child_exit.is_none() => { - child_exit = Some(result.map_err(|err| ToolError::Rejected(format!("wait for zsh fork command exit: {err}")))?); - } - stream = stream_rx.recv(), if stream_open => { - if let Some((output_stream, chunk)) = stream { - match output_stream { - ExecOutputStream::Stdout => stdout_bytes.extend_from_slice(&chunk), - ExecOutputStream::Stderr => stderr_bytes.extend_from_slice(&chunk), - } - session - .send_event( - turn, - EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { - call_id: call_id.to_string(), - stream: output_stream, - chunk, - }), - ) - .await; - } else { - stream_open = false; - } - } - accept_result = listener.accept(), if child_exit.is_none() => { - let (stream, _) = accept_result.map_err(|err| { - ToolError::Rejected(format!("failed to accept wrapper request: {err}")) - })?; - if self - .handle_wrapper_request(stream, req.justification.clone(), session, turn, call_id) - .await? - { - user_rejected = true; - } - } - _ = &mut expiration, if child_exit.is_none() => { - timed_out = true; - kill_child_process_group(&mut child).map_err(|err| { - ToolError::Rejected(format!("kill zsh fork command process group: {err}")) - })?; - child.start_kill().map_err(|err| { - ToolError::Rejected(format!("kill zsh fork command process: {err}")) - })?; - } - } - } - - let _ = std::fs::remove_file(&wrapper_socket_path); - - let status = child_exit.ok_or_else(|| { - ToolError::Rejected("zsh fork command did not return exit status".to_string()) - })?; - - if user_rejected { - return Err(ToolError::Rejected("rejected by user".to_string())); - } - - let stdout_text = crate::text_encoding::bytes_to_string_smart(&stdout_bytes); - let stderr_text = crate::text_encoding::bytes_to_string_smart(&stderr_bytes); - let output = ExecToolCallOutput { - exit_code: status.code().unwrap_or(-1), - stdout: crate::exec::StreamOutput::new(stdout_text.clone()), - stderr: crate::exec::StreamOutput::new(stderr_text.clone()), - aggregated_output: crate::exec::StreamOutput::new(format!( - "{stdout_text}{stderr_text}" - )), - duration: start.elapsed(), - timed_out, - }; - - Self::map_exec_result(req.sandbox, output) - } - - #[cfg(unix)] - async fn handle_wrapper_request( - &self, - mut stream: UnixStream, - approval_reason: Option, - session: &crate::codex::Session, - turn: &crate::codex::TurnContext, - call_id: &str, - ) -> Result { - let mut request_buf = Vec::new(); - stream.read_to_end(&mut request_buf).await.map_err(|err| { - ToolError::Rejected(format!("read wrapper request from socket: {err}")) - })?; - let request_line = String::from_utf8(request_buf).map_err(|err| { - ToolError::Rejected(format!("decode wrapper request as utf-8: {err}")) - })?; - let request = parse_wrapper_request_line(request_line.trim())?; - - let (request_id, file, argv, cwd) = match request { - WrapperIpcRequest::ExecRequest { - request_id, - file, - argv, - cwd, - } => (request_id, file, argv, cwd), - }; - - let command_for_approval = if argv.is_empty() { - vec![file.clone()] - } else { - argv.clone() - }; - - let approval_id = Uuid::new_v4().to_string(); - let decision = session - .request_command_approval( - turn, - call_id.to_string(), - Some(approval_id), - command_for_approval, - PathBuf::from(cwd), - approval_reason, - None, - None::, - ) - .await; - - let (action, reason, user_rejected) = match decision { - ReviewDecision::Approved - | ReviewDecision::ApprovedForSession - | ReviewDecision::ApprovedExecpolicyAmendment { .. } => { - (WrapperExecAction::Run, None, false) - } - ReviewDecision::Denied => ( - WrapperExecAction::Deny, - Some("command denied by host approval policy".to_string()), - true, - ), - ReviewDecision::Abort => ( - WrapperExecAction::Deny, - Some("command aborted by host approval policy".to_string()), - true, - ), - }; - - write_json_line( - &mut stream, - &WrapperIpcResponse::ExecResponse { - request_id, - action, - reason, - }, - ) - .await?; - - Ok(user_rejected) - } - - #[cfg(unix)] - fn map_exec_result( - sandbox: crate::exec::SandboxType, - output: ExecToolCallOutput, - ) -> Result { - if output.timed_out { - return Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Timeout { - output: Box::new(output), - }))); - } - - if crate::exec::is_likely_sandbox_denied(sandbox, &output) { - return Err(ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { - output: Box::new(output), - network_policy_decision: None, - }))); - } - - Ok(output) - } -} - -pub fn maybe_run_zsh_exec_wrapper_mode() -> anyhow::Result { - if std::env::var_os(ZSH_EXEC_WRAPPER_MODE_ENV_VAR).is_none() { - return Ok(false); - } - - run_exec_wrapper_mode()?; - Ok(true) -} - -fn run_exec_wrapper_mode() -> anyhow::Result<()> { - #[cfg(not(unix))] - { - anyhow::bail!("zsh exec wrapper mode is only supported on unix"); - } - - #[cfg(unix)] - { - use std::os::unix::net::UnixStream as StdUnixStream; - - let args: Vec = std::env::args().collect(); - if args.len() < 2 { - anyhow::bail!("exec wrapper mode requires target executable path"); - } - let file = args[1].clone(); - let argv = if args.len() > 2 { - args[2..].to_vec() - } else { - vec![file.clone()] - }; - let cwd = std::env::current_dir() - .context("resolve wrapper cwd")? - .to_string_lossy() - .to_string(); - let socket_path = std::env::var(ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR) - .context("missing wrapper socket path env var")?; - - let request_id = Uuid::new_v4().to_string(); - let request = WrapperIpcRequest::ExecRequest { - request_id: request_id.clone(), - file: file.clone(), - argv: argv.clone(), - cwd, - }; - let mut stream = StdUnixStream::connect(&socket_path) - .with_context(|| format!("connect to wrapper socket at {socket_path}"))?; - let encoded = serde_json::to_string(&request).context("serialize wrapper request")?; - stream - .write_all(encoded.as_bytes()) - .context("write wrapper request")?; - stream - .write_all(b"\n") - .context("write wrapper request newline")?; - stream - .shutdown(std::net::Shutdown::Write) - .context("shutdown wrapper write")?; - - let mut response_buf = String::new(); - stream - .read_to_string(&mut response_buf) - .context("read wrapper response")?; - let response: WrapperIpcResponse = - serde_json::from_str(response_buf.trim()).context("parse wrapper response")?; - - let (response_request_id, action, reason) = match response { - WrapperIpcResponse::ExecResponse { - request_id, - action, - reason, - } => (request_id, action, reason), - }; - if response_request_id != request_id { - anyhow::bail!( - "wrapper response request_id mismatch: expected {request_id}, got {response_request_id}" - ); - } - - if action == WrapperExecAction::Deny { - if let Some(reason) = reason { - tracing::warn!("execution denied: {reason}"); - } else { - tracing::warn!("execution denied"); - } - std::process::exit(1); - } - - let mut command = std::process::Command::new(&file); - if argv.len() > 1 { - command.args(&argv[1..]); - } - command.env_remove(ZSH_EXEC_WRAPPER_MODE_ENV_VAR); - command.env_remove(ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR); - command.env_remove(EXEC_WRAPPER_ENV_VAR); - let status = command.status().context("spawn wrapped executable")?; - std::process::exit(status.code().unwrap_or(1)); - } -} - -#[cfg(unix)] -fn parse_wrapper_request_line(request_line: &str) -> Result { - serde_json::from_str(request_line) - .map_err(|err| ToolError::Rejected(format!("parse wrapper request payload: {err}"))) -} - -#[cfg(unix)] -async fn write_json_line( - writer: &mut W, - message: &T, -) -> Result<(), ToolError> { - let encoded = serde_json::to_string(message) - .map_err(|err| ToolError::Rejected(format!("serialize wrapper message: {err}")))?; - tokio::io::AsyncWriteExt::write_all(writer, encoded.as_bytes()) - .await - .map_err(|err| ToolError::Rejected(format!("write wrapper message: {err}")))?; - tokio::io::AsyncWriteExt::write_all(writer, b"\n") - .await - .map_err(|err| ToolError::Rejected(format!("write wrapper newline: {err}")))?; - tokio::io::AsyncWriteExt::flush(writer) - .await - .map_err(|err| ToolError::Rejected(format!("flush wrapper message: {err}")))?; - Ok(()) -} - -#[cfg(all(test, unix))] -mod tests { - use super::*; - - #[test] - fn parse_wrapper_request_line_rejects_malformed_json() { - let err = parse_wrapper_request_line("this-is-not-json").unwrap_err(); - let ToolError::Rejected(message) = err else { - panic!("expected ToolError::Rejected"); - }; - assert!(message.starts_with("parse wrapper request payload:")); - } -} diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index a6a0721b98..74b252d97b 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -32,6 +32,8 @@ codex-core = { workspace = true } codex-execpolicy = { workspace = true } codex-protocol = { workspace = true } codex-shell-command = { workspace = true } + +[target.'cfg(unix)'.dependencies] codex-shell-escalation = { workspace = true } rmcp = { workspace = true, default-features = false, features = [ "auth", @@ -51,6 +53,7 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } shlex = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal"] } +tokio-util = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] } diff --git a/codex-rs/exec-server/src/unix.rs b/codex-rs/exec-server/src/unix.rs index 0e0355443b..2bd892fd64 100644 --- a/codex-rs/exec-server/src/unix.rs +++ b/codex-rs/exec-server/src/unix.rs @@ -67,7 +67,7 @@ use codex_execpolicy::Decision; use codex_execpolicy::Policy; use codex_execpolicy::RuleMatch; use codex_shell_command::is_dangerous_command::command_might_be_dangerous; -use codex_shell_escalation as shell_escalation; +use codex_shell_escalation::unix::escalate_client::run; use rmcp::ErrorData as McpError; use tokio::sync::RwLock; use tracing_subscriber::EnvFilter; @@ -160,7 +160,7 @@ pub async fn main_execve_wrapper() -> anyhow::Result<()> { .init(); let ExecveWrapperCli { file, argv } = ExecveWrapperCli::parse(); - let exit_code = shell_escalation::run(file, argv).await?; + let exit_code = run(file, argv).await?; std::process::exit(exit_code); } diff --git a/codex-rs/exec-server/src/unix/mcp.rs b/codex-rs/exec-server/src/unix/mcp.rs index 547d055c14..30dbc4d819 100644 --- a/codex-rs/exec-server/src/unix/mcp.rs +++ b/codex-rs/exec-server/src/unix/mcp.rs @@ -6,11 +6,19 @@ use anyhow::Context as _; use anyhow::Result; use codex_core::MCP_SANDBOX_STATE_CAPABILITY; use codex_core::MCP_SANDBOX_STATE_METHOD; -use codex_core::SandboxState; +use codex_core::SandboxState as CoreSandboxState; +use codex_core::exec::process_exec_tool_call; use codex_execpolicy::Policy; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::models::SandboxPermissions as ProtocolSandboxPermissions; use codex_protocol::protocol::SandboxPolicy; -use codex_shell_escalation::EscalationPolicyFactory; -use codex_shell_escalation::run_escalate_server; +use codex_shell_escalation::unix::escalate_server::EscalationPolicyFactory; +use codex_shell_escalation::unix::escalate_server::ExecParams as ShellExecParams; +use codex_shell_escalation::unix::escalate_server::ExecResult as ShellExecResult; +use codex_shell_escalation::unix::escalate_server::SandboxState as ShellEscalationSandboxState; +use codex_shell_escalation::unix::escalate_server::ShellCommandExecutor; +use codex_shell_escalation::unix::escalate_server::run_escalate_server; +use codex_shell_escalation::unix::stopwatch::Stopwatch; use rmcp::ErrorData as McpError; use rmcp::RoleServer; use rmcp::ServerHandler; @@ -27,7 +35,9 @@ use rmcp::tool_handler; use rmcp::tool_router; use rmcp::transport::stdio; use serde_json::json; +use std::collections::HashMap; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; use crate::unix::mcp_escalation_policy::McpEscalationPolicy; @@ -50,8 +60,8 @@ pub struct ExecResult { pub timed_out: bool, } -impl From for ExecResult { - fn from(result: codex_shell_escalation::ExecResult) -> Self { +impl From for ExecResult { + fn from(result: ShellExecResult) -> Self { Self { exit_code: result.exit_code, output: result.output, @@ -68,7 +78,7 @@ pub struct ExecTool { execve_wrapper: PathBuf, policy: Arc>, preserve_program_paths: bool, - sandbox_state: Arc>>, + sandbox_state: Arc>>, } #[derive(Debug, serde::Serialize, serde::Deserialize, rmcp::schemars::JsonSchema)] @@ -83,7 +93,7 @@ pub struct ExecParams { pub login: Option, } -impl From for codex_shell_escalation::ExecParams { +impl From for ShellExecParams { fn from(inner: ExecParams) -> Self { Self { command: inner.command, @@ -99,14 +109,51 @@ struct McpEscalationPolicyFactory { preserve_program_paths: bool, } +struct McpShellCommandExecutor; + +#[async_trait::async_trait] +impl ShellCommandExecutor for McpShellCommandExecutor { + async fn run( + &self, + command: Vec, + cwd: PathBuf, + env: HashMap, + cancel_rx: CancellationToken, + sandbox_state: &ShellEscalationSandboxState, + ) -> anyhow::Result { + let result = process_exec_tool_call( + codex_core::exec::ExecParams { + command, + cwd, + expiration: codex_core::exec::ExecExpiration::Cancellation(cancel_rx), + env, + network: None, + sandbox_permissions: ProtocolSandboxPermissions::UseDefault, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + justification: None, + arg0: None, + }, + &sandbox_state.sandbox_policy, + &sandbox_state.sandbox_cwd, + &sandbox_state.codex_linux_sandbox_exe, + sandbox_state.use_linux_sandbox_bwrap, + None, + ) + .await?; + + Ok(ShellExecResult { + exit_code: result.exit_code, + output: result.aggregated_output.text, + duration: result.duration, + timed_out: result.timed_out, + }) + } +} + impl EscalationPolicyFactory for McpEscalationPolicyFactory { type Policy = McpEscalationPolicy; - fn create_policy( - &self, - policy: Arc>, - stopwatch: codex_shell_escalation::Stopwatch, - ) -> Self::Policy { + fn create_policy(&self, policy: Arc>, stopwatch: Stopwatch) -> Self::Policy { McpEscalationPolicy::new( policy, self.context.clone(), @@ -151,15 +198,21 @@ impl ExecTool { .read() .await .clone() - .unwrap_or_else(|| SandboxState { + .unwrap_or_else(|| CoreSandboxState { sandbox_policy: SandboxPolicy::new_read_only_policy(), codex_linux_sandbox_exe: None, sandbox_cwd: PathBuf::from(¶ms.workdir), use_linux_sandbox_bwrap: false, }); + let shell_sandbox_state = ShellEscalationSandboxState { + sandbox_policy: sandbox_state.sandbox_policy.clone(), + codex_linux_sandbox_exe: sandbox_state.codex_linux_sandbox_exe.clone(), + sandbox_cwd: sandbox_state.sandbox_cwd.clone(), + use_linux_sandbox_bwrap: sandbox_state.use_linux_sandbox_bwrap, + }; let result = run_escalate_server( params.into(), - &sandbox_state, + &shell_sandbox_state, &self.bash_path, &self.execve_wrapper, self.policy.clone(), @@ -168,6 +221,7 @@ impl ExecTool { preserve_program_paths: self.preserve_program_paths, }, effective_timeout, + &McpShellCommandExecutor, ) .await .map_err(|e| McpError::internal_error(e.to_string(), None))?; @@ -236,7 +290,7 @@ impl ServerHandler for ExecTool { )); }; - let Ok(sandbox_state) = serde_json::from_value::(params.clone()) else { + let Ok(sandbox_state) = serde_json::from_value::(params.clone()) else { return Err(McpError::invalid_params( "failed to deserialize sandbox state".to_string(), Some(params), diff --git a/codex-rs/exec-server/src/unix/mcp_escalation_policy.rs b/codex-rs/exec-server/src/unix/mcp_escalation_policy.rs index 9863818261..73b2c34e42 100644 --- a/codex-rs/exec-server/src/unix/mcp_escalation_policy.rs +++ b/codex-rs/exec-server/src/unix/mcp_escalation_policy.rs @@ -2,9 +2,9 @@ use std::path::Path; use codex_core::sandboxing::SandboxPermissions; use codex_execpolicy::Policy; -use codex_shell_escalation::EscalateAction; -use codex_shell_escalation::EscalationPolicy; -use codex_shell_escalation::Stopwatch; +use codex_shell_escalation::unix::escalate_protocol::EscalateAction; +use codex_shell_escalation::unix::escalation_policy::EscalationPolicy; +use codex_shell_escalation::unix::stopwatch::Stopwatch; use rmcp::ErrorData as McpError; use rmcp::RoleServer; use rmcp::model::CreateElicitationRequestParams; diff --git a/codex-rs/shell-escalation/Cargo.toml b/codex-rs/shell-escalation/Cargo.toml index 49b8f10bd3..3cdd50250a 100644 --- a/codex-rs/shell-escalation/Cargo.toml +++ b/codex-rs/shell-escalation/Cargo.toml @@ -7,20 +7,21 @@ license.workspace = true [dependencies] anyhow = { workspace = true } async-trait = { workspace = true } -codex-core = { workspace = true } codex-execpolicy = { workspace = true } codex-protocol = { workspace = true } libc = { workspace = true } serde_json = { workspace = true } path-absolutize = { workspace = true } serde = { workspace = true, features = ["derive"] } -socket2 = { workspace = true } +socket2 = { workspace = true, features = ["all"] } tokio = { workspace = true, features = [ "io-std", + "net", "macros", "process", "rt-multi-thread", "signal", + "time", ] } tokio-util = { workspace = true } tracing = { workspace = true } diff --git a/codex-rs/shell-escalation/src/lib.rs b/codex-rs/shell-escalation/src/lib.rs index 555d0f89e0..fae4ef574b 100644 --- a/codex-rs/shell-escalation/src/lib.rs +++ b/codex-rs/shell-escalation/src/lib.rs @@ -1,21 +1,5 @@ #[cfg(unix)] -mod unix { - mod escalate_client; - mod escalate_protocol; - mod escalate_server; - mod escalation_policy; - mod socket; - mod stopwatch; - - pub use self::escalate_client::run; - pub use self::escalate_protocol::EscalateAction; - pub use self::escalate_server::EscalationPolicyFactory; - pub use self::escalate_server::ExecParams; - pub use self::escalate_server::ExecResult; - pub use self::escalate_server::run_escalate_server; - pub use self::escalation_policy::EscalationPolicy; - pub use self::stopwatch::Stopwatch; -} +pub mod unix; #[cfg(unix)] pub use unix::*; diff --git a/codex-rs/shell-escalation/src/unix/core_shell_escalation.rs b/codex-rs/shell-escalation/src/unix/core_shell_escalation.rs index 0be7af28fa..ca6bd347da 100644 --- a/codex-rs/shell-escalation/src/unix/core_shell_escalation.rs +++ b/codex-rs/shell-escalation/src/unix/core_shell_escalation.rs @@ -40,7 +40,7 @@ impl ShellPolicyFactory { } } -struct ShellEscalationPolicy { +pub struct ShellEscalationPolicy { provider: Arc, stopwatch: Stopwatch, } diff --git a/codex-rs/shell-escalation/src/unix/escalate_server.rs b/codex-rs/shell-escalation/src/unix/escalate_server.rs index 0ee5fc27c4..d437795bab 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_server.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_server.rs @@ -7,8 +7,8 @@ use std::sync::Arc; use std::time::Duration; use anyhow::Context as _; -use codex_core::SandboxState; use codex_execpolicy::Policy; +use codex_protocol::protocol::SandboxPolicy; use path_absolutize::Absolutize as _; use tokio::process::Command; use tokio::sync::RwLock; @@ -27,6 +27,26 @@ use crate::unix::socket::AsyncDatagramSocket; use crate::unix::socket::AsyncSocket; use crate::unix::stopwatch::Stopwatch; +#[derive(Debug, Clone)] +pub struct SandboxState { + pub sandbox_policy: SandboxPolicy, + pub codex_linux_sandbox_exe: Option, + pub sandbox_cwd: PathBuf, + pub use_linux_sandbox_bwrap: bool, +} + +#[async_trait::async_trait] +pub trait ShellCommandExecutor: Send + Sync { + async fn run( + &self, + command: Vec, + cwd: PathBuf, + env: HashMap, + cancel_rx: CancellationToken, + sandbox_state: &SandboxState, + ) -> anyhow::Result; +} + #[derive(Debug, serde::Deserialize, serde::Serialize)] pub struct ExecParams { /// The bash string to execute. @@ -71,11 +91,10 @@ impl EscalateServer { params: ExecParams, cancel_rx: CancellationToken, sandbox_state: &SandboxState, + command_executor: &dyn ShellCommandExecutor, ) -> anyhow::Result { let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?; let client_socket = escalate_client.into_inner(); - client_socket.set_cloexec(false)?; - let escalate_task = tokio::spawn(escalate_task(escalate_server, self.policy.clone())); let mut env = std::env::vars().collect::>(); env.insert( @@ -91,47 +110,27 @@ impl EscalateServer { self.execve_wrapper.to_string_lossy().to_string(), ); - let ExecParams { - command, - workdir, - timeout_ms: _, - login, - } = params; - let result = codex_core::exec::process_exec_tool_call( - codex_core::exec::ExecParams { - command: vec![ - self.bash_path.to_string_lossy().to_string(), - if login == Some(false) { - "-c".to_string() - } else { - "-lc".to_string() - }, - command, - ], - cwd: PathBuf::from(&workdir), - expiration: codex_core::exec::ExecExpiration::Cancellation(cancel_rx), - env, - network: None, - sandbox_permissions: codex_core::sandboxing::SandboxPermissions::UseDefault, - windows_sandbox_level: codex_protocol::config_types::WindowsSandboxLevel::Disabled, - justification: None, - arg0: None, + let command = vec![ + self.bash_path.to_string_lossy().to_string(), + if params.login == Some(false) { + "-c".to_string() + } else { + "-lc".to_string() }, - &sandbox_state.sandbox_policy, - &sandbox_state.sandbox_cwd, - &sandbox_state.codex_linux_sandbox_exe, - sandbox_state.use_linux_sandbox_bwrap, - None, - ) - .await?; + params.command, + ]; + let result = command_executor + .run( + command, + PathBuf::from(¶ms.workdir), + env, + cancel_rx, + sandbox_state, + ) + .await?; escalate_task.abort(); - Ok(ExecResult { - exit_code: result.exit_code, - output: result.aggregated_output.text, - duration: result.duration, - timed_out: result.timed_out, - }) + Ok(result) } } @@ -142,6 +141,7 @@ pub trait EscalationPolicyFactory { fn create_policy(&self, policy: Arc>, stopwatch: Stopwatch) -> Self::Policy; } +#[allow(clippy::too_many_arguments)] pub async fn run_escalate_server( exec_params: ExecParams, sandbox_state: &SandboxState, @@ -150,6 +150,7 @@ pub async fn run_escalate_server( policy: Arc>, escalation_policy_factory: impl EscalationPolicyFactory, effective_timeout: Duration, + command_executor: &dyn ShellCommandExecutor, ) -> anyhow::Result { let stopwatch = Stopwatch::new(effective_timeout); let cancel_token = stopwatch.cancellation_token(); @@ -160,7 +161,7 @@ pub async fn run_escalate_server( ); escalate_server - .exec(exec_params, cancel_token, sandbox_state) + .exec(exec_params, cancel_token, sandbox_state, command_executor) .await } @@ -272,6 +273,7 @@ async fn handle_escalate_session_with_policy( .await?; } } + Ok(()) } @@ -279,7 +281,6 @@ async fn handle_escalate_session_with_policy( mod tests { use super::*; use pretty_assertions::assert_eq; - use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; diff --git a/codex-rs/shell-escalation/src/unix/mod.rs b/codex-rs/shell-escalation/src/unix/mod.rs index 0ae7941da2..555bd7246e 100644 --- a/codex-rs/shell-escalation/src/unix/mod.rs +++ b/codex-rs/shell-escalation/src/unix/mod.rs @@ -1,7 +1,7 @@ +pub mod core_shell_escalation; pub mod escalate_client; pub mod escalate_protocol; pub mod escalate_server; pub mod escalation_policy; pub mod socket; -pub mod core_shell_escalation; pub mod stopwatch; diff --git a/codex-rs/shell-escalation/src/unix/socket.rs b/codex-rs/shell-escalation/src/unix/socket.rs index 35292367a6..fd3fd7cb0a 100644 --- a/codex-rs/shell-escalation/src/unix/socket.rs +++ b/codex-rs/shell-escalation/src/unix/socket.rs @@ -96,8 +96,8 @@ async fn read_frame_header( while filled < LENGTH_PREFIX_SIZE { let mut guard = async_socket.readable().await?; // The first read should come with a control message containing any FDs. - let result = if !captured_control { - guard.try_io(|inner| { + let read = if !captured_control { + match guard.try_io(|inner| { let mut bufs = [MaybeUninitSlice::new(&mut header[filled..])]; let (read, control_len) = { let mut msg = MsgHdrMut::new() @@ -109,16 +109,18 @@ async fn read_frame_header( control.truncate(control_len); captured_control = true; Ok(read) - }) + }) { + Ok(Ok(read)) => read, + Ok(Err(err)) => return Err(err), + Err(_would_block) => continue, + } } else { - guard.try_io(|inner| inner.get_ref().recv(&mut header[filled..])) + match guard.try_io(|inner| inner.get_ref().recv(&mut header[filled..])) { + Ok(Ok(read)) => read, + Ok(Err(err)) => return Err(err), + Err(_would_block) => continue, + } }; - let Ok(result) = result else { - // Would block, try again. - continue; - }; - - let read = result?; if read == 0 { return Err(std::io::Error::new( std::io::ErrorKind::UnexpectedEof, @@ -150,12 +152,11 @@ async fn read_frame_payload( let mut filled = 0; while filled < message_len { let mut guard = async_socket.readable().await?; - let result = guard.try_io(|inner| inner.get_ref().recv(&mut payload[filled..])); - let Ok(result) = result else { - // Would block, try again. - continue; + let read = match guard.try_io(|inner| inner.get_ref().recv(&mut payload[filled..])) { + Ok(Ok(read)) => read, + Ok(Err(err)) => return Err(err), + Err(_would_block) => continue, }; - let read = result?; if read == 0 { return Err(std::io::Error::new( std::io::ErrorKind::UnexpectedEof, @@ -261,7 +262,7 @@ impl AsyncSocket { } pub fn pair() -> std::io::Result<(AsyncSocket, AsyncSocket)> { - let (server, client) = Socket::pair(Domain::UNIX, Type::STREAM, None)?; + let (server, client) = Socket::pair_raw(Domain::UNIX, Type::STREAM, None)?; Ok((AsyncSocket::new(server)?, AsyncSocket::new(client)?)) } @@ -314,11 +315,11 @@ async fn send_stream_frame( let mut include_fds = !fds.is_empty(); while written < frame.len() { let mut guard = socket.writable().await?; - let result = guard.try_io(|inner| { - send_stream_chunk(inner.get_ref(), &frame[written..], fds, include_fds) - }); - let bytes_written = match result { - Ok(bytes_written) => bytes_written?, + let bytes_written = match guard + .try_io(|inner| send_stream_chunk(inner.get_ref(), &frame[written..], fds, include_fds)) + { + Ok(Ok(bytes_written)) => bytes_written, + Ok(Err(err)) => return Err(err), Err(_would_block) => continue, }; if bytes_written == 0 { @@ -370,7 +371,7 @@ impl AsyncDatagramSocket { } pub fn pair() -> std::io::Result<(Self, Self)> { - let (server, client) = Socket::pair(Domain::UNIX, Type::DGRAM, None)?; + let (server, client) = Socket::pair_raw(Domain::UNIX, Type::DGRAM, None)?; Ok((Self::new(server)?, Self::new(client)?)) } @@ -472,7 +473,7 @@ mod tests { #[test] fn send_datagram_bytes_rejects_excessive_fd_counts() -> std::io::Result<()> { - let (socket, _peer) = Socket::pair(Domain::UNIX, Type::DGRAM, None)?; + let (socket, _peer) = Socket::pair_raw(Domain::UNIX, Type::DGRAM, None)?; let fds = fd_list(MAX_FDS_PER_MESSAGE + 1)?; let err = send_datagram_bytes(&socket, b"hi", &fds).unwrap_err(); assert_eq!(std::io::ErrorKind::InvalidInput, err.kind()); @@ -481,7 +482,7 @@ mod tests { #[test] fn send_stream_chunk_rejects_excessive_fd_counts() -> std::io::Result<()> { - let (socket, _peer) = Socket::pair(Domain::UNIX, Type::STREAM, None)?; + let (socket, _peer) = Socket::pair_raw(Domain::UNIX, Type::STREAM, None)?; let fds = fd_list(MAX_FDS_PER_MESSAGE + 1)?; let err = send_stream_chunk(&socket, b"hello", &fds, true).unwrap_err(); assert_eq!(std::io::ErrorKind::InvalidInput, err.kind());