From 2e220af1f63e0dc8cd0282a283de9ec2ce5dfeb9 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 8 Sep 2026 20:45:42 +0000 Subject: [PATCH] Protect shell snapshots when credential brokerage is enabled (#43909) ## Why Shell startup files can introduce credentials or overwrite brokered dummy values. Snapshot capture and replay need to respect the command's sandbox and environment policy without persisting real credentials. ## What changed - Capture and validate brokered snapshots per command using its requested shell, working directory, login mode, sandbox, and environment policy. - Discover shell-initialized credentials and prepare snapshot exports for brokered replay. Preserve shell functions and aliases while preventing startup files from overwriting dummy credentials during replay. - Rebuild snapshots as credential broker state changes and retain each command's snapshot for the lifetime of its process. - Reject brokered POSIX shell execution when a protected snapshot cannot be created, and reject brokerage with `shell_zsh_fork`. Restore real credentials and remove brokerage markers for unsupported PowerShell and Cmd execution; escalated commands bypass protected snapshots. ## Testing Add regression coverage for credential discovery and redaction, environment filtering, readonly exports, shell function and heredoc replay, command working directories and login modes, sandbox enforcement, and missing or disabled snapshots. GitOrigin-RevId: 685356cb9e80f11cf69d053e4d239ee6ccf43a9e --- .../core/src/config/network_proxy_spec.rs | 4 + codex-rs/core/src/environment_selection.rs | 109 +++- codex-rs/core/src/session/mod.rs | 29 ++ codex-rs/core/src/session/session.rs | 55 +- codex-rs/core/src/session/tests.rs | 16 +- codex-rs/core/src/session/turn_context.rs | 49 +- codex-rs/core/src/shell_snapshot.rs | 424 ++++++++++++++-- codex-rs/core/src/shell_snapshot_sandbox.rs | 184 +++++++ codex-rs/core/src/shell_snapshot_tests.rs | 474 +++++++++++++++++- codex-rs/core/src/tasks/user_shell.rs | 11 +- codex-rs/core/src/tools/runtimes/mod.rs | 76 ++- codex-rs/core/src/tools/runtimes/mod_tests.rs | 145 +++++- .../core/src/tools/runtimes/unified_exec.rs | 142 +++++- codex-rs/core/src/unified_exec/process.rs | 4 + ...stdin_approval__tests__internal_grant.snap | 2 +- .../core/src/unified_exec/stdin_approval.rs | 2 +- .../src/unified_exec/stdin_approval_tests.rs | 33 +- codex-rs/core/tests/suite/approvals.rs | 374 ++++++++++++-- codex-rs/core/tests/suite/shell_snapshot.rs | 23 +- .../network-proxy/src/credential_broker.rs | 10 + .../src/credential_broker_tests.rs | 22 + codex-rs/network-proxy/src/proxy.rs | 11 + codex-rs/network-proxy/src/runtime.rs | 9 + .../src/shell_snapshot_render.rs | 3 + .../shell-command/src/shell_snapshot_tests.rs | 2 +- 25 files changed, 2080 insertions(+), 133 deletions(-) create mode 100644 codex-rs/core/src/shell_snapshot_sandbox.rs diff --git a/codex-rs/core/src/config/network_proxy_spec.rs b/codex-rs/core/src/config/network_proxy_spec.rs index c2235d6448..f802e402da 100644 --- a/codex-rs/core/src/config/network_proxy_spec.rs +++ b/codex-rs/core/src/config/network_proxy_spec.rs @@ -80,6 +80,10 @@ impl NetworkProxySpec { self.config.enabled } + pub(crate) fn credential_broker_enabled(&self) -> bool { + self.config.credential_broker && self.constraints.enabled != Some(false) + } + pub fn proxy_host_and_port(&self) -> String { host_and_port_from_network_addr(&self.config.proxy_url, /*default_port*/ 3128) } diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index aa473f3474..fd57ee6b3d 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -15,6 +15,7 @@ use codex_exec_server::ExecutorFileSystem; use codex_exec_server::SelectedCapabilityRootsStatus; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::config_types::ShellEnvironmentPolicy; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EnvironmentConfig; @@ -37,6 +38,7 @@ use crate::session::turn_context::ShellSnapshotTask; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; use crate::shell_snapshot::ShellSnapshot; +use crate::shell_snapshot::SnapshotCredentialBrokerState; /// Records whether a normalized config should follow later thread setting updates. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -142,6 +144,7 @@ struct ResolvedEnvironment { executor_platform_os: Option, temporary_directories: Option>, shell_snapshot: ShellSnapshotTask, + shell_snapshot_builder: ShellSnapshot, shell_snapshot_v2_supported: bool, installed_config: Option, } @@ -256,6 +259,13 @@ impl ThreadEnvironments { let selection = environment.selection; let config_origin = environment.config_origin; let selected_environment = Arc::clone(&environment.environment); + let inherited_snapshot = if !selected_environment.is_remote() + && shell_snapshot.should_rebuild_inherited() + { + futures::future::ready(None).boxed().shared() + } else { + environment.shell_snapshot + }; let resolution: TurnEnvironmentResolution = futures::future::ready(Ok(ResolvedEnvironment { environment: environment.environment, @@ -263,8 +273,11 @@ impl ThreadEnvironments { user_home_dir: environment.user_home_dir, executor_platform_os: environment.executor_platform_os, temporary_directories: environment.temporary_directories, - shell_snapshot: environment.shell_snapshot, - shell_snapshot_v2_supported: environment.shell_snapshot_v2_supported, + shell_snapshot: inherited_snapshot, + shell_snapshot_builder: shell_snapshot.clone(), + shell_snapshot_v2_supported: environment.shell_snapshot_v2_supported + && (selected_environment.is_remote() + || !shell_snapshot.should_rebuild_inherited()), installed_config: None, })) .boxed() @@ -292,6 +305,36 @@ impl ThreadEnvironments { } } + fn start_shell_snapshot_task( + shell_snapshot: ShellSnapshot, + environment: Arc, + cwd: PathUri, + shell: Option, + ) -> ShellSnapshotTask { + // Protected snapshots are captured by the command, using its sandbox. + if shell_snapshot.should_rebuild_inherited() { + return futures::future::ready(None).boxed().shared(); + } + let shell_snapshot = shell_snapshot + .build( + environment, + cwd, + shell, + /*allow_login_shell*/ true, + ShellEnvironmentPolicy::default(), + /*sandbox*/ None, + ) + .boxed() + .shared(); + drop(tokio::spawn( + shell_snapshot + .clone() + .in_current_span() + .with_current_subscriber(), + )); + shell_snapshot + } + pub(crate) fn update_selections( &self, environments: &[TurnEnvironmentSelection], @@ -482,6 +525,47 @@ impl ThreadEnvironments { self.environments.store(Arc::new(environments)); } + pub(crate) fn set_snapshot_credential_broker(&self, state: SnapshotCredentialBrokerState) { + if !self.shell_snapshot.set_credential_broker(state) { + return; + } + + let mut environments = Vec::clone(&self.environments.load()); + let mut changed = false; + for selected in &mut environments { + if !selected.environment.is_remote() + && let Some(Ok(resolved)) = selected.resolution.clone().now_or_never() + { + self.restart_shell_snapshot(selected, resolved); + changed = true; + } + } + if changed { + self.environments.store(Arc::new(environments)); + } + } + + fn restart_shell_snapshot( + &self, + selected: &mut SelectedTurnEnvironment, + resolved: ResolvedEnvironment, + ) { + let shell_snapshot = Self::start_shell_snapshot_task( + self.shell_snapshot.clone(), + Arc::clone(&resolved.environment), + selected.selection.cwd.clone(), + resolved.shell.clone(), + ); + selected.resolution = futures::future::ready(Ok(ResolvedEnvironment { + shell_snapshot, + shell_snapshot_v2_supported: cfg!(unix) + && !self.shell_snapshot.should_rebuild_inherited(), + ..resolved + })) + .boxed() + .shared(); + } + /// Combines persisted thread roots with installed attachment roots, keeping /// thread roots first and hiding attachments that are not ready yet. pub(crate) fn inspect_selected_capability_roots( @@ -667,13 +751,15 @@ impl ThreadEnvironments { cfg!(unix), ) }; - let task = shell_snapshot - .build(Arc::clone(&environment), selection.cwd, shell.clone()) - .boxed() - .shared(); - drop(tokio::spawn( - task.clone().in_current_span().with_current_subscriber(), - )); + let shell_snapshot_builder = shell_snapshot.clone(); + let task = Self::start_shell_snapshot_task( + shell_snapshot, + Arc::clone(&environment), + selection.cwd, + shell.clone(), + ); + let shell_snapshot_v2_supported = snapshot_v2 + && (environment.is_remote() || !shell_snapshot_builder.should_rebuild_inherited()); Ok(ResolvedEnvironment { environment, shell, @@ -681,7 +767,8 @@ impl ThreadEnvironments { executor_platform_os, temporary_directories: temporary_dirs, shell_snapshot: task, - shell_snapshot_v2_supported: snapshot_v2, + shell_snapshot_builder, + shell_snapshot_v2_supported, installed_config, }) } @@ -762,6 +849,8 @@ impl TurnEnvironmentState { ); turn_environment.executor_platform_os = environment.executor_platform_os; turn_environment.shell_snapshot = environment.shell_snapshot; + turn_environment.shell_snapshot_builder = + Some(Box::new(environment.shell_snapshot_builder)); turn_environment.shell_snapshot_v2_supported = environment.shell_snapshot_v2_supported; turn_environment.user_home_dir = environment.user_home_dir; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index bce3d7a512..f69af3f038 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -46,6 +46,7 @@ use crate::session::step_settings::ResolvedStepSettings; use crate::session::step_settings::StepSettings; use crate::session::turn_context::TurnEnvironment; use crate::session_prefix::format_inter_agent_completion_message; +use crate::shell_snapshot::SnapshotCredentialBrokerState; use crate::skills_load_input_from_config; use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context; use crate::turn_metadata::TurnMetadataState; @@ -1165,6 +1166,9 @@ impl Session { .cloned() else { self.services.network_proxy.store(None); + self.services + .turn_environments + .set_snapshot_credential_broker(SnapshotCredentialBrokerState::Inactive); return; }; @@ -1191,11 +1195,22 @@ impl Session { // listeners and must not be exposed as active managed proxy runtimes. if !spec.enabled() { self.services.network_proxy.store(None); + self.services + .turn_environments + .set_snapshot_credential_broker(SnapshotCredentialBrokerState::Inactive); return; } if let Some(started_proxy) = self.services.network_proxy.load_full() { if let Err(err) = spec.apply_to_started_proxy(started_proxy.as_ref()).await { warn!("failed to refresh managed network proxy for sandbox change: {err}"); + } else { + self.services + .turn_environments + .set_snapshot_credential_broker(if spec.credential_broker_enabled() { + SnapshotCredentialBrokerState::Ready(started_proxy.proxy()) + } else { + SnapshotCredentialBrokerState::Inactive + }); } return; } @@ -1214,11 +1229,25 @@ impl Session { .await { Ok((started_proxy, _session_network_proxy)) => { + if spec.credential_broker_enabled() { + self.services + .turn_environments + .set_snapshot_credential_broker(SnapshotCredentialBrokerState::Ready( + started_proxy.proxy(), + )); + } self.services .network_proxy .store(Some(Arc::new(started_proxy))); } Err(err) => { + self.services + .turn_environments + .set_snapshot_credential_broker(if spec.credential_broker_enabled() { + SnapshotCredentialBrokerState::Unavailable + } else { + SnapshotCredentialBrokerState::Inactive + }); warn!("failed to start managed network proxy for sandbox change: {err}"); } } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index bb50e2f558..ae95d17a37 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -14,6 +14,7 @@ use crate::hook_mcp_executor::CoreHookMcpExecutor; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; use crate::shell_snapshot::ShellSnapshot; +use crate::shell_snapshot::SnapshotCredentialBrokerState; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; use codex_http_client::ClientRouteClass; @@ -1091,13 +1092,10 @@ impl Session { }), }); } + let effective_config = config.config_layer_stack.effective_config(); let config_path = config.codex_home.join(CONFIG_TOML_FILE); if let Some(event) = unstable_features_warning_event( - config - .config_layer_stack - .effective_config() - .get("features") - .and_then(TomlValue::as_table), + effective_config.get("features").and_then(TomlValue::as_table), config.suppress_unstable_features_warning, &config.features, &config_path.display().to_string(), @@ -1217,7 +1215,27 @@ impl Session { } else { shell::default_user_shell() }; - let use_executor_shell_snapshots = config.features.enabled(Feature::ShellSnapshotV2) + let credential_broker_available = config.features.enabled(Feature::NetworkProxy) + && config + .config_layer_stack + .requirements() + .network + .as_ref() + .is_none_or(|network| network.value.enabled != Some(false)); + let credential_broker_configured = credential_broker_available + && effective_config + .get("features") + .and_then(|features| features.get("network_proxy")) + .and_then(|network_proxy| network_proxy.get("credential_broker")) + .and_then(TomlValue::as_bool) + .unwrap_or(false); + let credential_broker_active = credential_broker_configured + && config + .permissions + .network + .as_ref() + .is_some_and(crate::config::NetworkProxySpec::credential_broker_enabled); + let prefer_executor_shell_snapshots = config.features.enabled(Feature::ShellSnapshotV2) && config.features.enabled(Feature::ShellTool) && config.features.enabled(Feature::UnifiedExec) && matches!( @@ -1229,14 +1247,26 @@ impl Session { ), codex_tools::UnifiedExecShellMode::Direct ); + let use_executor_shell_snapshots = + prefer_executor_shell_snapshots && !credential_broker_active; let shell_snapshot = if config.features.enabled(Feature::ShellSnapshot) - && !use_executor_shell_snapshots + && (!use_executor_shell_snapshots || credential_broker_available) { + let snapshot_credential_broker = credential_broker_available.then(|| { + let state = if credential_broker_active { + SnapshotCredentialBrokerState::Starting + } else { + SnapshotCredentialBrokerState::Inactive + }; + watch::channel(state).0 + }); ShellSnapshot::new( config.codex_home.clone(), thread_id, session_telemetry.clone(), state_db_ctx.clone(), + snapshot_credential_broker, + prefer_executor_shell_snapshots, ) } else { ShellSnapshot::disabled() @@ -1355,6 +1385,17 @@ impl Session { } else { (None, None) }; + if let Some(network_proxy) = network_proxy.as_ref() + && config + .permissions + .network + .as_ref() + .is_some_and(crate::config::NetworkProxySpec::credential_broker_enabled) + { + turn_environments.set_snapshot_credential_broker( + SnapshotCredentialBrokerState::Ready(network_proxy.proxy()), + ); + } // Hooks and extensions share one stable thread-owned MCP runtime handle. let mcp_runtime = Arc::new(McpRuntime::empty( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 24ab858613..e262983e0e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5644,8 +5644,12 @@ async fn session_configuration_apply_permission_profile_accepts_direct_write_roo ); } +#[test_case::test_case(false; "ordinary_proxy")] +#[test_case::test_case(true; "credential_broker")] #[tokio::test] -async fn active_profile_update_rebuilds_network_proxy_config() -> std::io::Result<()> { +async fn active_profile_update_rebuilds_network_proxy_config( + credential_broker: bool, +) -> std::io::Result<()> { let codex_home = tempfile::tempdir().expect("create codex home"); let cwd = tempfile::tempdir().expect("create cwd"); let permissions = PermissionsToml { @@ -5690,7 +5694,14 @@ async fn active_profile_update_rebuilds_network_proxy_config() -> std::io::Resul ]), }; let base_config = ConfigToml { - features: Some(toml::from_str("network_proxy = true").expect("valid features")), + features: Some( + toml::from_str(if credential_broker { + "network_proxy = { enabled = true, credential_broker = true }" + } else { + "network_proxy = true" + }) + .expect("valid features"), + ), default_permissions: Some("locked-down".to_string()), permissions: Some(permissions), ..Default::default() @@ -5752,6 +5763,7 @@ async fn active_profile_update_rebuilds_network_proxy_config() -> std::io::Resul .expect("selected profile proxy should become the session proxy config"); assert_eq!(network.proxy_host_and_port(), "127.0.0.1:43128"); assert!(!network.socks_enabled()); + assert_eq!(network.credential_broker_enabled(), credential_broker); Ok(()) } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 8df4f1cab6..7e28ac8b51 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -6,7 +6,9 @@ use crate::config::TokenBudgetConfig; use crate::environment_selection::EnvironmentConfigOrigin; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::exec_policy::AllowPrefixRules; +use crate::shell_snapshot::ShellSnapshot; use crate::shell_snapshot::ShellSnapshotFile; +use crate::shell_snapshot::ShellSnapshotSandbox; use crate::tools::sandboxing::executor_windows_sandbox_level; use arc_swap::ArcSwap; use codex_core_plugins::PluginCommandAttribution; @@ -56,6 +58,7 @@ pub(crate) struct TurnEnvironment { /// OS reported by the selected executor; `None` for legacy executors. pub(crate) executor_platform_os: Option, pub(crate) shell_snapshot: ShellSnapshotTask, + pub(crate) shell_snapshot_builder: Option>, pub(crate) shell_snapshot_v2_supported: bool, } @@ -76,6 +79,7 @@ impl TurnEnvironment { shell, executor_platform_os: None, shell_snapshot: futures::future::ready(None).boxed().shared(), + shell_snapshot_builder: None, shell_snapshot_v2_supported: false, } } @@ -99,14 +103,47 @@ impl TurnEnvironment { config } - pub(crate) fn shell_snapshot(&self, cwd: &AbsolutePathBuf) -> Option { - if self.selection.cwd != PathUri::from_abs_path(cwd) { + pub(crate) async fn shell_snapshot( + &self, + cwd: &AbsolutePathBuf, + command: &[String], + shell: &shell::Shell, + config: &Config, + sandbox: Option, + ) -> Option> { + let credential_broker_enabled = config + .permissions + .network + .as_ref() + .is_some_and(|spec| spec.enabled() && spec.credential_broker_enabled()); + let cwd_matches_selection = self.selection.cwd == PathUri::from_abs_path(cwd); + if !config.features.enabled(Feature::ShellSnapshot) + || (!cwd_matches_selection && !credential_broker_enabled) + || command.len() < 3 + || !matches!(command[1].as_str(), "-lc" | "-c") + || (command[1] == "-c" && !credential_broker_enabled) + { return None; } - self.shell_snapshot - .peek()? - .as_deref() - .map(ShellSnapshotFile::path) + if credential_broker_enabled { + let sandbox = sandbox?; + self.shell_snapshot_builder + .as_ref()? + .as_ref() + .clone() + .build( + Arc::clone(&self.environment), + PathUri::from_abs_path(cwd), + Some(shell.clone()), + command[1] == "-lc", + self.shell_environment_policy().clone(), + Some(sandbox), + ) + .await + .filter(|snapshot| snapshot.is_brokered_for(self.shell_environment_policy())) + } else { + self.shell_snapshot.peek()?.clone() + } } pub(crate) fn cwd(&self) -> &PathUri { diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 4967172e34..1bc1aa07db 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::io::ErrorKind; use std::path::Path; use std::process::Stdio; @@ -9,29 +10,47 @@ use crate::StateDbHandle; use crate::rollout::list::find_thread_path_by_id_str; use crate::shell::Shell; use crate::shell::ShellType; -use crate::shell::get_shell; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use anyhow::bail; use codex_exec_server::Environment; +use codex_network_proxy::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; +use codex_network_proxy::NetworkProxy; +use codex_network_proxy::brokered_credential_dummy_env_keys; +use codex_network_proxy::brokered_credential_env_keys; +use codex_network_proxy::credential_broker_provider_context_env_keys; +use codex_network_proxy::is_credential_broker_provider_env_key; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; +use codex_protocol::config_types::ShellEnvironmentPolicy; +use codex_protocol::shell_environment::create_env_from_vars; use codex_shell_command::shell_snapshot::CapturedSnapshot; +use codex_shell_command::shell_snapshot::PreparedSnapshot; use codex_shell_command::shell_snapshot::SnapshotCaptureOptions; +use codex_shell_command::shell_snapshot::SnapshotCredentialEnvironment; use codex_shell_command::shell_snapshot::SnapshotStartup; +use codex_shell_command::shell_snapshot::prepare_snapshot_credentials; use codex_shell_command::shell_snapshot::snapshot_capture_script; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use tokio::fs; use tokio::process::Command; +use tokio::sync::watch; use tokio::time::timeout; use tracing::Instrument; use tracing::info_span; +#[path = "shell_snapshot_sandbox.rs"] +mod sandbox; + +pub(crate) use sandbox::ShellSnapshotSandbox; +pub(crate) use sandbox::snapshot_read_permissions; + #[derive(Clone)] pub(crate) struct ShellSnapshot { config: Option>, + credential_broker: Option>, } struct ShellSnapshotConfig { @@ -39,10 +58,33 @@ struct ShellSnapshotConfig { session_id: ThreadId, session_telemetry: SessionTelemetry, state_db: Option, + credential_broker: Option>, + prefer_executor_snapshots: bool, +} + +#[derive(Clone, PartialEq)] +pub(crate) enum SnapshotCredentialBrokerState { + Starting, + Inactive, + Unavailable, + Ready(NetworkProxy), } pub(crate) struct ShellSnapshotFile { path: AbsolutePathBuf, + credentials: Option, +} + +struct SnapshotCredentials { + network_proxy: NetworkProxy, + shell_environment_policy: ShellEnvironmentPolicy, + credential_env: HashMap, +} + +struct SnapshotCredentialBroker { + network_proxy: NetworkProxy, + shell_environment_policy: ShellEnvironmentPolicy, + allow_login_shell: bool, } const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); @@ -55,6 +97,8 @@ impl ShellSnapshot { session_id: ThreadId, session_telemetry: SessionTelemetry, state_db: Option, + credential_broker: Option>, + prefer_executor_snapshots: bool, ) -> Self { Self { config: Some(Arc::new(ShellSnapshotConfig { @@ -62,12 +106,31 @@ impl ShellSnapshot { session_id, session_telemetry, state_db, + credential_broker: credential_broker.as_ref().map(watch::Sender::subscribe), + prefer_executor_snapshots, })), + credential_broker, } } pub(crate) fn disabled() -> Self { - Self { config: None } + Self { + config: None, + credential_broker: None, + } + } + + pub(crate) fn set_credential_broker(&self, state: SnapshotCredentialBrokerState) -> bool { + self.credential_broker.as_ref().is_some_and(|sender| { + let previous = sender.send_replace(state.clone()); + previous != state + }) + } + + pub(crate) fn should_rebuild_inherited(&self) -> bool { + self.credential_broker.as_ref().is_some_and(|sender| { + !matches!(*sender.borrow(), SnapshotCredentialBrokerState::Inactive) + }) } pub(crate) async fn build( @@ -75,8 +138,11 @@ impl ShellSnapshot { environment: Arc, cwd: PathUri, shell: Option, + allow_login_shell: bool, + shell_environment_policy: ShellEnvironmentPolicy, + sandbox: Option, ) -> Option> { - let config = self.config.as_ref()?; + let config = Arc::clone(self.config.as_ref()?); if environment.is_remote() { return None; } @@ -85,16 +151,55 @@ impl ShellSnapshot { // TODO(anp): Migrate shell snapshot creation to accept PathUri and defer native // conversion to the spawned shell process. let cwd = cwd.to_abs_path().ok()?; - Self::build_for_cwd(Arc::clone(config), cwd, shell).await + drop(self); + Self::build_for_cwd( + config, + cwd, + shell, + allow_login_shell, + shell_environment_policy, + sandbox, + ) + .await } async fn build_for_cwd( config: Arc, cwd: AbsolutePathBuf, shell: Shell, + allow_login_shell: bool, + shell_environment_policy: ShellEnvironmentPolicy, + sandbox: Option, ) -> Option> { let snapshot_span = info_span!("shell_snapshot", thread_id = %config.session_id); async { + let credential_broker = if let Some(receiver) = config.credential_broker.as_ref() { + let mut receiver = receiver.clone(); + if matches!(&*receiver.borrow(), SnapshotCredentialBrokerState::Starting) + && receiver.changed().await.is_err() + { + return None; + } + let state = receiver.borrow().clone(); + match state { + SnapshotCredentialBrokerState::Starting + | SnapshotCredentialBrokerState::Unavailable => return None, + SnapshotCredentialBrokerState::Inactive if config.prefer_executor_snapshots => { + return None; + } + SnapshotCredentialBrokerState::Inactive => None, + SnapshotCredentialBrokerState::Ready(network_proxy) if sandbox.is_some() => { + Some(SnapshotCredentialBroker { + network_proxy, + shell_environment_policy, + allow_login_shell, + }) + } + SnapshotCredentialBrokerState::Ready(_) => return None, + } + } else { + None + }; let timer = config .session_telemetry .start_timer("codex.shell_snapshot.duration_ms", &[("version", "v1")]); @@ -104,6 +209,8 @@ impl ShellSnapshot { &cwd, &shell, config.state_db.clone(), + credential_broker, + sandbox.as_ref(), ) .await; let success_tag = if snapshot.is_ok() { "true" } else { "false" }; @@ -127,6 +234,8 @@ impl ShellSnapshot { session_cwd: &AbsolutePathBuf, shell: &Shell, state_db: Option, + credential_broker: Option, + sandbox: Option<&ShellSnapshotSandbox>, ) -> std::result::Result { // File to store the snapshot let extension = match shell.shell_type { @@ -156,19 +265,35 @@ impl ShellSnapshot { }); // Make the new snapshot. - if let Err(err) = write_shell_snapshot(shell.shell_type, &temp_path, session_cwd).await { + let credentials = write_shell_snapshot( + shell, + &temp_path, + session_cwd, + credential_broker.as_ref(), + sandbox, + ) + .await + .map_err(|err| { tracing::warn!( "Failed to create shell snapshot for {}: {err:?}", shell.name() ); - return Err("write_failed"); - } + "write_failed" + })?; tracing::info!( "Shell snapshot successfully created: {}", temp_path.display() ); - if let Err(err) = validate_snapshot(shell, &temp_path, session_cwd).await { + if let Err(err) = validate_snapshot( + shell, + &temp_path, + session_cwd, + credential_broker.as_ref(), + sandbox, + ) + .await + { tracing::error!("Shell snapshot validation failed: {err:?}"); remove_snapshot_file(&temp_path).await; return Err("validation_failed"); @@ -180,14 +305,48 @@ impl ShellSnapshot { return Err("write_failed"); } - Ok(ShellSnapshotFile { path }) + Ok(ShellSnapshotFile { path, credentials }) } } impl ShellSnapshotFile { + pub(crate) fn is_brokered_for( + &self, + shell_environment_policy: &ShellEnvironmentPolicy, + ) -> bool { + self.credentials.as_ref().is_some_and(|credentials| { + &credentials.shell_environment_policy == shell_environment_policy + }) + } + pub(crate) fn path(&self) -> AbsolutePathBuf { self.path.clone() } + + pub(crate) fn restore_credentials( + &self, + env: &mut HashMap, + shell_environment_policy: &ShellEnvironmentPolicy, + ) { + let Some(credentials) = self.credentials.as_ref() else { + return; + }; + + // This snapshot belongs to one command and was captured under this exact policy. + if &credentials.shell_environment_policy != shell_environment_policy { + return; + } + let mut snapshot_env = credentials.credential_env.clone(); + snapshot_env.insert( + CREDENTIAL_BROKER_ACTIVE_ENV_KEY.to_string(), + "1".to_string(), + ); + credentials + .network_proxy + .restore_brokered_credentials(&mut snapshot_env, &mut []); + snapshot_env.remove(CREDENTIAL_BROKER_ACTIVE_ENV_KEY); + env.extend(snapshot_env); + } } impl Drop for ShellSnapshotFile { @@ -202,17 +361,17 @@ impl Drop for ShellSnapshotFile { } async fn write_shell_snapshot( - shell_type: ShellType, + shell: &Shell, output_path: &AbsolutePathBuf, cwd: &AbsolutePathBuf, -) -> Result<()> { + credential_broker: Option<&SnapshotCredentialBroker>, + sandbox: Option<&ShellSnapshotSandbox>, +) -> Result> { + let shell_type = shell.shell_type; if shell_type == ShellType::PowerShell || shell_type == ShellType::Cmd { bail!("Shell snapshot not supported yet for {shell_type:?}"); } - let shell = - get_shell(shell_type).with_context(|| format!("No available shell for {shell_type:?}"))?; - - let snapshot = capture_snapshot(&shell, cwd).await?; + let (snapshot, credentials) = capture_snapshot(shell, cwd, credential_broker, sandbox).await?; if let Some(parent) = output_path.parent() { let parent_display = parent.display(); @@ -226,30 +385,170 @@ async fn write_shell_snapshot( .await .with_context(|| format!("Failed to write snapshot to {snapshot_path}"))?; - Ok(()) + Ok(credentials) } -async fn capture_snapshot(shell: &Shell, cwd: &AbsolutePathBuf) -> Result { +async fn capture_snapshot( + shell: &Shell, + cwd: &AbsolutePathBuf, + credential_broker: Option<&SnapshotCredentialBroker>, + sandbox: Option<&ShellSnapshotSandbox>, +) -> Result<(String, Option)> { let shell_type = shell.shell_type; + let shell_startup = if credential_broker.is_some_and(|broker| !broker.allow_login_shell) { + SnapshotStartup::NonInteractive + } else { + SnapshotStartup::Interactive + }; let script = snapshot_capture_script( shell_type, SnapshotCaptureOptions { - startup: SnapshotStartup::Interactive, + startup: shell_startup, declarations: true, - environment: false, + environment: credential_broker.is_some(), }, ) .ok_or_else(|| anyhow!("Shell snapshotting is not yet supported for {shell_type:?}"))?; - let captured = run_shell_script(shell, &script, cwd).await?; - CapturedSnapshot::parse(shell_type, captured.as_bytes()) - .map(|snapshot| snapshot.render_script()) - .ok_or_else(|| anyhow!("Invalid shell snapshot capture")) + let shell_mode = if credential_broker.is_none_or(|broker| broker.allow_login_shell) { + SnapshotShellMode::Login + } else { + SnapshotShellMode::NonLogin + }; + let raw_snapshot = run_script_with_timeout( + shell, + &script, + SNAPSHOT_TIMEOUT, + shell_mode, + cwd, + credential_broker, + sandbox, + ) + .await?; + let capture = CapturedSnapshot::parse(shell_type, raw_snapshot.as_bytes()) + .ok_or_else(|| anyhow!("invalid shell snapshot capture"))?; + let Some(credential_broker) = credential_broker else { + return Ok((capture.render_script(), None)); + }; + + let original_env = std::str::from_utf8(capture.environment)? + .split('\0') + .filter_map(|entry| entry.split_once('=')) + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect::>(); + let policy = &credential_broker.shell_environment_policy; + let inherited_env = create_env_from_vars(std::env::vars(), policy, /*thread_id*/ None); + let mut restored_env = original_env.clone(); + credential_broker + .network_proxy + .restore_brokered_credentials(&mut restored_env, &mut []); + let mut discovery_env = restored_env.clone(); + replace_provider_context_with_inherited( + &mut discovery_env, + &inherited_env, + credential_broker_provider_context_env_keys(), + ); + for (key, value) in &policy.r#set { + if !discovery_env.contains_key(key) + || credential_broker_provider_context_env_keys() + .any(|context_key| context_key.eq_ignore_ascii_case(key)) + { + discovery_env.insert(key.clone(), value.clone()); + } + } + credential_broker + .network_proxy + .apply_to_env(&mut discovery_env); + let brokered_keys = brokered_credential_dummy_env_keys(&discovery_env); + let mut env = create_env_from_vars( + restored_env + .iter() + .map(|(key, value)| (key.clone(), value.clone())), + policy, + /*thread_id*/ None, + ); + replace_provider_context_with_inherited( + &mut env, + &inherited_env, + credential_broker_provider_context_env_keys(), + ); + credential_broker.network_proxy.apply_to_env(&mut env); + let allowed_brokered_keys = brokered_credential_dummy_env_keys(&env); + let mut snapshot_env = env.clone(); + for key in credential_broker_provider_context_env_keys() { + if original_env.get(key) != env.get(key) { + snapshot_env.remove(key); + } + } + let PreparedSnapshot { + script: snapshot, .. + } = prepare_snapshot_credentials( + &capture, + SnapshotCredentialEnvironment { + original: &original_env, + restored: &restored_env, + configured: &policy.r#set, + discovered: &discovery_env, + allowed: &snapshot_env, + is_allowed_unset: &|key| { + create_env_from_vars( + std::iter::once((key.to_string(), String::new())), + policy, + /*thread_id*/ None, + ) + .contains_key(key) + }, + brokered_keys: &brokered_keys, + brokered_alias_keys: &[], + allowed_brokered_keys: &allowed_brokered_keys, + }, + |value| { + credential_broker + .network_proxy + .virtualize_brokered_text(value, &env) + }, + ) + .ok_or_else(|| anyhow!("shell snapshot contains a credential outside supported exports"))?; + + let credential_env = brokered_credential_env_keys(&env) + .map(str::to_string) + .chain(allowed_brokered_keys) + .filter_map(|key| env.get(&key).map(|value| (key, value.clone()))) + .collect(); + Ok(( + snapshot, + Some(SnapshotCredentials { + network_proxy: credential_broker.network_proxy.clone(), + shell_environment_policy: policy.clone(), + credential_env, + }), + )) +} + +fn replace_provider_context_with_inherited( + env: &mut HashMap, + inherited_env: &HashMap, + context_keys: impl Iterator, +) { + for key in context_keys { + if let Some(value) = inherited_env.get(key) { + env.insert(key.to_string(), value.clone()); + } + } +} + +#[derive(Clone, Copy)] +enum SnapshotShellMode<'a> { + Login, + NonLogin, + Validation(&'a AbsolutePathBuf), } async fn validate_snapshot( shell: &Shell, snapshot_path: &AbsolutePathBuf, cwd: &AbsolutePathBuf, + credential_broker: Option<&SnapshotCredentialBroker>, + sandbox: Option<&ShellSnapshotSandbox>, ) -> Result<()> { let snapshot_path_display = snapshot_path.display(); let script = format!("set -e; . \"{snapshot_path_display}\""); @@ -257,41 +556,90 @@ async fn validate_snapshot( shell, &script, SNAPSHOT_TIMEOUT, - /*use_login_shell*/ false, + SnapshotShellMode::Validation(snapshot_path), cwd, + credential_broker, + sandbox, ) .await .map(|_| ()) } -async fn run_shell_script(shell: &Shell, script: &str, cwd: &AbsolutePathBuf) -> Result { - run_script_with_timeout( - shell, - script, - SNAPSHOT_TIMEOUT, - /*use_login_shell*/ true, - cwd, - ) - .await -} - async fn run_script_with_timeout( shell: &Shell, script: &str, snapshot_timeout: Duration, - use_login_shell: bool, + shell_mode: SnapshotShellMode<'_>, cwd: &AbsolutePathBuf, + credential_broker: Option<&SnapshotCredentialBroker>, + sandbox: Option<&ShellSnapshotSandbox>, ) -> Result { - let args = shell.derive_exec_args(script, use_login_shell); + let suppress_startup_files = + credential_broker.is_some() && matches!(shell_mode, SnapshotShellMode::Validation(_)); + let mut args = shell.derive_exec_args(script, matches!(shell_mode, SnapshotShellMode::Login)); + if suppress_startup_files && shell.shell_type == ShellType::Zsh { + args[1] = "-fc".to_string(); + } let shell_name = shell.name(); + let mut prepared_env = None; + if let Some(credential_broker) = credential_broker { + let policy = &credential_broker.shell_environment_policy; + let mut inherited_policy = policy.clone(); + inherited_policy.r#set.clear(); + let mut env = + create_env_from_vars(std::env::vars(), &inherited_policy, /*thread_id*/ None); + env.extend( + policy + .r#set + .iter() + .filter(|(key, _)| { + (is_credential_broker_provider_env_key(key) + || matches!( + (shell.shell_type, key.as_str()), + (ShellType::Zsh, "ZDOTDIR") | (ShellType::Bash, "BASH_ENV") + )) + && (policy.include_only.is_empty() + || policy + .include_only + .iter() + .any(|pattern| pattern.matches(key))) + }) + .map(|(key, value)| (key.clone(), value.clone())), + ); + if suppress_startup_files { + env.remove("BASH_ENV"); + } + credential_broker.network_proxy.apply_to_env(&mut env); + prepared_env = Some(env); + } + if let Some(sandbox) = sandbox { + let snapshot_read_path = match shell_mode { + SnapshotShellMode::Validation(path) => Some(path), + SnapshotShellMode::Login | SnapshotShellMode::NonLogin => None, + }; + return sandbox + .run( + args, + cwd, + prepared_env.unwrap_or_else(|| std::env::vars().collect()), + snapshot_timeout, + shell_name, + snapshot_read_path, + ) + .await; + } // Handler is kept as guard to control the drop. The `mut` pattern is required because .args() // returns a ref of handler. let mut handler = Command::new(&args[0]); - codex_protocol::shell_environment::scrub_non_inheritable_env_vars(handler.as_std_mut()); handler.args(&args[1..]); handler.stdin(Stdio::null()); handler.current_dir(cwd); + if let Some(env) = prepared_env { + handler.env_clear(); + handler.envs(env); + } + codex_protocol::shell_environment::scrub_non_inheritable_env_vars(handler.as_std_mut()); #[cfg(unix)] unsafe { handler.pre_exec(|| { diff --git a/codex-rs/core/src/shell_snapshot_sandbox.rs b/codex-rs/core/src/shell_snapshot_sandbox.rs new file mode 100644 index 0000000000..2b0dbd1817 --- /dev/null +++ b/codex-rs/core/src/shell_snapshot_sandbox.rs @@ -0,0 +1,184 @@ +use crate::exec::ExecCapturePolicy; +use crate::exec::ExecExpiration; +use crate::sandboxing::ExecOptions; +use crate::sandboxing::ExecRequest; +use crate::sandboxing::execute_env; +use crate::tools::sandboxing::SandboxAttempt; +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use anyhow::bail; +use codex_network_proxy::NetworkProxy; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::models::AdditionalPermissionProfile; +use codex_protocol::models::FileSystemPermissions; +use codex_protocol::models::PermissionProfile; +use codex_sandboxing::SandboxCommand; +use codex_sandboxing::SandboxManager; +use codex_sandboxing::SandboxTransformRequest; +use codex_sandboxing::SandboxType; +use codex_sandboxing::policy_transforms::merge_permission_profiles; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use std::collections::HashMap; +use std::path::PathBuf; +use std::time::Duration; + +#[derive(Clone)] +pub(crate) struct ShellSnapshotSandbox { + sandbox: SandboxType, + sandbox_requested: bool, + permissions: PermissionProfile, + enforce_managed_network: bool, + sandbox_policy_cwd: PathUri, + workspace_roots: Vec, + codex_linux_sandbox_exe: Option, + use_legacy_landlock: bool, + windows_sandbox_level: WindowsSandboxLevel, + windows_sandbox_private_desktop: bool, + network: Option, + environment_id: String, + additional_permissions: Option, +} + +pub(crate) fn snapshot_read_permissions( + snapshot_path: &AbsolutePathBuf, + permissions: &PermissionProfile, + sandbox_cwd: &PathUri, +) -> Option { + if let Ok(cwd) = sandbox_cwd.to_abs_path() + && permissions + .file_system_sandbox_policy() + .can_read_local_path_with_cwd(snapshot_path.as_path(), cwd.as_path()) + { + return None; + } + Some(AdditionalPermissionProfile { + network: None, + file_system: Some(FileSystemPermissions::from_read_write_roots( + Some(vec![snapshot_path.clone()]), + /*write*/ None, + )), + }) +} + +impl ShellSnapshotSandbox { + pub(crate) fn new( + attempt: &SandboxAttempt<'_>, + network: Option<&NetworkProxy>, + environment_id: &str, + additional_permissions: Option<&AdditionalPermissionProfile>, + ) -> Self { + Self { + sandbox: attempt.sandbox, + sandbox_requested: attempt.sandbox_requested, + permissions: attempt.permissions.clone(), + enforce_managed_network: attempt.enforce_managed_network, + sandbox_policy_cwd: attempt.sandbox_cwd.clone(), + workspace_roots: attempt.workspace_roots.to_vec(), + codex_linux_sandbox_exe: attempt.codex_linux_sandbox_exe.cloned(), + use_legacy_landlock: attempt.use_legacy_landlock, + windows_sandbox_level: attempt.windows_sandbox_level, + windows_sandbox_private_desktop: attempt.windows_sandbox_private_desktop, + network: attempt.network_proxy(network).cloned(), + environment_id: environment_id.to_string(), + additional_permissions: additional_permissions.cloned(), + } + } + + pub(crate) async fn run( + &self, + args: Vec, + cwd: &AbsolutePathBuf, + mut env: HashMap, + snapshot_timeout: Duration, + shell_name: &str, + snapshot_read_path: Option<&AbsolutePathBuf>, + ) -> Result { + if self.sandbox_requested && self.sandbox == SandboxType::None { + bail!("shell snapshot sandbox cannot be enforced on this host"); + } + + let managed_network = if let Some(network) = self.network.as_ref() { + let prepared = network + .prepare_for_optional_environment(env, Some(&self.environment_id)) + .context("failed to prepare managed network for shell snapshot")?; + env = prepared.env; + Some(prepared.sandbox_context) + } else { + None + }; + let (program, args) = args + .split_first() + .ok_or_else(|| anyhow!("shell snapshot command is empty"))?; + let snapshot_permissions = snapshot_read_path.and_then(|path| { + snapshot_read_permissions( + path, + &self + .permissions + .clone() + .materialize_project_roots_with_path_uris(&self.workspace_roots), + &self.sandbox_policy_cwd, + ) + }); + let additional_permissions = merge_permission_profiles( + self.additional_permissions.as_ref(), + snapshot_permissions.as_ref(), + ); + let manager = SandboxManager::new(); + let request = manager + .transform(SandboxTransformRequest { + command: SandboxCommand { + program: program.clone().into(), + args: args.to_vec(), + cwd: PathUri::from_abs_path(cwd), + env, + managed_network, + additional_permissions, + }, + permissions: &self.permissions, + sandbox: self.sandbox, + enforce_managed_network: self.enforce_managed_network, + environment_id: Some(&self.environment_id), + network: self.network.as_ref(), + sandbox_policy_cwd: &self.sandbox_policy_cwd, + codex_linux_sandbox_exe: self.codex_linux_sandbox_exe.as_deref(), + use_legacy_landlock: self.use_legacy_landlock, + windows_sandbox_level: self.windows_sandbox_level, + windows_sandbox_private_desktop: self.windows_sandbox_private_desktop, + }) + .context("failed to prepare shell snapshot sandbox")?; + let workspace_roots = self + .workspace_roots + .iter() + .map(PathUri::to_abs_path) + .collect::>>() + .context("failed to resolve shell snapshot workspace roots")?; + let request = ExecRequest::from_sandbox_exec_request( + request, + ExecOptions { + expiration: ExecExpiration::Timeout(snapshot_timeout), + capture_policy: ExecCapturePolicy::FullBuffer, + }, + workspace_roots, + ) + .context("failed to prepare shell snapshot execution")?; + let output = tokio::time::timeout( + snapshot_timeout, + execute_env(request, /*stdout_stream*/ None), + ) + .await + .map_err(|_| anyhow!("Snapshot command timed out for {shell_name}"))? + .with_context(|| format!("Failed to execute sandboxed {shell_name}"))?; + + if output.exit_code != 0 { + bail!( + "Snapshot command exited with status {}: {}", + output.exit_code, + output.stderr.text + ); + } + + Ok(output.stdout.text) + } +} diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index 3c0995e139..28045a5c87 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -1,10 +1,22 @@ use super::*; +#[cfg(unix)] +use crate::config::NetworkProxySpec; +#[cfg(unix)] +use codex_network_proxy::NetworkProxyConfig; +#[cfg(unix)] +use codex_network_proxy::brokered_credential_binding_env_keys; +#[cfg(unix)] +use codex_protocol::config_types::EnvironmentVariablePattern; +#[cfg(unix)] +use codex_protocol::models::PermissionProfile; use core_test_support::PathBufExt; use core_test_support::PathExt; use pretty_assertions::assert_eq; #[cfg(unix)] use std::os::unix::ffi::OsStrExt; use std::path::PathBuf; +#[cfg(unix)] +use std::process::Command; #[cfg(target_os = "linux")] use std::process::Command as StdCommand; @@ -81,7 +93,16 @@ fn assert_posix_snapshot_sections(snapshot: &str) { async fn get_snapshot(shell_type: ShellType) -> Result { let dir = tempdir()?; let path = dir.path().join("snapshot.sh"); - write_shell_snapshot(shell_type, &path.abs(), &dir.path().abs()).await?; + let shell = crate::shell::get_shell(shell_type) + .with_context(|| format!("No available shell for {shell_type:?}"))?; + write_shell_snapshot( + &shell, + &path.abs(), + &dir.path().abs(), + /*credential_broker*/ None, + /*sandbox*/ None, + ) + .await?; let content = fs::read_to_string(&path).await?; Ok(content) } @@ -108,6 +129,443 @@ fn snapshot_file_name_parser_supports_legacy_and_suffixed_names() { ); } +#[cfg(unix)] +#[tokio::test] +async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() -> Result<()> { + let dir = tempdir()?; + std::fs::create_dir(dir.path().join(".codex"))?; + std::fs::write( + dir.path().join(".codex/startup.sh"), + "printf started > startup-ran\n", + )?; + let session_id = ThreadId::new(); + let session_telemetry = SessionTelemetry::new( + session_id, + "test", + "test", + /*account_id*/ None, + /*account_email*/ None, + /*auth_mode*/ None, + "test".to_string(), + /*log_user_prompts*/ false, + "test".to_string(), + codex_protocol::protocol::SessionSource::Cli, + ); + let shell = Shell { + shell_type: ShellType::Bash, + shell_path: PathBuf::from("/bin/bash"), + }; + let permission_profile = PermissionProfile::workspace_write(); + let mut network_config = NetworkProxyConfig::default(); + network_config.set_credential_broker_enabled(/*enabled*/ true); + let network_spec = NetworkProxySpec::from_config_and_constraints( + network_config, + /*requirements*/ None, + &permission_profile, + )?; + let started_proxy = network_spec + .start_proxy( + &permission_profile, + /*policy_decider*/ None, + /*blocked_request_observer*/ None, + /*enable_network_approval_flow*/ false, + Default::default(), + ) + .await?; + for (state, prefer_executor_snapshots, expect_snapshot) in [ + (SnapshotCredentialBrokerState::Unavailable, false, false), + (SnapshotCredentialBrokerState::Inactive, false, true), + (SnapshotCredentialBrokerState::Inactive, true, false), + ( + SnapshotCredentialBrokerState::Ready(started_proxy.proxy()), + false, + false, + ), + ] { + let (_sender, receiver) = watch::channel(state); + let config = Arc::new(ShellSnapshotConfig { + codex_home: dir.path().abs(), + session_id, + session_telemetry: session_telemetry.clone(), + state_db: None, + credential_broker: Some(receiver), + prefer_executor_snapshots, + }); + let snapshot = tokio::time::timeout( + SNAPSHOT_TIMEOUT, + ShellSnapshot::build_for_cwd( + config, + dir.path().abs(), + shell.clone(), + /*allow_login_shell*/ false, + ShellEnvironmentPolicy { + r#set: HashMap::from([( + "BASH_ENV".to_string(), + "./.codex/startup.sh".to_string(), + )]), + ..ShellEnvironmentPolicy::default() + }, + /*sandbox*/ None, + ), + ) + .await?; + + assert_eq!(snapshot.is_some(), expect_snapshot); + } + assert!(!dir.path().join("startup-ran").exists()); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test] +async fn snapshot_discovers_and_redacts_shell_initialized_credentials() -> Result<()> { + let dir = tempdir()?; + let startup = dir.path().join("startup.sh"); + std::fs::write( + &startup, + "export GH_TOKEN='ghp_shell_only_secret'\n\ + export AUTH_HEADER=\"Bearer $GH_TOKEN\"\n\ + declare -rx GITHUB_TOKEN='ghp_readonly_secret'\n\ + declare -rx HOMEBREW_GITHUB_API_TOKEN=\"$GITHUB_TOKEN\"\n\ + export GH_ENTERPRISE_TOKEN='ghp_enterprise_secret'\n\ + export GH_HOST='attacker.example'\n\ + export OPENAI_API_KEY='sk-proj-snapshot-secret'\n\ + export AUTH_BUNDLE=\"GitHub $GH_TOKEN\n\ + OpenAI $OPENAI_API_KEY\"\n\ + export OPENAI_BASE_URL='https://api.snapshot.example/v1'\n\ + export IDENTITY_SEEN=\"${OPENAI_IDENTITY_TOKEN_FILE-missing}\"\n\ + export EXCLUDED_PARENT_HOME=\"${HOME-missing}\"\n\ + export STARTUP_PATH_OVERRIDE_SEEN=\"${PATH%%:*}\"\n", + )?; + + let mut network_config = NetworkProxyConfig::default(); + network_config.set_credential_broker_enabled(/*enabled*/ true); + let permission_profile = PermissionProfile::workspace_write(); + let network_spec = NetworkProxySpec::from_config_and_constraints( + network_config, + /*requirements*/ None, + &permission_profile, + )?; + let started_proxy = network_spec + .start_proxy( + &permission_profile, + /*policy_decider*/ None, + /*blocked_request_observer*/ None, + /*enable_network_approval_flow*/ false, + Default::default(), + ) + .await?; + let network_proxy = started_proxy.proxy(); + + let shell = Shell { + shell_type: ShellType::Bash, + shell_path: PathBuf::from("/bin/bash"), + }; + let trusted_startup = ("BASH_ENV".to_string(), startup.display().to_string()); + let shell_environment_policy = ShellEnvironmentPolicy { + exclude: vec![EnvironmentVariablePattern::new_case_insensitive("HOME")], + r#set: HashMap::from([ + trusted_startup.clone(), + ("GH_HOST".to_string(), "github.example.com".to_string()), + ( + "OPENAI_IDENTITY_TOKEN_FILE".to_string(), + "identity-token-secret".to_string(), + ), + ( + "PATH".to_string(), + format!("/repository-controlled:{}", std::env::var("PATH")?), + ), + ]), + ..ShellEnvironmentPolicy::default() + }; + let credential_broker = SnapshotCredentialBroker { + network_proxy: network_proxy.clone(), + shell_environment_policy: shell_environment_policy.clone(), + allow_login_shell: false, + }; + let path = dir.path().join("snapshot.sh").abs(); + let credentials = write_shell_snapshot( + &shell, + &path, + &dir.path().abs(), + Some(&credential_broker), + /*sandbox*/ None, + ) + .await?; + let snapshot = fs::read_to_string(&path).await?; + + for secret in [ + "ghp_shell_only_secret", + "ghp_readonly_secret", + "ghp_enterprise_secret", + "sk-proj-snapshot-secret", + "identity-token-secret", + ] { + assert!(!snapshot.contains(secret), "snapshot exposed {secret}"); + } + assert!(!snapshot.contains("attacker.example")); + assert!(snapshot.contains("api.snapshot.example")); + assert!(snapshot.contains("IDENTITY_SEEN=\"missing\"")); + assert!(snapshot.contains("EXCLUDED_PARENT_HOME=\"missing\"")); + assert!(!snapshot.contains("STARTUP_PATH_OVERRIDE_SEEN=\"/repository-controlled\"")); + assert!(snapshot.contains("declare -rx HOMEBREW_GITHUB_API_TOKEN=\"${GITHUB_TOKEN-}\"")); + + validate_snapshot( + &shell, + &path, + &dir.path().abs(), + Some(&credential_broker), + /*sandbox*/ None, + ) + .await?; + + let validation_path = dir.path().join("validation.sh").abs(); + fs::write( + &validation_path, + "test \"${HOME-missing}\" = missing && \ + test \"${CODEX_NETWORK_PROXY_CREDENTIAL_BROKER_ACTIVE-}\" = 1\n", + ) + .await?; + validate_snapshot( + &shell, + &validation_path, + &dir.path().abs(), + Some(&credential_broker), + /*sandbox*/ None, + ) + .await?; + + let filtered_startup_broker = SnapshotCredentialBroker { + network_proxy: network_proxy.clone(), + shell_environment_policy: ShellEnvironmentPolicy { + r#set: HashMap::from([trusted_startup.clone()]), + include_only: vec![ + EnvironmentVariablePattern::new_case_insensitive("PATH"), + EnvironmentVariablePattern::new_case_insensitive("EXCLUDED_PARENT_HOME"), + ], + ..ShellEnvironmentPolicy::default() + }, + allow_login_shell: false, + }; + let filtered_startup_path = dir.path().join("filtered-startup-snapshot.sh").abs(); + write_shell_snapshot( + &shell, + &filtered_startup_path, + &dir.path().abs(), + Some(&filtered_startup_broker), + /*sandbox*/ None, + ) + .await?; + assert!( + !fs::read_to_string(&filtered_startup_path) + .await? + .contains("EXCLUDED_PARENT_HOME") + ); + + let inherited_secret = "ghp_inherited_enterprise_secret"; + let inherited_context = HashMap::from([( + "GH_HOST".to_string(), + "github.inherited.example".to_string(), + )]); + let mut inherited_discovery_env = HashMap::from([ + ( + "GH_ENTERPRISE_TOKEN".to_string(), + inherited_secret.to_string(), + ), + ("GH_HOST".to_string(), "attacker.example".to_string()), + ]); + replace_provider_context_with_inherited( + &mut inherited_discovery_env, + &inherited_context, + credential_broker_provider_context_env_keys(), + ); + assert_eq!( + inherited_discovery_env.get("GH_HOST").map(String::as_str), + Some("github.inherited.example") + ); + network_proxy.apply_to_env(&mut inherited_discovery_env); + let mut inherited_allowed_env = HashMap::from([( + "GH_ENTERPRISE_TOKEN".to_string(), + inherited_secret.to_string(), + )]); + replace_provider_context_with_inherited( + &mut inherited_allowed_env, + &inherited_context, + brokered_credential_binding_env_keys(&inherited_discovery_env), + ); + network_proxy.apply_to_env(&mut inherited_allowed_env); + assert_eq!( + inherited_allowed_env.get("GH_HOST").map(String::as_str), + Some("github.inherited.example") + ); + assert_ne!( + inherited_allowed_env + .get("GH_ENTERPRISE_TOKEN") + .map(String::as_str), + Some(inherited_secret) + ); + let mut inherited_snapshot = format!("export GH_ENTERPRISE_TOKEN={inherited_secret}\n"); + assert!( + network_proxy.virtualize_brokered_text(&mut inherited_snapshot, &inherited_allowed_env) + ); + assert!(!inherited_snapshot.contains(inherited_secret)); + + let snapshot_file = ShellSnapshotFile { path, credentials }; + let mut env = HashMap::from([("GH_HOST".to_string(), "github.example.com".to_string())]); + snapshot_file.restore_credentials(&mut env, &shell_environment_policy); + for (key, value) in [ + ("GH_TOKEN", "ghp_shell_only_secret"), + ("GITHUB_TOKEN", "ghp_readonly_secret"), + ("GH_ENTERPRISE_TOKEN", "ghp_enterprise_secret"), + ("OPENAI_API_KEY", "sk-proj-snapshot-secret"), + ] { + assert_eq!(env.get(key).map(String::as_str), Some(value), "{key}"); + } + assert_eq!( + env.get("OPENAI_BASE_URL").map(String::as_str), + Some("https://api.snapshot.example/v1") + ); + + let unbrokered_replay = Command::new("/bin/bash") + .arg("-c") + .arg(". \"$1\" && printf '%s\\n%s' \"$AUTH_HEADER\" \"$AUTH_BUNDLE\"") + .arg("snapshot") + .arg(snapshot_file.path().as_path()) + .env_clear() + .envs(&env) + .output()?; + assert!(unbrokered_replay.status.success()); + assert_eq!( + String::from_utf8(unbrokered_replay.stdout)?, + "Bearer ghp_shell_only_secret\nGitHub ghp_shell_only_secret\nOpenAI sk-proj-snapshot-secret" + ); + + network_proxy.apply_to_env(&mut env); + assert_ne!(env["GH_TOKEN"], "ghp_shell_only_secret"); + assert_ne!(env["GITHUB_TOKEN"], "ghp_readonly_secret"); + assert_ne!(env["GH_ENTERPRISE_TOKEN"], "ghp_enterprise_secret"); + assert_ne!(env["OPENAI_API_KEY"], "sk-proj-snapshot-secret"); + let replay = Command::new("/bin/bash") + .arg("-c") + .arg( + ". \"$1\" && printf '%s\\n%s\\n%s' \"$HOMEBREW_GITHUB_API_TOKEN\" \"$AUTH_HEADER\" \"$AUTH_BUNDLE\"", + ) + .arg("snapshot") + .arg(snapshot_file.path().as_path()) + .env_clear() + .envs(&env) + .output()?; + assert!( + replay.status.success(), + "snapshot replay failed: {replay:?}" + ); + assert_eq!( + String::from_utf8(replay.stdout)?, + format!( + "{}\nBearer {}\nGitHub {}\nOpenAI {}", + env["GITHUB_TOKEN"], env["GH_TOKEN"], env["GH_TOKEN"], env["OPENAI_API_KEY"] + ) + ); + + let filtered_policy = ShellEnvironmentPolicy { + include_only: vec![EnvironmentVariablePattern::new_case_insensitive("GH_HOST")], + ..ShellEnvironmentPolicy::default() + }; + let mut filtered_env = HashMap::new(); + snapshot_file.restore_credentials(&mut filtered_env, &filtered_policy); + assert!(filtered_env.is_empty()); + + let partially_filtered_credential_broker = SnapshotCredentialBroker { + network_proxy: network_proxy.clone(), + shell_environment_policy: ShellEnvironmentPolicy { + r#set: HashMap::from([trusted_startup.clone()]), + include_only: vec![ + EnvironmentVariablePattern::new_case_insensitive("BASH_ENV"), + EnvironmentVariablePattern::new_case_insensitive("GH_TOKEN"), + EnvironmentVariablePattern::new_case_insensitive("AUTH_HEADER"), + EnvironmentVariablePattern::new_case_insensitive("AUTH_BUNDLE"), + ], + ..ShellEnvironmentPolicy::default() + }, + allow_login_shell: false, + }; + for (github_token, openai_api_key) in [ + ("ghp_shell_only_secret", "sk-proj-snapshot-secret"), + (env["GH_TOKEN"].as_str(), env["OPENAI_API_KEY"].as_str()), + ] { + std::fs::write( + &startup, + format!( + "export GH_TOKEN='{github_token}'\n\ + export GITHUB_TOKEN=\"$GH_TOKEN\"\n\ + export AUTH_HEADER=\"Bearer $GH_TOKEN\"\n\ + export OPENAI_API_KEY='{openai_api_key}'\n\ + export AUTH_BUNDLE=\"GitHub $GH_TOKEN\n\ + OpenAI $OPENAI_API_KEY\"\n\ + unset OPENAI_API_KEY\n" + ), + )?; + let partially_filtered_path = dir.path().join("partially-filtered-snapshot.sh").abs(); + write_shell_snapshot( + &shell, + &partially_filtered_path, + &dir.path().abs(), + Some(&partially_filtered_credential_broker), + /*sandbox*/ None, + ) + .await?; + let filtered_replay = Command::new("/bin/bash") + .arg("-c") + .arg(". \"$1\" && printf '%s\\n%s' \"$AUTH_HEADER\" \"${AUTH_BUNDLE-unset}\"") + .arg("snapshot") + .arg(partially_filtered_path.as_path()) + .env_clear() + .env("GH_TOKEN", "ghp_filtered_dummy") + .output()?; + assert!(filtered_replay.status.success()); + assert_eq!( + String::from_utf8(filtered_replay.stdout)?, + "Bearer ghp_filtered_dummy\nunset" + ); + } + + std::fs::write( + &startup, + "export GH_TOKEN='ghp_hidden_alias_secret'\n\ + export HIDDEN_HEADER=\"Bearer $GH_TOKEN\"\n\ + unset GH_TOKEN GITHUB_ENTERPRISE_TOKEN\n", + )?; + let mut previously_brokered_env = HashMap::from([( + "GH_TOKEN".to_string(), + "ghp_hidden_alias_secret".to_string(), + )]); + network_proxy.apply_to_env(&mut previously_brokered_env); + let inherited_credential_broker = SnapshotCredentialBroker { + network_proxy, + shell_environment_policy: ShellEnvironmentPolicy { + r#set: HashMap::from([trusted_startup]), + ..ShellEnvironmentPolicy::default() + }, + allow_login_shell: false, + }; + let inherited_path = dir.path().join("inherited-snapshot.sh").abs(); + write_shell_snapshot( + &shell, + &inherited_path, + &dir.path().abs(), + Some(&inherited_credential_broker), + /*sandbox*/ None, + ) + .await?; + assert!( + !fs::read_to_string(inherited_path) + .await? + .contains("ghp_hidden_alias_secret") + ); + + Ok(()) +} + #[cfg(unix)] #[tokio::test] async fn try_create_creates_and_deletes_snapshot_file() -> Result<()> { @@ -123,6 +581,8 @@ async fn try_create_creates_and_deletes_snapshot_file() -> Result<()> { &dir.path().abs(), &shell, /*state_db*/ None, + /*credential_broker*/ None, + /*sandbox*/ None, ) .await .expect("snapshot should be created"); @@ -152,6 +612,8 @@ async fn try_create_uses_distinct_generation_paths() -> Result<()> { &dir.path().abs(), &shell, /*state_db*/ None, + /*credential_broker*/ None, + /*sandbox*/ None, ) .await .expect("initial snapshot should be created"); @@ -161,6 +623,8 @@ async fn try_create_uses_distinct_generation_paths() -> Result<()> { &dir.path().abs(), &shell, /*state_db*/ None, + /*credential_broker*/ None, + /*sandbox*/ None, ) .await .expect("refreshed snapshot should be created"); @@ -218,8 +682,10 @@ async fn snapshot_shell_does_not_inherit_stdin() -> Result<()> { &shell, &script, Duration::from_secs(2), - /*use_login_shell*/ true, + SnapshotShellMode::Login, &home, + /*credential_broker*/ None, + /*sandbox*/ None, ) .await .context("run snapshot command")?; @@ -261,8 +727,10 @@ async fn timed_out_snapshot_shell_is_terminated() -> Result<()> { &shell, &script, Duration::from_secs(1), - /*use_login_shell*/ true, + SnapshotShellMode::Login, &dir.path().abs(), + /*credential_broker*/ None, + /*sandbox*/ None, ) .await .expect_err("snapshot shell should time out"); diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index ad7097126e..052fb4a6fc 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -163,7 +163,16 @@ pub(crate) async fn execute_user_shell_command( .await; return; }; - let shell_snapshot_location = turn_environment.shell_snapshot(&cwd); + let shell_snapshot = turn_environment + .shell_snapshot( + &cwd, + &display_command, + environment_shell, + &turn_context.config, + /*sandbox*/ None, + ) + .await; + let shell_snapshot_location = shell_snapshot.as_ref().map(|snapshot| snapshot.path()); let shell_environment_policy = turn_environment.shell_environment_policy(); let mut exec_env_map = create_env(shell_environment_policy, Some(session.thread_id)); inject_session_env(&mut exec_env_map, session.session_id()); diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index 9f4b2063b2..b2c0300ac9 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -17,6 +17,7 @@ use codex_core_plugins::PLUGIN_METRICS_OUTPUT_ENV_VAR; use codex_install_context::InstallContext; #[cfg(target_os = "macos")] use codex_network_proxy::CODEX_PROXY_GIT_SSH_COMMAND_MARKER; +use codex_network_proxy::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; use codex_network_proxy::CUSTOM_CA_ENV_KEYS; use codex_network_proxy::PROXY_ACTIVE_ENV_KEY; use codex_network_proxy::PROXY_ENV_KEYS; @@ -200,8 +201,22 @@ fn prepare_powershell_command_for_elevated_windows_sandbox_with_fallback( command } +pub(crate) fn prepare_brokered_shell_snapshot_env( + env: &mut HashMap, + shell_snapshot: Option<&AbsolutePathBuf>, +) { + if shell_snapshot.is_some() + && env + .get(CREDENTIAL_BROKER_ACTIVE_ENV_KEY) + .is_some_and(|active| active == "1") + { + env.remove("BASH_ENV"); + } +} + /// POSIX-only helper: for commands produced by `Shell::derive_exec_args` -/// for Bash/Zsh/sh of the form `[shell_path, "-lc", "