openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-25 18:28:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve complete shell snapshot exports through filtering and replay (#43907)

Browse Source 
## 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
This commit is contained in:
Winston Howes
2026-09-08 20:41:01 +00:00
committed by copyberry
parent f046cf35df
commit 1530f828cb
11 changed files with 1472 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/shell-command/src/shell_snapshot.rs
+99 -216
View File
@@ -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<String> {
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<String> {
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,
}
codex-rs/shell-command/src/shell_snapshot_capture.rs
+335
View File
@@ -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<String> {
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<CapturedStartupEnvironment<'a>>,
pub(super) state: &'a str,
pub(super) aliases: &'a str,
pub(super) exports: Vec<CapturedExport<'a>>,
/// 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<Self> {
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)
}
codex-rs/shell-command/src/shell_snapshot_credentials.rs
+216
View File
@@ -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<String, String>,
pub restored: &'a HashMap<String, String>,
pub configured: &'a HashMap<String, String>,
pub discovered: &'a HashMap<String, String>,
pub allowed: &'a HashMap<String, String>,
/// 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<String, String>,
pub rejected_alias_keys: Vec<String>,
}
/// 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<PreparedSnapshot> {
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::<Vec<_>>();
(!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::<HashMap<_, _>>();
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::<Vec<_>>();
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,
})
}
codex-rs/shell-command/src/shell_snapshot_exports.rs
+81
View File
@@ -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'")
}
codex-rs/shell-command/src/shell_snapshot_render.rs
+69
View File
@@ -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<ValuePart>,
}
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<String> {
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<String> {
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)
}
codex-rs/shell-command/src/shell_snapshot_tests.rs
+546 -8
View File
@@ -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<String> {
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::<Vec<_>>()
};
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<PreparedSnapshot> {
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<String> {
let captured =
CapturedSnapshot::parse(shell_type, source.as_bytes()).context("invalid native capture")?;
Ok(captured.render_script())
}