diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index 359bc770f2..b8a3e61311 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -28,6 +28,9 @@ use codex_network_proxy::PROXY_GIT_SSH_COMMAND_ENV_KEY; pub(crate) use codex_network_proxy::is_managed_proxy_env_var; pub(crate) use codex_network_proxy::strip_managed_proxy_env; use codex_protocol::shell_environment::is_non_inheritable_env_var; +use codex_shell_command::shell_snapshot::SNAPSHOT_BROKERED_UNSET_ENV_PREFIX; +use codex_shell_command::shell_snapshot::SNAPSHOT_BROKERED_VALUE_ENV_PREFIX; +use codex_shell_command::shell_snapshot::build_brokered_credential_exports; use codex_shell_command::shell_snapshot::posix_env_path_expansion_function; use codex_utils_absolute_path::AbsolutePathBuf; use std::collections::HashMap; @@ -39,8 +42,6 @@ pub(crate) mod unified_exec; const SNAPSHOT_ORIGINAL_BASH_ENV_ENV_KEY: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_ORIGINAL_BASH_ENV"; const SNAPSHOT_ORIGINAL_POSIX_ENV_ENV_KEY: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_ORIGINAL_POSIX_ENV"; const SNAPSHOT_ORIGINAL_ZDOTDIR_ENV_KEY: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_ORIGINAL_ZDOTDIR"; -const SNAPSHOT_BROKERED_VALUE_ENV_PREFIX: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_BROKERED_VALUE_"; -const SNAPSHOT_BROKERED_UNSET_ENV_PREFIX: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_BROKERED_UNSET_"; pub(crate) fn exec_env_for_sandbox_permissions( env: &HashMap, @@ -517,60 +518,6 @@ unset __CODEX_SNAPSHOT_ORIGINAL_ENV_SET __CODEX_SNAPSHOT_ORIGINAL_ENV \ rewritten } -fn build_brokered_credential_exports(env: &HashMap, remove_copies: bool) -> String { - let mut value_copies = env - .keys() - .filter_map(|copy_key| { - let key = copy_key.strip_prefix(SNAPSHOT_BROKERED_VALUE_ENV_PREFIX)?; - is_valid_shell_variable_name(key).then_some((key, copy_key)) - }) - .collect::>(); - value_copies.sort_unstable(); - let value_exports = value_copies - .into_iter() - .map(|(key, copy_key)| { - let restore = format!( - "if [ -z \"${{{copy_key}+x}}\" ]; then exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ] && [ \"${{{key}}}\" != \"${{{copy_key}}}\" ]; then export {key}=\"${{{copy_key}}}\" || exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ] && [ \"${{{key}}}\" != \"${{{copy_key}}}\" ]; then exit 1; fi" - ); - if remove_copies { - format!("{restore}\nunset {copy_key} || exit 1") - } else { - restore - } - }) - .collect::>() - .join("\n"); - let mut unset_markers = env - .keys() - .filter_map(|marker_key| { - let key = marker_key.strip_prefix(SNAPSHOT_BROKERED_UNSET_ENV_PREFIX)?; - is_valid_shell_variable_name(key).then_some((key, marker_key)) - }) - .collect::>(); - unset_markers.sort_unstable(); - let unset_exports = unset_markers - .into_iter() - .map(|(key, marker_key)| { - let restore = format!( - "if [ -z \"${{{marker_key}+x}}\" ]; then exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ]; then unset {key} || exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ]; then exit 1; fi" - ); - if remove_copies { - format!("{restore}\nunset {marker_key} || exit 1") - } else { - restore - } - }) - .collect::>() - .join("\n"); - let exports = join_shell_blocks([value_exports, unset_exports]); - if exports.is_empty() { - return exports; - } - format!( - "case $- in\n *x*) __CODEX_SNAPSHOT_BROKER_XTRACE=1; set +x ;;\n *) __CODEX_SNAPSHOT_BROKER_XTRACE= ;;\nesac\n{exports}\nif [ -n \"$__CODEX_SNAPSHOT_BROKER_XTRACE\" ]; then\n unset __CODEX_SNAPSHOT_BROKER_XTRACE\n set -x\nelse\n unset __CODEX_SNAPSHOT_BROKER_XTRACE\nfi" - ) -} - fn build_override_exports( explicit_env_overrides: &HashMap, restore_even_when_absent: &[&str], diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index cde30c959d..a227c94a9e 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -723,12 +723,16 @@ async fn warn_if_remote_credential_masking_unavailable( format!( "Credential masking is unavailable for this network proxy configuration on environment `{environment_id}`. This command will run with unmasked credentials." ) - } else if req.shell.shell.shell_type != ShellType::Bash { + } else if !matches!(req.shell.shell.shell_type, ShellType::Bash | ShellType::Zsh) { format!( "Credential masking is unavailable for shell `{}` on environment `{environment_id}`. This command will run with unmasked credentials.", req.shell.shell.name(), ) - } else if req.shell.use_login_shell || req.shell_snapshot.is_some() { + } else if match req.shell.shell.shell_type { + ShellType::Bash => req.shell.use_login_shell || req.shell_snapshot.is_some(), + ShellType::Zsh => !req.shell.use_login_shell || req.shell_snapshot.is_none(), + ShellType::Sh | ShellType::PowerShell | ShellType::Cmd => true, + } { format!( "Credential masking is unavailable for this shell startup on environment `{environment_id}`. This command will run with unmasked credentials." ) diff --git a/codex-rs/core/tests/suite/scenarios_credential_masking.rs b/codex-rs/core/tests/suite/scenarios_credential_masking.rs index bdc9d748cf..3bd9396de4 100644 --- a/codex-rs/core/tests/suite/scenarios_credential_masking.rs +++ b/codex-rs/core/tests/suite/scenarios_credential_masking.rs @@ -22,10 +22,12 @@ async fn remote_credential_masking_request_history( ) -> Result<()> { skip_if_no_network!(Ok(())); skip_if_sandbox!(Ok(())); - if codex_core::shell::default_user_shell().name() != "bash" { - eprintln!("skipping credential masking: this scenario requires a Bash executor"); + let shell = codex_core::shell::default_user_shell(); + if !matches!(shell.name(), "bash" | "zsh") { + eprintln!("skipping credential masking: this scenario requires a Bash or Zsh executor"); return Ok(()); } + let use_zsh_snapshot = shell.name() == "zsh" && !configured_provider; let home = Arc::new(TempDir::new()?); if let Some(warning) = codex_sandboxing::system_bwrap_warning(&PermissionProfile::workspace_write(), home.path()) @@ -86,9 +88,15 @@ credential_masking = true ) .expect("write credential masking config"); }) - .with_config(|config| { + .with_config(move |config| { configure_scenario_catalog(config); config.project_doc_max_bytes = 0; + if use_zsh_snapshot { + config + .features + .enable(Feature::ShellSnapshotV2) + .expect("enable shell snapshots"); + } }) .with_model_info_override("gpt-5.5", |model| { model.tool_mode = Some(ToolMode::Direct); @@ -118,7 +126,7 @@ credential_masking = true &json!({ "cmd": command, "environment_id": "remote", - "login": use_login_shell, + "login": use_login_shell || use_zsh_snapshot, "yield_time_ms": 1000, }) .to_string(), @@ -195,6 +203,11 @@ credential_masking = true .replace_all(&snapshot, replacement) .into_owned(); } - insta::assert_snapshot!("remote_credential_masking", snapshot); + let snapshot_name = if use_zsh_snapshot { + "remote_credential_masking_zsh" + } else { + "remote_credential_masking" + }; + insta::assert_snapshot!(snapshot_name, snapshot); Ok(()) } diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__credential_masking__remote_credential_masking_zsh.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__credential_masking__remote_credential_masking_zsh.snap new file mode 100644 index 0000000000..443d337a08 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__credential_masking__remote_credential_masking_zsh.snap @@ -0,0 +1,24 @@ +--- +source: core/tests/suite/scenarios_credential_masking.rs +expression: snapshot +--- +Scenario: A remote executor masks its credential before the command output reaches the model. + +## Window 1 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Print OPENAI_API_KEY from the executor environment. +-- request 2 (turn) -- +04:function_call/exec_command:{"cmd":"printf '%s' \\"$OPENAI_API_KEY\\"","environment_id":"remote","login":true,"yield_time_ms":1000} +05:function_call_output:Chunk ID: + Wall time: seconds + Process exited with code 0 + Original token count: + Output: + diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index b745e4e7a6..b246f87837 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -266,7 +266,7 @@ impl EnvironmentInfo { let shell = codex_shell_command::shell_detect::default_user_shell(); let credential_masking = cfg!(any(target_os = "linux", target_os = "macos")) - && shell.shell_type == ShellType::Bash; + && matches!(shell.shell_type, ShellType::Bash | ShellType::Zsh); Self { shell: shell.into(), executor_version: unknown_executor_version(), diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index f3f174e998..fced427a23 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -385,6 +385,7 @@ impl LocalProcess { &mut prepared, &self.inner.telemetry, CapturePurpose::Execution, + &self.inner.credential_brokers, ) .await?; if prepared.command.is_empty() { @@ -853,6 +854,7 @@ impl ExecBackend for LocalProcess { &mut prepared, &self.inner.telemetry, CapturePurpose::Prewarm, + &self.inner.credential_brokers, ) .await .map(|_| ()) diff --git a/codex-rs/exec-server/src/process_sandbox.rs b/codex-rs/exec-server/src/process_sandbox.rs index 74efe6c352..b9e51913d3 100644 --- a/codex-rs/exec-server/src/process_sandbox.rs +++ b/codex-rs/exec-server/src/process_sandbox.rs @@ -70,6 +70,16 @@ struct PreparedWindowsSandboxRequest { } impl PreparedExecRequest { + #[cfg(unix)] + pub(crate) fn credential_snapshot_proxy(&self, params: &ExecParams) -> Option<&NetworkProxy> { + if !params.network_proxy.as_ref()?.proxy.credential_masking + || !supports_credential_snapshot(params) + { + return None; + } + self.network_proxy_handle.as_ref()?.proxy() + } + pub(crate) fn windows_sandbox_spawn_request(&self) -> Option> { self.windows_sandbox .as_ref() @@ -85,6 +95,18 @@ impl PreparedExecRequest { } } +fn supports_credential_snapshot(params: &ExecParams) -> bool { + cfg!(any(target_os = "linux", target_os = "macos")) + && params.arg0.is_none() + && params.shell_snapshot.as_ref().is_some_and(|snapshot| { + snapshot.shell.name == "zsh" + && matches!(params.argv.as_slice(), [program, flag, _] + if program == &snapshot.shell.path + && flag == "-lc" + && Path::new(program).file_name().is_some_and(|name| name == "zsh")) + }) +} + #[tracing::instrument( name = "codex.exec_server.process_prepare_sandbox", skip_all, @@ -456,7 +478,10 @@ async fn prepare_managed_network( internal_error(format!("failed to prepare executor network proxy: {err}")) })?; if credential_masking { - if cfg!(any(target_os = "linux", target_os = "macos")) + if supports_credential_snapshot(params) { + // Requested snapshots trust startup profiles before masking their captured state. + proxy.restore_and_disable_brokered_credentials(&mut prepared.env, &mut []); + } else if cfg!(any(target_os = "linux", target_os = "macos")) && supported_shell && params.arg0.is_none() && params.shell_snapshot.is_none() diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index dfb02ff92a..55a831c7cb 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -13,6 +13,7 @@ use std::time::Duration; use codex_exec_server_protocol::JSONRPCErrorError; use codex_network_proxy::PROXY_ACTIVE_ENV_KEY; +use codex_network_proxy::RemoteCredentialBrokerCache; use codex_network_proxy::strip_managed_proxy_env; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_protocol::shell_environment; @@ -54,6 +55,9 @@ const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); const SNAPSHOT_RETRY_BACKOFF: Duration = Duration::from_secs(1); const MAX_SNAPSHOT_ATTEMPTS: usize = 3; +#[path = "shell_snapshot_credentials.rs"] +mod credentials; + #[derive(Clone, Copy)] enum SnapshotReplay { File, @@ -106,7 +110,9 @@ impl ShellSnapshotCache { prepared: &mut PreparedExecRequest, telemetry: &ExecServerTelemetry, purpose: CapturePurpose, + credential_brokers: &RemoteCredentialBrokerCache, ) -> Result, JSONRPCErrorError> { + let protected = prepared.credential_snapshot_proxy(params).is_some(); let Some(request) = params.shell_snapshot.as_ref() else { return Ok(None); }; @@ -136,7 +142,10 @@ impl ShellSnapshotCache { }; let wait_started_at = std::time::Instant::now(); - let (snapshot, attempt) = { + let (snapshot, attempt) = if protected { + // Protected captures are prepared for this launch and never enter the cache. + (Arc::new(OnceCell::new()), 1) + } else { let mut entries = self.entries.lock().await; let position = entries.iter().position(|entry| { &entry.request == request @@ -195,13 +204,12 @@ impl ShellSnapshotCache { CapturePurpose::Prewarm => "prewarm", }; let started_at = std::time::Instant::now(); - let result = - capture_snapshot(params, prepared, shell_type) - .await - .map(|mut snapshot| { - snapshot.prewarmed = prewarmed; - snapshot - }); + let result = capture_snapshot(params, prepared, shell_type, credential_brokers) + .await + .map(|mut snapshot| { + snapshot.prewarmed = prewarmed; + snapshot + }); telemetry.shell_snapshot_captured( started_at.elapsed(), result.as_ref().map(|_| ()).map_err(|(reason, _)| *reason), @@ -212,26 +220,37 @@ impl ShellSnapshotCache { ("sandbox", prepared.sandbox.as_metric_tag()), ], ); - result.map_err(|(_, error)| error) + result }; - let snapshot = match purpose { - CapturePurpose::Execution => { - snapshot - .get_or_init(|| async { - availability = "on_demand"; - capture.await.map_err(|err| { - tracing::warn!("failed to capture shell snapshot: {err:?}"); - Instant::now() + SNAPSHOT_RETRY_BACKOFF + let uncached; + let snapshot = if protected { + uncached = match capture.await { + Err(("credential_proxy", error)) => return Err(error), + result => result.map_err(|_| Instant::now()), + }; + &uncached + } else { + match purpose { + CapturePurpose::Execution => { + snapshot + .get_or_init(|| async { + availability = "on_demand"; + capture.await.map_err(|(_, err)| { + tracing::warn!("failed to capture shell snapshot: {err:?}"); + Instant::now() + SNAPSHOT_RETRY_BACKOFF + }) }) - }) - .await - } - CapturePurpose::Prewarm => { - // Leave the cell uninitialized on failure: a waiting real command - // can capture immediately, without spending its retry budget. - snapshot - .get_or_try_init(|| async { capture.await.map(Ok) }) - .await? + .await + } + CapturePurpose::Prewarm => { + // Leave the cell uninitialized on failure: a waiting real command + // can capture immediately, without spending its retry budget. + snapshot + .get_or_try_init(|| async { + capture.await.map(Ok).map_err(|(_, error)| error) + }) + .await? + } } }; let wait = wait_started_at.elapsed(); @@ -239,6 +258,9 @@ impl ShellSnapshotCache { return Ok(None); } let Ok(snapshot) = snapshot else { + if protected { + credentials::restore_fallback(prepared); + } telemetry.shell_snapshot_command(wait, availability, "fallback"); return Ok(None); }; @@ -262,6 +284,9 @@ impl ShellSnapshotCache { "cannot prepare shell snapshot transport; using normal startup" ); telemetry.shell_snapshot_command(wait, availability, "fallback"); + if protected { + credentials::restore_fallback(prepared); + } return Ok(None); } } @@ -279,13 +304,17 @@ impl ShellSnapshotCache { ) }) .collect::>(); - prepared.env.extend( - snapshot - .environment - .iter() - .map(|(name, value)| (name.clone(), value.clone())), - ); - prepared.env.extend(request_overrides); + if protected { + prepared.env.clone_from(&snapshot.environment); + } else { + prepared.env.extend( + snapshot + .environment + .iter() + .map(|(name, value)| (name.clone(), value.clone())), + ); + prepared.env.extend(request_overrides); + } apply_exec_metadata(&mut prepared.env, params.metadata.as_ref()); prepared .env @@ -309,7 +338,13 @@ impl ShellSnapshotCache { } let state_expansion = state_variables .iter() - .map(|name| format!("${{{name}}}")) + .map(|name| { + if protected { + format!("${{{name}:?missing shell snapshot chunk}}") + } else { + format!("${{{name}}}") + } + }) .collect::(); let state_variables = state_variables.join(" "); format!("eval \"unset {state_variables}\n{state_expansion}\"") @@ -324,10 +359,17 @@ impl ShellSnapshotCache { ShellType::PowerShell | ShellType::Cmd => unreachable!(), }; prepared.command[shell_start + 1] = shell_flag.to_string(); - prepared.command[shell_start + 2] = format!( - "{startup}if ! {restore} >/dev/null; then printf 'failed to restore shell snapshot\\n' >&2; fi\n{}", - params.argv[2] - ); + prepared.command[shell_start + 2] = if protected { + format!( + "{startup}if ! {restore} >/dev/null 2>&1; then printf 'failed to restore shell snapshot\\n' >&2; exit 1; fi\n{}", + params.argv[2] + ) + } else { + format!( + "{startup}if ! {restore} >/dev/null; then printf 'failed to restore shell snapshot\\n' >&2; fi\n{}", + params.argv[2] + ) + }; telemetry.shell_snapshot_command(wait, availability, "used"); Ok(reader) @@ -343,12 +385,14 @@ async fn capture_snapshot( params: &ExecParams, prepared: &PreparedExecRequest, shell_type: ShellType, + credential_brokers: &RemoteCredentialBrokerCache, ) -> CaptureResult { + let proxy = prepared.credential_snapshot_proxy(params); let mut script = snapshot_source_capture_script( shell_type, SnapshotCaptureOptions { startup: SnapshotStartup::Interactive, - declarations: false, + declarations: proxy.is_some(), environment: true, }, ) @@ -358,6 +402,16 @@ async fn capture_snapshot( invalid_params("unsupported shell snapshot script".to_string()), ) })?; + if proxy.is_some() { + // Readonly exports cannot be restored after the global zshenv runs again. + script.push_str( + r#" +for __codex_snapshot_key in ${(f)"$(builtin typeset +x)"}; do + case "${(tP)__codex_snapshot_key}" in *readonly*) exit 74 ;; esac +done +"#, + ); + } // Probe under the exact capture sandbox, including its /proc fallback and // read policy. Keep cached env replay when descriptor paths cannot be opened. let probe = if shell_type == ShellType::Sh { @@ -506,7 +560,11 @@ async fn capture_snapshot( } SnapshotReplay::Environment }; - let snapshot = parse_snapshot(shell_type, captured, params.env_policy.as_ref(), replay)?; + let snapshot = if let Some(proxy) = proxy { + credentials::prepare(params, captured, replay, proxy, credential_brokers).await? + } else { + parse_snapshot(shell_type, captured, params.env_policy.as_ref(), replay)? + }; child.preserve_helpers(); Ok(snapshot) } diff --git a/codex-rs/exec-server/src/shell_snapshot_credentials.rs b/codex-rs/exec-server/src/shell_snapshot_credentials.rs new file mode 100644 index 0000000000..8d5a17aa99 --- /dev/null +++ b/codex-rs/exec-server/src/shell_snapshot_credentials.rs @@ -0,0 +1,301 @@ +//! Prepare one uncached login capture before creating any replay transport. +//! Startup profiles are trusted with real credentials. Raw capture bytes stay executor-local; +//! unsafe executable state falls back to normal startup. + +use std::collections::BTreeSet; +use std::collections::HashMap; + +use codex_network_proxy::CredentialMaskingFailureMode; +use codex_network_proxy::NetworkProxy; +use codex_network_proxy::NetworkProxyHandle; +use codex_network_proxy::RemoteCredentialBrokerCache; +use codex_network_proxy::brokered_credential_value_env_keys; +use codex_network_proxy::is_managed_proxy_env_var; +use codex_protocol::config_types::ShellEnvironmentPolicyInherit; +use codex_protocol::shell_environment; +use codex_shell_command::shell_detect::ShellType; +use codex_shell_command::shell_snapshot::CapturedSnapshot; +use codex_shell_command::shell_snapshot::SNAPSHOT_BROKERED_UNSET_ENV_PREFIX; +use codex_shell_command::shell_snapshot::SNAPSHOT_BROKERED_VALUE_ENV_PREFIX; +use codex_shell_command::shell_snapshot::SnapshotCredentialEnvironment; +use codex_shell_command::shell_snapshot::build_brokered_credential_exports; +use codex_shell_command::shell_snapshot::prepare_snapshot_credentials; + +use super::CaptureResult; +use super::MAX_FILE_SNAPSHOT_BYTES; +use super::MAX_SNAPSHOT_BYTES; +use super::SnapshotReplay; +use super::parse_snapshot; +use crate::local_process::shell_environment_policy; +use crate::process_sandbox::PreparedExecRequest; +use crate::protocol::ExecParams; +use crate::rpc::internal_error; + +const CONFIGURED_VALUE_PREFIX: &str = "__CODEX_SHELL_SNAPSHOT_CONFIGURED_"; + +pub(super) fn restore_fallback(prepared: &mut PreparedExecRequest) { + if let Some(proxy) = prepared + .network_proxy_handle + .as_ref() + .and_then(NetworkProxyHandle::proxy) + { + for value in prepared.env.values_mut() { + proxy.restore_brokered_text(value); + } + proxy.restore_and_disable_brokered_credentials(&mut prepared.env, &mut []); + let warning = "Credential masking is unavailable for this shell snapshot; running with unmasked credentials."; + tracing::warn!("{warning}"); + if let Some(script) = prepared.command.last_mut() { + *script = format!("\\builtin printf '%s\\n' '{warning}' >&2\n{script}"); + } + } +} + +pub(super) async fn prepare( + params: &ExecParams, + output: &[u8], + replay: SnapshotReplay, + proxy: &NetworkProxy, + credential_brokers: &RemoteCredentialBrokerCache, +) -> CaptureResult { + let unsupported = || { + ( + "credential_snapshot", + internal_error("shell snapshot cannot safely replay credentials".to_string()), + ) + }; + let captured = CapturedSnapshot::parse(ShellType::Zsh, output).ok_or_else(unsupported)?; + let original = captured + .environment + .split(|byte| *byte == 0) + .filter(|entry| !entry.is_empty()) + .map(|entry| { + let (key, value) = std::str::from_utf8(entry).ok()?.split_once('=')?; + Some((key.to_string(), value.to_string())) + }) + .collect::>>() + .ok_or_else(unsupported)?; + let mut snapshot = parse_snapshot(ShellType::Zsh, output, params.env_policy.as_ref(), replay)?; + let mut policy = params + .env_policy + .as_ref() + .map(shell_environment_policy) + .unwrap_or_default(); + policy.inherit = ShellEnvironmentPolicyInherit::All; + let mut configured = shell_environment::create_env_from_vars( + std::iter::empty(), + &policy, + /*thread_id*/ None, + ); + configured.extend(params.env.clone()); + configured.retain(|key, _| { + !shell_environment::is_non_inheritable_env_var(key) + && ![ + crate::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR, + shell_environment::CODEX_THREAD_ID_ENV_VAR, + shell_environment::CODEX_TOOL_CALL_ID_ENV_VAR, + ] + .contains(&key.as_str()) + && !key.starts_with(SNAPSHOT_BROKERED_VALUE_ENV_PREFIX) + && !key.starts_with(SNAPSHOT_BROKERED_UNSET_ENV_PREFIX) + && !key.starts_with(CONFIGURED_VALUE_PREFIX) + }); + let mut restored = original.clone(); + for value in restored.values_mut() { + proxy.restore_brokered_text(value); + } + let mut context = restored.clone(); + context.extend(configured.clone()); + proxy + .bind_credential_broker_environment(credential_brokers, &context) + .await + .map_err(|_| { + ( + "credential_proxy", + internal_error("failed to bind startup credential destinations".to_string()), + ) + })?; + // Saved placeholders may belong to the final destination context selected above. + for value in restored + .values_mut() + .chain(snapshot.environment.values_mut()) + { + proxy.restore_brokered_text(value); + } + let mut discovered = restored.clone(); + let provider_context_keys = proxy + .credential_broker_environment(&discovered) + .provider_context_keys; + for (key, value) in &configured { + if !discovered.contains_key(key) || provider_context_keys.contains(key) { + discovered.insert(key.clone(), value.clone()); + } + } + proxy.apply_to_env_for_snapshot(&mut discovered); + let discovered_metadata = proxy.credential_broker_environment(&discovered); + snapshot.environment.extend(configured.clone()); + snapshot.environment.retain(|key, _| { + !key.starts_with(SNAPSHOT_BROKERED_VALUE_ENV_PREFIX) + && !key.starts_with(SNAPSHOT_BROKERED_UNSET_ENV_PREFIX) + && !key.starts_with(CONFIGURED_VALUE_PREFIX) + }); + proxy.apply_to_env(&mut snapshot.environment); + proxy + .handle_unmasked_credentials( + &mut snapshot.environment, + CredentialMaskingFailureMode::Warn, + ) + .map_err(|_| { + ( + "credential_proxy", + internal_error("failed to prepare startup credentials".to_string()), + ) + })?; + let allowed_metadata = proxy.credential_broker_environment(&snapshot.environment); + let mut credential_keys = discovered_metadata.credential_keys.clone(); + credential_keys.extend( + restored + .iter() + .filter(|(key, value)| { + !value.is_empty() + && discovered_metadata.provider_keys.contains(key) + && !discovered_metadata.provider_context_keys.contains(key) + }) + .map(|(key, _)| key.clone()), + ); + let mut allowed_keys = allowed_metadata.credential_keys; + allowed_keys.extend( + credential_keys + .iter() + .filter(|key| { + snapshot + .environment + .get(*key) + .is_some_and(|value| !value.is_empty()) + }) + .cloned(), + ); + // Proxy state belongs to the launcher. Linux rewrites its endpoints inside each + // sandbox network namespace after snapshot preparation. + let managed_proxy_keys = snapshot + .environment + .iter() + .filter(|(key, value)| is_managed_proxy_env_var(key, value)) + .map(|(key, _)| key.clone()) + .collect::>(); + let is_managed_proxy_key = + |key: &str| managed_proxy_keys.contains(key) || is_managed_proxy_env_var(key, ""); + let mut replay_environment = snapshot.environment.clone(); + replay_environment.retain(|key, _| !is_managed_proxy_key(key)); + let prepared = prepare_snapshot_credentials( + &captured, + SnapshotCredentialEnvironment { + original: &original, + restored: &restored, + configured: &configured, + discovered: &discovered, + allowed: &replay_environment, + is_allowed_unset: &|key| { + !is_managed_proxy_key(key) + && (params.env_policy.is_none() + || shell_environment::create_env_from_vars( + [(key.to_string(), String::new())], + &policy, + /*thread_id*/ None, + ) + .contains_key(key)) + }, + brokered_keys: &credential_keys, + brokered_alias_keys: &[], + allowed_brokered_keys: &allowed_keys, + }, + |text| proxy.virtualize_brokered_text(text, &snapshot.environment), + ) + .ok_or_else(unsupported)?; + if !prepared.rejected_alias_keys.is_empty() { + return Err(unsupported()); + } + snapshot.environment.extend(prepared.aliases); + let keys = brokered_credential_value_env_keys(&discovered) + .into_iter() + .chain(credential_keys) + .chain(discovered_metadata.context_keys) + .chain(discovered_metadata.provider_keys) + .filter(|key| !is_managed_proxy_key(key)) + .collect::>(); + snapshot.environment.retain(|key, _| { + !key.starts_with(SNAPSHOT_BROKERED_VALUE_ENV_PREFIX) + && !key.starts_with(SNAPSHOT_BROKERED_UNSET_ENV_PREFIX) + }); + let mut exports = String::new(); + for key in keys { + if key.is_empty() + || !key.bytes().enumerate().all(|(index, byte)| { + byte == b'_' || byte.is_ascii_alphabetic() || index != 0 && byte.is_ascii_digit() + }) + { + return Err(unsupported()); + } + if let Some(value) = snapshot.environment.get(&key).cloned() { + let copy = format!("{SNAPSHOT_BROKERED_VALUE_ENV_PREFIX}{key}"); + snapshot.environment.insert(copy.clone(), value); + // Login replay must also restore values cleared by the global zshenv. + exports.push_str(&format!( + "\\builtin export {key}=\"${{{copy}}}\" || exit 1\n" + )); + } else { + snapshot.environment.insert( + format!("{SNAPSHOT_BROKERED_UNSET_ENV_PREFIX}{key}"), + "1".to_string(), + ); + exports.push_str(&format!("\\builtin unset {key} || exit 1\n")); + } + } + // Restore explicit overrides after captured declarations, including empty values. + let mut overrides = String::new(); + for (index, key) in configured + .keys() + .filter(|key| !is_managed_proxy_key(key)) + .collect::>() + .into_iter() + .enumerate() + { + if key.is_empty() + || !key.bytes().enumerate().all(|(index, byte)| { + byte == b'_' || byte.is_ascii_alphabetic() || index != 0 && byte.is_ascii_digit() + }) + { + continue; + } + if let Some(value) = snapshot.environment.get(key).cloned() { + let copy = format!("{CONFIGURED_VALUE_PREFIX}{index}"); + snapshot.environment.insert(copy.clone(), value); + overrides.push_str(&format!( + "if [ -z \"${{{copy}+x}}\" ]; then exit 1; fi\n\\builtin export {key}=\"${{{copy}}}\" || exit 1\n\\builtin unset {copy} || exit 1\n" + )); + } + } + // Keep replay guards in the bounded snapshot transport, not the -c argument. + let before = + build_brokered_credential_exports(&snapshot.environment, /*remove_copies*/ false); + let after = + build_brokered_credential_exports(&snapshot.environment, /*remove_copies*/ true); + snapshot.state = format!( + "{before}\n{exports}{}\n{overrides}\n\\builtin test \"$?\" -eq 0 || exit 1\n{after}\n", + prepared.script + ); + let environment_bytes = snapshot + .environment + .iter() + .map(|(key, value)| key.len() + value.len() + 2) + .sum::(); + let state_limit = if snapshot.file_source { + MAX_FILE_SNAPSHOT_BYTES + } else { + MAX_SNAPSHOT_BYTES.saturating_sub(environment_bytes) + }; + if environment_bytes > MAX_SNAPSHOT_BYTES || snapshot.state.len() > state_limit { + return Err(unsupported()); + } + Ok(snapshot) +} diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index 66dc90e23e..7f5298fcb5 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -123,10 +123,11 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( } else { CapturePurpose::Execution }; + let credential_brokers = codex_network_proxy::RemoteCredentialBrokerCache::default(); let (first, second) = tokio::join!( biased; - cache.prepare(¶ms, &mut prepared, &telemetry, purpose), - cache.prepare(¶ms, &mut concurrent, &telemetry, CapturePurpose::Execution), + cache.prepare(¶ms, &mut prepared, &telemetry, purpose, &credential_brokers), + cache.prepare(¶ms, &mut concurrent, &telemetry, CapturePurpose::Execution, &credential_brokers), ); let first = if prewarming { first.expect_err("prewarm must report failure without caching it"); @@ -162,6 +163,7 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( &mut prepared, &telemetry, CapturePurpose::Execution, + &credential_brokers, ) .await .expect("capture must stay cached during backoff"); diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 3c74a7e176..7e82dfc912 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -495,21 +495,43 @@ async fn exec_metadata_overrides_policy_and_cached_snapshot( } #[cfg(unix)] -#[test_case(false, "bash", &["-lc", "-lc"]; "unmasked")] -#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "bash", &["-c", "-c", "-lc"]; "masked_bash"))] -#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "bash", &["-lc", "-lc"]; "bash_profile_fallback"))] -#[cfg_attr(target_os = "macos", test_case(true, "zsh", &["-lc", "-lc"]; "zsh_profile_gateway_fallback"))] +#[derive(Clone, Copy)] +enum CredentialStartup { + Inherited, + Profile, + UnsupportedProfile, + OrdinaryLogin, +} + +#[cfg(unix)] +#[test_case(false, "bash", &["-lc", "-lc"], CredentialStartup::Inherited; "unmasked")] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "bash", &["-c", "-c", "-lc"], CredentialStartup::Inherited; "masked_bash"))] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "bash", &["-lc", "-lc"], CredentialStartup::Inherited; "bash_profile_fallback"))] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "zsh", &["-lc", "-lc"], CredentialStartup::OrdinaryLogin; "zsh_profile_gateway_fallback"))] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "zsh", &["-lc", "-lc"], CredentialStartup::Inherited; "masked_zsh_inherited_credentials"))] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "zsh", &["-lc", "-lc"], CredentialStartup::Profile; "masked_zsh_profile_credentials"))] +#[cfg_attr(any(target_os = "linux", target_os = "macos"), test_case(true, "zsh", &["-lc", "-lc"], CredentialStartup::UnsupportedProfile; "zsh_unsupported_profile_fallback"))] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(remote_exec_server)] async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( credential_masking: bool, shell_name: &str, shell_args: &[&str], + credential_startup: CredentialStartup, ) -> Result<()> { const API_KEY: &str = "sk-proj-executor-only-0000000000000000000000000000000000"; + const PROFILE_GH_TOKEN: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + const REQUEST_GH_TOKEN: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + const CUSTOM_HOST: &str = "chatgpt.com"; let shell_path = format!("/bin/{shell_name}"); - let api_host = if shell_name == "zsh" { - "chatgpt.com" + if shell_name == "zsh" && !std::path::Path::new(&shell_path).exists() { + eprintln!("skipping zsh test: {shell_path} is unavailable"); + return Ok(()); + } + let unset_destinations = + shell_name == "zsh" && matches!(credential_startup, CredentialStartup::Inherited); + let api_host = if shell_name == "zsh" && !unset_destinations { + CUSTOM_HOST } else { "api.openai.com" }; @@ -524,16 +546,31 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( return Ok(()); } let upstream = credential_masking - .then(|| tls_proxy::TlsInterceptingProxy::start(shell_args.len())) + .then(|| { + tls_proxy::TlsInterceptingProxy::start( + shell_args.len() * if unset_destinations { 2 } else { 1 }, + ) + }) .transpose()?; let mut executor_env = vec![("OPENAI_API_KEY", API_KEY.to_string())]; if let Some(upstream) = &upstream { let ca_path = home.path().join("upstream-ca.pem"); std::fs::write(&ca_path, &upstream.ca_cert_pem)?; + let inherited_api_host = if unset_destinations { + CUSTOM_HOST + } else { + "api.openai.com" + }; executor_env.extend([ ("AUTH_ALIAS", format!("Bearer {API_KEY}")), - ("OPENAI_BASE_URL", "https://api.openai.com/v1".to_string()), - ("GH_HOST", String::new()), + ( + "OPENAI_BASE_URL", + format!("https://{inherited_api_host}/v1"), + ), + ( + "GH_HOST", + if unset_destinations { CUSTOM_HOST } else { "" }.to_string(), + ), ( "GH_ENTERPRISE_TOKEN", "executor-unbound-enterprise-token".to_string(), @@ -551,6 +588,14 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( ("no_proxy", String::new()), ]); } + if matches!(credential_startup, CredentialStartup::Profile) { + executor_env.extend([("GH_TOKEN", String::new()), ("GH_ALIAS", String::new())]); + for (key, value) in &mut executor_env { + if matches!(*key, "OPENAI_API_KEY" | "AUTH_ALIAS") { + value.clear(); + } + } + } let mut server = common::exec_server::exec_server_with_env(executor_env, &[]).await?; let environment = Environment::create_for_tests(Some(server.websocket_url().to_string()))?; let backend = environment.get_exec_backend(); @@ -561,26 +606,56 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( &profile_bin.join("profile_command"), "#!/bin/sh\nprintf helper\n", )?; + let shell_path = if shell_name == "zsh" { + // Observe the launcher's proxy before replay and simulate a global zshenv export. + let path = home.path().join("zsh"); + codex_utils_cargo_bin::write_executable( + &path, + "#!/bin/sh\nexport GITHUB_TOKEN=global-test-token\nexport TEST_LIVE_HTTP_PROXY=\"$HTTP_PROXY\"\nexec /bin/zsh \"$@\"\n", + )?; + path.to_string_lossy().into_owned() + } else { + shell_path + }; let profile = "printf '%s\\n' \"$HTTP_PROXY\" >> \"$HOME/captures\"\ntest \"$CODEX_NETWORK_PROXY_ACTIVE\" = 1 || exit 41\nexport PROFILE_ALLOWED=profile\nprofile_helper() { profile_command; }\n"; std::fs::write(home.path().join(".bashrc"), profile)?; std::fs::write( home.path().join(".bash_profile"), "export PATH=\"$HOME/profile-bin:$PATH\"\n", )?; + let profile_credentials = match credential_startup { + CredentialStartup::Inherited | CredentialStartup::OrdinaryLogin => String::new(), + CredentialStartup::Profile => format!( + "export OPENAI_API_KEY={API_KEY}\nexport AUTH_ALIAS=\"Bearer $OPENAI_API_KEY\"\nexport GH_TOKEN={PROFILE_GH_TOKEN}\nexport GH_ALIAS=\"Bearer $GH_TOKEN\"\nfor index in {{1..256}}; do export \"AUTH_SIZE_$index=Bearer $OPENAI_API_KEY\"; done\nunset index\n" + ), + CredentialStartup::UnsupportedProfile => { + format!("credential_literal() {{ printf '%s' '{API_KEY}'; }}\n") + } + }; + let profile_destinations = if unset_destinations { + "unset OPENAI_BASE_URL GH_HOST\n".to_string() + } else { + format!("export OPENAI_BASE_URL=https://{api_host}/v1\n") + }; std::fs::write( home.path().join(".zshenv"), format!( - "export PATH=\"$HOME/profile-bin:$PATH\"\nexport OPENAI_BASE_URL=https://{api_host}/v1\n{profile}" + "unset GITHUB_TOKEN\nexport PATH=\"$HOME/profile-bin:$PATH\"\n{profile_destinations}typeset -ux FOO=profile\nexport PROFILE_OVERRIDE=profile\nprofile_local=local\ncd \"$HOME/profile-bin\"\numask 027\n{profile_credentials}test \"$OPENAI_API_KEY\" = '{API_KEY}' || exit 48\n{profile}" ), )?; + std::fs::write(home.path().join(".zshrc"), ": > \"$HOME/zshrc-loaded\"\n")?; let mut policy = ExecEnvPolicy { inherit: ShellEnvironmentPolicyInherit::All, ignore_default_excludes: true, exclude: Vec::new(), - r#set: HashMap::from([( - "HOME".to_string(), - home.path().to_string_lossy().into_owned(), - )]), + r#set: HashMap::from([ + ( + "HOME".to_string(), + home.path().to_string_lossy().into_owned(), + ), + ("DEBUG_FLAG".to_string(), "debug".to_string()), + ("PROFILE_OVERRIDE".to_string(), "policy".to_string()), + ]), include_only: vec![ "HOME".to_string(), "PATH".to_string(), @@ -591,7 +666,9 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( if credential_masking { policy.include_only.extend( [ + "OPENAI_BASE_URL", "AUTH_ALIAS", + "GITHUB_TOKEN", "GH_HOST", "GH_ENTERPRISE_TOKEN", "CURL_CA_BUNDLE", @@ -599,25 +676,70 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( .map(str::to_string), ); } + if matches!(credential_startup, CredentialStartup::Profile) { + policy + .include_only + .extend(["GH_TOKEN", "GH_ALIAS", "AUTH_SIZE_*"].map(str::to_string)); + } let mut proxy_settings = NetworkProxyConfig { enabled: true, ..NetworkProxyConfig::default() }; - proxy_settings.set_allowed_domains(vec![api_host.to_string()]); + let mut allowed_domains = vec![api_host.to_string()]; + if unset_destinations { + allowed_domains.push(CUSTOM_HOST.to_string()); + } + proxy_settings.set_allowed_domains(allowed_domains); let mut proxy_config = RemoteNetworkProxyConfig::from_effective_config(&proxy_settings)?; proxy_config.credential_masking = credential_masking; - let mut proxy_addresses = Vec::new(); let mut saved_api_key = None; for (attempt, &shell_args) in shell_args.iter().enumerate() { - let use_snapshot = shell_name == "bash" && shell_args == "-lc"; - let expect_masking = credential_masking && shell_name == "bash" && !use_snapshot; + let use_snapshot = + shell_args == "-lc" && !matches!(credential_startup, CredentialStartup::OrdinaryLogin); + let expect_snapshot_fallback = + matches!(credential_startup, CredentialStartup::UnsupportedProfile) + || matches!(credential_startup, CredentialStartup::Profile) && attempt == 1; + let expect_masking = credential_masking + && (shell_name == "bash" && !use_snapshot + || shell_name == "zsh" && use_snapshot && !expect_snapshot_fallback); let expect_profile = use_snapshot || shell_name == "zsh"; + let request_foo = if attempt == 0 { "request" } else { "" }; + let mut env = HashMap::from([("FOO".to_string(), request_foo.to_string())]); + if shell_name == "zsh" && expect_masking { + env.insert("HTTP_PROXY".to_string(), "http://127.0.0.1:1".to_string()); + } + if matches!(credential_startup, CredentialStartup::Profile) { + let (key, value) = if attempt == 0 { + ("GH_TOKEN", REQUEST_GH_TOKEN) + } else { + ("OPENAI_API_KEY", "") + }; + env.insert(key.to_string(), value.to_string()); + } let command = if expect_profile { "profile_helper; printf '|%s|%s|%s' \"$PROFILE_ALLOWED\" \"$CODEX_NETWORK_PROXY_ACTIVE\" \"$HTTP_PROXY\"" } else { "printf 'masked|%s|%s' \"$CODEX_NETWORK_PROXY_ACTIVE\" \"$HTTP_PROXY\"" }; + let command = if shell_name == "zsh" && expect_masking { + format!( + "test \"$HTTP_PROXY\" = \"$TEST_LIVE_HTTP_PROXY\" && test \"$HTTPS_PROXY\" = \"$TEST_LIVE_HTTP_PROXY\" || exit 51; test \"${{DEBUG_FLAG+x}}\" = '' && test \"${{GITHUB_TOKEN+x}}\" = '' && test \"$PROFILE_OVERRIDE\" = policy && test \"${{FOO-unset}}\" = '{request_foo}' && test -e \"$HOME/zshrc-loaded\" || exit 47; {command}" + ) + } else if matches!(credential_startup, CredentialStartup::OrdinaryLogin) { + format!( + "test \"$profile_local\" = local && test \"$PWD\" = \"$HOME/profile-bin\" && test \"$(umask -S)\" = 'u=rwx,g=rx,o=' && test ! -e \"$HOME/zshrc-loaded\" || exit 49; {command}" + ) + } else { + command.to_string() + }; + let command = if matches!(credential_startup, CredentialStartup::Profile) && attempt == 0 { + format!( + "test -n \"$GH_TOKEN\" && test \"$GH_ALIAS\" = \"Bearer $GH_TOKEN\" && test \"$AUTH_SIZE_1\" = \"Bearer $OPENAI_API_KEY\" && test \"$AUTH_SIZE_256\" = \"Bearer $OPENAI_API_KEY\" || exit 50; case \"$GH_TOKEN|$GH_ALIAS\" in *'{PROFILE_GH_TOKEN}'*|*'{REQUEST_GH_TOKEN}'*) exit 50 ;; esac; {command}" + ) + } else { + command + }; let command = if credential_masking { format!( "{command}; test \"$GH_ENTERPRISE_TOKEN\" = executor-unbound-enterprise-token || exit 45; test \"$AUTH_ALIAS\" = \"Bearer $OPENAI_API_KEY\" || exit 42; printf '|%s' \"$OPENAI_API_KEY\"; curl --silent --show-error --fail --max-time 5 --noproxy '' --config \"$HOME/curl.conf\" \"${{OPENAI_BASE_URL:-https://api.openai.com/v1}}/models\" >/dev/null" @@ -625,6 +747,13 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( } else { format!("{command}; test \"$OPENAI_API_KEY\" = '{API_KEY}' || exit 46") }; + let command = if unset_destinations { + format!( + "test \"${{OPENAI_BASE_URL+x}}\" = '' && test \"${{GH_HOST+x}}\" = '' || exit 52; {command} && curl --silent --show-error --fail --max-time 5 --noproxy '' --config \"$HOME/curl.conf\" https://{CUSTOM_HOST}/v1/models >/dev/null" + ) + } else { + command + }; let command = if credential_masking && attempt == 0 { format!( "printf 'header = \"Authorization: %s\"\\n' \"$AUTH_ALIAS\" > \"$HOME/curl.conf\" || exit 44; {command}" @@ -632,8 +761,9 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( } else { command }; - let started = backend - .start(ExecParams { + let started = timeout( + Duration::from_secs(30), + backend.start(ExecParams { metadata: Default::default(), process_id: ProcessId::from(format!("managed-snapshot-{attempt}")), argv: vec![shell_path.clone(), shell_args.to_string(), command], @@ -646,7 +776,7 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( path: shell_path.clone(), }, }), - env: HashMap::new(), + env, tty: false, pipe_stdin: false, arg0: None, @@ -664,11 +794,20 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( format!("managed-snapshot-{attempt}"), ), ), - }) - .await?; + }), + ) + .await??; let (stdout, stderr, status, closed) = collect_process_output_from_events(started.process).await?; - assert_eq!((stderr, status, closed), (String::new(), Some(0), true)); + let expected_stderr = if expect_snapshot_fallback { + "Credential masking is unavailable for this shell snapshot; running with unmasked credentials.\n" + } else { + "" + }; + assert_eq!( + (stderr.as_str(), status, closed), + (expected_stderr, Some(0), true) + ); let prefix = if expect_profile { "helper|profile|1|" } else { @@ -706,26 +845,46 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( Some(&vec![format!("Bearer {API_KEY}")]), ) ); + if unset_destinations { + let request = upstream.next_request()?; + assert_eq!( + ( + request.connect_authority.as_str(), + request.request_line.as_str(), + request.headers.get("authorization"), + ), + ( + format!("{CUSTOM_HOST}:443").as_str(), + "GET /v1/models HTTP/1.1", + Some(&vec![format!("Bearer {saved_api_key}")]), + ) + ); + } proxy_address } else { proxy_state }; assert!(proxy_address.starts_with("http://127.0.0.1:")); - proxy_addresses.push(proxy_address.to_string()); } - let expected_captures = if shell_name == "zsh" { - Some(format!("{}\n", proxy_addresses.join("\n"))) - } else { - shell_args - .iter() - .position(|flag| *flag == "-lc") - .map(|attempt| format!("{}\n", proxy_addresses[attempt])) - }; - if let Some(expected_captures) = expected_captures { - assert_eq!( - std::fs::read_to_string(home.path().join("captures"))?, - expected_captures + let expected_capture_count = + if matches!(credential_startup, CredentialStartup::UnsupportedProfile) { + Some(shell_args.len() * 2) + } else if matches!(credential_startup, CredentialStartup::Profile) { + Some(shell_args.len() + 1) + } else if shell_name == "zsh" { + Some(shell_args.len()) + } else { + shell_args.iter().position(|flag| *flag == "-lc").map(|_| 1) + }; + if let Some(expected_capture_count) = expected_capture_count { + // Capture and replay can have different proxy ports in separate Linux namespaces. + let captures = std::fs::read_to_string(home.path().join("captures"))?; + assert_eq!(captures.lines().count(), expected_capture_count); + assert!( + captures + .lines() + .all(|address| address.starts_with("http://127.0.0.1:")) ); } else { assert!(!home.path().join("captures").exists()); diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index ac69847ae4..7b968a5ba4 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -927,6 +927,18 @@ impl NetworkProxy { self.state.current_cfg().await } + /// Selects cached credentials for the startup environment without changing listeners or policy. + /// Restore any existing placeholders before binding, then prepare the environment for execution. + pub async fn bind_credential_broker_environment( + &self, + credential_brokers: &crate::RemoteCredentialBrokerCache, + environment: &HashMap, + ) -> Result<()> { + self.state + .bind_credential_broker_environment(credential_brokers, environment) + .await + } + /// Revision of credential configuration, excluding ordinary network policy changes. pub fn credential_broker_config_revision(&self) -> u64 { self.state.credential_broker_config_revision() @@ -1253,7 +1265,8 @@ impl NetworkProxy { let environment_id = self .execution_scope .as_ref() - .map(|scope| scope.environment_id.as_str()); + .map(|scope| scope.environment_id.as_str()) + .or(self.state.environment_id()); self.state .virtualize_snapshot_credentials(env, environment_id); } @@ -1646,6 +1659,7 @@ impl NetworkProxy { *active_route = Some(Arc::clone(&route)); drop(active_route); return Ok(NetworkProxyHandle::windows_shared( + self.clone(), route, Arc::clone(&windows_runtime.active_route), Arc::clone(&self.environment_proxies), @@ -1720,6 +1734,7 @@ impl NetworkProxy { Ok(NetworkProxyHandle { runtime: Some(ProxyRuntime::Listeners(listeners)), + proxy: Some(self.clone()), environment_proxies: Some(Arc::clone(&self.environment_proxies)), #[cfg(target_os = "windows")] windows_active_route: None, @@ -1729,6 +1744,7 @@ impl NetworkProxy { pub struct NetworkProxyHandle { runtime: Option, + proxy: Option, environment_proxies: Option>>>, #[cfg(target_os = "windows")] windows_active_route: Option>>>>, @@ -1738,6 +1754,7 @@ impl NetworkProxyHandle { fn noop() -> Self { Self { runtime: None, + proxy: None, environment_proxies: None, #[cfg(target_os = "windows")] windows_active_route: None, @@ -1746,17 +1763,24 @@ impl NetworkProxyHandle { #[cfg(target_os = "windows")] fn windows_shared( + proxy: NetworkProxy, route: Arc, active_route: Arc>>>, environment_proxies: Arc>>, ) -> Self { Self { runtime: Some(ProxyRuntime::SharedIngress { route }), + proxy: Some(proxy), environment_proxies: Some(environment_proxies), windows_active_route: Some(active_route), } } + /// Returns the running proxy. Disabled proxy handles return `None`. + pub fn proxy(&self) -> Option<&NetworkProxy> { + self.proxy.as_ref() + } + pub async fn wait(mut self) -> Result<()> { let result = match self.runtime.as_mut() { Some(runtime) => runtime.wait().await, @@ -1857,6 +1881,7 @@ mod tests { }); let handle = NetworkProxyHandle { runtime: Some(ProxyRuntime::Listeners(listeners)), + proxy: None, environment_proxies: Some(Arc::new(Mutex::new(HashMap::new()))), #[cfg(target_os = "windows")] windows_active_route: None, diff --git a/codex-rs/network-proxy/src/remote_config_tests.rs b/codex-rs/network-proxy/src/remote_config_tests.rs index e589868be7..0379aff882 100644 --- a/codex-rs/network-proxy/src/remote_config_tests.rs +++ b/codex-rs/network-proxy/src/remote_config_tests.rs @@ -226,8 +226,8 @@ fn launch_config_materializes_audit_and_execution_attribution() { assert_eq!(state.execution_id().as_deref(), Some("execution-1")); } -#[test] -fn remote_credential_mappings_follow_the_environment_and_broker_config() { +#[tokio::test] +async fn remote_credential_mappings_follow_the_environment_and_broker_config() { const TOKEN: &str = "ghp_abcdefghijklmnopqrstuvwxyz0123456789"; let brokers = crate::RemoteCredentialBrokerCache::default(); let launch = |environment_id: &str, @@ -283,13 +283,20 @@ fn remote_credential_mappings_follow_the_environment_and_broker_config() { "thread", NetworkMode::Full, /*allow_local_binding*/ false, - &next_env, + &env, ); + let listener = next.clone(); + next.bind_credential_broker_environment(&brokers, &next_env) + .await + .expect("bind startup credential context"); next.virtualize_child_credentials(&mut next_env); let mut headers = rama_http::HeaderMap::from_iter([( rama_http::header::AUTHORIZATION, rama_http::HeaderValue::from_str(&format!("Bearer {placeholder}")).unwrap(), )]); + let mut listener_headers = headers.clone(); + listener.inject_request_credentials("github.enterprise.example", &mut listener_headers); + assert_eq!(listener_headers, headers); first.inject_request_credentials("github.enterprise.example", &mut headers); assert_eq!( headers[rama_http::header::AUTHORIZATION].to_str().unwrap(), diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 63d50400b7..5f518b71f2 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -3,6 +3,7 @@ use crate::config::NetworkMode; use crate::config::NetworkProxyConfig; use crate::config::ValidatedUnixSocketPath; use crate::credential_broker::CredentialBroker; +use crate::credential_broker::env_value; use crate::domain_matcher::DomainPatternSet; use crate::mitm::MitmState; use crate::mitm_hook::HookEvaluation; @@ -274,7 +275,8 @@ pub struct NetworkProxyState { pub(crate) policy_audit_observer: Option, pub(crate) launch_span_context: Option, pub(crate) process_log_metadata: crate::NetworkProxyProcessLogMetadata, - credential_broker: CredentialBroker, + /// Shared by this proxy's listeners; rebinding leaves other proxies' cached brokers intact. + credential_broker: Arc>, audit_metadata: NetworkProxyAuditMetadata, execution_attributions: Arc>>, environment_id: Option>, @@ -367,7 +369,7 @@ impl NetworkProxyState { ) }; if let Some(broker) = cached_broker { - proxy_state.credential_broker = broker; + proxy_state.credential_broker = Arc::new(std::sync::RwLock::new(broker)); } Ok(proxy_state) } @@ -420,7 +422,7 @@ impl NetworkProxyState { let credential_broker = CredentialBroker::new(state.config.credential_broker); credential_broker.configure(&local_binding_policy.resolved_config(&state.config)); Self { - credential_broker, + credential_broker: Arc::new(std::sync::RwLock::new(credential_broker)), state: Arc::new(RwLock::new(state)), local_binding_policy, proxy_private_ips_via_upstream: false, @@ -519,6 +521,45 @@ impl NetworkProxyState { &self.audit_metadata } + fn credential_broker(&self) -> CredentialBroker { + self.credential_broker + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + pub(crate) async fn bind_credential_broker_environment( + &self, + credential_brokers: &crate::RemoteCredentialBrokerCache, + environment: &HashMap, + ) -> Result<()> { + self.reload_if_needed().await?; + let mut state = self.state.write().await; + let mut config = self.local_binding_policy.resolved_config(&state.config); + config.configure_credential_broker_environment(environment); + // Captures are complete environments: absent hints are removals, not launch fallbacks. + config.credential_broker_context = self + .credential_broker_environment(environment) + .provider_context_keys + .into_iter() + .filter_map(|key| env_value(environment, &key).map(|value| (key, value.to_string()))) + .collect::>() + .into(); + let broker = credential_brokers + .broker_if_enabled( + self.environment_id(), + self.audit_metadata.conversation_id.as_deref(), + &config, + ) + .context("credential brokerage is disabled for the startup environment")?; + state.config.credential_broker_context = config.credential_broker_context; + *self + .credential_broker + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = broker; + Ok(()) + } + pub fn virtualize_child_credentials(&self, env: &mut HashMap) { self.virtualize_child_credentials_for_environment(env, self.environment_id()); } @@ -531,10 +572,13 @@ impl NetworkProxyState { let parent_env = std::env::vars_os() .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))) .collect(); - self.credential_broker - .discover_parent_credentials_for_environment(&parent_env, env, environment_id); - self.credential_broker - .virtualize_child_env_for_environment(env, environment_id); + let credential_broker = self.credential_broker(); + credential_broker.discover_parent_credentials_for_environment( + &parent_env, + env, + environment_id, + ); + credential_broker.virtualize_child_env_for_environment(env, environment_id); } pub(crate) fn restore_child_credentials( @@ -542,7 +586,7 @@ impl NetworkProxyState { env: &mut HashMap, command: &mut [String], ) { - self.credential_broker.restore_child_env(env, command); + self.credential_broker().restore_child_env(env, command); } pub(crate) fn virtualize_snapshot_credentials( @@ -550,7 +594,7 @@ impl NetworkProxyState { env: &mut HashMap, environment_id: Option<&str>, ) { - self.credential_broker + self.credential_broker() .virtualize_snapshot_env(env, environment_id); } @@ -561,7 +605,7 @@ impl NetworkProxyState { snapshot_value: &str, environment_id: Option<&str>, ) -> bool { - self.credential_broker + self.credential_broker() .child_alias_matches(key, value, snapshot_value, environment_id) } @@ -570,7 +614,7 @@ impl NetworkProxyState { env: &mut HashMap, command: &mut [String], ) { - self.credential_broker + self.credential_broker() .restore_and_disable_child_env(env, command); } @@ -579,7 +623,7 @@ impl NetworkProxyState { text: &mut String, env: &HashMap, ) -> bool { - self.credential_broker.virtualize_text(text, env) + self.credential_broker().virtualize_text(text, env) } pub(crate) fn handle_unmasked_credentials( @@ -587,7 +631,7 @@ impl NetworkProxyState { env: &mut HashMap, mode: crate::CredentialMaskingFailureMode, ) -> anyhow::Result<()> { - self.credential_broker + self.credential_broker() .handle_unmasked_credentials(env, mode) } @@ -595,7 +639,7 @@ impl NetworkProxyState { &self, env: &HashMap, ) -> crate::CredentialBrokerEnvironment { - self.credential_broker.environment(env) + self.credential_broker().environment(env) } pub(crate) fn credential_broker_environment_for_text( @@ -603,7 +647,7 @@ impl NetworkProxyState { text: &str, env: &HashMap, ) -> crate::CredentialBrokerEnvironment { - self.credential_broker.environment_for_text(text, env) + self.credential_broker().environment_for_text(text, env) } pub(crate) fn credential_broker_source_matches_text( @@ -612,7 +656,7 @@ impl NetworkProxyState { source_value: &str, text: &str, ) -> bool { - self.credential_broker + self.credential_broker() .source_matches_text(source, source_value, text) } @@ -623,16 +667,20 @@ impl NetworkProxyState { source_env: &HashMap, is_allowed: impl Fn(&str) -> bool, ) -> bool { - self.credential_broker - .provider_sources_allowed(value, virtualized, source_env, is_allowed) + self.credential_broker().provider_sources_allowed( + value, + virtualized, + source_env, + is_allowed, + ) } pub(crate) fn restore_brokered_text(&self, text: &mut String) -> bool { - self.credential_broker.restore_text(text) + self.credential_broker().restore_text(text) } pub fn inject_request_credentials(&self, host: &str, headers: &mut rama_http::HeaderMap) { - self.credential_broker + self.credential_broker() .inject_request_headers_for_environment(host, headers, self.environment_id()); } @@ -651,7 +699,7 @@ impl NetworkProxyState { } pub(crate) fn credential_broker_config_revision(&self) -> u64 { - self.credential_broker.config_revision() + self.credential_broker().config_revision() } pub(crate) async fn current_cfg_with_brokerage_provenance( @@ -688,7 +736,7 @@ impl NetworkProxyState { Ok(mut new_state) => { { let mut guard = self.state.write().await; - self.credential_broker + self.credential_broker() .configure(&self.local_binding_policy.resolved_config(&new_state.config)); // Log policy diffs without dumping potentially sensitive config values. log_policy_changes(&guard.config, &new_state.config); @@ -710,7 +758,7 @@ impl NetworkProxyState { pub async fn replace_config_state(&self, mut new_state: ConfigState) -> Result<()> { self.reload_if_needed().await?; let mut guard = self.state.write().await; - self.credential_broker + self.credential_broker() .configure(&self.local_binding_policy.resolved_config(&new_state.config)); log_policy_changes(&guard.config, &new_state.config); new_state.blocked = guard.blocked.clone(); @@ -1016,8 +1064,9 @@ impl NetworkProxyState { } }; + let credential_broker = self.credential_broker(); if credential_destinations.iter().any(|destination| { - self.credential_broker.request_matches_hooked_host_alias( + credential_broker.request_matches_hooked_host_alias( destination, req.headers(), &hooked_hosts, @@ -1041,7 +1090,7 @@ impl NetworkProxyState { let guard = self.state.read().await; guard.mitm_hooks.contains_key(&normalized_host) }; - let protocols = self.credential_broker.host_protocols_for_environment( + let protocols = self.credential_broker().host_protocols_for_environment( &normalized_host, port, self.environment_id(), @@ -1134,7 +1183,7 @@ impl NetworkProxyState { Some(mut new_state) => { { let mut guard = self.state.write().await; - self.credential_broker + self.credential_broker() .configure(&self.local_binding_policy.resolved_config(&new_state.config)); log_policy_changes(&guard.config, &new_state.config); new_state.blocked = guard.blocked.clone(); diff --git a/codex-rs/shell-command/src/shell_snapshot.rs b/codex-rs/shell-command/src/shell_snapshot.rs index a1fa1740a1..726be2c6f6 100644 --- a/codex-rs/shell-command/src/shell_snapshot.rs +++ b/codex-rs/shell-command/src/shell_snapshot.rs @@ -10,6 +10,8 @@ mod exports; mod literals; #[path = "shell_snapshot_render.rs"] mod render; +#[path = "shell_snapshot_restore.rs"] +mod restore; pub use capture::CapturedSnapshot; pub use capture::CapturedStartupEnvironment; @@ -19,6 +21,9 @@ pub use capture::snapshot_source_capture_script; pub use credentials::PreparedSnapshot; pub use credentials::SnapshotCredentialEnvironment; pub use credentials::prepare_snapshot_credentials; +pub use restore::SNAPSHOT_BROKERED_UNSET_ENV_PREFIX; +pub use restore::SNAPSHOT_BROKERED_VALUE_ENV_PREFIX; +pub use restore::build_brokered_credential_exports; use capture::BASH_SH_SNAPSHOT_HEADER; diff --git a/codex-rs/shell-command/src/shell_snapshot_credentials.rs b/codex-rs/shell-command/src/shell_snapshot_credentials.rs index 8a74c3cc89..2c11044750 100644 --- a/codex-rs/shell-command/src/shell_snapshot_credentials.rs +++ b/codex-rs/shell-command/src/shell_snapshot_credentials.rs @@ -8,6 +8,7 @@ use super::render; use super::render::Export; use super::render::Value; use super::render::ValuePart; +use crate::shell_detect::ShellType; use std::collections::HashMap; /// Environment views used to render protected credential references into a shell snapshot. @@ -106,7 +107,9 @@ pub fn prepare_snapshot_credentials( let next_credential = credential_keys .iter() .filter_map(|credential_key| { - let dummy = discovered.get(*credential_key)?; + let dummy = discovered + .get(*credential_key) + .filter(|dummy| !dummy.is_empty())?; let real = real_credential_value(credential_key)?; let allowed_key = allowed_brokered_keys .iter() @@ -296,11 +299,24 @@ pub fn prepare_snapshot_credentials( return Some(Export::Captured(line)); } return Some(Export::Assignment { - declaration, + declaration: if captured.shell_type == ShellType::Zsh + && configured.contains_key(key) + { + format!("export {key}").into() + } else { + declaration.into() + }, value: assignment, }); } + // Explicit scalar overrides must not inherit captured coercion attributes. + if captured.shell_type == ShellType::Zsh + && configured.contains_key(key) + && matches!(export.value, ExportValue::Plain) + { + return None; + } Some(Export::Captured(line)) }) .collect::>(); diff --git a/codex-rs/shell-command/src/shell_snapshot_render.rs b/codex-rs/shell-command/src/shell_snapshot_render.rs index 1d8e48ccf2..e4e06638ae 100644 --- a/codex-rs/shell-command/src/shell_snapshot_render.rs +++ b/codex-rs/shell-command/src/shell_snapshot_render.rs @@ -4,6 +4,7 @@ use super::capture::CapturedSnapshot; use codex_protocol::config_types::ShellEnvironmentPolicy; use codex_protocol::shell_environment::create_env_from_vars; +use std::borrow::Cow; impl CapturedSnapshot<'_> { /// Render native functions, options, and aliases; the executor restores environment separately. @@ -45,7 +46,7 @@ pub(super) enum ValuePart { pub(super) enum Export<'a> { Captured(&'a str), Assignment { - declaration: &'a str, + declaration: Cow<'a, str>, value: Value, }, Array { diff --git a/codex-rs/shell-command/src/shell_snapshot_restore.rs b/codex-rs/shell-command/src/shell_snapshot_restore.rs new file mode 100644 index 0000000000..0c7dcb7250 --- /dev/null +++ b/codex-rs/shell-command/src/shell_snapshot_restore.rs @@ -0,0 +1,80 @@ +//! Restore brokered values after shell startup and snapshot replay. + +use std::collections::HashMap; + +pub const SNAPSHOT_BROKERED_VALUE_ENV_PREFIX: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_BROKERED_VALUE_"; +pub const SNAPSHOT_BROKERED_UNSET_ENV_PREFIX: &str = "CODEX_NETWORK_PROXY_SNAPSHOT_BROKERED_UNSET_"; + +/// Restore protected credential copies and unsets after shell startup. +/// Keep the copies for another replay stage, or remove them before the user command. +pub fn build_brokered_credential_exports( + env: &HashMap, + remove_copies: bool, +) -> String { + let mut value_copies = env + .keys() + .filter_map(|copy_key| { + let key = copy_key.strip_prefix(SNAPSHOT_BROKERED_VALUE_ENV_PREFIX)?; + is_valid_shell_variable_name(key).then_some((key, copy_key)) + }) + .collect::>(); + value_copies.sort_unstable(); + let value_exports = value_copies + .into_iter() + .map(|(key, copy_key)| { + let restore = format!( + "if [ -z \"${{{copy_key}+x}}\" ]; then exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ] && [ \"${{{key}}}\" != \"${{{copy_key}}}\" ]; then export {key}=\"${{{copy_key}}}\" || exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ] && [ \"${{{key}}}\" != \"${{{copy_key}}}\" ]; then exit 1; fi" + ); + if remove_copies { + format!("{restore}\nunset {copy_key} || exit 1") + } else { + restore + } + }) + .collect::>() + .join("\n"); + let mut unset_markers = env + .keys() + .filter_map(|marker_key| { + let key = marker_key.strip_prefix(SNAPSHOT_BROKERED_UNSET_ENV_PREFIX)?; + is_valid_shell_variable_name(key).then_some((key, marker_key)) + }) + .collect::>(); + unset_markers.sort_unstable(); + let unset_exports = unset_markers + .into_iter() + .map(|(key, marker_key)| { + let restore = format!( + "if [ -z \"${{{marker_key}+x}}\" ]; then exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ]; then unset {key} || exit 1; fi\nif [ -n \"${{{key}+x}}\" ] && [ -n \"${{{key}}}\" ]; then exit 1; fi" + ); + if remove_copies { + format!("{restore}\nunset {marker_key} || exit 1") + } else { + restore + } + }) + .collect::>() + .join("\n"); + let exports = [value_exports, unset_exports] + .into_iter() + .filter(|block| !block.is_empty()) + .collect::>() + .join("\n"); + if exports.is_empty() { + return exports; + } + format!( + "case $- in\n *x*) __CODEX_SNAPSHOT_BROKER_XTRACE=1; set +x ;;\n *) __CODEX_SNAPSHOT_BROKER_XTRACE= ;;\nesac\n{exports}\nif [ -n \"$__CODEX_SNAPSHOT_BROKER_XTRACE\" ]; then\n unset __CODEX_SNAPSHOT_BROKER_XTRACE\n set -x\nelse\n unset __CODEX_SNAPSHOT_BROKER_XTRACE\nfi" + ) +} + +fn is_valid_shell_variable_name(name: &str) -> bool { + let mut chars = name.chars(); + let Some(first) = chars.next() else { + return false; + }; + if !(first == '_' || first.is_ascii_alphabetic()) { + return false; + } + chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) +}