github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 08:16:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support remote credential masking with Zsh shell snapshots (#54432)

Browse Source 
## Why

Remote credential masking previously supported only Bash commands without login startup or shell snapshots, leaving Zsh commands with unmasked credentials.

## What changed

- Enable masking for Zsh login commands using shell snapshots on Linux and macOS.
- Capture trusted startup profiles per launch, bind the credential broker to the resulting environment, and replace credentials and supported aliases with placeholders before snapshot replay.
- Preserve explicit environment overrides, credential unsets, and launcher-managed proxy settings during replay.
- Warn and fall back to unmasked startup when a snapshot cannot safely replay credentials; stop execution if protected snapshot restoration fails.

## Testing

Extend remote execution tests for inherited and profile-defined credentials, overrides, destination changes, and unsupported-profile fallback. Add a Zsh scenario snapshot showing a credential placeholder in command output.

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