From a548463b78773802b00b0a226b89ff48876a4e1e Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Wed, 9 Sep 2026 03:28:59 +0000 Subject: [PATCH] Handle copied credentials in the broker and shell snapshots (#44038) ## Why Shell startup can copy a credential into another variable and unset its original source. These aliases need credential brokering without losing source restrictions or mistaking ordinary shell content for credentials. ## What changed - Discover supported GitHub and OpenAI credentials embedded in environment values, even without their canonical variables. Preserve source ownership so unbound enterprise tokens cannot acquire a default host binding. - Redact unregistered supported credentials during text virtualization, distinguish adjacent credentials, and avoid matching unrelated provider prefixes and common hashed paths. - Expose helpers for checking allowed credential sources and restoring known dummy credentials in trusted text for fail-open execution. - Render allowed credential aliases as snapshot exports and replace known credential values in shell state, preserving exported functions and credential-shaped function names. ## Testing Add regression coverage for copied aliases, enterprise host binding, source filtering, adjacent tokens, path false positives, and dummy restoration. Add a Bash snapshot test that preserves exported functions while replacing credentials in heredoc content. GitOrigin-RevId: f1c3d531405a96ac6bfa3882a11b74643fa6f123 --- .../network-proxy/src/credential_broker.rs | 219 ++++++-- .../src/credential_broker/matching.rs | 315 +++++++++++ .../src/credential_broker/providers.rs | 4 + .../src/credential_broker/providers/github.rs | 4 + .../src/credential_broker/providers/openai.rs | 16 +- .../src/credential_broker_tests.rs | 500 +++++++++++++++++- codex-rs/network-proxy/src/lib.rs | 1 + codex-rs/network-proxy/src/proxy.rs | 5 + codex-rs/network-proxy/src/runtime.rs | 4 + .../src/shell_snapshot_credentials.rs | 25 +- .../src/shell_snapshot_render.rs | 4 + .../shell-command/src/shell_snapshot_tests.rs | 82 +++ 12 files changed, 1126 insertions(+), 53 deletions(-) create mode 100644 codex-rs/network-proxy/src/credential_broker/matching.rs diff --git a/codex-rs/network-proxy/src/credential_broker.rs b/codex-rs/network-proxy/src/credential_broker.rs index 04e9b8d5cf..10722da569 100644 --- a/codex-rs/network-proxy/src/credential_broker.rs +++ b/codex-rs/network-proxy/src/credential_broker.rs @@ -1,3 +1,4 @@ +mod matching; mod providers; use crate::config::NetworkProxyConfig; @@ -22,9 +23,15 @@ struct CredentialBrokerState { enabled: bool, openai_api_host: Option, credentials: Vec, + credential_owners: Vec, credential_aliases: Vec, } +struct CredentialOwner { + env_var: String, + real_value: String, +} + struct CredentialRecord { env_var: String, provider: &'static providers::CredentialProvider, @@ -91,6 +98,7 @@ impl CredentialBroker { if state.enabled != config.credential_broker { state.enabled = config.credential_broker; state.credentials.clear(); + state.credential_owners.clear(); state.credential_aliases.clear(); } if state.openai_api_host != config.credential_broker_openai_host { @@ -116,6 +124,7 @@ impl CredentialBroker { if !state.enabled { return; } + state.observe_credential_owners(parent_env); for provider in providers::credential_providers() { for source in provider.sources() { @@ -154,6 +163,7 @@ impl CredentialBroker { return; } set_env_value(env, CREDENTIAL_BROKER_ACTIVE_ENV_KEY, "1".to_string()); + state.observe_credential_owners(env); for provider in providers::credential_providers() { for source in provider.sources() { @@ -167,6 +177,71 @@ impl CredentialBroker { } } } + for provider in providers::credential_providers() { + let Some((source, host_binding)) = provider.sources().iter().rev().find_map(|source| { + (source.host_binding)(env, state.openai_api_host.as_deref()) + .map(|binding| (source, binding)) + }) else { + continue; + }; + for (key, value) in env.iter() { + if key.eq_ignore_ascii_case("PATH") || key.to_ascii_uppercase().ends_with("_PATH") { + continue; + } + for prefix in provider.credential_prefixes { + for (start, _) in value.match_indices(prefix) { + let credential = + matching::builtin_credential_candidate(provider, value, start); + if credential.len() < provider.minimum_credential_len + || matching::is_operational_path_match( + value, + start, + start + credential.len(), + ) + || provider.ignored_credential_prefixes.iter().any(|ignored| { + credential.starts_with(ignored) + && provider + .credential_watermark + .is_none_or(|watermark| !credential.contains(watermark)) + }) + || state.is_dummy_value(credential) + || source.binding_env_vars.is_empty() + && state.credential_owners.iter().any(|existing| { + !source + .env_vars + .iter() + .any(|key| env_key_matches(key, &existing.env_var)) + && existing.real_value == credential + }) + || state.credentials.iter().any(|existing| { + std::ptr::eq(existing.provider, provider) + && existing.real_value == credential + }) + || provider.sources().iter().any(|candidate| { + candidate + .env_vars + .iter() + .any(|key| env_value(env, key) == Some(credential)) + && (candidate.host_binding)( + env, + state.openai_api_host.as_deref(), + ) + .is_none() + }) + || provider.request_header_value(credential).is_none() + { + continue; + } + state.register( + source.env_vars[0], + provider, + host_binding.clone(), + credential, + ); + } + } + } + } let credentials = prioritized_credentials(&state, env); let mut credential_aliases = Vec::new(); for (key, value) in env.iter_mut() { @@ -281,46 +356,32 @@ impl CredentialBroker { } pub(crate) fn virtualize_text(&self, text: &mut String, env: &HashMap) -> bool { + let state = self.read_state(); + matching::virtualize_text(&state, text, env) + } + + pub(crate) fn restore_text(&self, text: &mut String) -> bool { let state = self.read_state(); if !state.enabled { - return true; + return false; } - let allowed_keys = brokered_credential_dummy_env_keys(env); - let credentials = prioritized_credentials(&state, env); - let mut allowed = true; - for credential in &credentials { - let contains_real = text.as_str() == credential.real_value - || credential.real_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH - && text.contains(&credential.real_value); - let contains_dummy = text.as_str() == credential.dummy_value - || credential.dummy_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH - && text.contains(&credential.dummy_value); - if !contains_real && !contains_dummy { - continue; - } - - let replacement = credentials.iter().copied().find(|candidate| { - std::ptr::eq(candidate.provider, credential.provider) - && candidate.real_value == credential.real_value - && allowed_keys.iter().any(|key| { - env_key_matches(key, &candidate.env_var) - && env_value(env, key) == Some(candidate.dummy_value.as_str()) - }) - }); - if replacement.is_none() { - allowed = false; - } - let replacement = replacement.map_or("", |candidate| candidate.dummy_value.as_str()); - if contains_real { - *text = text.replace(&credential.real_value, replacement); - } - if contains_dummy && credential.dummy_value != replacement { - *text = text.replace(&credential.dummy_value, replacement); + let mut credentials = state.credentials.iter().collect::>(); + credentials + .sort_unstable_by_key(|credential| std::cmp::Reverse(credential.dummy_value.len())); + let mut restored = false; + for credential in credentials { + if text.as_str() == credential.dummy_value { + text.clone_from(&credential.real_value); + restored = true; + } else if credential.dummy_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH + && text.contains(&credential.dummy_value) + { + *text = text.replace(&credential.dummy_value, &credential.real_value); + restored = true; } } - - allowed + restored } pub(crate) fn inject_request_headers(&self, host: &str, headers: &mut HeaderMap) { @@ -403,6 +464,32 @@ fn brokerable_credential_value<'a>( } impl CredentialBrokerState { + fn observe_credential_owners(&mut self, env: &HashMap) { + // Ownership is known even when a destination is missing or invalid. + for provider in providers::credential_providers() { + for env_var in provider.sources().iter().flat_map(|source| source.env_vars) { + if let Some(real_value) = + brokerable_credential_value(env, self, env_var, provider).map(str::to_string) + { + self.remember_credential_owner(env_var, &real_value); + } + } + } + } + + fn remember_credential_owner(&mut self, env_var: &str, real_value: &str) { + if !self + .credential_owners + .iter() + .any(|owner| env_key_matches(&owner.env_var, env_var) && owner.real_value == real_value) + { + self.credential_owners.push(CredentialOwner { + env_var: env_var.to_string(), + real_value: real_value.to_string(), + }); + } + } + fn register( &mut self, env_var: &str, @@ -410,6 +497,7 @@ impl CredentialBrokerState { host_binding: providers::CredentialHostBinding, real_value: &str, ) -> String { + self.remember_credential_owner(env_var, real_value); if let Some(existing) = self.credentials.iter().find(|credential| { credential.env_var == env_var && std::ptr::eq(credential.provider, provider) @@ -588,6 +676,69 @@ pub fn credential_broker_provider_context_env_keys() -> impl Iterator, + is_allowed: impl Fn(&str) -> bool, +) -> bool { + let mut recognized = false; + let allowed = providers::credential_providers() + .filter(move |provider| { + provider.credential_prefixes.iter().any(|prefix| { + value.match_indices(*prefix).any(|(start, _)| { + matching::recognized_credential_match(provider, value, virtualized, start) + .is_some() + }) + }) + }) + .all(|provider| { + recognized = true; + let actual_sources = provider + .sources() + .iter() + .flat_map(|source| source.env_vars.iter().copied()) + .filter(|source| { + env_value(source_env, source).is_some_and(|source_value| { + source_value.len() >= provider.minimum_credential_len + && value.contains(source_value) + }) + }) + .collect::>(); + let unattributed = provider.credential_prefixes.iter().any(|prefix| { + value.match_indices(*prefix).any(|(start, _)| { + matching::recognized_credential_match(provider, value, virtualized, start) + .is_some_and(|credential| { + !actual_sources.iter().any(|source| { + env_value(source_env, source).is_some_and(|source_value| { + credential == source_value + || credential + .strip_prefix(source_value) + .is_some_and(|suffix| suffix.starts_with(['_', '-'])) + }) + }) + }) + }) + }); + if actual_sources.is_empty() || unattributed { + provider + .sources() + .iter() + .flat_map(|source| source.env_vars.iter().copied()) + .all(&is_allowed) + } else { + actual_sources.iter().all(|source| { + actual_sources.iter().any(|equivalent| { + env_value(source_env, source) == env_value(source_env, equivalent) + && is_allowed(equivalent) + }) + }) + } + }); + allowed && recognized +} + /// Returns whether an environment key belongs to a supported credential provider. pub fn is_credential_broker_provider_env_key(key: &str) -> bool { providers::credential_providers().any(|provider| { diff --git a/codex-rs/network-proxy/src/credential_broker/matching.rs b/codex-rs/network-proxy/src/credential_broker/matching.rs new file mode 100644 index 0000000000..328805bfa3 --- /dev/null +++ b/codex-rs/network-proxy/src/credential_broker/matching.rs @@ -0,0 +1,315 @@ +use super::CredentialBrokerState; +use super::MIN_EMBEDDED_CREDENTIAL_LENGTH; +use super::brokered_credential_dummy_env_keys; +use super::env_key_matches; +use super::env_value; +use super::prioritized_credentials; +use super::providers; +use std::collections::HashMap; +use std::path::Path; +use url::Position; +use url::Url; + +pub(super) fn virtualize_text( + state: &CredentialBrokerState, + text: &mut String, + env: &HashMap, +) -> bool { + if !state.enabled { + return true; + } + + let allowed_keys = brokered_credential_dummy_env_keys(env); + let credentials = prioritized_credentials(state, env); + let mut allowed = true; + for credential in &credentials { + let contains_real = text.as_str() == credential.real_value + || credential.real_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH + && text.contains(&credential.real_value); + let contains_dummy = text.as_str() == credential.dummy_value + || credential.dummy_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH + && text.contains(&credential.dummy_value); + if !contains_real && !contains_dummy { + continue; + } + + let replacement = credentials.iter().copied().find(|candidate| { + std::ptr::eq(candidate.provider, credential.provider) + && candidate.real_value == credential.real_value + && (allowed_keys.iter().any(|key| { + env_key_matches(key, &candidate.env_var) + && env_value(env, key) == Some(candidate.dummy_value.as_str()) + }) || state.credential_aliases.iter().any(|alias| { + env_value(env, &alias.env_var) == Some(alias.dummy_value.as_str()) + && (alias.dummy_value == candidate.dummy_value + || candidate.dummy_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH + && alias.dummy_value.contains(&candidate.dummy_value)) + && !state.credentials.iter().any(|other| { + other.dummy_value != candidate.dummy_value + && (alias.dummy_value == other.dummy_value + || other.dummy_value.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH + && alias.dummy_value.contains(&other.dummy_value)) + }) + })) + }); + if replacement.is_none() { + allowed = false; + } + let replacement = replacement.map_or("", |candidate| candidate.dummy_value.as_str()); + if contains_real { + *text = text.replace(&credential.real_value, replacement); + } + if contains_dummy && credential.dummy_value != replacement { + *text = text.replace(&credential.dummy_value, replacement); + } + } + + // Startup can copy a supported credential before unsetting its source variable. + for provider in providers::credential_providers() { + for prefix in provider.credential_prefixes { + let mut offset = 0; + while let Some(position) = text[offset..].find(prefix) { + let start = offset + position; + if let Some(length) = state + .credentials + .iter() + .flat_map(|credential| [&credential.real_value, &credential.dummy_value]) + .filter(|credential| credential.len() >= MIN_EMBEDDED_CREDENTIAL_LENGTH) + .filter(|credential| text[start..].starts_with(credential.as_str())) + .map(String::len) + .max() + { + offset = start + length; + continue; + } + let candidate = builtin_credential_candidate(provider, text, start); + let length = candidate.len(); + if provider + .ignored_credential_prefixes + .iter() + .any(|prefix| candidate.starts_with(prefix)) + && provider + .credential_watermark + .is_none_or(|watermark| !candidate.contains(watermark)) + { + let known_length = env + .values() + .filter(|known| { + known.len() >= provider.minimum_credential_len + && candidate + .strip_prefix(known.as_str()) + .is_some_and(|suffix| { + provider.credential_prefixes.iter().any(|prefix| { + suffix.match_indices(prefix).any(|(offset, _)| { + suffix.len() - offset + >= provider.minimum_credential_len + }) + }) + }) + }) + .map(String::len) + .min(); + let embedded_supported = provider + .credential_prefixes + .iter() + .flat_map(|prefix| candidate.match_indices(prefix)) + .filter_map(|(offset, _)| { + (offset > 0 + && candidate.len() - offset >= provider.minimum_credential_len + && !provider + .ignored_credential_prefixes + .iter() + .any(|ignored| candidate[offset..].starts_with(ignored))) + .then_some(offset) + }) + .min(); + offset = start + + known_length + .into_iter() + .chain(embedded_supported) + .min() + .unwrap_or(length); + continue; + } + let end = start + length; + let credential = &text[start..end]; + let ignored_credential_match = + ignored_credential_match(provider, text, start, credential); + if length >= provider.minimum_credential_len && !ignored_credential_match { + text.replace_range(start..end, ""); + allowed = false; + } else { + offset = start + + if ignored_credential_match { + prefix.len() + } else { + length + }; + } + } + } + } + + allowed +} + +pub(super) fn is_operational_path_match(text: &str, start: usize, end: usize) -> bool { + let is_value_boundary = |character: char| { + character.is_ascii_whitespace() || matches!(character, '"' | '\'' | '=' | '`') + }; + let value_start = text[..start] + .rfind(is_value_boundary) + .map_or(0, |index| index + 1); + let value_end = text[end..] + .find(is_value_boundary) + .map_or(text.len(), |index| end + index); + let value = &text[value_start..value_end]; + if let Ok(url) = Url::parse(value) + && url.has_host() + { + let relative_start = start - value_start; + let relative_end = end - value_start; + return relative_start >= url[..Position::BeforeHost].len() + && relative_end <= url[..Position::AfterPort].len(); + } + + let relative_start = start - value_start; + let relative_end = end - value_start; + if !value[..relative_start].contains(['/', '\\']) + && !value[relative_end..].contains(['/', '\\']) + { + return false; + } + + let path = Path::new(value); + path.has_root() || path.components().count() > 1 || value.contains('\\') +} + +fn ignored_credential_match( + provider: &providers::CredentialProvider, + text: &str, + start: usize, + credential: &str, +) -> bool { + provider + .ignored_credential_prefixes + .iter() + .any(|prefix| credential.starts_with(prefix)) + && provider + .credential_watermark + .is_none_or(|watermark| !credential.contains(watermark)) + || credential.starts_with("sk-") + && provider + .credential_watermark + .is_none_or(|watermark| !credential.contains(watermark)) + && credential[3..].split(['-', '_']).any(|segment| { + segment.len() == 64 + && segment.bytes().all(|byte| byte.is_ascii_hexdigit()) + && credential + .find(segment) + .is_some_and(|offset| offset < provider.minimum_credential_len) + }) + && text[..start] + .rsplit(|character: char| { + character.is_ascii() && !character.is_ascii_alphanumeric() + }) + .next() + .is_some_and(|word| { + ((1..=3).contains(&word.len()) + || word + .get(word.len().saturating_sub(2)..) + .is_some_and(|suffix| { + suffix.eq_ignore_ascii_case("di") + || suffix.eq_ignore_ascii_case("ta") + && !word.eq_ignore_ascii_case("data") + }) + || credential.strip_prefix("sk-").is_some_and(|hash| { + hash.len() == 64 && hash.bytes().all(|byte| byte.is_ascii_hexdigit()) + })) + && word.bytes().all(|byte| byte.is_ascii_alphabetic()) + && text[..start] + .rsplit_once(['/', '\\']) + .is_some_and(|(_, component)| { + component.chars().all(|character| { + !character.is_ascii() + || character.is_ascii_alphanumeric() + || matches!(character, '_' | '-') + }) + }) + }) +} + +pub(super) fn recognized_credential_match<'a>( + provider: &providers::CredentialProvider, + value: &'a str, + virtualized: &str, + start: usize, +) -> Option<&'a str> { + let credential = builtin_credential_candidate(provider, value, start); + let length = credential.len(); + let enclosing_start = value.as_bytes()[..start] + .iter() + .rposition(|byte| !byte.is_ascii_alphanumeric() && !matches!(byte, b'_' | b'-')) + .map_or(0, |offset| offset + 1); + let enclosing = &value[enclosing_start..start + length]; + (length >= provider.minimum_credential_len || !virtualized.contains(credential)) + .then_some(credential) + .filter(|credential| !ignored_credential_match(provider, value, start, credential)) + .filter(|_| { + !provider.ignored_credential_prefixes.iter().any(|prefix| { + enclosing.match_indices(prefix).any(|(offset, _)| { + let ignored = &enclosing[offset..]; + offset <= start - enclosing_start + && ignored_credential_match( + provider, + value, + enclosing_start + offset, + ignored, + ) + && virtualized.contains(ignored) + }) + }) + }) +} + +pub(super) fn builtin_credential_candidate<'a>( + provider: &providers::CredentialProvider, + value: &'a str, + start: usize, +) -> &'a str { + let mut length = value.as_bytes()[start..] + .iter() + .take_while(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + .count(); + let candidate = &value[start..start + length]; + let current_prefix_len = provider + .credential_prefixes + .iter() + .filter(|prefix| candidate.starts_with(**prefix)) + .map(|prefix| prefix.len()) + .max() + .unwrap_or(0); + if let Some(separator) = providers::credential_providers() + .flat_map(|candidate_provider| { + candidate_provider + .credential_prefixes + .iter() + .map(move |prefix| (prefix, candidate_provider.minimum_credential_len)) + }) + .filter_map(|(candidate_prefix, minimum_length)| { + candidate[current_prefix_len..] + .match_indices(*candidate_prefix) + .find_map(|(offset, _)| { + let offset = current_prefix_len + offset; + (matches!(candidate.as_bytes()[offset - 1], b'_' | b'-') + && offset > provider.minimum_credential_len + && candidate.len() - offset >= minimum_length) + .then_some(offset - 1) + }) + }) + .min() + { + length = separator; + } + &value[start..start + length] +} diff --git a/codex-rs/network-proxy/src/credential_broker/providers.rs b/codex-rs/network-proxy/src/credential_broker/providers.rs index 88657f5ee3..be22d60481 100644 --- a/codex-rs/network-proxy/src/credential_broker/providers.rs +++ b/codex-rs/network-proxy/src/credential_broker/providers.rs @@ -16,6 +16,10 @@ type TranslateRequestHeader = fn(&HeaderMap, &str, &str) -> Option; /// identities when deduplicating credential records. pub(super) struct CredentialProvider { pub(super) context_env_vars: &'static [&'static str], + pub(super) credential_prefixes: &'static [&'static str], + pub(super) ignored_credential_prefixes: &'static [&'static str], + pub(super) credential_watermark: Option<&'static str>, + pub(super) minimum_credential_len: usize, sources: &'static [CredentialSource], pub(super) reset_on_configuration_change: bool, dummy_value: fn(&str) -> String, diff --git a/codex-rs/network-proxy/src/credential_broker/providers/github.rs b/codex-rs/network-proxy/src/credential_broker/providers/github.rs index e972ede28a..f538a8bfde 100644 --- a/codex-rs/network-proxy/src/credential_broker/providers/github.rs +++ b/codex-rs/network-proxy/src/credential_broker/providers/github.rs @@ -21,6 +21,10 @@ const GITHUB_CLOUD_HOST_SUFFIXES: &[&str] = &[".ghe.com"]; pub(super) static PROVIDER: CredentialProvider = CredentialProvider { context_env_vars: &[GH_HOST_ENV_VAR], + credential_prefixes: GITHUB_TOKEN_PREFIXES, + ignored_credential_prefixes: &[], + credential_watermark: None, + minimum_credential_len: GITHUB_TOKEN_MIN_LEN, sources: &[ CredentialSource { env_vars: GITHUB_CLOUD_TOKEN_ENV_VARS, diff --git a/codex-rs/network-proxy/src/credential_broker/providers/openai.rs b/codex-rs/network-proxy/src/credential_broker/providers/openai.rs index 716d3b3d4e..5e0178b1e2 100644 --- a/codex-rs/network-proxy/src/credential_broker/providers/openai.rs +++ b/codex-rs/network-proxy/src/credential_broker/providers/openai.rs @@ -10,12 +10,17 @@ use rama_http::header::AUTHORIZATION; use std::collections::HashMap; const OPENAI_API_KEY_ENV_VARS: &[&str] = &["OPENAI_API_KEY"]; +const OPENAI_API_KEY_PREFIXES: &[&str] = &["sk-proj-", "sk-svcacct-", "sk-admin-", "sk-"]; const OPENAI_BASE_URL_ENV_VAR: &str = "OPENAI_BASE_URL"; const OPENAI_API_KEY_MIN_LEN: usize = 51; const OPENAI_API_HOST: &str = "api.openai.com"; pub(super) static PROVIDER: CredentialProvider = CredentialProvider { context_env_vars: &[OPENAI_BASE_URL_ENV_VAR], + credential_prefixes: OPENAI_API_KEY_PREFIXES, + ignored_credential_prefixes: &["sk-ant-", "sk-or-"], + credential_watermark: Some("T3BlbkFJ"), + minimum_credential_len: OPENAI_API_KEY_MIN_LEN, sources: &[CredentialSource { env_vars: OPENAI_API_KEY_ENV_VARS, binding_env_vars: &[OPENAI_BASE_URL_ENV_VAR], @@ -78,10 +83,9 @@ fn host_binding( } fn openai_api_key_prefix(value: &str) -> &str { - let Some(suffix) = value.strip_prefix("sk-") else { - return "sk-"; - }; - suffix - .find('-') - .map_or("sk-", |separator| &value[..separator + 4]) + OPENAI_API_KEY_PREFIXES + .iter() + .copied() + .find(|prefix| value.starts_with(prefix)) + .unwrap_or("sk-") } diff --git a/codex-rs/network-proxy/src/credential_broker_tests.rs b/codex-rs/network-proxy/src/credential_broker_tests.rs index dc18da691c..d9155987a3 100644 --- a/codex-rs/network-proxy/src/credential_broker_tests.rs +++ b/codex-rs/network-proxy/src/credential_broker_tests.rs @@ -86,6 +86,260 @@ fn virtualize_child_env_replaces_supported_credentials() { let mut excluded_dummies = format!("{github_dummy}\n{openai_dummy}"); assert!(!broker.virtualize_text(&mut excluded_dummies, &filtered_env)); assert_eq!(excluded_dummies, format!("{github_dummy}\n")); + let unknown_github_token = "ghp_0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefgh"; + let unknown_openai_key = "sk-proj-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefgh"; + let unknown_legacy_openai_key = format!("sk-{}", "a".repeat(48)); + let mut unregistered = format!( + "{unknown_github_token}\n{unknown_openai_key}\n{unknown_legacy_openai_key}\nghp_x sk-proj-x" + ); + assert!(!broker.virtualize_text(&mut unregistered, &env)); + assert_eq!(unregistered, "\n\n\nghp_x sk-proj-x"); + for key in [ + "sk-ant-api03-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno", + "sk-ant-oat01-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno", + "sk-or-v1-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno", + ] { + let mut unrelated_provider = key.to_string(); + assert!(broker.virtualize_text(&mut unrelated_provider, &env)); + assert_eq!(unrelated_provider, key); + assert!(!credential_broker_provider_sources_allowed( + key, + key, + &HashMap::new(), + |_| true, + )); + } + let mut embedded_tokens = format!("v1{unknown_github_token}\n/opt/ta{unknown_openai_key}/bin"); + assert!(!broker.virtualize_text(&mut embedded_tokens, &env)); + assert_eq!(embedded_tokens, "v1\n/opt/ta/bin"); + let qualified_legacy = format!( + "sk-{}-{}T3BlbkFJ{}", + "a".repeat(20), + "b".repeat(19), + "c".repeat(20) + ); + let legacy_broker = CredentialBroker::new(/*enabled*/ true); + let mut legacy_env = env_map([("OPENAI_API_KEY", qualified_legacy.as_str())]); + legacy_broker.virtualize_child_env(&mut legacy_env); + assert_credential_shape(&qualified_legacy, &legacy_env["OPENAI_API_KEY"], "sk-"); + let unmarked_legacy = format!("sk-{}-{}", "a".repeat(15), "b".repeat(35)); + let mut unmarked_legacy_alias = format!("Bearer {unmarked_legacy}"); + assert!(!broker.virtualize_text(&mut unmarked_legacy_alias, &env)); + assert_eq!(unmarked_legacy_alias, "Bearer "); + let collision = format!( + "sk-{}-sk-{}T3BlbkFJ{}", + "a".repeat(20), + "b".repeat(16), + "c".repeat(20) + ); + let mut collision_alias = format!("Bearer {collision}"); + assert!(!broker.virtualize_text(&mut collision_alias, &env)); + assert_eq!(collision_alias, "Bearer "); + let unrelated_collision = format!("sk-ant-api03-{}-sk-{}", "a".repeat(40), "b".repeat(48)); + let mut redacted_collision = unrelated_collision.clone(); + assert!(!broker.virtualize_text(&mut redacted_collision, &env)); + assert_eq!( + redacted_collision, + format!("sk-ant-api03-{}-", "a".repeat(40)) + ); + assert!(credential_broker_provider_sources_allowed( + &unrelated_collision, + &redacted_collision, + &HashMap::new(), + |_| true, + )); + let unrelated = "sk-ant-api03-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno"; + let mut known_provider_env = env.clone(); + known_provider_env.insert("ANTHROPIC_API_KEY".to_string(), unrelated.to_string()); + for separator in ["_", "__", "--", "-_", "_-", "", "_openai_", "_Bearer_"] { + let mut known_adjacent = format!("{unrelated}{separator}{unknown_legacy_openai_key}"); + assert!(!broker.virtualize_text(&mut known_adjacent, &known_provider_env)); + assert_eq!(known_adjacent, format!("{unrelated}{separator}")); + } + let first_bundle = format!("{unrelated}_{unknown_legacy_openai_key}"); + known_provider_env.insert("FIRST_BUNDLE".to_string(), first_bundle.clone()); + let openrouter = "sk-or-v1-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno"; + for ignored in [unrelated, openrouter] { + for separator in ["_", "-", "__", "--", "", "_openai_", "_Bearer_"] { + let mut unknown_adjacent = format!("{ignored}{separator}{unknown_legacy_openai_key}"); + assert!(!broker.virtualize_text(&mut unknown_adjacent, &env)); + assert_eq!(unknown_adjacent, format!("{ignored}{separator}")); + } + } + let mut nested_bundle = format!("{first_bundle}_{openrouter}"); + assert!(!broker.virtualize_text(&mut nested_bundle, &known_provider_env)); + assert_eq!(nested_bundle, format!("{unrelated}__{openrouter}")); + let mixed_providers = format!("{unknown_github_token}_{unrelated_collision}"); + let mut virtualized_providers = mixed_providers.clone(); + assert!(!broker.virtualize_text(&mut virtualized_providers, &env)); + assert!(!credential_broker_provider_sources_allowed( + &mixed_providers, + &virtualized_providers, + &HashMap::new(), + |source| source != "OPENAI_API_KEY", + )); + let mixed_credentials = + format!("{unknown_github_token}_{unknown_legacy_openai_key}_{unrelated}"); + let mut virtualized_credentials = mixed_credentials.clone(); + assert!(!broker.virtualize_text(&mut virtualized_credentials, &env)); + assert!(!credential_broker_provider_sources_allowed( + &mixed_credentials, + &virtualized_credentials, + &HashMap::new(), + |source| source != "OPENAI_API_KEY", + )); + let equivalent_sources = env_map([ + ("GH_TOKEN", unknown_github_token), + ("GITHUB_TOKEN", unknown_github_token), + ]); + assert!(credential_broker_provider_sources_allowed( + unknown_github_token, + "", + &equivalent_sources, + |source| source == "GH_TOKEN", + )); + let distinct_github_token = "ghp_abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH"; + let distinct_sources = env_map([ + ("GH_TOKEN", unknown_github_token), + ("GITHUB_TOKEN", distinct_github_token), + ]); + assert!(!credential_broker_provider_sources_allowed( + &format!("{unknown_github_token}\n{distinct_github_token}"), + "", + &distinct_sources, + |source| source == "GH_TOKEN", + )); + for unrelated in [ + unrelated, + "sk-or-v1-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmno", + ] { + for separator in ['_', '-'] { + for (mut adjacent, expected) in [ + ( + format!("{unrelated}{separator}{qualified_legacy}"), + format!("{unrelated}{separator}"), + ), + ( + format!("{unknown_legacy_openai_key}{separator}{unrelated}"), + format!("{separator}{unrelated}"), + ), + ] { + assert!(!broker.virtualize_text(&mut adjacent, &env)); + assert_eq!(adjacent, expected); + } + } + } + let hash = "a".repeat(64); + let copied_openai_key = format!("sk-proj-{hash}"); + let mut copied_credential = format!("ta{copied_openai_key}"); + assert!(!broker.virtualize_text(&mut copied_credential, &env)); + assert_eq!(copied_credential, "ta"); + let mut path_embedded_credentials = + format!("/prefix/{unknown_legacy_openai_key}:/next\n/prefix/sk-proj-{hash}"); + assert!(!broker.virtualize_text(&mut path_embedded_credentials, &env)); + assert_eq!(path_embedded_credentials, "/prefix/:/next\n/prefix/"); + for separator in ['-', '_'] { + let mut embedded_path = format!("/workspace/token{separator}sk-proj-{hash}/bin"); + assert!(!broker.virtualize_text(&mut embedded_path, &env)); + assert_eq!(embedded_path, format!("/workspace/token{separator}/bin")); + } + let mut adjacent_path = format!("/workspace/tokensk-proj-{hash}/bin"); + assert!(!broker.virtualize_text(&mut adjacent_path, &env)); + assert_eq!(adjacent_path, "/workspace/token/bin"); + let mut word_adjacent_path = format!("/workspace/datask-proj-{hash}/bin"); + assert!(!broker.virtualize_text(&mut word_adjacent_path, &env)); + assert_eq!(word_adjacent_path, "/workspace/data/bin"); + for prefix in ["a", "di", "ma", "ri", "bri"] { + let mut embedded_credential = format!("Bearer {prefix}sk-{hash}"); + assert!(!broker.virtualize_text(&mut embedded_credential, &env)); + assert_eq!(embedded_credential, format!("Bearer {prefix}")); + } + for component in ["cafe\u{e9}task", "cafe\u{301}task"] { + let mut unicode_adjacent_path = format!("/workspace/{component}-proj-{hash}/bin"); + assert!(!broker.virtualize_text(&mut unicode_adjacent_path, &env)); + assert_eq!( + unicode_adjacent_path, + format!("/workspace/{}/bin", component.strip_suffix("sk").unwrap()) + ); + } + let mut hashed_credential_path = + format!("/workspace/task-proj-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefgh-{hash}/bin"); + assert!(!broker.virtualize_text(&mut hashed_credential_path, &env)); + assert_eq!(hashed_credential_path, "/workspace/ta/bin"); + let mut watermarked_path = format!("/workspace/tokensk-proj-{hash}-T3BlbkFJsuffix/bin"); + assert!(!broker.virtualize_text(&mut watermarked_path, &env)); + assert_eq!(watermarked_path, "/workspace/token/bin"); + for suffix in ["build-", "", "proj-", "admin-", "svcacct-"] { + let path_value = + format!("/task-{suffix}{hash}:/task-{suffix}{hash}/x:/task-{suffix}{hash}"); + for path in [ + path_value, + format!("declare -x PATH=\"/task-{suffix}{hash}-build:/task-{suffix}{hash}-release\""), + format!("export -UT PATH path=(/task-{suffix}{hash} /usr/bin)"), + format!("alias activate='source /task-{suffix}{hash}/bin/activate'"), + format!(r"C:\task-{suffix}{hash}\Scripts"), + ] { + let mut virtualized_path = path.clone(); + assert!(broker.virtualize_text(&mut virtualized_path, &env)); + assert_eq!(virtualized_path, path); + } + } + for component in [ + "my_task", + "flask", + "disk", + "mask", + "risk", + "brisk", + "subtask", + "mytask", + "devtask", + "multitask", + "buildtask", + "mydisk", + "harddisk", + "MY_Task", + "Flask", + "cafe\u{e9}_task", + "cafe\u{301}_task", + ] { + for suffix in ["", "proj-", "admin-", "svcacct-"] { + for path in [ + format!("/workspace/{component}-{suffix}{hash}/bin"), + format!("VIRTUAL_ENV=/workspace/{component}-{suffix}{hash}"), + format!( + "alias activate='source /workspace/{component}-{suffix}{hash}/bin/activate'" + ), + format!(r"C:\\workspace\\{component}-{suffix}{hash}\\Scripts"), + ] { + let mut virtualized_path = path.clone(); + assert!(broker.virtualize_text(&mut virtualized_path, &env)); + assert_eq!(virtualized_path, path); + } + } + } + let registered_hex_credential = format!("sk-{hash}"); + let registered_broker = CredentialBroker::new(/*enabled*/ true); + let mut registered_env = env_map([("OPENAI_API_KEY", registered_hex_credential.as_str())]); + registered_broker.virtualize_child_env(&mut registered_env); + let registered_dummy = ®istered_env["OPENAI_API_KEY"]; + let mut credential_path = format!("/workspace/multita{registered_hex_credential}/bin"); + assert!(registered_broker.virtualize_text(&mut credential_path, ®istered_env)); + assert_eq!( + credential_path, + format!("/workspace/multita{registered_dummy}/bin") + ); + for (key, placeholder, credential) in [ + ("GH_TOKEN", "ghp_", unknown_github_token), + ("OPENAI_API_KEY", "sk-", unknown_legacy_openai_key.as_str()), + ] { + let broker = CredentialBroker::new(/*enabled*/ true); + let mut env = env_map([(key, placeholder)]); + broker.virtualize_child_env(&mut env); + let mut alias = format!("Bearer {credential}"); + assert!(!broker.virtualize_text(&mut alias, &env)); + assert!(!alias.contains(credential)); + } let mut command = vec![ format!("Authorization: Bearer {github_dummy}"), format!("Authorization: Bearer {openai_dummy}"), @@ -242,6 +496,45 @@ fn virtualize_child_env_preserves_live_dummy_mappings() { assert_eq!(distinct_credentials["GH_TOKEN"], secondary_dummy); } +#[test] +fn unbound_enterprise_aliases_retain_source_ownership() { + let token = "ghp_abcdefghijklmnopqrstuvwxyz0123456789"; + let real_header = format!("Bearer {token}"); + for source in ["GH_ENTERPRISE_TOKEN", "GITHUB_ENTERPRISE_TOKEN"] { + for host in [None, Some("")] { + for parent_discovery in [false, true] { + let broker = CredentialBroker::new(/*enabled*/ true); + let mut parent = env_map([(source, token)]); + if let Some(host) = host { + parent.insert("GH_HOST".to_string(), host.to_string()); + } + let mut child = env_map([("AUTH_HEADER", real_header.as_str())]); + if parent_discovery { + broker.discover_parent_credentials(&parent, &child); + } else { + broker.virtualize_child_env(&mut parent); + } + broker.virtualize_child_env(&mut child); + assert_eq!(child["AUTH_HEADER"], real_header, "{source}, {host:?}"); + assert!(broker.read_state().credentials.is_empty()); + + // A later explicit source and destination can still register normally. + child.insert(source.to_string(), token.to_string()); + child.insert("GH_HOST".to_string(), "enterprise.example".to_string()); + broker.virtualize_child_env(&mut child); + let dummy = &child[source]; + assert_ne!(dummy, token); + let mut headers = headers_with_bearer(dummy); + let original = headers.clone(); + broker.inject_request_headers("api.github.com", &mut headers); + assert_eq!(headers, original); + broker.inject_request_headers("enterprise.example", &mut headers); + assert_eq!(authorization(&headers), Some(real_header.as_str())); + } + } + } +} + #[test] fn virtualize_child_env_replaces_aliases_of_filtered_parent_credentials() { let broker = CredentialBroker::new(/*enabled*/ true); @@ -264,6 +557,10 @@ fn virtualize_child_env_replaces_aliases_of_filtered_parent_credentials() { assert_eq!(child_env["AUTH_HEADER"], format!("Bearer {dummy}")); assert!(!child_env.contains_key("GH_TOKEN")); + let mut virtualized_alias = child_env["AUTH_HEADER"].clone(); + assert!(broker.virtualize_text(&mut virtualized_alias, &child_env)); + assert_eq!(virtualized_alias, child_env["AUTH_HEADER"]); + let mut headers = headers_with_bearer(&dummy); broker.inject_request_headers("api.github.com", &mut headers); assert_eq!(authorization(&headers), Some(authorization_header.as_str())); @@ -272,6 +569,139 @@ fn virtualize_child_env_replaces_aliases_of_filtered_parent_credentials() { assert_eq!(child_env["HOMEBREW_GITHUB_API_TOKEN"], github_token); assert_eq!(child_env["AUTH_HEADER"], authorization_header); assert!(!child_env.contains_key("GH_TOKEN")); + + let openai_token = "sk-proj-abcdefghijklmnopqrstuvwxyz1234567890"; + let mixed_bundle = format!("GitHub {github_token}\nOpenAI {openai_token}"); + let mut mixed_env = env_map([ + ("GH_TOKEN", github_token), + ("OPENAI_API_KEY", openai_token), + ("AUTH_BUNDLE", mixed_bundle.as_str()), + ]); + broker.virtualize_child_env(&mut mixed_env); + let mut mixed_alias = mixed_env["AUTH_BUNDLE"].clone(); + let excluded_dummy = mixed_env.remove("OPENAI_API_KEY").expect("OpenAI dummy"); + assert!(!broker.virtualize_text(&mut mixed_alias, &mixed_env)); + assert!(!mixed_alias.contains(&excluded_dummy)); +} + +#[test] +fn virtualize_child_env_preserves_paths_unless_the_credential_is_known() { + let broker = CredentialBroker::new(/*enabled*/ true); + let real = format!("sk-proj-{}", "a".repeat(64)); + let path = format!("/workspace/my_ta{real}/bin"); + let mut env = env_map([("VIRTUAL_ENV", &path)]); + + broker.virtualize_child_env(&mut env); + + assert_eq!( + env, + env_map([ + ("VIRTUAL_ENV", &path), + (CREDENTIAL_BROKER_ACTIVE_ENV_KEY, "1"), + (BROKERED_CREDENTIALS_ENV_KEY, "[]"), + ]) + ); + + let copied = format!("ta{real}"); + let mut copied_env = env_map([("COPIED", &copied)]); + broker.virtualize_child_env(&mut copied_env); + let dummy = copied_env["COPIED"].strip_prefix("ta").unwrap(); + assert_credential_shape(&real, dummy, "sk-proj-"); + let mut headers = headers_with_bearer(dummy); + broker.inject_request_headers("api.openai.com", &mut headers); + assert_eq!( + authorization(&headers), + Some(format!("Bearer {real}").as_str()) + ); + + broker.virtualize_child_env(&mut env); + assert_eq!(env["VIRTUAL_ENV"], path.replace(&real, dummy)); +} + +#[test] +fn virtualize_child_env_discovers_credentials_without_canonical_variables() { + for (token, canonical_key, host) in [ + ( + "ghp_abcdefghijklmnopqrstuvwxyz1234567890", + "GH_TOKEN", + "api.github.com", + ), + ( + "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", + "OPENAI_API_KEY", + "api.openai.com", + ), + ] { + let broker = CredentialBroker::new(/*enabled*/ true); + let authorization_header = format!("Bearer {token}"); + let mut env = env_map([("AUTH_HEADER", authorization_header.as_str())]); + + broker.virtualize_child_env(&mut env); + + let dummy_header = &env["AUTH_HEADER"]; + assert_ne!(dummy_header, &authorization_header); + assert!(!env.contains_key(canonical_key)); + let mut headers = headers_with_authorization(dummy_header); + broker.inject_request_headers(host, &mut headers); + assert_eq!(authorization(&headers), Some(authorization_header.as_str())); + + let mut snapshot = format!("export AUTH_HEADER='{authorization_header}'"); + assert!(broker.virtualize_text(&mut snapshot, &env)); + assert_eq!(snapshot, format!("export AUTH_HEADER='{dummy_header}'")); + + let broker = CredentialBroker::new(/*enabled*/ true); + let mut env = env_map([ + (canonical_key, "another-canonical-credential"), + ("AUTH_HEADER", authorization_header.as_str()), + ]); + broker.virtualize_child_env(&mut env); + assert_ne!(env["AUTH_HEADER"], authorization_header); + let mut headers = headers_with_authorization(&env["AUTH_HEADER"]); + broker.inject_request_headers(host, &mut headers); + assert_eq!(authorization(&headers), Some(authorization_header.as_str())); + } +} + +#[test] +fn virtualize_child_env_keeps_adjacent_provider_credentials_separate() { + let broker = CredentialBroker::new(/*enabled*/ true); + let github_token = "ghp_abcdefghijklmnopqrstuvwxyz1234567890ABCD"; + let openai_token = "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + let real_bundle = format!("{github_token}_{openai_token}"); + let mut env = env_map([("AUTH_BUNDLE", real_bundle.as_str())]); + + broker.virtualize_child_env(&mut env); + + let dummy_bundle = env["AUTH_BUNDLE"].clone(); + assert!(!dummy_bundle.contains(github_token)); + assert!(!dummy_bundle.contains(openai_token)); + let github_dummy = &dummy_bundle[..github_token.len()]; + let openai_dummy = &dummy_bundle[github_token.len() + 1..]; + + let mut github_headers = headers_with_bearer(github_dummy); + broker.inject_request_headers("api.github.com", &mut github_headers); + assert_eq!( + authorization(&github_headers), + Some(format!("Bearer {github_token}").as_str()) + ); + + let mut openai_headers = headers_with_bearer(openai_dummy); + broker.inject_request_headers("api.openai.com", &mut openai_headers); + assert_eq!( + authorization(&openai_headers), + Some(format!("Bearer {openai_token}").as_str()) + ); + + let mut bundled_headers = headers_with_bearer(&dummy_bundle); + broker.inject_request_headers("api.github.com", &mut bundled_headers); + assert_eq!( + authorization(&bundled_headers), + Some(format!("Bearer {dummy_bundle}").as_str()) + ); + + let mut restored_bundle = dummy_bundle; + assert!(broker.restore_text(&mut restored_bundle)); + assert_eq!(restored_bundle, real_bundle); } #[test] @@ -362,24 +792,80 @@ fn child_without_dummy_cannot_use_previous_child_credential() { } #[test] -fn virtualize_child_env_preserves_unbound_enterprise_token() { +fn virtualize_child_env_keeps_unbound_enterprise_token_out_of_persisted_text() { let broker = CredentialBroker::new(/*enabled*/ true); - let mut env = env_map([("GH_ENTERPRISE_TOKEN", "ghp-enterprise-real")]); + let token = "ghp_abcdefghijklmnopqrstuvwxyz1234567890"; + let authorization_header = format!("Bearer {token}"); + let mut env = env_map([ + ("GH_ENTERPRISE_TOKEN", token), + ("AUTH_HEADER", authorization_header.as_str()), + ]); broker.virtualize_child_env(&mut env); - let inert_token = "ghp_abcdefghijklmnopqrstuvwxyz1234567890"; - let mut headers = headers_with_bearer(inert_token); + assert_eq!(env["GH_ENTERPRISE_TOKEN"], token); + assert_eq!(env["AUTH_HEADER"], authorization_header); + for alias in [ + format!("export GH_ENTERPRISE_TOKEN={token}"), + format!("export AUTH_HEADER='Bearer {token}_suffix'"), + format!("export AUTH_HEADER='Bearer {token}-suffix'"), + ] { + let mut persisted = alias; + assert!(!broker.virtualize_text(&mut persisted, &env)); + assert!(!persisted.contains(token)); + } + let distinct_token = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789abcdefghijkl"; + let mut adjacent = format!("{token}_{distinct_token}"); + assert!(!broker.virtualize_text(&mut adjacent, &env)); + assert_eq!(adjacent, "_"); + + let mut truncated_env = env_map([("GH_ENTERPRISE_TOKEN", "ghp_abcdefghijkl")]); + broker.virtualize_child_env(&mut truncated_env); + let mut hidden = token.to_string(); + assert!(!broker.virtualize_text(&mut hidden, &truncated_env)); + assert!(hidden.is_empty()); + assert!(!credential_broker_provider_sources_allowed( + token, + "", + &truncated_env, + |source| source != "GH_TOKEN", + )); + assert!(!credential_broker_provider_sources_allowed( + token, + "", + &HashMap::new(), + |source| source != "GH_TOKEN", + )); + + let fine_grained = "github_pat_11AA0bbCC_abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH"; + let mut truncated_env = env_map([( + "GH_ENTERPRISE_TOKEN", + "github_pat_11AA0bbCC_abcdefghijklmnopqrs", + )]); + broker.virtualize_child_env(&mut truncated_env); + let mut hidden = fine_grained.to_string(); + assert!(!broker.virtualize_text(&mut hidden, &truncated_env)); + assert!(hidden.is_empty()); + assert!(!credential_broker_provider_sources_allowed( + fine_grained, + "", + &env_map([("GH_TOKEN", "github_pat_11AA0bbCC")]), + |source| source == "GH_TOKEN", + )); + let mut headers = headers_with_bearer(token); broker.inject_request_headers("attacker.example", &mut headers); - assert_eq!(env["GH_ENTERPRISE_TOKEN"], "ghp-enterprise-real"); - assert_eq!(headers, headers_with_bearer(inert_token)); + assert_eq!(env["GH_ENTERPRISE_TOKEN"], token); + assert_eq!(headers, headers_with_bearer(token)); assert!(!broker.host_requires_mitm("attacker.example")); env.insert("GH_HOST".to_string(), "github.example.com".to_string()); broker.virtualize_child_env(&mut env); let mut headers = headers_with_bearer(&env["GH_ENTERPRISE_TOKEN"]); broker.inject_request_headers("github.example.com", &mut headers); - assert_eq!(authorization(&headers), Some("Bearer ghp-enterprise-real")); + assert_eq!( + authorization(&headers), + Some(format!("Bearer {token}").as_str()) + ); } #[test] diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 38ac654d46..f71a334aed 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -48,6 +48,7 @@ pub use credential_broker::brokered_credential_binding_env_keys; pub use credential_broker::brokered_credential_dummy_env_keys; pub use credential_broker::brokered_credential_env_keys; pub use credential_broker::credential_broker_provider_context_env_keys; +pub use credential_broker::credential_broker_provider_sources_allowed; pub use credential_broker::is_credential_broker_provider_env_key; pub use environment_policy::EnvironmentNetworkPolicy; pub use mitm_hook::InjectedHeaderConfig; diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index aab1c44a9e..d002b6fe0f 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -1153,6 +1153,11 @@ impl NetworkProxy { self.state.virtualize_brokered_text(text, env) } + /// Restores known dummy credentials in trusted text captured for fail-open execution. + pub fn restore_brokered_text(&self, text: &mut String) -> bool { + self.state.restore_brokered_text(text) + } + pub fn apply_to_env_for_environment( &self, env: &mut HashMap, diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 3e38030ab3..4cfd2bdd87 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -477,6 +477,10 @@ impl NetworkProxyState { self.credential_broker.virtualize_text(text, env) } + pub(crate) fn restore_brokered_text(&self, text: &mut String) -> bool { + self.credential_broker.restore_text(text) + } + pub fn inject_request_credentials(&self, host: &str, headers: &mut rama_http::HeaderMap) { self.credential_broker.inject_request_headers(host, headers); } diff --git a/codex-rs/shell-command/src/shell_snapshot_credentials.rs b/codex-rs/shell-command/src/shell_snapshot_credentials.rs index 23fe061c25..d2cf0054f4 100644 --- a/codex-rs/shell-command/src/shell_snapshot_credentials.rs +++ b/codex-rs/shell-command/src/shell_snapshot_credentials.rs @@ -136,7 +136,7 @@ pub fn prepare_snapshot_credentials( let mut alias_values = HashMap::new(); let mut rejected_alias_keys = Vec::new(); let mut invalid_export = false; - let exports = captured + let mut exports = captured .exports .iter() .filter_map(|export| { @@ -185,7 +185,6 @@ pub fn prepare_snapshot_credentials( Some(Export::Captured(line)) }) .collect::>(); - let mut snapshot = render::render(captured.state, captured.aliases, &exports)?; if invalid_export { return None; } @@ -193,19 +192,33 @@ pub fn prepare_snapshot_credentials( 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 + && let Some((assignment, value)) = allowed .get(key) .and_then(|value| credential_alias_assignment(value, &credential_keys)) { - // Native capture can deliberately omit readonly exports; retain only their metadata. + exports.push(Export::Alias { + key, + value: assignment, + }); alias_values.insert(key.to_string(), value); } else { rejected_alias_keys.push(key.to_string()); } } - if !virtualize_text(&mut snapshot) { - return None; + let mut snapshot = render::render(captured.state, captured.aliases, &exports)?; + let mut credential_values = original + .values() + .chain(restored.values()) + .chain(configured.values()) + .collect::>(); + credential_values.sort_unstable_by_key(|value| (std::cmp::Reverse(value.len()), *value)); + credential_values.dedup(); + for value in credential_values { + let mut replacement = value.clone(); + if virtualize_text(&mut replacement) && replacement != *value { + snapshot = snapshot.replace(value, &replacement); + } } Some(PreparedSnapshot { diff --git a/codex-rs/shell-command/src/shell_snapshot_render.rs b/codex-rs/shell-command/src/shell_snapshot_render.rs index 61a4db789a..1881988bdf 100644 --- a/codex-rs/shell-command/src/shell_snapshot_render.rs +++ b/codex-rs/shell-command/src/shell_snapshot_render.rs @@ -32,6 +32,7 @@ pub(super) enum ValuePart { pub(super) enum Export<'a> { Captured(&'a str), + Alias { key: &'a str, value: Value }, Assignment { declaration: &'a str, value: Value }, } @@ -63,6 +64,9 @@ pub(super) fn render(state: &str, aliases: &str, exports: &[Export<'_>]) -> Opti for export in exports { match export { Export::Captured(source) => output.push_str(source), + Export::Alias { key, value } => { + output.push_str(&format!("export {key}={}\n", value.render()?)); + } Export::Assignment { declaration, value } => { output.push_str(&format!("{declaration}={}\n", value.render()?)); } diff --git a/codex-rs/shell-command/src/shell_snapshot_tests.rs b/codex-rs/shell-command/src/shell_snapshot_tests.rs index abd359b3a1..bd68bd34c5 100644 --- a/codex-rs/shell-command/src/shell_snapshot_tests.rs +++ b/codex-rs/shell-command/src/shell_snapshot_tests.rs @@ -698,3 +698,85 @@ fn captured_script(shell_type: ShellType, source: &str) -> Result { CapturedSnapshot::parse(shell_type, source.as_bytes()).context("invalid native capture")?; Ok(captured.render_script()) } + +#[test] +fn brokered_bash_snapshot_preserves_exported_functions_and_redacts_heredocs() -> Result<()> { + let dir = tempdir()?; + let startup = dir.path().join("startup.sh"); + let real = "ghp_abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH"; + let dummy = "ghp_brokered_dummy_abcdefghijklmnopqrstuvwxyz0123"; + let credential_shaped_function = "ghp_0123456789abcdefghijklmnopqrstuvwxyzABCD"; + std::fs::write( + &startup, + format!( + "helper() {{\n cat <<'EOF'\nexport LEGACY_SETTING=production\nAuthorization: {real}\nEOF\n}}\n{credential_shaped_function}() {{ printf UNEXPECTED_EXECUTION; }}\nexport -f helper\nexport GH_TOKEN='{real}'\n" + ), + )?; + let output = Command::new("/bin/bash") + .arg("-c") + .arg( + snapshot_capture_script( + ShellType::Bash, + SnapshotCaptureOptions { + startup: SnapshotStartup::NonInteractive, + ..CAPTURE_ALL + }, + ) + .expect("bash supports snapshots"), + ) + .env_clear() + .env("HOME", dir.path()) + .env("PATH", "/usr/bin:/bin") + .env("BASH_ENV", &startup) + .output()?; + assert!(output.status.success()); + + let mut snapshot = String::from_utf8(output.stdout)?; + let original = HashMap::from([("GH_TOKEN".to_string(), real.to_string())]); + let discovered = HashMap::from([("GH_TOKEN".to_string(), dummy.to_string())]); + let brokered_keys = vec!["GH_TOKEN".to_string()]; + rewrite_snapshot_credentials( + ShellType::Bash, + &mut snapshot, + SnapshotCredentialEnvironment { + original: &original, + restored: &original, + configured: &HashMap::new(), + discovered: &discovered, + allowed: &discovered, + is_allowed_unset: &|_| false, + brokered_keys: &brokered_keys, + brokered_alias_keys: &[], + allowed_brokered_keys: &brokered_keys, + }, + |text| { + if text.contains(credential_shaped_function) { + *text = text.replace(credential_shaped_function, ""); + return false; + } + *text = text.replace(real, dummy); + true + }, + ); + assert!(!snapshot.contains(real)); + assert!(snapshot.contains(credential_shaped_function)); + let path = dir.path().join("snapshot.sh"); + std::fs::write(&path, snapshot)?; + + let restored = Command::new("/bin/bash") + .arg("-c") + .arg(". \"$1\"; exec /bin/bash -c helper") + .arg("snapshot") + .arg(&path) + .env_clear() + .env("PATH", "/usr/bin:/bin") + .env("GH_TOKEN", dummy) + .output()?; + assert!(restored.status.success()); + assert_eq!( + String::from_utf8(restored.stdout)?, + format!("export LEGACY_SETTING=production\nAuthorization: {dummy}\n") + ); + + Ok(()) +}