From 1530f828cbaea015bc0fc53c0486e2f889a677f3 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 8 Sep 2026 20:12:02 +0000 Subject: [PATCH] Preserve complete shell snapshot exports through filtering and replay (#43907) ## Why Line-based export parsing can truncate multiline values or mistake their contents for declarations. Snapshot replay also needs to restore Bash options before parsing functions that use extended glob syntax. ## What changed - Capture shell state, aliases, and complete export records with NUL-delimited boundaries, separating capture, credential preparation, and rendering. - Apply credential and environment policy to whole exports, preserving multiline values and making credential aliases follow allowed overrides. - Restore Bash `shopt` options before functions, including for Bash running as `sh`. - Support non-evaluating POSIX `ENV` path expansion and preserve native environment metadata and unset exports. - Move core and exec-server snapshot consumers to the shared capture parser. Allow capture overhead while enforcing the exec-server's 512 KiB state-and-environment limit before filtering. ## Testing Add regression coverage for complete capture records, multiline export filtering and replay, credential alias overrides, `ENV` expansion, and `PATH` export state. Extend exec-server tests to cover Bash-backed `sh`, `extglob`, `nocasematch`, `set -u`, and large environments. GitOrigin-RevId: 7589ab4d137529a395dcee10b5162718de24a621 --- codex-rs/core/src/shell_snapshot.rs | 33 +- codex-rs/core/src/shell_snapshot_tests.rs | 23 +- codex-rs/exec-server/src/shell_snapshot.rs | 71 ++- .../exec-server/src/shell_snapshot_tests.rs | 30 +- codex-rs/exec-server/tests/exec_process.rs | 45 +- codex-rs/shell-command/src/shell_snapshot.rs | 315 ++++------ .../src/shell_snapshot_capture.rs | 335 +++++++++++ .../src/shell_snapshot_credentials.rs | 216 +++++++ .../src/shell_snapshot_exports.rs | 81 +++ .../src/shell_snapshot_render.rs | 69 +++ .../shell-command/src/shell_snapshot_tests.rs | 554 +++++++++++++++++- 11 files changed, 1472 insertions(+), 300 deletions(-) create mode 100644 codex-rs/shell-command/src/shell_snapshot_capture.rs create mode 100644 codex-rs/shell-command/src/shell_snapshot_credentials.rs create mode 100644 codex-rs/shell-command/src/shell_snapshot_exports.rs create mode 100644 codex-rs/shell-command/src/shell_snapshot_render.rs diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 2e0d734a98..4967172e34 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -17,7 +17,10 @@ use anyhow::bail; use codex_exec_server::Environment; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; -use codex_shell_command::shell_snapshot::snapshot_script; +use codex_shell_command::shell_snapshot::CapturedSnapshot; +use codex_shell_command::shell_snapshot::SnapshotCaptureOptions; +use codex_shell_command::shell_snapshot::SnapshotStartup; +use codex_shell_command::shell_snapshot::snapshot_capture_script; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use tokio::fs; @@ -209,8 +212,7 @@ async fn write_shell_snapshot( let shell = get_shell(shell_type).with_context(|| format!("No available shell for {shell_type:?}"))?; - let raw_snapshot = capture_snapshot(&shell, cwd).await?; - let snapshot = strip_snapshot_preamble(&raw_snapshot)?; + let snapshot = capture_snapshot(&shell, cwd).await?; if let Some(parent) = output_path.parent() { let parent_display = parent.display(); @@ -229,18 +231,19 @@ async fn write_shell_snapshot( async fn capture_snapshot(shell: &Shell, cwd: &AbsolutePathBuf) -> Result { let shell_type = shell.shell_type; - let script = snapshot_script(shell_type) - .ok_or_else(|| anyhow!("Shell snapshotting is not yet supported for {shell_type:?}"))?; - run_shell_script(shell, &script, cwd).await -} - -fn strip_snapshot_preamble(snapshot: &str) -> Result { - let marker = "# Snapshot file"; - let Some(start) = snapshot.find(marker) else { - bail!("Snapshot output missing marker {marker}"); - }; - - Ok(snapshot[start..].to_string()) + let script = snapshot_capture_script( + shell_type, + SnapshotCaptureOptions { + startup: SnapshotStartup::Interactive, + declarations: true, + environment: false, + }, + ) + .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")) } async fn validate_snapshot( diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index f27e27d446..3c0995e139 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -86,19 +86,6 @@ async fn get_snapshot(shell_type: ShellType) -> Result { Ok(content) } -#[test] -fn strip_snapshot_preamble_removes_leading_output() { - let snapshot = "noise\n# Snapshot file\nexport PATH=/bin\n"; - let cleaned = strip_snapshot_preamble(snapshot).expect("snapshot marker exists"); - assert_eq!(cleaned, "# Snapshot file\nexport PATH=/bin\n"); -} - -#[test] -fn strip_snapshot_preamble_requires_marker() { - let result = strip_snapshot_preamble("missing header"); - assert!(result.is_err()); -} - #[test] fn snapshot_file_name_parser_supports_legacy_and_suffixed_names() { let session_id = "019cf82b-6a62-7700-bbbd-46909794ef89"; @@ -217,7 +204,15 @@ async fn snapshot_shell_does_not_inherit_stdin() -> Result<()> { let home_display = home.display(); let script = format!( "HOME=\"{home_display}\"; export HOME; {}", - snapshot_script(ShellType::Bash).expect("bash supports snapshots") + snapshot_capture_script( + ShellType::Bash, + SnapshotCaptureOptions { + startup: SnapshotStartup::Interactive, + declarations: true, + environment: false, + } + ) + .expect("bash supports snapshots") ); let output = run_script_with_timeout( &shell, diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 7254c1bda5..e848f6d82c 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -10,7 +10,10 @@ use codex_network_proxy::strip_managed_proxy_env; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_protocol::shell_environment; use codex_shell_command::shell_detect::ShellType; -use codex_shell_command::shell_snapshot::snapshot_state_and_environment_script; +use codex_shell_command::shell_snapshot::CapturedSnapshot; +use codex_shell_command::shell_snapshot::SnapshotCaptureOptions; +use codex_shell_command::shell_snapshot::SnapshotStartup; +use codex_shell_command::shell_snapshot::snapshot_capture_script; use codex_utils_path_uri::PathUri; use tokio::io::AsyncReadExt; use tokio::process::Command; @@ -30,6 +33,8 @@ use crate::telemetry::ExecServerTelemetry; const MAX_CACHED_SNAPSHOTS: usize = 16; const MAX_SNAPSHOT_BYTES: usize = 512 * 1024; +// Capture also includes quoted export records and an optional pre-startup environment. +const MAX_SNAPSHOT_CAPTURE_BYTES: usize = 8 * MAX_SNAPSHOT_BYTES; const MAX_SNAPSHOT_ENV_VALUE_BYTES: usize = 60 * 1024; const MAX_SNAPSHOT_SCOPE_BYTES: usize = 256; const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); @@ -255,7 +260,15 @@ async fn capture_snapshot( prepared: &PreparedExecRequest, shell_type: ShellType, ) -> CaptureResult { - let script = snapshot_state_and_environment_script(shell_type).ok_or_else(|| { + let script = snapshot_capture_script( + shell_type, + SnapshotCaptureOptions { + startup: SnapshotStartup::Interactive, + declarations: false, + environment: true, + }, + ) + .ok_or_else(|| { ( "unsupported_shell", invalid_params("unsupported shell snapshot script".to_string()), @@ -299,7 +312,7 @@ async fn capture_snapshot( let capture = async { let mut output = Vec::new(); stdout - .take((MAX_SNAPSHOT_BYTES + 1) as u64) + .take((MAX_SNAPSHOT_CAPTURE_BYTES + 1) as u64) .read_to_end(&mut output) .await .map_err(|err| { @@ -308,10 +321,12 @@ async fn capture_snapshot( internal_error(format!("cannot read shell snapshot: {err}")), ) })?; - if output.len() > MAX_SNAPSHOT_BYTES { + if output.len() > MAX_SNAPSHOT_CAPTURE_BYTES { return Err(( "too_large", - internal_error(format!("shell snapshot exceeds {MAX_SNAPSHOT_BYTES} bytes")), + internal_error(format!( + "shell snapshot capture exceeds {MAX_SNAPSHOT_CAPTURE_BYTES} bytes" + )), )); } let status = child.wait().await.map_err(|err| { @@ -337,35 +352,30 @@ async fn capture_snapshot( ) })??; - parse_snapshot(&output, params.env_policy.as_ref()) + parse_snapshot(shell_type, &output, params.env_policy.as_ref()) } -fn parse_snapshot(output: &[u8], env_policy: Option<&ExecEnvPolicy>) -> CaptureResult { - let separator = output.iter().position(|byte| *byte == 0).ok_or_else(|| { +fn parse_snapshot( + shell_type: ShellType, + output: &[u8], + env_policy: Option<&ExecEnvPolicy>, +) -> CaptureResult { + let captured = CapturedSnapshot::parse(shell_type, output).ok_or_else(|| { ( - "missing_environment", - internal_error("shell snapshot is missing its environment".to_string()), - ) - })?; - let state = &output[..separator]; - let marker = b"# Snapshot file"; - let start = state - .windows(marker.len()) - .position(|window| window == marker) - .ok_or_else(|| { - ( - "missing_state", - internal_error("shell snapshot is missing its state marker".to_string()), - ) - })?; - let state = std::str::from_utf8(&state[start..]).map_err(|err| { - ( - "invalid_utf8", - internal_error(format!("shell snapshot state is not UTF-8: {err}")), + "invalid_capture", + internal_error("invalid shell snapshot capture".to_string()), ) })?; + let state = captured.render_state(); + if state.len().saturating_add(captured.environment.len()) > MAX_SNAPSHOT_BYTES { + return Err(( + "too_large", + internal_error(format!("shell snapshot exceeds {MAX_SNAPSHOT_BYTES} bytes")), + )); + } - let mut environment = output[separator + 1..] + let mut environment = captured + .environment .split(|byte| *byte == 0) .filter(|entry| !entry.is_empty()) .filter_map(|entry| { @@ -388,10 +398,7 @@ fn parse_snapshot(output: &[u8], env_policy: Option<&ExecEnvPolicy>) -> CaptureR environment.remove("OLDPWD"); environment.retain(|name, _| !shell_environment::is_non_inheritable_env_var(name)); - Ok(ShellSnapshot { - state: state.to_string(), - environment, - }) + Ok(ShellSnapshot { state, environment }) } #[cfg(test)] diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index 74e5c63cf6..cdab0f7508 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -12,6 +12,7 @@ use test_case::test_case; use super::CapturePurpose; use super::MAX_SNAPSHOT_ATTEMPTS; +use super::MAX_SNAPSHOT_BYTES; use super::SNAPSHOT_RETRY_BACKOFF; use super::ShellSnapshotCache; use super::parse_snapshot; @@ -228,6 +229,20 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( Ok(()) } +#[test] +fn snapshot_size_limit_counts_state_and_environment_before_filtering() { + let half = "x".repeat(MAX_SNAPSHOT_BYTES / 2); + let oversized = format!("# Snapshot file\n# {half}\n\0\0\0FILTERED={half}\0"); + let policy = ExecEnvPolicy { + inherit: ShellEnvironmentPolicyInherit::All, + ignore_default_excludes: false, + exclude: vec!["FILTERED".to_string()], + r#set: HashMap::new(), + include_only: Vec::new(), + }; + assert!(parse_snapshot(ShellType::Bash, oversized.as_bytes(), Some(&policy)).is_err()); +} + #[test] fn snapshot_filters_profile_exports_after_capture() { let policy = ExecEnvPolicy { @@ -238,7 +253,8 @@ fn snapshot_filters_profile_exports_after_capture() { include_only: vec!["PROFILE_*".to_string()], }; let snapshot = parse_snapshot( - b"profile noise\n# Snapshot file\nfunction profile_helper() { :; }\n\0PROFILE_ALLOWED=profile\0PROFILE_DENIED=denied\0PROFILE_SECRET=secret\0PWD=/tmp\0", + ShellType::Bash, + b"profile \xff noise\n# Snapshot file\nfunction profile_helper() { :; }\n\0alias profile_alias='profile_helper'\n\0PROFILE_DENIED\0export PROFILE_DENIED=denied\n\0NON_UTF8\0export NON_UTF8='\xff'\n\0\0PROFILE_ALLOWED=profile\0PROFILE_DENIED=denied\0PROFILE_SECRET=secret\0PWD=/tmp\0NON_UTF8=\xff\0", Some(&policy), ) .expect("snapshot should parse"); @@ -249,7 +265,7 @@ fn snapshot_filters_profile_exports_after_capture() { ); assert_eq!( snapshot.state, - "# Snapshot file\nfunction profile_helper() { :; }\n" + "# Snapshot file\nfunction profile_helper() { :; }\nalias profile_alias='profile_helper'\n" ); } @@ -267,7 +283,8 @@ fn snapshot_preserves_profile_exports_with_restrictive_inheritance() { include_only: Vec::new(), }; let snapshot = parse_snapshot( - b"# Snapshot file\n\0PROFILE_ALLOWED=profile\0SDKROOT=/sdk\0PROFILE_SECRET=secret\0PROFILE_DENIED=denied\0", + ShellType::Bash, + b"# Snapshot file\n\0\0\0PROFILE_ALLOWED=profile\0SDKROOT=/sdk\0PROFILE_SECRET=secret\0PROFILE_DENIED=denied\0", Some(&policy), ) .expect("snapshot should parse"); @@ -300,10 +317,11 @@ fn snapshot_caches_only_unmanaged_proxy_state() { ]), ), ] { - let output = format!("# Snapshot file\n\0{exports}"); - let snapshot = - parse_snapshot(output.as_bytes(), /*env_policy*/ None).expect("snapshot should parse"); + let output = format!("# Snapshot file\n\0\0\0{exports}"); + let snapshot = parse_snapshot(ShellType::Bash, output.as_bytes(), /*env_policy*/ None) + .expect("snapshot should parse"); assert_eq!(snapshot.environment, expected); } } +use codex_shell_command::shell_detect::ShellType; diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index c7bb503849..f1df5d78ae 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -1,7 +1,7 @@ mod common; use std::collections::HashMap; -#[cfg(target_os = "linux")] +#[cfg(unix)] use std::os::unix::fs::PermissionsExt; use std::sync::Arc; @@ -181,6 +181,7 @@ async fn codex_home_symlink_opt_out_respects_host_config_and_scope() -> Result<( #[test_case(true, true, false, false, "bash"; "remote_tty")] #[test_case(true, false, true, false, "bash"; "remote_sandbox")] #[test_case(false, false, false, false, "sh"; "local_sh_pipe")] +#[test_case(false, false, false, false, "bash-sh"; "local_bash_backed_sh")] #[test_case(false, false, false, true, "bash"; "local_bash_env")] #[test_case(true, false, false, true, "bash"; "remote_bash_env")] #[cfg_attr( @@ -219,6 +220,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( "bash" if automatic_startup => ("/bin/bash", ".bash-env"), "bash" => ("/bin/bash", ".bashrc"), "sh" => ("/bin/sh", ".snapshot-env"), + "bash-sh" => ("/bin/bash", ".snapshot-env"), "zsh" if automatic_startup => ("/bin/zsh", ".zshenv"), "zsh" => ("/bin/zsh", ".zshrc"), name => anyhow::bail!("unsupported test shell {name}"), @@ -226,6 +228,14 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( let profile_path = home.path().join(profile_name); let profile_path_entry = home.path().join("profile-bin"); let runtime_path_entry = home.path().join("runtime-bin"); + std::fs::create_dir(&profile_path_entry)?; + let wc = profile_path_entry.join("wc"); + std::fs::write( + &wc, + "#!/bin/sh\nprintf x >> \"$HOME/tool-captures\"\nexec /usr/bin/wc \"$@\"\n", + )?; + std::fs::set_permissions(&wc, std::fs::Permissions::from_mode(0o755))?; + let posix_shell = matches!(shell_name, "sh" | "bash-sh"); let padding = if !use_remote && !tty && shell_name == "bash" { format!( "snapshot_padding() {{ printf '%s' '{}'; }}\n", @@ -234,7 +244,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( } else { String::new() }; - let shadowed_builtins = if shell_name == "sh" { + let shadowed_builtins = if posix_shell { "" } else { "unset() { exit 41; }\nbuiltin() { :; }\n" @@ -242,7 +252,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( std::fs::write( &profile_path, format!( - "printf x >> \"$HOME/captures\"\nexport PATH=\"$HOME/profile-bin:/usr/bin:/bin\"\nexport PROFILE_ALLOWED=profile\nexport PROFILE_SECRET=secret\nexport PROFILE_DENIED=denied\nprofile_helper() {{ printf helper; }}\n{shadowed_builtins}{padding}" + "printf x >> \"$HOME/captures\"\nexport PATH=\"$HOME/profile-bin:/usr/bin:/bin\"\nexport PROFILE_ALLOWED=profile\nexport PROFILE_SECRET=secret\nexport PROFILE_DENIED=denied\nprofile_helper() {{ printf helper; }}\nif [ -n \"${{BASH_VERSION-}}\" ]; then\n shopt -s extglob nocasematch\n eval 'profile_helper() {{ case $1 in @(foo|bar)*) printf helper ;; *) return 1 ;; esac; }}'\nfi\nset -u\n{shadowed_builtins}{padding}" ), )?; if shell_name == "zsh" && automatic_startup { @@ -255,11 +265,15 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( "HOME".to_string(), home.path().to_string_lossy().into_owned(), )]); - if shell_name == "sh" { + if posix_shell { configured_environment.insert( "ENV".to_string(), - profile_path.to_string_lossy().into_owned(), + "${XDG_CONFIG_HOME:-$HOME}/.snapshot-env".to_string(), ); + // Keep coverage for large values alongside the many-small-entry case below. + for index in 0..3 { + configured_environment.insert(format!("PROFILE_SDK_{index}"), "x".repeat(60 * 1024)); + } } if shell_name == "bash" && automatic_startup { configured_environment.insert( @@ -267,6 +281,13 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( profile_path.to_string_lossy().into_owned(), ); } + // Many small entries exercise capture overhead separately from the byte limit above. + let many_entries = !use_remote && !tty && !automatic_startup; + if many_entries { + configured_environment.extend( + (0..1_000).map(|index| (format!("PROFILE_ENTRY_{index}"), format!("value-{index}"))), + ); + } let policy = ExecEnvPolicy { inherit: ShellEnvironmentPolicyInherit::All, ignore_default_excludes: false, @@ -283,10 +304,15 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( let (command_prefix, expected_prefix) = if shell_name == "sh" { ("", "") } else { - ("profile_helper; ", "helper") + ("profile_helper FOObar; ", "helper") + }; + let entry_check = if many_entries { + "[ \"${PROFILE_ENTRY_999-missing}\" = value-999 ] || exit 43; " + } else { + "" }; let command = format!( - "export PATH='{}':\"$PATH\"; {command_prefix}printf '|%s|%s|%s|%s|%s|%s' \"$PROFILE_ALLOWED\" \"${{PROFILE_SECRET-missing}}\" \"${{PROFILE_DENIED-missing}}\" \"$PATH\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_0-missing}}\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_1-missing}}\"", + "case $- in *u*) ;; *) exit 42 ;; esac; {entry_check}export PATH='{}':\"$PATH\"; {command_prefix}printf '|%s|%s|%s|%s|%s|%s' \"$PROFILE_ALLOWED\" \"${{PROFILE_SECRET-missing}}\" \"${{PROFILE_DENIED-missing}}\" \"$PATH\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_0-missing}}\" \"${{__CODEX_SHELL_SNAPSHOT_STATE_1-missing}}\"", runtime_path_entry.display(), ); let expected_stdout = format!( @@ -307,14 +333,14 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( shell_snapshot: Some(ShellSnapshotRequest { scope_id: "attachment-1".to_string(), shell: ShellInfo { - name: shell_name.to_string(), + name: if posix_shell { "sh" } else { shell_name }.to_string(), path: shell_path.to_string(), }, }), env: HashMap::new(), tty, pipe_stdin: false, - arg0: None, + arg0: (shell_name == "bash-sh").then(|| "sh".to_string()), sandbox: (use_sandbox && attempt == 0).then(|| { FileSystemSandboxContext::from_permission_profile_with_cwd( PermissionProfile::read_only(), @@ -335,6 +361,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( } assert_eq!(std::fs::read_to_string(home.path().join("captures"))?, "x"); + assert!(!std::fs::read(home.path().join("tool-captures"))?.is_empty()); if let Some(server) = context._server { assert!(!server.codex_home().join("shell_snapshots").exists()); } diff --git a/codex-rs/shell-command/src/shell_snapshot.rs b/codex-rs/shell-command/src/shell_snapshot.rs index d0cf01c3b3..228029cbd1 100644 --- a/codex-rs/shell-command/src/shell_snapshot.rs +++ b/codex-rs/shell-command/src/shell_snapshot.rs @@ -1,225 +1,108 @@ -use crate::shell_detect::ShellType; +//! Capture, prepare credential-safe state, and render restorable shell snapshots. + +#[path = "shell_snapshot_capture.rs"] +mod capture; +#[path = "shell_snapshot_credentials.rs"] +mod credentials; +#[path = "shell_snapshot_exports.rs"] +mod exports; +#[path = "shell_snapshot_render.rs"] +mod render; + +pub use capture::CapturedSnapshot; +pub use capture::CapturedStartupEnvironment; +pub use capture::SnapshotCaptureOptions; +pub use capture::snapshot_capture_script; +pub use credentials::PreparedSnapshot; +pub use credentials::SnapshotCredentialEnvironment; +pub use credentials::prepare_snapshot_credentials; #[cfg(all(test, unix))] #[path = "shell_snapshot_tests.rs"] mod tests; -const EXCLUDED_EXPORT_VARS: &[&str] = &["PWD", "OLDPWD"]; -const EXPORT_CAPTURE_MARKER: &str = "# Capture exported variables"; - -/// Returns the shell-native script used to capture restorable shell state. +/// Returns the POSIX shell helper used to resolve supported `ENV` startup paths. /// -/// Command Prompt does not expose the POSIX or PowerShell state required by -/// this representation and therefore does not support snapshots. -pub fn snapshot_script(shell_type: ShellType) -> Option { - match shell_type { - ShellType::Zsh => Some(zsh_snapshot_script()), - ShellType::Bash => Some(bash_snapshot_script()), - ShellType::Sh => Some(sh_snapshot_script()), - ShellType::PowerShell => Some(powershell_snapshot_script().to_string()), - ShellType::Cmd => None, - } +/// The helper deliberately supports only non-evaluating path forms. Unsupported +/// shell expressions are returned unchanged rather than executed. +pub fn posix_env_path_expansion_function() -> &'static str { + r#"__codex_snapshot_expand_env() ( + set +u + __codex_snapshot_getenv() { + # Preserve the value separately from lookup status and its formatting newline. + __codex_env_expanded=$( + if command -v printenv >/dev/null 2>&1; then + printenv "$1" + elif [ "$1" = PATH ]; then + [ "${PATH+x}" = x ] && printf '%s\n' "$PATH" + else + command -p printenv "$1" + fi && printf '.' + ) || return 1 + __codex_env_expanded=${__codex_env_expanded%?} + __codex_env_expanded=${__codex_env_expanded%?} + } + __codex_env_file=$1 + case "$__codex_env_file" in + '~/'*) __codex_env_file="${HOME-}/${__codex_env_file#*/}" ;; + '${PATH%%:*}') __codex_env_file="${PATH%%:*}" ;; + '${PATH%%:*}/'*) __codex_env_file="${PATH%%:*}/${__codex_env_file#*/}" ;; + '${'*) + __codex_env_body=${__codex_env_file#\$\{} + case "$__codex_env_body" in + *\}*) + __codex_env_name=${__codex_env_body%%\}*} + __codex_env_suffix=${__codex_env_body#*\}} + case "$__codex_env_name" in + *:-*) + __codex_env_default=${__codex_env_name#*:-} + __codex_env_name=${__codex_env_name%%:-*} + __codex_env_has_default=1 + ;; + *) __codex_env_has_default= ;; + esac + case "$__codex_env_name" in + ''|[0-9]*|*[!A-Za-z0-9_]*) ;; + *) + case "$__codex_env_suffix" in + ''|/*) + if __codex_snapshot_getenv "$__codex_env_name" 2>/dev/null && + { [ -n "$__codex_env_expanded" ] || [ -z "$__codex_env_has_default" ]; }; then + __codex_env_file="$__codex_env_expanded$__codex_env_suffix" + elif [ -n "$__codex_env_has_default" ]; then + __codex_env_default=$(__codex_snapshot_expand_env "$__codex_env_default") + __codex_env_file="$__codex_env_default$__codex_env_suffix" + fi + ;; + esac + ;; + esac + ;; + esac + ;; + '$'*) + __codex_env_name=${__codex_env_file%%/*} + __codex_env_name=${__codex_env_name#\$} + case "$__codex_env_name" in + ''|[0-9]*|*[!A-Za-z0-9_]*) ;; + *) + if __codex_snapshot_getenv "$__codex_env_name" 2>/dev/null; then + if [ "$__codex_env_file" = "\$$__codex_env_name" ]; then + __codex_env_file=$__codex_env_expanded + else + __codex_env_file="$__codex_env_expanded/${__codex_env_file#*/}" + fi + fi + ;; + esac + ;; + esac + printf '%s' "$__codex_env_file" +)"# } -/// Captures shell state and a separate NUL-delimited exported environment. -/// -/// Keeping exports outside the restorable script lets executors apply their -/// environment policy after shell profiles have run. -pub fn snapshot_state_and_environment_script(shell_type: ShellType) -> Option { - let script = snapshot_script(shell_type)?; - let (state, _) = script.split_once(EXPORT_CAPTURE_MARKER)?; - Some(format!("{state}printf '\\0'\n/usr/bin/env -0\n")) -} - -fn excluded_exports_regex() -> String { - EXCLUDED_EXPORT_VARS.join("|") -} - -fn zsh_snapshot_script() -> String { - let excluded = excluded_exports_regex(); - let script = r##"if [[ -n "$ZDOTDIR" ]]; then - rc="$ZDOTDIR/.zshrc" -else - rc="$HOME/.zshrc" -fi -[[ -r "$rc" ]] && . "$rc" -print '# Snapshot file' -print '# Unset all aliases to avoid conflicts with functions' -print 'unalias -a 2>/dev/null || true' -print '# Functions' -functions -print '' -setopt_count=$(setopt | wc -l | tr -d ' ') -print "# setopts $setopt_count" -setopt | sed 's/^/setopt /' -print '' -alias_count=$(alias -L | wc -l | tr -d ' ') -print "# aliases $alias_count" -alias -L -print '' -# Capture exported variables -export_lines=$(export -p | awk ' -/^(export|declare -x|typeset -x) / { - line=$0 - name=line - sub(/^(export|declare -x|typeset -x) /, "", name) - if (name ~ /^-[A-Za-z]*r[A-Za-z]* /) { - next - } - if (name ~ /^-[A-Za-z]*T[A-Za-z]* /) { - sub(/^-[A-Za-z]*T[A-Za-z]* /, "", name) - sub(/ [A-Za-z_][A-Za-z0-9_]*=.*/, "", name) - } - sub(/=.*/, "", name) - if (name ~ /^(EXCLUDED_EXPORTS)$/) { - next - } - if (name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) { - print line - } -}') -export_count=$(printf '%s\n' "$export_lines" | sed '/^$/d' | wc -l | tr -d ' ') -print "# exports $export_count" -if [[ -n "$export_lines" ]]; then - print -r -- "$export_lines" -fi -"##; - script.replace("EXCLUDED_EXPORTS", &excluded) -} - -fn bash_snapshot_script() -> String { - let excluded = excluded_exports_regex(); - let script = r##"if [ -z "$BASH_ENV" ] && [ -r "$HOME/.bashrc" ]; then - . "$HOME/.bashrc" -fi -echo '# Snapshot file' -echo '# Unset all aliases to avoid conflicts with functions' -unalias -a 2>/dev/null || true -echo '# Functions' -declare -f -echo '' -bash_opts=$(set -o | awk '$2=="on"{print $1}') -bash_opt_count=$(printf '%s\n' "$bash_opts" | sed '/^$/d' | wc -l | tr -d ' ') -echo "# setopts $bash_opt_count" -if [ -n "$bash_opts" ]; then - printf 'set -o %s\n' $bash_opts -fi -echo '' -alias_count=$(alias -p | wc -l | tr -d ' ') -echo "# aliases $alias_count" -alias -p -echo '' -# Capture exported variables -export_lines=$( - while IFS= read -r name; do - if [[ "$name" =~ ^(EXCLUDED_EXPORTS)$ ]]; then - continue - fi - if [[ ! "$name" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]]; then - continue - fi - declare -xp "$name" 2>/dev/null || true - done < <(compgen -e) -) -export_count=$(printf '%s\n' "$export_lines" | sed '/^$/d' | wc -l | tr -d ' ') -echo "# exports $export_count" -if [ -n "$export_lines" ]; then - printf '%s\n' "$export_lines" -fi -"##; - script.replace("EXCLUDED_EXPORTS", &excluded) -} - -fn sh_snapshot_script() -> String { - let excluded = excluded_exports_regex(); - let script = r##"if [ -n "$ENV" ] && [ -r "$ENV" ]; then - . "$ENV" -fi -echo '# Snapshot file' -echo '# Unset all aliases to avoid conflicts with functions' -unalias -a 2>/dev/null || true -echo '# Functions' -if command -v typeset >/dev/null 2>&1; then - typeset -f -elif command -v declare >/dev/null 2>&1; then - declare -f -fi -echo '' -if set -o >/dev/null 2>&1; then - sh_opts=$(set -o | awk '$2=="on"{print $1}') - sh_opt_count=$(printf '%s\n' "$sh_opts" | sed '/^$/d' | wc -l | tr -d ' ') - echo "# setopts $sh_opt_count" - if [ -n "$sh_opts" ]; then - printf 'set -o %s\n' $sh_opts - fi -else - echo '# setopts 0' -fi -echo '' -if alias >/dev/null 2>&1; then - alias_count=$(alias | wc -l | tr -d ' ') - echo "# aliases $alias_count" - alias - echo '' -else - echo '# aliases 0' -fi -# Capture exported variables -if export -p >/dev/null 2>&1; then - export_lines=$(export -p | awk ' -/^(export|declare -x|typeset -x) / { - line=$0 - name=line - sub(/^(export|declare -x|typeset -x) /, "", name) - sub(/=.*/, "", name) - if (name ~ /^(EXCLUDED_EXPORTS)$/) { - next - } - if (name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) { - print line - } -}') - export_count=$(printf '%s\n' "$export_lines" | sed '/^$/d' | wc -l | tr -d ' ') - echo "# exports $export_count" - if [ -n "$export_lines" ]; then - printf '%s\n' "$export_lines" - fi -else - export_count=$(env | sort | awk -F= '$1 ~ /^[A-Za-z_][A-Za-z0-9_]*$/ { count++ } END { print count }') - echo "# exports $export_count" - env | sort | while IFS='=' read -r key value; do - case "$key" in - ""|[0-9]*|*[!A-Za-z0-9_]*|EXCLUDED_EXPORTS) continue ;; - esac - escaped=$(printf "%s" "$value" | sed "s/'/'\"'\"'/g") - printf "export %s='%s'\n" "$key" "$escaped" - done -fi -"##; - script.replace("EXCLUDED_EXPORTS", &excluded) -} - -fn powershell_snapshot_script() -> &'static str { - r##"$ErrorActionPreference = 'Stop' -Write-Output '# Snapshot file' -Write-Output '# Unset all aliases to avoid conflicts with functions' -Write-Output 'Remove-Item Alias:* -ErrorAction SilentlyContinue' -Write-Output '# Functions' -Get-ChildItem Function: | ForEach-Object { - "function {0} {{`n{1}`n}}" -f $_.Name, $_.Definition -} -Write-Output '' -$aliases = Get-Alias -Write-Output ("# aliases " + $aliases.Count) -$aliases | ForEach-Object { - "Set-Alias -Name {0} -Value {1}" -f $_.Name, $_.Definition -} -Write-Output '' -$envVars = Get-ChildItem Env: -Write-Output ("# exports " + $envVars.Count) -$envVars | ForEach-Object { - $escaped = $_.Value -replace "'", "''" - "`$env:{0}='{1}'" -f $_.Name, $escaped -} -"## +#[derive(Clone, Copy)] +pub enum SnapshotStartup { + Interactive, + NonInteractive, } diff --git a/codex-rs/shell-command/src/shell_snapshot_capture.rs b/codex-rs/shell-command/src/shell_snapshot_capture.rs new file mode 100644 index 0000000000..ecf8b4eef7 --- /dev/null +++ b/codex-rs/shell-command/src/shell_snapshot_capture.rs @@ -0,0 +1,335 @@ +//! Capture native shell state and decode complete export records without applying credential policy. +//! Retain native declaration flags and reject incomplete records. + +use super::SnapshotStartup; +use super::exports; +use super::posix_env_path_expansion_function; +use crate::shell_detect::ShellType; +use std::borrow::Cow; + +const SNAPSHOT_COMMAND_HELPER: &str = r#"__codex_snapshot_command() { + if command -v "$1" >/dev/null 2>&1; then + "$@" + else + command -p "$@" + fi +}"#; + +const SNAPSHOT_ENVIRONMENT: &str = r#"if command -v env >/dev/null 2>&1; then + "env" -0 +else + # Resolve the fallback separately: Bash's command -p changes the utility's PATH. + "$(PATH="$(command -p getconf PATH)" command -v env)" -0 +fi +"#; + +/// Selects startup behavior and the records needed by a snapshot consumer. +/// Shell state and aliases are always captured. Credential preparation needs both +/// declarations and environment values; executor replay needs only the environment. +#[derive(Clone, Copy)] +pub struct SnapshotCaptureOptions { + pub startup: SnapshotStartup, + pub declarations: bool, + pub environment: bool, +} + +/// Returns a script that captures native state, aliases, and the requested export records. +/// +/// Startup behavior is independent of how the captured records will be rendered. +/// PowerShell and Command Prompt snapshots are not supported by the snapshot callers. +pub fn snapshot_capture_script( + shell_type: ShellType, + options: SnapshotCaptureOptions, +) -> Option { + let script = match shell_type { + ShellType::Zsh => zsh_snapshot_script(options.startup), + ShellType::Bash => bash_snapshot_script(options.startup), + ShellType::Sh => sh_snapshot_script(options.startup), + ShellType::PowerShell | ShellType::Cmd => return None, + }; + let declarations = if options.declarations { + exports::script(shell_type) + } else { + String::new() + }; + Some( + script + .replace("SNAPSHOT_EXPORTS", &declarations) + .replace( + "SNAPSHOT_DECLARATION_ENVIRONMENT", + if options.declarations { SNAPSHOT_ENVIRONMENT } else { "" }, + ) + .replace("SNAPSHOT_COMMAND_HELPER", SNAPSHOT_COMMAND_HELPER) + .replace( + "SNAPSHOT_STARTUP_ENVIRONMENT", + if options.declarations && options.environment { + "printf '\\0CODEX_SNAPSHOT_POSIX_STARTUP\\0%s\\0' \"$__codex_env_file\"\n SNAPSHOT_ENVIRONMENT\n printf '\\0'" + } else { + "" + }, + ) + .replace( + "SNAPSHOT_ENVIRONMENT", + if options.environment { + SNAPSHOT_ENVIRONMENT + } else { + "" + }, + ), + ) +} + +fn zsh_snapshot_script(shell_startup: SnapshotStartup) -> String { + let startup = match shell_startup { + SnapshotStartup::Interactive => { + r##"if [[ -n "${ZDOTDIR-}" ]]; then + rc="$ZDOTDIR/.zshrc" +elif [[ -n "${HOME-}" ]]; then + rc="$HOME/.zshrc" +else + rc= +fi +[[ -r "$rc" ]] && . "$rc" +"## + } + SnapshotStartup::NonInteractive => "", + }; + let script = r##"print '# Snapshot file' +print '# Unset all aliases to avoid conflicts with functions' +print 'unalias -a 2>/dev/null || true' +print '# Functions' +functions +print '' +SNAPSHOT_COMMAND_HELPER +setopt_count=$(setopt | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +print "# setopts $setopt_count" +setopt | __codex_snapshot_command sed 's/^/setopt /' +print '' +printf '\0' +alias_count=$(alias -L | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +print "# aliases $alias_count" +alias -L +print '' +printf '\0' +SNAPSHOT_EXPORTS +printf '\0' +SNAPSHOT_ENVIRONMENT +"##; + format!("{startup}{script}") +} + +fn bash_snapshot_script(shell_startup: SnapshotStartup) -> String { + let startup = match shell_startup { + SnapshotStartup::Interactive => { + r##"if [ -z "${BASH_ENV-}" ] && [ -n "${HOME-}" ] && [ -r "$HOME/.bashrc" ]; then + . "$HOME/.bashrc" +fi +"## + } + SnapshotStartup::NonInteractive => "", + }; + let script = r##"echo '# Snapshot file' +echo '# Unset all aliases to avoid conflicts with functions' +echo 'unalias -a 2>/dev/null || true' +shopt -p || true +echo '# Functions' +declare -f +echo '' +SNAPSHOT_COMMAND_HELPER +bash_opts=$(set -o | __codex_snapshot_command awk '$2=="on"{print $1}') +bash_opt_count=$(printf '%s\n' "$bash_opts" | __codex_snapshot_command sed '/^$/d' | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +echo "# setopts $bash_opt_count" +if [ -n "$bash_opts" ]; then + printf 'set -o %s\n' $bash_opts +fi +echo '' +printf '\0' +alias_count=$(alias -p | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') +echo "# aliases $alias_count" +alias -p +echo '' +printf '\0' +SNAPSHOT_EXPORTS +printf '\0' +SNAPSHOT_ENVIRONMENT +"##; + format!("{startup}{script}") +} + +fn sh_snapshot_script(shell_startup: SnapshotStartup) -> String { + let startup = match shell_startup { + SnapshotStartup::Interactive => [ + posix_env_path_expansion_function(), + r##"if [ -n "${ENV-}" ]; then + __codex_env_file=$(__codex_snapshot_expand_env "$ENV") + if [ -r "$__codex_env_file" ] && [ ! -d "$__codex_env_file" ]; then + SNAPSHOT_STARTUP_ENVIRONMENT + case "$__codex_env_file" in + /*) . "$__codex_env_file" ;; + *) . "./$__codex_env_file" ;; + esac + fi + command unset __codex_env_file +fi +command unset -f __codex_snapshot_expand_env +"##, + ] + .join("\n"), + SnapshotStartup::NonInteractive => String::new(), + }; + let script = r##"echo '# Snapshot file' +if [ -n "${BASH_VERSINFO-}" ]; then + echo '# Bash-backed sh' + shopt -p || true +fi +echo '# Unset all aliases to avoid conflicts with functions' +unalias -a 2>/dev/null || true +echo '# Functions' +if command -v typeset >/dev/null 2>&1; then + typeset -f +elif command -v declare >/dev/null 2>&1; then + declare -f +fi +echo '' +SNAPSHOT_COMMAND_HELPER +if set -o >/dev/null 2>&1; then + sh_opts=$(set -o | __codex_snapshot_command awk '$2=="on"{print $1}') + sh_opt_count=$(printf '%s\n' "$sh_opts" | __codex_snapshot_command sed '/^$/d' | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') + echo "# setopts $sh_opt_count" + if [ -n "$sh_opts" ]; then + printf 'set -o %s\n' $sh_opts + fi +else + echo '# setopts 0' +fi +echo '' +printf '\0' +if alias >/dev/null 2>&1; then + alias_count=$(alias | __codex_snapshot_command wc -l | __codex_snapshot_command tr -d ' ') + echo "# aliases $alias_count" + alias + echo '' +else + echo '# aliases 0' +fi +printf '\0' +SNAPSHOT_EXPORTS +printf '\0' +SNAPSHOT_DECLARATION_ENVIRONMENT +printf '\0' +SNAPSHOT_ENVIRONMENT +"##; + format!("{startup}{script}") +} + +/// Native shell sections captured before applying credential or environment policy. +/// +/// Function and alias declarations remain shell source. The shell supplies each export's +/// name and boundary; only tied arrays need syntax decoding. No captured code is evaluated. +pub struct CapturedSnapshot<'a> { + /// Resolved POSIX startup file and its environment immediately before sourcing it. + pub startup_environment: Option>, + pub(super) state: &'a str, + pub(super) aliases: &'a str, + pub(super) exports: Vec>, + /// NUL-delimited environment entries from the same shell invocation. + pub environment: &'a [u8], +} + +/// Optional pre-startup state captured alongside the final snapshot records. +pub struct CapturedStartupEnvironment<'a> { + pub path: &'a str, + pub environment: &'a [u8], +} + +pub(super) struct CapturedExport<'a> { + pub source: Cow<'a, str>, + pub key: &'a str, +} + +impl<'a> CapturedSnapshot<'a> { + /// Decode the local capture format emitted by `snapshot_capture_script`. + /// Keep raw environment bytes so callers can discard non-UTF-8 entries independently. + pub fn parse(shell_type: ShellType, captured: &'a [u8]) -> Option { + if !matches!(shell_type, ShellType::Bash | ShellType::Zsh | ShellType::Sh) { + return None; + } + let mut remaining = captured; + let startup_marker = b"\0CODEX_SNAPSHOT_POSIX_STARTUP\0"; + let startup_environment = if shell_type == ShellType::Sh + && let Some(start) = remaining + .windows(startup_marker.len()) + .position(|part| part == startup_marker) + { + remaining = &remaining[start + startup_marker.len()..]; + let path = std::str::from_utf8(take_record(&mut remaining)?).ok()?; + let environment = remaining; + while !take_record(&mut remaining)?.is_empty() {} + let environment = &environment[..environment.len() - remaining.len() - 1]; + Some(CapturedStartupEnvironment { path, environment }) + } else { + None + }; + // state NUL aliases NUL (name NUL declaration NUL)* NUL environment. + // POSIX sh additionally captures declaration values as a NUL-delimited environment. + let state = take_record(&mut remaining)?; + let marker = b"# Snapshot file"; + let start = state + .windows(marker.len()) + .position(|part| part == marker)?; + let state = std::str::from_utf8(&state[start..]).ok()?; + let aliases = std::str::from_utf8(take_record(&mut remaining)?).ok()?; + let mut records = Vec::new(); + loop { + let key = std::str::from_utf8(take_record(&mut remaining)?).ok()?; + if key.is_empty() { + if shell_type == ShellType::Sh { + loop { + let entry = take_record(&mut remaining)?; + if entry.is_empty() { + break; + } + let equals = entry.iter().position(|byte| *byte == b'=')?; + let Ok(key) = std::str::from_utf8(&entry[..equals]) else { + continue; + }; + if key.is_empty() + || matches!(key, "PWD" | "OLDPWD") + || !key.bytes().enumerate().all(|(index, byte)| { + byte == b'_' + || byte.is_ascii_alphabetic() + || index > 0 && byte.is_ascii_digit() + }) + { + continue; + } + let value = String::from_utf8_lossy(&entry[equals + 1..]); + records.push(CapturedExport { + source: Cow::Owned(format!( + "export {key}={}\n", + shlex::try_quote(&value).ok()? + )), + key, + }); + } + } + return Some(Self { + startup_environment, + state, + aliases, + exports: records, + environment: remaining, + }); + } + let source = String::from_utf8_lossy(take_record(&mut remaining)?); + records.push(CapturedExport { source, key }); + } + } +} + +fn take_record<'a>(remaining: &mut &'a [u8]) -> Option<&'a [u8]> { + let boundary = remaining.iter().position(|byte| *byte == 0)?; + let record = &remaining[..boundary]; + *remaining = &remaining[boundary + 1..]; + Some(record) +} diff --git a/codex-rs/shell-command/src/shell_snapshot_credentials.rs b/codex-rs/shell-command/src/shell_snapshot_credentials.rs new file mode 100644 index 0000000000..23fe061c25 --- /dev/null +++ b/codex-rs/shell-command/src/shell_snapshot_credentials.rs @@ -0,0 +1,216 @@ +//! Apply credential and environment policy to captured exports, producing typed render values. +//! Native declarations are rendered only after applying the supplied credential policy. + +use super::capture::CapturedSnapshot; +use super::render; +use super::render::Export; +use super::render::Value; +use super::render::ValuePart; +use std::collections::HashMap; + +/// Environment views used to render protected credential references into a shell snapshot. +pub struct SnapshotCredentialEnvironment<'a> { + pub original: &'a HashMap, + pub restored: &'a HashMap, + pub configured: &'a HashMap, + pub discovered: &'a HashMap, + pub allowed: &'a HashMap, + /// Applies environment policy to export-only names absent from the captured environment. + pub is_allowed_unset: &'a dyn Fn(&str) -> bool, + pub brokered_keys: &'a [String], + pub brokered_alias_keys: &'a [String], + pub allowed_brokered_keys: &'a [String], +} + +/// Credential-checked replay script and aliases to restore after shell startup. +pub struct PreparedSnapshot { + pub script: String, + pub aliases: HashMap, + pub rejected_alias_keys: Vec, +} + +/// Prepare a replay script from captured state without changing the captured data. +/// +/// Credential policy is applied before rendering. A rejected capture produces no script, +/// when the assembled source still contains a recognized credential. +pub fn prepare_snapshot_credentials( + captured: &CapturedSnapshot<'_>, + environment: SnapshotCredentialEnvironment<'_>, + mut virtualize_text: impl FnMut(&mut String) -> bool, +) -> Option { + let SnapshotCredentialEnvironment { + original, + restored, + configured, + discovered, + allowed, + is_allowed_unset, + brokered_keys, + brokered_alias_keys, + allowed_brokered_keys, + } = environment; + let real_credential_value = |key: &str| restored.get(key).or_else(|| configured.get(key)); + let mut credential_aliases = original + .iter() + .filter(|(key, _)| !brokered_keys.contains(key) && !brokered_alias_keys.contains(key)) + .filter_map(|(key, value)| { + let credential_keys = brokered_keys + .iter() + .filter(|credential_key| { + real_credential_value(credential_key).is_some_and(|real| { + value == real || real.len() >= 16 && value.contains(real) + }) || discovered.get(*credential_key).is_some_and(|dummy| { + value == dummy || dummy.len() >= 16 && value.contains(dummy) + }) + }) + .map(String::as_str) + .collect::>(); + (!credential_keys.is_empty()).then_some((key.as_str(), credential_keys)) + }) + .filter(|(key, _)| { + key.bytes().enumerate().all(|(index, byte)| { + byte == b'_' || byte.is_ascii_alphabetic() || (index != 0 && byte.is_ascii_digit()) + }) + }) + .collect::>(); + let credential_alias_assignment = |value: &str, credential_keys: &[&str]| { + let mut remaining = value; + let mut assignment = Value::default(); + let mut normalized = String::new(); + while !remaining.is_empty() { + let next_credential = credential_keys + .iter() + .filter_map(|credential_key| { + let dummy = discovered.get(*credential_key)?; + let real = real_credential_value(credential_key)?; + let allowed_key = allowed_brokered_keys + .iter() + .find(|allowed_key| allowed_key.as_str() == *credential_key) + .or_else(|| { + allowed_brokered_keys.iter().find(|allowed_key| { + real_credential_value(allowed_key) == Some(real) + }) + })?; + remaining + .find(dummy) + .map(|index| (index, allowed_key.as_str(), dummy.as_str())) + }) + .min_by_key(|(index, _, _)| *index); + let Some((index, credential_key, dummy)) = next_credential else { + assignment + .parts + .push(ValuePart::Literal(remaining.to_string())); + normalized.push_str(remaining); + break; + }; + if index != 0 { + assignment + .parts + .push(ValuePart::Literal(remaining[..index].to_string())); + normalized.push_str(&remaining[..index]); + } + assignment.parts.push(ValuePart::Credential { + key: credential_key.to_string(), + }); + normalized.push_str(allowed.get(credential_key)?); + remaining = &remaining[index + dummy.len()..]; + } + Some((assignment, normalized)) + }; + let credential_alias_is_allowed = |credential_keys: &[&str]| { + credential_keys.iter().all(|credential_key| { + real_credential_value(credential_key).is_some_and(|real| { + allowed_brokered_keys + .iter() + .any(|allowed_key| real_credential_value(allowed_key) == Some(real)) + }) + }) + }; + let is_disallowed_credential_alias = + |key: &str, virtualize_text: &mut dyn FnMut(&mut String) -> bool| { + original.get(key).is_some_and(|value| { + let mut virtualized = value.clone(); + !virtualize_text(&mut virtualized) + }) + }; + let mut alias_values = HashMap::new(); + let mut rejected_alias_keys = Vec::new(); + let mut invalid_export = false; + let exports = captured + .exports + .iter() + .filter_map(|export| { + let line = export.source.as_ref(); + let key = export.key; + + if (!allowed.contains_key(key) + && (original.contains_key(key) || !is_allowed_unset(key))) + || brokered_keys + .iter() + .any(|credential_key| credential_key == key) + || brokered_alias_keys + .iter() + .any(|credential_key| credential_key == key) + || is_disallowed_credential_alias(key, &mut virtualize_text) + { + return None; + } + + if let Some(credential_keys) = credential_aliases.remove(key) { + if !credential_alias_is_allowed(&credential_keys) { + rejected_alias_keys.push(key.to_string()); + return None; + } + let declaration = line + .split_once('=') + .map_or(line.trim_end(), |(prefix, _)| prefix); + let (assignment, value) = + credential_alias_assignment(allowed.get(key)?, &credential_keys)?; + alias_values.insert(key.to_string(), value); + if line + .split_whitespace() + .skip(/*n*/ 1) + .take_while(|word| word.starts_with('-')) + .any(|flags| flags.contains('T')) + { + invalid_export = true; + return Some(Export::Captured(line)); + } + return Some(Export::Assignment { + declaration, + value: assignment, + }); + } + + Some(Export::Captured(line)) + }) + .collect::>(); + let mut snapshot = render::render(captured.state, captured.aliases, &exports)?; + if invalid_export { + return None; + } + + for (key, credential_keys) in credential_aliases { + if !is_disallowed_credential_alias(key, &mut virtualize_text) + && credential_alias_is_allowed(&credential_keys) + && let Some((_, value)) = allowed + .get(key) + .and_then(|value| credential_alias_assignment(value, &credential_keys)) + { + // Native capture can deliberately omit readonly exports; retain only their metadata. + alias_values.insert(key.to_string(), value); + } else { + rejected_alias_keys.push(key.to_string()); + } + } + + if !virtualize_text(&mut snapshot) { + return None; + } + + Some(PreparedSnapshot { + script: snapshot, + aliases: alias_values, + rejected_alias_keys, + }) +} diff --git a/codex-rs/shell-command/src/shell_snapshot_exports.rs b/codex-rs/shell-command/src/shell_snapshot_exports.rs new file mode 100644 index 0000000000..491ba09375 --- /dev/null +++ b/codex-rs/shell-command/src/shell_snapshot_exports.rs @@ -0,0 +1,81 @@ +//! Capture native export declarations and unset POSIX exports. +//! Set POSIX values are captured in bulk for quoting in Rust. + +use crate::shell_detect::ShellType; + +pub(super) fn script(shell_type: ShellType) -> String { + let script = match shell_type { + ShellType::Bash => { + r#"( + while IFS= read -r __codex_snapshot_export_name; do + case "$__codex_snapshot_export_name" in + ""|[0-9]*|*[!A-Za-z0-9_]*|PWD|OLDPWD) continue ;; + esac + RECORD_START + declare -xp "$__codex_snapshot_export_name" 2>/dev/null || true + RECORD_END + done < <(compgen -e) +) +"# + } + ShellType::Zsh => { + r#"( + unsetopt rcquotes + # The bundled Zsh does not include the zsh/parameter module. + for __codex_snapshot_export_name in ${(f)"$(typeset +x)"}; do + case "$__codex_snapshot_export_name" in + ""|[0-9]*|*[!A-Za-z0-9_]*|PWD|OLDPWD) continue ;; + esac + case "${(tP)__codex_snapshot_export_name}" in + *readonly*) continue ;; + *export*) ;; + *) continue ;; + esac + RECORD_START + typeset -xp "$__codex_snapshot_export_name" + RECORD_END + done +) +"# + } + ShellType::Sh => { + r#"if export -p >/dev/null 2>&1; then +export -p | __codex_snapshot_command awk ' +/^(export|declare -x|typeset -x) [A-Za-z_][A-Za-z0-9_]*$/ { + name=$0 + sub(/^(export|declare -x|typeset -x) /, "", name) + if (name ~ /^(PWD|OLDPWD)$/) { + next + } + if (!seen[name]++) { + print name + } +}' | while IFS= read -r __codex_snapshot_export_name; do + # Only the validated identifier enters eval. Set values come from the bulk environment. + if eval '[ "${'"$__codex_snapshot_export_name"'+x}" = x ]'; then + continue + fi + # Only preserve an unset export if the native shell confirms it already exists. + if ( + set +a + set -- "$__codex_snapshot_export_name" "$(export -p)" + export "$1" + [ "$2" = "$(export -p)" ] + ); then + RECORD_START + printf 'export %s\n' "$__codex_snapshot_export_name" + RECORD_END + fi +done +fi +"# + } + ShellType::PowerShell | ShellType::Cmd => return String::new(), + }; + script + .replace( + "RECORD_START", + "printf '%s\\0' \"$__codex_snapshot_export_name\"", + ) + .replace("RECORD_END", "printf '\\0'") +} diff --git a/codex-rs/shell-command/src/shell_snapshot_render.rs b/codex-rs/shell-command/src/shell_snapshot_render.rs new file mode 100644 index 0000000000..77814806c1 --- /dev/null +++ b/codex-rs/shell-command/src/shell_snapshot_render.rs @@ -0,0 +1,69 @@ +//! Serialize prepared snapshot values without inspecting credentials or executing shell code. +//! Native declarations stay unchanged unless the credential stage supplies a replacement. + +use super::capture::CapturedSnapshot; + +impl CapturedSnapshot<'_> { + /// Render native functions, options, and aliases; the executor restores environment separately. + pub fn render_state(&self) -> String { + format!("{}{}", self.state, self.aliases) + } + + /// Render a complete replay script without applying credential or environment policy. + pub fn render_script(&self) -> String { + let mut script = self.render_state(); + script.push_str("# exports (native declarations)\n"); + for export in &self.exports { + script.push_str(&export.source); + } + script + } +} + +#[derive(Default)] +pub(super) struct Value { + pub parts: Vec, +} + +pub(super) enum ValuePart { + Literal(String), + Credential { key: String }, +} + +pub(super) enum Export<'a> { + Captured(&'a str), + Assignment { declaration: &'a str, value: Value }, +} + +impl Value { + fn render(&self) -> Option { + let mut output = String::new(); + for part in &self.parts { + match part { + ValuePart::Literal(value) => { + output.push_str(shlex::try_quote(value).ok()?.as_ref()) + } + ValuePart::Credential { key } => { + output.push_str(&format!("\"${{{key}-}}\"")); + } + } + } + if output.is_empty() { + output.push_str("''"); + } + Some(output) + } +} + +pub(super) fn render(state: &str, aliases: &str, exports: &[Export<'_>]) -> Option { + let mut output = format!("{state}{aliases}"); + for export in exports { + match export { + Export::Captured(source) => output.push_str(source), + Export::Assignment { declaration, value } => { + output.push_str(&format!("{declaration}={}\n", value.render()?)); + } + } + } + Some(output) +} diff --git a/codex-rs/shell-command/src/shell_snapshot_tests.rs b/codex-rs/shell-command/src/shell_snapshot_tests.rs index b280c54e34..a73d4fc558 100644 --- a/codex-rs/shell-command/src/shell_snapshot_tests.rs +++ b/codex-rs/shell-command/src/shell_snapshot_tests.rs @@ -1,11 +1,90 @@ -use super::snapshot_script; +use super::CapturedSnapshot; +use super::PreparedSnapshot; +use super::SnapshotCaptureOptions; +use super::SnapshotCredentialEnvironment; +use super::SnapshotStartup; +use super::prepare_snapshot_credentials; +use super::snapshot_capture_script; use crate::shell_detect::ShellType; +use anyhow::Context; use anyhow::Result; -#[cfg(target_os = "macos")] use pretty_assertions::assert_eq; +use std::collections::HashMap; use std::process::Command; use tempfile::tempdir; +const CAPTURE_ALL: SnapshotCaptureOptions = SnapshotCaptureOptions { + startup: SnapshotStartup::Interactive, + declarations: true, + environment: true, +}; + +fn snapshot_script(shell_type: ShellType) -> Option { + snapshot_capture_script(shell_type, CAPTURE_ALL) +} + +#[test] +fn snapshot_capture_requires_marker_and_complete_records() { + for captured in [ + "missing header\0\0\0", + "# Snapshot file\n\0", + "# Snapshot file\n\0\0NAME\0export NAME=value\n", + ] { + assert!(CapturedSnapshot::parse(ShellType::Bash, captured.as_bytes()).is_none()); + } +} + +#[test] +fn snapshot_capture_selects_declarations_and_environment() -> Result<()> { + let dir = tempdir()?; + for (shell, shell_type) in [ + ("/bin/bash", ShellType::Bash), + ("/bin/sh", ShellType::Sh), + ("/bin/zsh", ShellType::Zsh), + ] { + if !std::path::Path::new(shell).exists() { + continue; + } + for (declarations, environment) in + [(true, false), (false, true), (true, true), (false, false)] + { + let script = snapshot_capture_script( + shell_type, + SnapshotCaptureOptions { + startup: SnapshotStartup::NonInteractive, + declarations, + environment, + }, + ) + .unwrap(); + let output = Command::new(shell) + .args(["-c", &script]) + .env_clear() + .env("HOME", dir.path()) + .env("PATH", "/usr/bin:/bin") + .env("APP_SETTING", "value") + .output()?; + assert!(output.status.success(), "{shell}: {output:?}"); + let captured = CapturedSnapshot::parse(shell_type, &output.stdout).unwrap(); + assert_eq!( + ( + captured + .exports + .iter() + .any(|export| export.key == "APP_SETTING"), + captured + .environment + .split(|byte| *byte == 0) + .any(|entry| entry == b"APP_SETTING=value"), + ), + (declarations, environment), + "{shell}" + ); + } + } + Ok(()) +} + #[test] fn bash_snapshot_filters_invalid_exports() -> Result<()> { let output = Command::new("/bin/bash") @@ -20,7 +99,7 @@ fn bash_snapshot_filters_invalid_exports() -> Result<()> { assert!(output.status.success()); - let stdout = String::from_utf8_lossy(&output.stdout); + let stdout = captured_script(ShellType::Bash, &String::from_utf8(output.stdout)?)?; assert!(stdout.contains("VALID_NAME")); assert!(!stdout.contains("PWD=/tmp/stale")); assert!(!stdout.contains("NEXTEST_BIN_EXE_codex-write-config-schema")); @@ -41,7 +120,7 @@ fn bash_snapshot_preserves_multiline_exports() -> Result<()> { assert!(output.status.success()); - let stdout = String::from_utf8_lossy(&output.stdout); + let stdout = captured_script(ShellType::Bash, &String::from_utf8(output.stdout)?)?; assert!( stdout.contains("MULTILINE_CERT=") || stdout.contains("MULTILINE_CERT"), "snapshot should include the multiline export name" @@ -68,6 +147,429 @@ fn bash_snapshot_preserves_multiline_exports() -> Result<()> { Ok(()) } +#[test] +fn snapshot_environment_preserves_native_path_exports_and_shell_level() -> Result<()> { + let dir = tempdir()?; + for (shell, shell_type, setup) in [ + ("/bin/bash", ShellType::Bash, "export -n PATH"), + ("/bin/zsh", ShellType::Zsh, "typeset +x PATH"), + ] { + if !std::path::Path::new(shell).exists() { + continue; + } + let run = |script: &str| { + // Keep Bash's final-command exec optimization from decrementing SHLVL. + Command::new(shell) + .args(["-c", &format!("{setup}\n{script}\n:")]) + .env_clear() + .env("HOME", dir.path()) + .env("PATH", "/usr/bin:/bin") + .env("SHLVL", "3") + .output() + }; + let native = run("/usr/bin/env -0")?; + let script = snapshot_capture_script( + shell_type, + SnapshotCaptureOptions { + startup: SnapshotStartup::NonInteractive, + ..CAPTURE_ALL + }, + ) + .unwrap(); + let output = run(&script)?; + assert!(native.status.success() && output.status.success()); + let captured = CapturedSnapshot::parse(shell_type, &output.stdout).unwrap(); + let metadata = |environment: &[u8]| { + environment + .split(|byte| *byte == 0) + .filter(|entry| entry.starts_with(b"PATH=") || entry.starts_with(b"SHLVL=")) + .map(<[u8]>::to_vec) + .collect::>() + }; + assert_eq!( + metadata(captured.environment), + metadata(&native.stdout), + "{shell}" + ); + } + Ok(()) +} + +#[test] +fn sh_declarations_preserve_path_export_state() -> Result<()> { + let dir = tempdir()?; + let snapshot_path = dir.path().join("snapshot.sh"); + let script = snapshot_capture_script( + ShellType::Sh, + SnapshotCaptureOptions { + startup: SnapshotStartup::NonInteractive, + ..CAPTURE_ALL + }, + ) + .expect("sh supports snapshots"); + for shell in [ + "/bin/sh", + #[cfg(target_os = "macos")] + "/bin/dash", + ] { + if !std::path::Path::new(shell).exists() { + continue; + } + for (setup, expected) in [ + ("PATH=/enterprise/bin; export PATH", "x:/enterprise/bin"), + ("PATH=; export PATH", "x:"), + ("unset PATH; export PATH", ":"), + ("unset PATH; PATH=/enterprise/bin", ":"), + ] { + let captured = Command::new(shell) + .args(["-uc", &format!("{setup}\n{script}")]) + .env_clear() + .output()?; + assert!(captured.status.success(), "capture failed: {captured:?}"); + std::fs::write( + &snapshot_path, + captured_script(ShellType::Sh, &String::from_utf8(captured.stdout)?)?, + )?; + let restored = Command::new(shell) + .args([ + "-c", + "unset PATH; . \"$1\"; printf '%s:%s' \"${PATH+x}\" \"${PATH-}\"", + "snapshot", + ]) + .arg(&snapshot_path) + .env_clear() + .output()?; + assert!(restored.status.success(), "replay failed: {restored:?}"); + assert_eq!(restored.stdout, expected.as_bytes(), "{setup}"); + } + } + Ok(()) +} + +#[test] +fn posix_startup_path_expansion_preserves_supported_forms() -> Result<()> { + let script = format!( + "{}\n__codex_snapshot_expand_env \"$ENV\"", + super::posix_env_path_expansion_function() + ); + let output = Command::new("/bin/sh") + .arg("-uc") + .arg(&script) + .env("ENV", "$MISSING") + .output()?; + + assert!(output.status.success(), "snapshot failed: {output:?}"); + assert_eq!(output.stdout, b"$MISSING"); + + let dir = tempdir()?; + let startup = dir.path().join("startup.sh"); + for (env_value, expansion_env) in [ + ( + startup.display().to_string(), + Vec::<(String, String)>::new(), + ), + ( + "$STARTUP_FILE".to_string(), + vec![("STARTUP_FILE".to_string(), startup.display().to_string())], + ), + ( + "${STARTUP_FILE}".to_string(), + vec![("STARTUP_FILE".to_string(), startup.display().to_string())], + ), + ( + "${STARTUP_DIR}/startup.sh".to_string(), + vec![("STARTUP_DIR".to_string(), dir.path().display().to_string())], + ), + ( + "${PATH%%:*}/startup.sh".to_string(), + vec![( + "PATH".to_string(), + format!("{}:/usr/bin:/bin", dir.path().display()), + )], + ), + ( + "$PATH/startup.sh".to_string(), + vec![("PATH".to_string(), dir.path().display().to_string())], + ), + ( + "${PATH}/startup.sh".to_string(), + vec![("PATH".to_string(), dir.path().display().to_string())], + ), + ( + format!("$PATH{}", startup.display()), + vec![("PATH".to_string(), String::new())], + ), + ( + format!("${{PATH}}{}", startup.display()), + vec![("PATH".to_string(), String::new())], + ), + ( + "~/startup.sh".to_string(), + vec![("HOME".to_string(), dir.path().display().to_string())], + ), + ] { + let output = Command::new("/bin/sh") + .arg("-uc") + .arg(&script) + .env("ENV", &env_value) + .envs(expansion_env) + .output()?; + + assert!(output.status.success(), "snapshot failed: {output:?}"); + assert_eq!( + String::from_utf8(output.stdout)?, + startup.display().to_string(), + "ENV={env_value}" + ); + } + + let default_startup_dir = dir.path().join(".config"); + for xdg_config_home in [ + None, + Some(std::ffi::OsStr::new("")), + Some(default_startup_dir.as_os_str()), + ] { + let mut command = Command::new("/bin/sh"); + command + .arg("-uc") + .arg(&script) + .env("HOME", dir.path()) + .env("ENV", "${XDG_CONFIG_HOME:-$HOME/.config}/startup.sh"); + if let Some(xdg_config_home) = xdg_config_home { + command.env("XDG_CONFIG_HOME", xdg_config_home); + } else { + command.env_remove("XDG_CONFIG_HOME"); + } + let output = command.output()?; + + assert!(output.status.success(), "snapshot failed: {output:?}"); + assert_eq!( + String::from_utf8(output.stdout)?, + default_startup_dir.join("startup.sh").display().to_string() + ); + } + + let injected_marker = dir.path().join("env-injected"); + let injected_env = format!( + "\"; touch '{}'; __codex_env_file=\"", + injected_marker.display() + ); + let output = Command::new("/bin/sh") + .arg("-uc") + .arg(&script) + .env("ENV", injected_env) + .output()?; + assert!(output.status.success(), "snapshot failed: {output:?}"); + assert!( + !injected_marker.exists(), + "snapshot interpreted ENV as shell source" + ); + + Ok(()) +} + +#[test] +fn brokered_snapshots_filter_complete_multiline_exports() -> Result<()> { + let dir = tempdir()?; + for (shell, shell_type) in [ + ("/bin/bash", ShellType::Bash), + ("/bin/sh", ShellType::Sh), + ("/bin/dash", ShellType::Sh), + ("/bin/zsh", ShellType::Zsh), + ] { + if !std::path::Path::new(shell).exists() { + continue; + } + for value in [ + "before\nexport NOT_A_REAL_VARIABLE=value\nafter", + "before'\ndeclare -x NOT_A_REAL_VARIABLE=\"quoted\"\ntypeset OTHER=value\nafter\\", + "before\t\n# exports 99\nexport NOT_A_REAL_VARIABLE=value\nafter", + "certificate\nwith trailing newlines\n\n", + "before\nexport NOT_A_REAL_VARIABLE\nafter", + "before\nexport name=not_the_actual_value\nafter", + ] { + let original = HashMap::from([ + ("APP_SCRIPT".to_string(), value.to_string()), + ("EXCLUDED_SCRIPT".to_string(), value.to_string()), + ("ZZZ_AFTER".to_string(), "kept".to_string()), + ("name".to_string(), "original name".to_string()), + ("value".to_string(), "original value".to_string()), + ("escaped".to_string(), "original escaped".to_string()), + ]); + let mut allowed = original.clone(); + allowed.remove("EXCLUDED_SCRIPT"); + let function_value = "before\n# exports 1\nexport NOT_A_REAL_VARIABLE=value\nafter"; + let capture_script = snapshot_capture_script( + shell_type, + SnapshotCaptureOptions { + startup: SnapshotStartup::NonInteractive, + ..CAPTURE_ALL + }, + ) + .expect("shell supports snapshots"); + let capture = if matches!(shell_type, ShellType::Sh) { + capture_script + } else { + format!("emit_source() {{ cat <<'EOF'\n{function_value}\nEOF\n}}\n{capture_script}") + }; + let capture = if matches!(shell_type, ShellType::Sh | ShellType::Zsh) { + format!("export APP_SETTING EXCLUDED_SETTING\n{capture}") + } else { + capture + }; + let captured = Command::new(shell) + .args(["-c", &capture]) + .env_clear() + .envs(&original) + .env("PATH", "/usr/bin:/bin") + .output()?; + assert!(captured.status.success(), "{shell}: {captured:?}"); + let mut snapshot = String::from_utf8(captured.stdout)?; + let records = CapturedSnapshot::parse(shell_type, snapshot.as_bytes()).unwrap(); + assert_eq!( + records + .exports + .iter() + .filter(|export| export.key == "name") + .count(), + 1, + "{shell}: duplicate export record" + ); + if matches!(shell_type, ShellType::Sh) { + let raw_path = dir.path().join("raw-snapshot.sh"); + std::fs::write(&raw_path, captured_script(shell_type, &snapshot)?)?; + let replayed = Command::new(shell) + .args(["-c", ". \"$1\"; printf '%s\\n' \"${APP_SETTING-unset}\"; APP_SETTING=production; /usr/bin/printenv APP_SETTING; NOT_A_REAL_VARIABLE=unexpected; /usr/bin/printenv NOT_A_REAL_VARIABLE", "snapshot"]) + .arg(&raw_path) + .env_clear() + .output()?; + assert_eq!( + replayed.stdout, b"unset\nproduction\n", + "{shell}: {snapshot}" + ); + assert_eq!(replayed.status.code(), Some(1)); + } + rewrite_snapshot_credentials( + shell_type, + &mut snapshot, + SnapshotCredentialEnvironment { + original: &original, + restored: &original, + configured: &HashMap::new(), + discovered: &allowed, + allowed: &allowed, + is_allowed_unset: &|key| key == "APP_SETTING", + brokered_keys: &[], + brokered_alias_keys: &[], + allowed_brokered_keys: &[], + }, + |_| true, + ); + let path = dir.path().join("snapshot.sh"); + std::fs::write(&path, snapshot)?; + let replayed = Command::new(shell) + .args(["-c", ". \"$1\"; printf '%s\\0%s\\0%s\\0%s\\0%s\\0%s' \"$APP_SCRIPT\" \"${EXCLUDED_SCRIPT-unset}\" \"$ZZZ_AFTER\" \"$name\" \"$value\" \"$escaped\"; if command -v emit_source >/dev/null; then emit_source; fi; APP_SETTING=production; EXCLUDED_SETTING=denied; /usr/bin/printenv APP_SETTING; /usr/bin/printenv EXCLUDED_SETTING || :", "snapshot"]) + .arg(&path) + .env_clear() + .output()?; + assert!(replayed.status.success(), "{shell}: {replayed:?}"); + let function_value = if matches!(shell_type, ShellType::Sh) { + String::new() + } else { + format!("{function_value}\n") + }; + let exported_setting = if matches!(shell_type, ShellType::Sh | ShellType::Zsh) { + "production\n" + } else { + "" + }; + assert_eq!( + replayed.stdout, + format!("{value}\0unset\0kept\0original name\0original value\0original escaped{function_value}{exported_setting}").as_bytes(), + "{shell}" + ); + } + } + Ok(()) +} + +#[test] +fn brokered_snapshot_alias_values_follow_allowed_credential_overrides() { + let real = "ghp_abcdefghijklmnopqrstuvwxyz0123456789"; + let discovered_dummy = "ghp_discovered_dummy_abcdefghijklmnopqr"; + let configured_dummy = "ghp_configured_dummy_abcdefghijklmnopqr"; + let original = HashMap::from([ + ("GH_TOKEN".to_string(), real.to_string()), + ("AUTH_HEADER".to_string(), format!("Bearer {real}")), + ]); + let discovered = HashMap::from([("GH_TOKEN".to_string(), discovered_dummy.to_string())]); + let brokered_keys = vec!["GH_TOKEN".to_string()]; + for (source_allowed, alias, expected) in [ + ( + true, + format!("Bearer {discovered_dummy}"), + Some(( + format!("Bearer {configured_dummy}"), + r#"export AUTH_HEADER='Bearer '"${GH_TOKEN-}""#, + )), + ), + ( + true, + String::new(), + Some((String::new(), "export AUTH_HEADER=''")), + ), + (false, format!("Bearer {discovered_dummy}"), None), + ] { + let mut allowed = HashMap::from([ + ("GH_TOKEN".to_string(), configured_dummy.to_string()), + ("AUTH_HEADER".to_string(), alias), + ]); + let allowed_brokered_keys = if source_allowed { + brokered_keys.as_slice() + } else { + allowed.remove("GH_TOKEN"); + &[] + }; + let mut snapshot = + format!("# Snapshot file\n\0\0AUTH_HEADER\0export AUTH_HEADER='Bearer {real}'\n\0\0"); + let prepared = rewrite_snapshot_credentials( + ShellType::Bash, + &mut snapshot, + SnapshotCredentialEnvironment { + original: &original, + restored: &original, + configured: &HashMap::new(), + discovered: &discovered, + allowed: &allowed, + is_allowed_unset: &|_| false, + brokered_keys: &brokered_keys, + brokered_alias_keys: &[], + allowed_brokered_keys, + }, + |text| { + *text = text.replace(real, discovered_dummy); + true + }, + ); + let expected = match expected { + Some((value, assignment)) => ( + format!("# Snapshot file\n{assignment}\n"), + HashMap::from([("AUTH_HEADER".to_string(), value)]), + Vec::new(), + ), + None => ( + "# Snapshot file\n".to_string(), + HashMap::new(), + vec!["AUTH_HEADER".to_string()], + ), + }; + assert_eq!( + prepared.map(|rewrite| (rewrite.script, rewrite.aliases, rewrite.rejected_alias_keys)), + Some(expected), + ); + } +} + #[cfg(target_os = "macos")] #[test] fn zsh_snapshot_restores_tied_path() -> Result<()> { @@ -98,7 +600,8 @@ fn zsh_snapshot_restores_tied_path() -> Result<()> { assert!(snapshot.status.success()); let snapshot_path = dir.path().join("snapshot.sh"); - std::fs::write(&snapshot_path, &snapshot.stdout)?; + let snapshot = captured_script(ShellType::Zsh, &String::from_utf8(snapshot.stdout)?)?; + std::fs::write(&snapshot_path, &snapshot)?; let restored = Command::new("/bin/zsh") .arg("-f") @@ -115,7 +618,6 @@ fn zsh_snapshot_restores_tied_path() -> Result<()> { expected_path ); - let snapshot = String::from_utf8(snapshot.stdout)?; assert!( snapshot .lines() @@ -133,7 +635,11 @@ fn zsh_snapshot_restores_tied_path() -> Result<()> { .env("ZDOTDIR", dir.path()) .output()?; assert!(readonly_snapshot.status.success()); - std::fs::write(&snapshot_path, &readonly_snapshot.stdout)?; + let readonly_snapshot = captured_script( + ShellType::Zsh, + &String::from_utf8(readonly_snapshot.stdout)?, + )?; + std::fs::write(&snapshot_path, &readonly_snapshot)?; let readonly_restored = Command::new("/bin/zsh") .arg("-f") @@ -150,7 +656,6 @@ fn zsh_snapshot_restores_tied_path() -> Result<()> { "/codex-path:/usr/bin:/bin" ); - let readonly_snapshot = String::from_utf8(readonly_snapshot.stdout)?; assert!( !readonly_snapshot .lines() @@ -160,3 +665,36 @@ fn zsh_snapshot_restores_tied_path() -> Result<()> { Ok(()) } + +// Literal-only fixtures are shell source rather than capture output. Give those fixtures +// empty alias/export sections; real shell fixtures exercise the full framed capture. +fn rewrite_snapshot_credentials( + shell_type: ShellType, + snapshot: &mut String, + environment: SnapshotCredentialEnvironment<'_>, + virtualize_text: impl FnMut(&mut String) -> bool, +) -> Option { + let source_only = !snapshot.contains('\0'); + let framed = if source_only { + format!("# Snapshot file\n{snapshot}\0\0\0") + } else { + snapshot.clone() + }; + let capture = CapturedSnapshot::parse(shell_type, framed.as_bytes())?; + let prepared = prepare_snapshot_credentials(&capture, environment, virtualize_text)?; + *snapshot = if source_only { + prepared + .script + .strip_prefix("# Snapshot file\n")? + .to_string() + } else { + prepared.script.clone() + }; + Some(prepared) +} + +fn captured_script(shell_type: ShellType, source: &str) -> Result { + let captured = + CapturedSnapshot::parse(shell_type, source.as_bytes()).context("invalid native capture")?; + Ok(captured.render_script()) +}