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", "