diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a8c52ef917..72eca6f3f3 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4488,7 +4488,6 @@ dependencies = [ name = "codex-workload-identity-providers" version = "0.0.0" dependencies = [ - "anyhow", "codex-workload-identity", "codex-workload-identity-aws", "codex-workload-identity-azure", @@ -4496,7 +4495,6 @@ dependencies = [ "codex-workload-identity-github-actions", "codex-workload-identity-spiffe", "reqwest 0.12.28", - "tempfile", "tokio", ] @@ -4504,7 +4502,6 @@ dependencies = [ name = "codex-workload-identity-spiffe" version = "0.0.0" dependencies = [ - "anyhow", "codex-workload-identity", "spiffe", "tokio", diff --git a/codex-rs/codex-mcp/src/catalog.rs b/codex-rs/codex-mcp/src/catalog.rs index 3a752c82ba..d46697175c 100644 --- a/codex-rs/codex-mcp/src/catalog.rs +++ b/codex-rs/codex-mcp/src/catalog.rs @@ -4,6 +4,7 @@ use std::collections::BTreeSet; use std::collections::HashMap; use codex_config::McpServerConfig; +use codex_config::McpServerTransportConfig; /// Plugin identity retained with an MCP registration for tool attribution. #[derive(Clone, Debug, PartialEq, Eq)] @@ -228,6 +229,38 @@ impl McpCatalogBuilder { .push(CatalogAction::Register(Box::new(registration))); } + /// Retains environment variables accepted by `retain` across every pending registration. + pub fn retain_environment_variables(&mut self, retain: impl Fn(&str) -> bool) { + for action in &mut self.actions { + let CatalogAction::Register(registration) = action else { + continue; + }; + match &mut registration.config.transport { + McpServerTransportConfig::Stdio { env, env_vars, .. } => { + if let Some(env) = env { + env.retain(|name, _| retain(name)); + } + env_vars.retain(|env_var| retain(env_var.name())); + } + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + env_http_headers, + .. + } => { + if bearer_token_env_var + .as_deref() + .is_some_and(|name| !retain(name)) + { + *bearer_token_env_var = None; + } + if let Some(env_http_headers) = env_http_headers { + env_http_headers.retain(|_, env_var| retain(env_var)); + } + } + } + } + } + /// Applies the legacy name-scoped disabled veto after source resolution. pub fn disable(&mut self, name: String) { self.disabled_server_names.insert(name); diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index 696876cc40..cfb26b660e 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -4,6 +4,7 @@ use std::time::Duration; use codex_config::AppToolApproval; use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::McpServerConfig; +use codex_config::McpServerEnvVar; use codex_config::McpServerToolConfig; use codex_config::McpServerTransportConfig; use pretty_assertions::assert_eq; @@ -49,6 +50,76 @@ fn plugin(plugin_id: &str) -> McpPluginAttribution { McpPluginAttribution::new(plugin_id.to_string(), plugin_id.to_string()) } +#[test] +fn environment_variable_filter_applies_to_every_pending_registration() { + let mut stdio = server("https://unused.example/mcp"); + stdio.transport = McpServerTransportConfig::Stdio { + command: "mcp-server".to_string(), + args: Vec::new(), + env: Some(HashMap::from([ + ("SECRET".to_string(), "secret".to_string()), + ("SAFE".to_string(), "safe".to_string()), + ])), + env_vars: vec![ + McpServerEnvVar::from("SECRET"), + McpServerEnvVar::from("SAFE"), + ], + cwd: None, + }; + let mut http = server("https://http.example/mcp"); + http.transport = McpServerTransportConfig::StreamableHttp { + url: "https://http.example/mcp".to_string(), + bearer_token_env_var: Some("secret".to_string()), + http_headers: None, + env_http_headers: Some(HashMap::from([ + ("X-Secret".to_string(), "SECRET".to_string()), + ("X-Safe".to_string(), "SAFE".to_string()), + ])), + }; + let mut expected_stdio = stdio.clone(); + expected_stdio.transport = McpServerTransportConfig::Stdio { + command: "mcp-server".to_string(), + args: Vec::new(), + env: Some(HashMap::from([("SAFE".to_string(), "safe".to_string())])), + env_vars: vec![McpServerEnvVar::from("SAFE")], + cwd: None, + }; + let mut expected_http = http.clone(); + expected_http.transport = McpServerTransportConfig::StreamableHttp { + url: "https://http.example/mcp".to_string(), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: Some(HashMap::from([("X-Safe".to_string(), "SAFE".to_string())])), + }; + + let mut builder = ResolvedMcpCatalog::builder(); + builder.register(McpServerRegistration::from_plugin( + "stdio".to_string(), + plugin("plugin@test"), + /*plugin_order*/ 0, + stdio, + )); + builder.register(McpServerRegistration::from_extension( + "http".to_string(), + "extension", + /*contribution_order*/ 0, + http, + )); + builder.retain_environment_variables(|name| !name.eq_ignore_ascii_case("SECRET")); + + let catalog = builder.build(); + assert_eq!( + catalog + .server("stdio") + .map(super::ResolvedMcpServer::config), + Some(&expected_stdio) + ); + assert_eq!( + catalog.server("http").map(super::ResolvedMcpServer::config), + Some(&expected_http) + ); +} + fn plugin_source(plugin_id: &str) -> McpServerSource { McpServerSource::Plugin(plugin(plugin_id)) } diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 7038fad22d..0a21b333c6 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -86,10 +86,12 @@ use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::permissions::ReadDenyMatcher; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::NetworkAccess; use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::SandboxPolicy; +use codex_workload_identity::CredentialSourceConfig; use serde::Deserialize; use tempfile::tempdir; @@ -218,8 +220,8 @@ X-Safe = "SAFE_VAR" let plugins_manager = PluginsManager::new(codex_home.path().to_path_buf()); let mcp_config = config.to_mcp_config(&plugins_manager).await; - let local = mcp_config - .configured_mcp_servers + let configured_mcp_servers = mcp_config.mcp_server_catalog.configured_servers(); + let local = configured_mcp_servers .get("local") .expect("local MCP config"); let McpServerTransportConfig::Stdio { env, env_vars, .. } = &local.transport else { @@ -234,10 +236,7 @@ X-Safe = "SAFE_VAR" ); assert_eq!(env_vars, &vec![McpServerEnvVar::from("SAFE_VAR")]); - let http = mcp_config - .configured_mcp_servers - .get("http") - .expect("HTTP MCP config"); + let http = configured_mcp_servers.get("http").expect("HTTP MCP config"); let McpServerTransportConfig::StreamableHttp { bearer_token_env_var, env_http_headers, @@ -257,6 +256,159 @@ X-Safe = "SAFE_VAR" Ok(()) } +#[tokio::test] +async fn workload_assertion_cli_override_is_excluded_from_model_processes() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let cwd = TempDir::new()?; + std::fs::write( + codex_home.path().join(CONFIG_TOML_FILE), + r#" +[workload_identity] +identity_provider_id = "idp_example" +identity_provider_mapping_id = "idpm_example" +audience = "openai-audience" + +[workload_identity.credential_source] +type = "azure" +"#, + )?; + let config = ConfigBuilder::without_managed_config_for_tests() + .codex_home(codex_home.path().to_path_buf()) + .cli_overrides(vec![ + ( + "workload_identity.credential_source.type".to_string(), + toml::Value::String("environment".to_string()), + ), + ( + "workload_identity.credential_source.variable".to_string(), + toml::Value::String("CODEX_WIF_SUBJECT_TOKEN".to_string()), + ), + ]) + .harness_overrides(ConfigOverrides { + cwd: Some(cwd.path().to_path_buf()), + ..Default::default() + }) + .build() + .await?; + + let shell_env = codex_protocol::shell_environment::create_env_from_vars( + [( + "CODEX_WIF_SUBJECT_TOKEN".to_string(), + "secret.assertion".to_string(), + )], + &config.permissions.shell_environment_policy, + /*thread_id*/ None, + ); + assert_eq!(shell_env, HashMap::new()); + Ok(()) +} + +#[tokio::test] +async fn workload_credential_file_is_denied_even_with_danger_full_access() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let token_file = codex_home.path().join("azure-token"); + let cfg = ConfigToml { + workload_identity: Some(WorkloadIdentityConfig { + identity_provider_id: "idp_example".to_string(), + identity_provider_mapping_id: "idpm_example".to_string(), + audience: "openai-audience".to_string(), + token_url: "https://auth.openai.com/oauth/token".to_string(), + credential_source: CredentialSourceConfig::Azure { + token_file: Some(token_file.clone()), + }, + }), + ..Default::default() + }; + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides { + sandbox_mode: Some(SandboxMode::DangerFullAccess), + ..Default::default() + }, + codex_home.abs(), + ) + .await?; + + let file_system_policy = config.permissions.file_system_sandbox_policy(); + let matcher = ReadDenyMatcher::new(&file_system_policy, config.cwd.as_path()) + .expect("workload credential should install a deny-read matcher"); + assert!(matcher.is_read_denied(&token_file)); + assert_eq!( + config.permissions.permission_profile().enforcement(), + SandboxEnforcement::Managed + ); + Ok(()) +} + +#[tokio::test] +async fn workload_identity_rejects_forced_api_key_login() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let cfg = ConfigToml { + forced_login_method: Some(ForcedLoginMethod::Api), + workload_identity: Some(WorkloadIdentityConfig { + identity_provider_id: "idp_example".to_string(), + identity_provider_mapping_id: "idpm_example".to_string(), + audience: "openai-audience".to_string(), + token_url: "https://auth.openai.com/oauth/token".to_string(), + credential_source: CredentialSourceConfig::Environment { + variable: "WORKLOAD_IDENTITY_TOKEN".to_string(), + }, + }), + ..Default::default() + }; + + let error = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await + .expect_err("forced API key login must reject ChatGPT workload identity"); + + assert!(error.to_string().contains( + "`workload_identity` requires ChatGPT auth but `forced_login_method` requires API key auth" + )); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test] +async fn rotating_workload_credential_denies_its_secret_directory() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let secret_dir = codex_home.path().join("projected-secret"); + let version_dir = secret_dir.join("version"); + std::fs::create_dir_all(&version_dir)?; + std::fs::write(version_dir.join("token"), "subject-token")?; + let token_file = secret_dir.join("token"); + std::os::unix::fs::symlink("version/token", &token_file)?; + + let cfg = ConfigToml { + workload_identity: Some(WorkloadIdentityConfig { + identity_provider_id: "idp_example".to_string(), + identity_provider_mapping_id: "idpm_example".to_string(), + audience: "openai-audience".to_string(), + token_url: "https://auth.openai.com/oauth/token".to_string(), + credential_source: CredentialSourceConfig::Azure { + token_file: Some(token_file.clone()), + }, + }), + ..Default::default() + }; + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await?; + + let file_system_policy = config.permissions.file_system_sandbox_policy(); + let matcher = ReadDenyMatcher::new(&file_system_policy, config.cwd.as_path()) + .expect("workload credential should install a deny-read matcher"); + assert!(matcher.is_read_denied(&token_file)); + assert!(matcher.is_read_denied(&secret_dir.join("next-version/token"))); + Ok(()) +} + async fn derive_legacy_sandbox_policy_for_test( cfg: &ConfigToml, sandbox_mode_override: Option, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 8e8e1a8d4a..bf175f63c6 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -69,6 +69,7 @@ use codex_features::NetworkProxyConfigToml; use codex_git_utils::resolve_root_git_project_for_trust; use codex_install_context::InstallContext; use codex_login::AuthManagerConfig; +use codex_mcp::McpCatalogBuilder; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; use codex_mcp::McpServerRegistration; @@ -1487,43 +1488,7 @@ impl Config { server.clone(), )); } - - let sensitive_environment_variables = self - .workload_identity - .as_ref() - .map(|workload_identity| { - workload_identity - .credential_source - .sensitive_environment_variables() - }) - .unwrap_or_default(); - let is_sensitive = |name: &str| { - sensitive_environment_variables - .iter() - .any(|sensitive| name.eq_ignore_ascii_case(sensitive)) - }; - for server in configured_mcp_servers.values_mut() { - match &mut server.transport { - McpServerTransportConfig::Stdio { env, env_vars, .. } => { - if let Some(env) = env { - env.retain(|name, _| !is_sensitive(name)); - } - env_vars.retain(|env_var| !is_sensitive(env_var.name())); - } - McpServerTransportConfig::StreamableHttp { - bearer_token_env_var, - env_http_headers, - .. - } => { - if bearer_token_env_var.as_deref().is_some_and(&is_sensitive) { - *bearer_token_env_var = None; - } - if let Some(env_http_headers) = env_http_headers { - env_http_headers.retain(|_, env_var| !is_sensitive(env_var)); - } - } - } - } + self.remove_workload_identity_environment_variables(&mut catalog); McpConfig { chatgpt_base_url: self.chatgpt_base_url.clone(), @@ -1556,6 +1521,26 @@ impl Config { } } + pub(crate) fn remove_workload_identity_environment_variables( + &self, + catalog: &mut McpCatalogBuilder, + ) { + let sensitive_environment_variables = self + .workload_identity + .as_ref() + .map(|workload_identity| { + workload_identity + .credential_source + .sensitive_environment_variables() + }) + .unwrap_or_default(); + catalog.retain_environment_variables(|name| { + !sensitive_environment_variables + .iter() + .any(|sensitive| name.eq_ignore_ascii_case(sensitive)) + }); + } + pub(crate) fn prefix_mcp_tool_names(&self) -> bool { !self.features.enabled(Feature::NonPrefixedMcpToolNames) } @@ -2360,6 +2345,47 @@ fn apply_managed_filesystem_constraints( } } +fn apply_workload_identity_filesystem_constraints( + file_system_sandbox_policy: &mut FileSystemSandboxPolicy, + workload_identity: &WorkloadIdentityConfig, +) { + let entries = workload_identity + .credential_source + .credential_file_paths() + .into_iter() + .filter_map(workload_identity_credential_deny_path) + .map(|path| codex_protocol::permissions::FileSystemSandboxEntry { + path: codex_protocol::permissions::FileSystemPath::Path { path }, + access: codex_protocol::permissions::FileSystemAccessMode::Deny, + }) + .collect::>(); + if entries.is_empty() { + return; + } + + file_system_sandbox_policy + .preserve_deny_read_restrictions_from(&FileSystemSandboxPolicy::restricted(entries)); +} + +fn workload_identity_credential_deny_path(path: PathBuf) -> Option { + let path = AbsolutePathBuf::relative_to_current_dir(path).ok()?; + let Some(parent) = path.as_path().parent() else { + return Some(path); + }; + let parent = std::fs::canonicalize(parent).unwrap_or_else(|_| parent.to_path_buf()); + + // Projected service-account tokens are symlinks whose targets rotate. Mask + // the containing secret directory so a rotation cannot expose a new target. + let deny_path = if std::fs::symlink_metadata(path.as_path()) + .is_ok_and(|metadata| metadata.file_type().is_symlink()) + { + parent + } else { + parent.join(path.as_path().file_name()?) + }; + AbsolutePathBuf::from_absolute_path(deny_path).ok() +} + /// Optional overrides for user configuration (e.g., from CLI flags). #[derive(Default, Debug, Clone)] pub struct ConfigOverrides { @@ -3560,6 +3586,21 @@ impl Config { filesystem_requirements, ); } + if let Some(workload_identity) = cfg.workload_identity.as_ref() { + if cfg.forced_login_method == Some(ForcedLoginMethod::Api) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "`workload_identity` requires ChatGPT auth but `forced_login_method` requires API key auth", + )); + } + workload_identity.validate().map_err(|error| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, error.to_string()) + })?; + apply_workload_identity_filesystem_constraints( + &mut effective_file_system_sandbox_policy, + workload_identity, + ); + } let effective_file_system_sandbox_policy = effective_file_system_sandbox_policy .with_additional_readable_roots(resolved_cwd.as_path(), &helper_readable_roots); let effective_permission_profile = PermissionProfile::from_runtime_permissions_with_enforcement( @@ -3577,11 +3618,6 @@ impl Config { profile_workspace_roots, ) .map_err(std::io::Error::from)?; - if let Some(workload_identity) = cfg.workload_identity.as_ref() { - workload_identity.validate().map_err(|error| { - std::io::Error::new(std::io::ErrorKind::InvalidInput, error.to_string()) - })?; - } let otel = otel::resolve_config(cfg.otel.unwrap_or_default(), &mut startup_warnings); let config = Self { model, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 90e6b01e75..867362f967 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -169,6 +169,7 @@ impl McpManager { } => catalog.remove_extension(name, contributor_id, contribution_order), } } + config.remove_workload_identity_environment_variables(&mut catalog); let catalog = catalog.build(); for conflict in catalog.conflicts() { tracing::warn!( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f733bb8920..3e08a8934d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -617,12 +617,7 @@ impl Session { auth.as_ref(), ) .await; - Ok::<_, std::io::Error>(( - auth, - mcp_servers, - auth_statuses, - tool_plugin_provenance, - )) + Ok::<_, std::io::Error>((auth, mcp_servers, auth_statuses, tool_plugin_provenance)) } .instrument(info_span!( "session_init.auth_mcp", @@ -630,11 +625,8 @@ impl Session { )); // Join all independent futures. - let ( - thread_persistence_result, - state_db_ctx, - auth_and_mcp_result, - ) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); + let (thread_persistence_result, state_db_ctx, auth_and_mcp_result) = + tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let (auth, mcp_servers, auth_statuses, tool_plugin_provenance) = auth_and_mcp_result?; let mut live_thread_init = @@ -810,10 +802,23 @@ impl Session { shell::default_user_shell() }; let shell_snapshot = if config.features.enabled(Feature::ShellSnapshot) { + let excluded_environment_variables = config + .workload_identity + .as_ref() + .map(|workload_identity| { + workload_identity + .credential_source + .sensitive_environment_variables() + .into_iter() + .map(str::to_string) + .collect() + }) + .unwrap_or_default(); ShellSnapshot::new( config.codex_home.clone(), thread_id, session_telemetry.clone(), + excluded_environment_variables, state_db_ctx.clone(), ) } else { diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 6a95093003..c5d929185d 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -33,6 +33,7 @@ struct ShellSnapshotConfig { codex_home: AbsolutePathBuf, session_id: ThreadId, session_telemetry: SessionTelemetry, + excluded_environment_variables: Vec, state_db: Option, } @@ -50,6 +51,7 @@ impl ShellSnapshot { codex_home: AbsolutePathBuf, session_id: ThreadId, session_telemetry: SessionTelemetry, + excluded_environment_variables: Vec, state_db: Option, ) -> Self { Self { @@ -57,6 +59,7 @@ impl ShellSnapshot { codex_home, session_id, session_telemetry, + excluded_environment_variables, state_db, })), } @@ -97,6 +100,7 @@ impl ShellSnapshot { config.session_id, &cwd, &shell, + &config.excluded_environment_variables, config.state_db.clone(), ) .await; @@ -120,6 +124,7 @@ impl ShellSnapshot { session_id: ThreadId, session_cwd: &AbsolutePathBuf, shell: &Shell, + excluded_environment_variables: &[String], state_db: Option, ) -> std::result::Result { // File to store the snapshot @@ -150,7 +155,14 @@ impl ShellSnapshot { }); // Make the new snapshot. - if let Err(err) = write_shell_snapshot(shell.shell_type, &temp_path, session_cwd).await { + if let Err(err) = write_shell_snapshot( + shell.shell_type, + &temp_path, + session_cwd, + excluded_environment_variables, + ) + .await + { tracing::warn!( "Failed to create shell snapshot for {}: {err:?}", shell.name() @@ -162,7 +174,14 @@ impl ShellSnapshot { temp_path.display() ); - if let Err(err) = validate_snapshot(shell, &temp_path, session_cwd).await { + if let Err(err) = validate_snapshot( + shell, + &temp_path, + session_cwd, + excluded_environment_variables, + ) + .await + { tracing::error!("Shell snapshot validation failed: {err:?}"); remove_snapshot_file(&temp_path).await; return Err("validation_failed"); @@ -199,6 +218,7 @@ async fn write_shell_snapshot( shell_type: ShellType, output_path: &AbsolutePathBuf, cwd: &AbsolutePathBuf, + excluded_environment_variables: &[String], ) -> Result<()> { if shell_type == ShellType::PowerShell || shell_type == ShellType::Cmd { bail!("Shell snapshot not supported yet for {shell_type:?}"); @@ -206,7 +226,7 @@ async fn write_shell_snapshot( let shell = get_shell(shell_type, /*path*/ None) .with_context(|| format!("No available shell for {shell_type:?}"))?; - let raw_snapshot = capture_snapshot(&shell, cwd).await?; + let raw_snapshot = capture_snapshot(&shell, cwd, excluded_environment_variables).await?; let snapshot = strip_snapshot_preamble(&raw_snapshot)?; if let Some(parent) = output_path.parent() { @@ -224,13 +244,49 @@ async fn write_shell_snapshot( Ok(()) } -async fn capture_snapshot(shell: &Shell, cwd: &AbsolutePathBuf) -> Result { +async fn capture_snapshot( + shell: &Shell, + cwd: &AbsolutePathBuf, + excluded_environment_variables: &[String], +) -> Result { let shell_type = shell.shell_type; match shell_type { - ShellType::Zsh => run_shell_script(shell, &zsh_snapshot_script(), cwd).await, - ShellType::Bash => run_shell_script(shell, &bash_snapshot_script(), cwd).await, - ShellType::Sh => run_shell_script(shell, &sh_snapshot_script(), cwd).await, - ShellType::PowerShell => run_shell_script(shell, powershell_snapshot_script(), cwd).await, + ShellType::Zsh => { + run_shell_script( + shell, + &zsh_snapshot_script(excluded_environment_variables), + cwd, + excluded_environment_variables, + ) + .await + } + ShellType::Bash => { + run_shell_script( + shell, + &bash_snapshot_script(excluded_environment_variables), + cwd, + excluded_environment_variables, + ) + .await + } + ShellType::Sh => { + run_shell_script( + shell, + &sh_snapshot_script(excluded_environment_variables), + cwd, + excluded_environment_variables, + ) + .await + } + ShellType::PowerShell => { + run_shell_script( + shell, + powershell_snapshot_script(), + cwd, + excluded_environment_variables, + ) + .await + } ShellType::Cmd => bail!("Shell snapshotting is not yet supported for {shell_type:?}"), } } @@ -248,6 +304,7 @@ async fn validate_snapshot( shell: &Shell, snapshot_path: &AbsolutePathBuf, cwd: &AbsolutePathBuf, + excluded_environment_variables: &[String], ) -> Result<()> { let snapshot_path_display = snapshot_path.display(); let script = format!("set -e; . \"{snapshot_path_display}\""); @@ -257,18 +314,25 @@ async fn validate_snapshot( SNAPSHOT_TIMEOUT, /*use_login_shell*/ false, cwd, + excluded_environment_variables, ) .await .map(|_| ()) } -async fn run_shell_script(shell: &Shell, script: &str, cwd: &AbsolutePathBuf) -> Result { +async fn run_shell_script( + shell: &Shell, + script: &str, + cwd: &AbsolutePathBuf, + excluded_environment_variables: &[String], +) -> Result { run_script_with_timeout( shell, script, SNAPSHOT_TIMEOUT, /*use_login_shell*/ true, cwd, + excluded_environment_variables, ) .await } @@ -279,6 +343,7 @@ async fn run_script_with_timeout( snapshot_timeout: Duration, use_login_shell: bool, cwd: &AbsolutePathBuf, + excluded_environment_variables: &[String], ) -> Result { let args = shell.derive_exec_args(script, use_login_shell); let shell_name = shell.name(); @@ -289,6 +354,7 @@ async fn run_script_with_timeout( handler.args(&args[1..]); handler.stdin(Stdio::null()); handler.current_dir(cwd); + remove_environment_variables(&mut handler, excluded_environment_variables); #[cfg(unix)] unsafe { handler.pre_exec(|| { @@ -311,12 +377,24 @@ async fn run_script_with_timeout( Ok(String::from_utf8_lossy(&output.stdout).into_owned()) } -fn excluded_exports_regex() -> String { - EXCLUDED_EXPORT_VARS.join("|") +fn remove_environment_variables(command: &mut Command, variables: &[String]) { + for variable in variables { + command.env_remove(variable); + } } -fn zsh_snapshot_script() -> String { - let excluded = excluded_exports_regex(); +fn excluded_exports_regex(excluded_environment_variables: &[String]) -> String { + EXCLUDED_EXPORT_VARS + .iter() + .copied() + .chain(excluded_environment_variables.iter().map(String::as_str)) + .map(regex_lite::escape) + .collect::>() + .join("|") +} + +fn zsh_snapshot_script(excluded_environment_variables: &[String]) -> String { + let excluded = excluded_exports_regex(excluded_environment_variables); let script = r##"if [[ -n "$ZDOTDIR" ]]; then rc="$ZDOTDIR/.zshrc" else @@ -359,8 +437,8 @@ fi script.replace("EXCLUDED_EXPORTS", &excluded) } -fn bash_snapshot_script() -> String { - let excluded = excluded_exports_regex(); +fn bash_snapshot_script(excluded_environment_variables: &[String]) -> String { + let excluded = excluded_exports_regex(excluded_environment_variables); let script = r##"if [ -z "$BASH_ENV" ] && [ -r "$HOME/.bashrc" ]; then . "$HOME/.bashrc" fi @@ -401,8 +479,8 @@ fi script.replace("EXCLUDED_EXPORTS", &excluded) } -fn sh_snapshot_script() -> String { - let excluded = excluded_exports_regex(); +fn sh_snapshot_script(excluded_environment_variables: &[String]) -> String { + let excluded = excluded_exports_regex(excluded_environment_variables); let script = r##"if [ -n "$ENV" ] && [ -r "$ENV" ]; then . "$ENV" fi diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index dcba3e2459..7a2d4f8d4d 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -83,7 +83,13 @@ fn assert_posix_snapshot_sections(snapshot: &str) { async fn get_snapshot(shell_type: ShellType) -> Result { let dir = tempdir()?; let path = dir.path().join("snapshot.sh"); - write_shell_snapshot(shell_type, &path.abs(), &dir.path().abs()).await?; + write_shell_snapshot( + shell_type, + &path.abs(), + &dir.path().abs(), + /*excluded_environment_variables*/ &[], + ) + .await?; let content = fs::read_to_string(&path).await?; Ok(content) } @@ -128,7 +134,7 @@ fn snapshot_file_name_parser_supports_legacy_and_suffixed_names() { fn bash_snapshot_filters_invalid_exports() -> Result<()> { let output = Command::new("/bin/bash") .arg("-c") - .arg(bash_snapshot_script()) + .arg(bash_snapshot_script(&[])) .env("BASH_ENV", "/dev/null") .env("VALID_NAME", "ok") .env("PWD", "/tmp/stale") @@ -147,13 +153,34 @@ fn bash_snapshot_filters_invalid_exports() -> Result<()> { Ok(()) } +#[cfg(unix)] +#[test] +fn bash_snapshot_excludes_configured_exports() -> Result<()> { + let output = Command::new("/bin/bash") + .arg("-c") + .arg(bash_snapshot_script(&[ + "CODEX_WIF_SUBJECT_TOKEN".to_string() + ])) + .env("BASH_ENV", "/dev/null") + .env("CODEX_WIF_SUBJECT_TOKEN", "secret.assertion") + .env("SAFE_VAR", "safe") + .output()?; + + assert!(output.status.success()); + + let stdout = String::from_utf8_lossy(&output.stdout); + assert!(!stdout.contains("CODEX_WIF_SUBJECT_TOKEN")); + assert!(stdout.contains("SAFE_VAR")); + Ok(()) +} + #[cfg(unix)] #[test] fn bash_snapshot_preserves_multiline_exports() -> Result<()> { let multiline_cert = "-----BEGIN CERTIFICATE-----\nabc\n-----END CERTIFICATE-----"; let output = Command::new("/bin/bash") .arg("-c") - .arg(bash_snapshot_script()) + .arg(bash_snapshot_script(&[])) .env("BASH_ENV", "/dev/null") .env("MULTILINE_CERT", multiline_cert) .output()?; @@ -201,6 +228,7 @@ async fn try_create_creates_and_deletes_snapshot_file() -> Result<()> { ThreadId::new(), &dir.path().abs(), &shell, + /*excluded_environment_variables*/ &[], /*state_db*/ None, ) .await @@ -230,6 +258,7 @@ async fn try_create_uses_distinct_generation_paths() -> Result<()> { session_id, &dir.path().abs(), &shell, + /*excluded_environment_variables*/ &[], /*state_db*/ None, ) .await @@ -239,6 +268,7 @@ async fn try_create_uses_distinct_generation_paths() -> Result<()> { session_id, &dir.path().abs(), &shell, + /*excluded_environment_variables*/ &[], /*state_db*/ None, ) .await @@ -283,7 +313,7 @@ async fn snapshot_shell_does_not_inherit_stdin() -> Result<()> { let home_display = home.display(); let script = format!( "HOME=\"{home_display}\"; export HOME; {}", - bash_snapshot_script() + bash_snapshot_script(&[]) ); let output = run_script_with_timeout( &shell, @@ -291,6 +321,7 @@ async fn snapshot_shell_does_not_inherit_stdin() -> Result<()> { Duration::from_secs(2), /*use_login_shell*/ true, &home, + /*excluded_environment_variables*/ &[], ) .await .context("run snapshot command")?; @@ -334,6 +365,7 @@ async fn timed_out_snapshot_shell_is_terminated() -> Result<()> { Duration::from_secs(1), /*use_login_shell*/ true, &dir.path().abs(), + /*excluded_environment_variables*/ &[], ) .await .expect_err("snapshot shell should time out"); @@ -368,6 +400,22 @@ async fn timed_out_snapshot_shell_is_terminated() -> Result<()> { Ok(()) } +#[test] +fn snapshot_command_removes_configured_environment_variables() { + let variable = "CODEX_WIF_SNAPSHOT_ASSERTION"; + let mut command = tokio::process::Command::new("true"); + command.env(variable, "secret.assertion"); + + remove_environment_variables(&mut command, &[variable.to_string()]); + + assert!( + command + .as_std() + .get_envs() + .any(|(name, value)| name == variable && value.is_none()) + ); +} + #[cfg(target_os = "macos")] #[tokio::test] async fn macos_zsh_snapshot_includes_sections() -> Result<()> { diff --git a/codex-rs/hooks/src/engine/mod_tests.rs b/codex-rs/hooks/src/engine/mod_tests.rs index c12cc4a809..f92293fc6f 100644 --- a/codex-rs/hooks/src/engine/mod_tests.rs +++ b/codex-rs/hooks/src/engine/mod_tests.rs @@ -1206,6 +1206,7 @@ fn profile_user_layers_load_shared_hooks_json_once() { CommandShell { program: String::new(), args: Vec::new(), + excluded_environment_variables: Vec::new(), }, ); @@ -1280,6 +1281,7 @@ fn malformed_hooks_json_is_reported_as_startup_warning() { CommandShell { program: String::new(), args: Vec::new(), + excluded_environment_variables: Vec::new(), }, ); diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index f6efc5698d..378162fe6b 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -4,6 +4,21 @@ version.workspace = true edition.workspace = true license.workspace = true +[features] +default = ["workload-identity-azure"] +workload-identity-all-providers = [ + "workload-identity-aws", + "workload-identity-azure", + "workload-identity-gcp", + "workload-identity-github-actions", + "workload-identity-spiffe", +] +workload-identity-aws = ["codex-workload-identity-providers/aws"] +workload-identity-azure = ["codex-workload-identity-providers/azure"] +workload-identity-gcp = ["codex-workload-identity-providers/gcp"] +workload-identity-github-actions = ["codex-workload-identity-providers/github-actions"] +workload-identity-spiffe = ["codex-workload-identity-providers/spiffe"] + [lints] workspace = true @@ -21,13 +36,7 @@ codex-protocol = { workspace = true } codex-secrets = { workspace = true } codex-terminal-detection = { workspace = true } codex-utils-template = { workspace = true } -codex-workload-identity-providers = { workspace = true, features = [ - "aws", - "azure", - "gcp", - "github-actions", - "spiffe", -] } +codex-workload-identity-providers = { workspace = true } once_cell = { workspace = true } os_info = { workspace = true } rand = { workspace = true } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index e2865364ff..4ad7cf51e5 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -549,6 +549,19 @@ fn external_auth_tokens_without_chatgpt_metadata_cannot_seed_chatgpt_auth() { ); } +#[test] +fn external_auth_tokens_debug_redacts_access_token() { + let tokens = ExternalAuthTokens::chatgpt( + "secret-access-token", + "workspace-id", + Some("team".to_string()), + ); + + let debug = format!("{tokens:?}"); + assert!(!debug.contains("secret-access-token")); + assert!(debug.contains("[REDACTED]")); +} + struct TestExternalChatgptAuth { initial: ExternalAuthTokens, refreshed: ExternalAuthTokens, @@ -557,40 +570,44 @@ struct TestExternalChatgptAuth { struct FailingExternalChatgptAuth; -#[async_trait] impl ExternalAuth for FailingExternalChatgptAuth { fn auth_mode(&self) -> AuthMode { AuthMode::Chatgpt } - async fn resolve(&self) -> std::io::Result> { - Err(std::io::Error::other("external token exchange failed")) + fn requires_successful_resolution(&self) -> bool { + true } - async fn refresh( + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + Box::pin(async { Err(std::io::Error::other("external token exchange failed")) }) + } + + fn refresh( &self, _context: ExternalAuthRefreshContext, - ) -> std::io::Result { - Err(std::io::Error::other("external token exchange failed")) + ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { + Box::pin(async { Err(std::io::Error::other("external token exchange failed")) }) } } -#[async_trait] impl ExternalAuth for TestExternalChatgptAuth { fn auth_mode(&self) -> AuthMode { AuthMode::Chatgpt } - async fn resolve(&self) -> std::io::Result> { - Ok(Some(self.initial.clone())) + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + Box::pin(async { Ok(Some(self.initial.clone())) }) } - async fn refresh( + fn refresh( &self, _context: ExternalAuthRefreshContext, - ) -> std::io::Result { - self.refresh_count.fetch_add(1, Ordering::Relaxed); - Ok(self.refreshed.clone()) + ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { + Box::pin(async { + self.refresh_count.fetch_add(1, Ordering::Relaxed); + Ok(self.refreshed.clone()) + }) } } @@ -622,6 +639,7 @@ async fn external_chatgpt_auth_resolves_and_refreshes_without_writing_auth_file( /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), ) .await; manager.set_external_auth(external.clone()); @@ -649,6 +667,50 @@ async fn external_chatgpt_auth_resolves_and_refreshes_without_writing_auth_file( assert!(!get_auth_file(codex_home.path()).exists()); } +#[tokio::test] +async fn external_chatgpt_auth_allows_host_to_change_account_during_refresh() { + let codex_home = tempdir().expect("create temp codex home"); + let token = fake_jwt_for_auth_file_params(&AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("enterprise".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), + }) + .expect("build test access token"); + let external = Arc::new(TestExternalChatgptAuth { + initial: ExternalAuthTokens::chatgpt( + token.clone(), + WORKSPACE_ID_ALLOWED, + Some("enterprise".to_string()), + ), + refreshed: ExternalAuthTokens::chatgpt( + token, + WORKSPACE_ID_DISALLOWED, + Some("enterprise".to_string()), + ), + refresh_count: AtomicUsize::new(0), + }); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + ) + .await; + manager.set_external_auth(external); + manager.auth().await.expect("initial auth should resolve"); + + manager + .refresh_token() + .await + .expect("host-managed account change should remain supported"); + + assert_eq!( + manager.auth_cached().and_then(|auth| auth.get_account_id()), + Some(WORKSPACE_ID_DISALLOWED.to_string()) + ); +} + #[tokio::test] async fn auth_result_preserves_external_auth_resolution_errors() { let manager = AuthManager::shared( @@ -656,6 +718,7 @@ async fn auth_result_preserves_external_auth_resolution_errors() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), ) .await; manager.set_external_auth(Arc::new(FailingExternalChatgptAuth)); @@ -714,6 +777,7 @@ async fn external_bearer_only_auth_manager_returns_none_when_command_fails() { let script = ProviderAuthScript::new_failing().unwrap(); let manager = AuthManager::external_bearer_only(script.auth_config()); + assert_eq!(manager.auth_result().await.unwrap(), None); assert_eq!(manager.auth().await, None); } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 209d8e84cf..d6e2868e06 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -123,12 +123,22 @@ pub enum RefreshTokenError { Transient(#[from] std::io::Error), } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct ExternalAuthTokens { pub access_token: String, pub chatgpt_metadata: Option, } +impl Debug for ExternalAuthTokens { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ExternalAuthTokens") + .field("access_token", &"[REDACTED]") + .field("chatgpt_metadata", &self.chatgpt_metadata) + .finish() + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ExternalAuthChatgptMetadata { pub account_id: String, @@ -181,6 +191,11 @@ pub trait ExternalAuth: Send + Sync { /// Indicates which top-level auth mode this external provider supplies. fn auth_mode(&self) -> AuthMode; + /// Whether a resolution failure must prevent the caller from using fallback auth. + fn requires_successful_resolution(&self) -> bool { + false + } + /// Returns cached or immediately available auth, if this provider can resolve it synchronously /// from the caller's perspective. fn resolve(&self) -> ExternalAuthFuture<'_, Option> { @@ -1700,8 +1715,17 @@ impl AuthManager { /// Resolves current auth while preserving errors from explicitly configured external auth. pub async fn auth_result(&self) -> std::io::Result> { - if self.has_external_auth() { - return self.resolve_external_auth().await; + if let Some(external_auth) = self.external_auth() { + match self.resolve_external_auth(&external_auth).await { + Ok(Some(auth)) => return Ok(Some(auth)), + Ok(None) => {} + Err(error) if external_auth.requires_successful_resolution() => { + return Err(error); + } + Err(error) => { + tracing::error!("Failed to resolve optional external auth: {error}"); + } + } } let Some(auth) = self.auth_cached() else { @@ -1932,7 +1956,7 @@ impl AuthManager { if let Some(workload_identity) = config.workload_identity() { let client = build_workload_identity_client( workload_identity, - CLIENT_ID, + oauth_client_id(), build_reqwest_client(), ); auth_manager.set_external_auth(Arc::new(WorkloadIdentityExternalAuth::new(client))); @@ -1961,11 +1985,10 @@ impl AuthManager { self.external_auth_mode() == Some(AuthMode::ApiKey) } - async fn resolve_external_auth(&self) -> std::io::Result> { - let Some(external_auth) = self.external_auth() else { - return Ok(None); - }; - + async fn resolve_external_auth( + &self, + external_auth: &Arc, + ) -> std::io::Result> { match external_auth.resolve().await { Ok(Some(tokens)) if external_auth.auth_mode() == AuthMode::ApiKey => { Ok(Some(CodexAuth::from_api_key(&tokens.access_token))) @@ -1975,7 +1998,7 @@ impl AuthManager { self.set_cached_auth(Some(auth.clone())); Ok(Some(auth)) } - Ok(None) => Ok(self.auth_cached()), + Ok(None) => Ok(None), Err(error) => Err(error), } } @@ -2003,6 +2026,7 @@ impl AuthManager { auth_dot_json, AuthCredentialsStoreMode::Ephemeral, self.chatgpt_base_url.as_deref(), + self.keyring_backend_kind, ) .await } @@ -2216,31 +2240,11 @@ impl AuthManager { if external_auth.auth_mode() == AuthMode::ApiKey { return Ok(()); } - let Some(chatgpt_metadata) = refreshed.chatgpt_metadata() else { - return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresh did not return ChatGPT metadata", - ))); - }; - if let Some(expected_workspace_ids) = forced_chatgpt_workspace_id.as_deref() - && !expected_workspace_ids.contains(&chatgpt_metadata.account_id) - { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth refresh returned workspace {:?}, expected one of {:?}", - chatgpt_metadata.account_id, expected_workspace_ids, - ), - ))); - } - let auth_dot_json = - AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - ) - .map_err(RefreshTokenError::Transient)?; - self.reload().await; + let auth = self + .external_chatgpt_auth(&refreshed) + .await + .map_err(RefreshTokenError::Transient)?; + self.set_cached_auth(Some(auth)); Ok(()) } diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs index bfc065e155..955c45ccf8 100644 --- a/codex-rs/login/src/auth/workload_identity.rs +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -1,8 +1,8 @@ -use async_trait::async_trait; use codex_app_server_protocol::AuthMode; use codex_workload_identity_providers::ConfiguredWorkloadIdentityClient; use super::ExternalAuth; +use super::ExternalAuthFuture; use super::ExternalAuthRefreshContext; use super::ExternalAuthTokens; @@ -31,20 +31,23 @@ impl WorkloadIdentityExternalAuth { } } -#[async_trait] impl ExternalAuth for WorkloadIdentityExternalAuth { fn auth_mode(&self) -> AuthMode { AuthMode::Chatgpt } - async fn resolve(&self) -> std::io::Result> { - self.tokens(false).await.map(Some) + fn requires_successful_resolution(&self) -> bool { + true } - async fn refresh( + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + Box::pin(async move { self.tokens(false).await.map(Some) }) + } + + fn refresh( &self, _context: ExternalAuthRefreshContext, - ) -> std::io::Result { - self.tokens(true).await + ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { + Box::pin(async move { self.tokens(true).await }) } } diff --git a/codex-rs/workload-identity-azure/src/lib.rs b/codex-rs/workload-identity-azure/src/lib.rs index ef98c30d69..62362ea81d 100644 --- a/codex-rs/workload-identity-azure/src/lib.rs +++ b/codex-rs/workload-identity-azure/src/lib.rs @@ -9,28 +9,32 @@ const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; #[derive(Clone, Debug)] pub struct AzureSubjectTokenProvider { - source: Option, + token_file: Option, } impl AzureSubjectTokenProvider { pub fn new(token_file: Option) -> Self { let token_file = token_file .or_else(|| std::env::var_os(AZURE_FEDERATED_TOKEN_FILE_ENV).map(PathBuf::from)); - Self { - source: token_file.map(|path| FileSubjectTokenSource::for_source("azure", path)), - } + Self { token_file } } } impl SubjectTokenProvider for AzureSubjectTokenProvider { async fn subject_token(&self) -> Result { - match &self.source { - Some(source) => source.subject_token().await, - None => Err(SubjectTokenError::MissingPrerequisite { - provider: "azure", - prerequisite: AZURE_FEDERATED_TOKEN_FILE_ENV.to_string(), - }), + let token_file = + self.token_file + .as_ref() + .ok_or_else(|| SubjectTokenError::MissingPrerequisite { + provider: "azure", + prerequisite: AZURE_FEDERATED_TOKEN_FILE_ENV.to_string(), + })?; + if !token_file.is_absolute() { + return Err(SubjectTokenError::InvalidConfiguration { provider: "azure" }); } + FileSubjectTokenSource::for_source("azure", token_file.clone()) + .subject_token() + .await } } @@ -57,4 +61,14 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn rejects_relative_token_path() { + let source = AzureSubjectTokenProvider::new(Some(PathBuf::from("azure-token"))); + + assert!(matches!( + source.subject_token().await, + Err(SubjectTokenError::InvalidConfiguration { provider: "azure" }) + )); + } } diff --git a/codex-rs/workload-identity-gcp/src/lib.rs b/codex-rs/workload-identity-gcp/src/lib.rs index e76be1f5a3..0ed9ace4e3 100644 --- a/codex-rs/workload-identity-gcp/src/lib.rs +++ b/codex-rs/workload-identity-gcp/src/lib.rs @@ -1,9 +1,11 @@ use std::time::Duration; +use codex_workload_identity::BoundedResponseBodyError; use codex_workload_identity::MAX_SUBJECT_TOKEN_BYTES; use codex_workload_identity::SubjectToken; use codex_workload_identity::SubjectTokenError; use codex_workload_identity::SubjectTokenProvider; +use codex_workload_identity::read_bounded_response_body; use reqwest::header::HeaderValue; use url::Url; @@ -83,14 +85,17 @@ impl SubjectTokenProvider for GcpSubjectTokenProvider { if !response.status().is_success() { return Err(SubjectTokenError::InvalidResponse { provider: "gcp" }); } - let body = response - .bytes() + let body = read_bounded_response_body(response, MAX_SUBJECT_TOKEN_BYTES) .await - .map_err(|_| SubjectTokenError::InvalidResponse { provider: "gcp" })?; - if body.len() > MAX_SUBJECT_TOKEN_BYTES { - return Err(SubjectTokenError::TooLarge { provider: "gcp" }); - } - let body = String::from_utf8(body.to_vec()) + .map_err(|error| match error { + BoundedResponseBodyError::Request(_) => { + SubjectTokenError::InvalidResponse { provider: "gcp" } + } + BoundedResponseBodyError::TooLarge => { + SubjectTokenError::TooLarge { provider: "gcp" } + } + })?; + let body = String::from_utf8(body) .map_err(|_| SubjectTokenError::InvalidResponse { provider: "gcp" })?; SubjectToken::jwt(body, "gcp") } diff --git a/codex-rs/workload-identity-github-actions/Cargo.toml b/codex-rs/workload-identity-github-actions/Cargo.toml index 4b83fb7e30..1814dc2405 100644 --- a/codex-rs/workload-identity-github-actions/Cargo.toml +++ b/codex-rs/workload-identity-github-actions/Cargo.toml @@ -11,12 +11,12 @@ workspace = true codex-workload-identity = { workspace = true } reqwest = { workspace = true, features = ["json"] } serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } url = { workspace = true } [dev-dependencies] anyhow = { workspace = true } pretty_assertions = { workspace = true } -serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } wiremock = { workspace = true } diff --git a/codex-rs/workload-identity-github-actions/src/lib.rs b/codex-rs/workload-identity-github-actions/src/lib.rs index d4c4c7a20e..649f976d73 100644 --- a/codex-rs/workload-identity-github-actions/src/lib.rs +++ b/codex-rs/workload-identity-github-actions/src/lib.rs @@ -1,8 +1,11 @@ use std::time::Duration; +use codex_workload_identity::BoundedResponseBodyError; +use codex_workload_identity::MAX_SUBJECT_TOKEN_BYTES; use codex_workload_identity::SubjectToken; use codex_workload_identity::SubjectTokenError; use codex_workload_identity::SubjectTokenProvider; +use codex_workload_identity::read_bounded_response_body; use serde::Deserialize; use url::Host; use url::Url; @@ -53,7 +56,11 @@ impl GithubActionsSubjectTokenProvider { Host::Ipv4(address) => address.is_loopback(), Host::Ipv6(address) => address.is_loopback(), }); - if !github_https && !loopback_http { + if (!github_https && !loopback_http) + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { return Err(SubjectTokenError::InvalidConfiguration { provider: "github_actions", }); @@ -107,7 +114,17 @@ impl SubjectTokenProvider for GithubActionsSubjectTokenProvider { provider: "github_actions", }); } - let response = response.json::().await.map_err(|_| { + let body = read_bounded_response_body(response, MAX_SUBJECT_TOKEN_BYTES) + .await + .map_err(|error| match error { + BoundedResponseBodyError::Request(_) => SubjectTokenError::InvalidResponse { + provider: "github_actions", + }, + BoundedResponseBodyError::TooLarge => SubjectTokenError::TooLarge { + provider: "github_actions", + }, + })?; + let response = serde_json::from_slice::(&body).map_err(|_| { SubjectTokenError::InvalidResponse { provider: "github_actions", } @@ -182,4 +199,26 @@ mod tests { assert_eq!(audiences, vec!["openai-audience"]); Ok(()) } + + #[test] + fn rejects_request_url_user_info_and_fragments() { + for request_url in [ + "https://user@actions.githubusercontent.com/token", + "https://actions.githubusercontent.com/token#fragment", + ] { + let source = GithubActionsSubjectTokenProvider { + request_url: Some(request_url.to_string()), + request_token: Some("runner-request-secret".to_string()), + audience: "openai-audience".to_string(), + http: reqwest::Client::new(), + }; + + assert!(matches!( + source.request_url(), + Err(SubjectTokenError::InvalidConfiguration { + provider: "github_actions" + }) + )); + } + } } diff --git a/codex-rs/workload-identity-providers/BUILD.bazel b/codex-rs/workload-identity-providers/BUILD.bazel index e394cabde7..11f8716ce4 100644 --- a/codex-rs/workload-identity-providers/BUILD.bazel +++ b/codex-rs/workload-identity-providers/BUILD.bazel @@ -2,11 +2,6 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "workload-identity-providers", + crate_features = ["azure"], crate_name = "codex_workload_identity_providers", - crate_features = [ - "azure", - "gcp", - "github-actions", - "spiffe", - ], ) diff --git a/codex-rs/workload-identity-providers/Cargo.toml b/codex-rs/workload-identity-providers/Cargo.toml index bd70e0baa8..d6286ab536 100644 --- a/codex-rs/workload-identity-providers/Cargo.toml +++ b/codex-rs/workload-identity-providers/Cargo.toml @@ -25,9 +25,7 @@ codex-workload-identity-spiffe = { workspace = true, optional = true } reqwest = { workspace = true } [dev-dependencies] -anyhow = { workspace = true } -tempfile = { workspace = true } -tokio = { workspace = true, features = ["fs", "macros", "rt-multi-thread"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } [lib] doctest = false diff --git a/codex-rs/workload-identity-spiffe/Cargo.toml b/codex-rs/workload-identity-spiffe/Cargo.toml index 2d8d135063..7137e1ecdc 100644 --- a/codex-rs/workload-identity-spiffe/Cargo.toml +++ b/codex-rs/workload-identity-spiffe/Cargo.toml @@ -12,7 +12,6 @@ codex-workload-identity = { workspace = true } spiffe = { workspace = true, features = ["workload-api-jwt"] } [dev-dependencies] -anyhow = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } [lib] diff --git a/codex-rs/workload-identity-spiffe/src/lib.rs b/codex-rs/workload-identity-spiffe/src/lib.rs index a9fe6b3212..6025ff6798 100644 --- a/codex-rs/workload-identity-spiffe/src/lib.rs +++ b/codex-rs/workload-identity-spiffe/src/lib.rs @@ -3,6 +3,7 @@ use codex_workload_identity::SubjectTokenError; use codex_workload_identity::SubjectTokenProvider; use spiffe::SpiffeId; use spiffe::WorkloadApiClient; +use spiffe::transport::Endpoint; const SPIFFE_ENDPOINT_SOCKET_ENV: &str = "SPIFFE_ENDPOINT_SOCKET"; @@ -38,13 +39,18 @@ impl SubjectTokenProvider for SpiffeSubjectTokenProvider { provider: "spiffe", prerequisite: SPIFFE_ENDPOINT_SOCKET_ENV.to_string(), })?; + let endpoint = Endpoint::parse(endpoint_socket) + .map_err(|_| SubjectTokenError::InvalidConfiguration { provider: "spiffe" })?; + if !matches!(&endpoint, Endpoint::Unix(_)) { + return Err(SubjectTokenError::InvalidConfiguration { provider: "spiffe" }); + } let spiffe_id = self .spiffe_id .as_deref() .map(str::parse::) .transpose() .map_err(|_| SubjectTokenError::InvalidConfiguration { provider: "spiffe" })?; - let client = WorkloadApiClient::connect_to(endpoint_socket) + let client = WorkloadApiClient::connect(endpoint) .await .map_err(|_| SubjectTokenError::Unavailable { provider: "spiffe" })?; let jwt_svid = client @@ -72,4 +78,18 @@ mod tests { Err(SubjectTokenError::InvalidConfiguration { provider: "spiffe" }) )); } + + #[tokio::test] + async fn rejects_tcp_endpoint() { + let source = SpiffeSubjectTokenProvider::new( + Some("tcp://127.0.0.1:8081".to_string()), + None, + "openai-audience".to_string(), + ); + + assert!(matches!( + source.subject_token().await, + Err(SubjectTokenError::InvalidConfiguration { provider: "spiffe" }) + )); + } } diff --git a/codex-rs/workload-identity/Cargo.toml b/codex-rs/workload-identity/Cargo.toml index c88638f38f..74ff3ab995 100644 --- a/codex-rs/workload-identity/Cargo.toml +++ b/codex-rs/workload-identity/Cargo.toml @@ -11,14 +11,14 @@ workspace = true reqwest = { workspace = true, features = ["json"] } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true, features = ["fs", "sync"] } +tokio = { workspace = true, features = ["fs", "io-util", "sync"] } url = { workspace = true } [dev-dependencies] anyhow = { workspace = true } pretty_assertions = { workspace = true } -serde_json = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } wiremock = { workspace = true } diff --git a/codex-rs/workload-identity/src/client.rs b/codex-rs/workload-identity/src/client.rs index c84d768852..0e5de67911 100644 --- a/codex-rs/workload-identity/src/client.rs +++ b/codex-rs/workload-identity/src/client.rs @@ -2,9 +2,11 @@ use std::sync::Mutex; use std::time::Duration; use std::time::Instant; +use crate::BoundedResponseBodyError; use crate::SubjectTokenError; use crate::SubjectTokenProvider; use crate::WorkloadIdentityConfig; +use crate::read_bounded_response_body; use reqwest::StatusCode; use serde::Deserialize; use serde::Serialize; @@ -16,10 +18,12 @@ const ACCESS_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:access_token"; const MAX_REFRESH_LEAD: Duration = Duration::from_secs(5 * 60); const TOKEN_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(10); const FAILED_EXCHANGE_RETRY_DELAY: Duration = Duration::from_secs(5); +const MAX_TOKEN_RESPONSE_BYTES: usize = 64 * 1024; #[derive(Clone, PartialEq, Eq)] pub struct WorkloadIdentityAccessToken { pub access_token: String, + pub user_id: String, pub chatgpt_account_id: String, pub chatgpt_plan_type: Option, } @@ -29,6 +33,7 @@ impl std::fmt::Debug for WorkloadIdentityAccessToken { formatter .debug_struct("WorkloadIdentityAccessToken") .field("access_token", &"[REDACTED]") + .field("user_id", &self.user_id) .field("chatgpt_account_id", &self.chatgpt_account_id) .field("chatgpt_plan_type", &self.chatgpt_plan_type) .finish() @@ -47,12 +52,20 @@ pub enum WorkloadIdentityError { EmptyAccessToken, #[error("workload identity token exchange returned an empty ChatGPT account ID")] EmptyAccountId, + #[error("workload identity token exchange returned an empty user ID")] + EmptyUserId, #[error("workload identity token exchange returned a token with no usable lifetime")] InvalidLifetime, #[error("workload identity token exchange returned an unexpected token type")] UnexpectedTokenType, + #[error("workload identity token exchange returned an invalid response")] + InvalidResponse, + #[error("workload identity token exchange returned an oversized response")] + ResponseTooLarge, #[error("workload identity token exchange is unavailable")] ExchangeUnavailable, + #[error("workload identity token exchange changed the mapped principal during refresh")] + PrincipalChanged, #[error("{0}")] RecentFailure(String), } @@ -71,6 +84,7 @@ struct CachedExchangeFailure { struct CacheState { token: Option, failure: Option, + principal: Option<(String, String)>, } pub struct WorkloadIdentityClient { @@ -186,16 +200,22 @@ where .send() .await?; let status = response.status(); + let body = read_bounded_response_body(response, MAX_TOKEN_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedResponseBodyError::Request(error) => WorkloadIdentityError::Request(error), + BoundedResponseBodyError::TooLarge => WorkloadIdentityError::ResponseTooLarge, + })?; if !status.is_success() { - let message = response - .json::() - .await + let message = serde_json::from_slice::(&body) .ok() - .and_then(|response| response.error) + .and_then(TokenExchangeErrorResponse::code) + .filter(|code| is_safe_error_code(code)) .unwrap_or_else(|| "token endpoint rejected the request".to_string()); return Err(WorkloadIdentityError::Rejected { status, message }); } - let response: TokenExchangeResponse = response.json().await?; + let response: TokenExchangeResponse = + serde_json::from_slice(&body).map_err(|_| WorkloadIdentityError::InvalidResponse)?; if !response.token_type.eq_ignore_ascii_case("bearer") || response.issued_token_type != ACCESS_TOKEN_TYPE { @@ -207,12 +227,16 @@ where if response.chatgpt_account_id.trim().is_empty() { return Err(WorkloadIdentityError::EmptyAccountId); } + if response.user_id.trim().is_empty() { + return Err(WorkloadIdentityError::EmptyUserId); + } if response.expires_in == 0 { return Err(WorkloadIdentityError::InvalidLifetime); } let token = WorkloadIdentityAccessToken { access_token: response.access_token, + user_id: response.user_id, chatgpt_account_id: response.chatgpt_account_id, chatgpt_plan_type: response.chatgpt_plan_type, }; @@ -223,16 +247,24 @@ where let refresh_at = Instant::now() .checked_add(lifetime.saturating_sub(refresh_lead)) .ok_or(WorkloadIdentityError::InvalidLifetime)?; - *self + let mut cache = self .cache .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = CacheState { - failure: None, - token: Some(CachedAccessToken { - token: token.clone(), - refresh_at, - }), - }; + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some((user_id, account_id)) = cache.principal.as_ref() + && (user_id != &token.user_id || account_id != &token.chatgpt_account_id) + { + cache.token = None; + return Err(WorkloadIdentityError::PrincipalChanged); + } + cache + .principal + .get_or_insert_with(|| (token.user_id.clone(), token.chatgpt_account_id.clone())); + cache.failure = None; + cache.token = Some(CachedAccessToken { + token: token.clone(), + refresh_at, + }); Ok(token) } } @@ -254,6 +286,7 @@ struct TokenExchangeResponse { issued_token_type: String, token_type: String, expires_in: u64, + user_id: String, chatgpt_account_id: String, #[serde(default)] chatgpt_plan_type: Option, @@ -262,7 +295,31 @@ struct TokenExchangeResponse { #[derive(Deserialize)] struct TokenExchangeErrorResponse { #[serde(default)] - error: Option, + error: Option, +} + +impl TokenExchangeErrorResponse { + fn code(self) -> Option { + match self.error? { + TokenExchangeError::Code(code) => Some(code), + TokenExchangeError::Details { code } => code, + } + } +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum TokenExchangeError { + Code(String), + Details { code: Option }, +} + +fn is_safe_error_code(code: &str) -> bool { + !code.is_empty() + && code.len() <= 64 + && code + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) } #[cfg(test)] diff --git a/codex-rs/workload-identity/src/client_tests.rs b/codex-rs/workload-identity/src/client_tests.rs index 860abbf2cc..68dcd05586 100644 --- a/codex-rs/workload-identity/src/client_tests.rs +++ b/codex-rs/workload-identity/src/client_tests.rs @@ -54,6 +54,7 @@ fn exchange_response(access_token: &str) -> ResponseTemplate { "issued_token_type": ACCESS_TOKEN_TYPE, "token_type": "Bearer", "expires_in": 3600, + "user_id": "user_example", "chatgpt_account_id": "workspace_example", "chatgpt_plan_type": "enterprise", })) @@ -73,6 +74,7 @@ async fn exchanges_azure_subject_token_and_caches_access_token() -> anyhow::Resu let expected = WorkloadIdentityAccessToken { access_token: "first.access.token".to_string(), + user_id: "user_example".to_string(), chatgpt_account_id: "workspace_example".to_string(), chatgpt_plan_type: Some("enterprise".to_string()), }; @@ -144,3 +146,109 @@ async fn forced_refresh_performs_a_new_exchange() -> anyhow::Result<()> { client.refresh().await?; Ok(()) } + +#[tokio::test] +async fn refresh_rejects_a_changed_principal() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(exchange_response("first.access.token")) + .expect(1) + .mount(&server) + .await; + let client = test_client(&server).await?; + client.resolve().await?; + + server.reset().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "second.access.token", + "issued_token_type": ACCESS_TOKEN_TYPE, + "token_type": "Bearer", + "expires_in": 3600, + "user_id": "different_user", + "chatgpt_account_id": "workspace_example", + }))) + .expect(1) + .mount(&server) + .await; + + assert!(matches!( + client.refresh().await, + Err(WorkloadIdentityError::PrincipalChanged) + )); + assert!(client.cache.lock().expect("cache lock").token.is_none()); + Ok(()) +} + +#[tokio::test] +async fn reports_safe_nested_error_code() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "identity_provider_mapping_not_found", + "message": "untrusted diagnostic text", + }, + }))) + .expect(1) + .mount(&server) + .await; + let client = test_client(&server).await?; + + assert_eq!( + client + .resolve() + .await + .expect_err("exchange should fail") + .to_string(), + "workload identity token exchange was rejected with HTTP 400 Bad Request: identity_provider_mapping_not_found" + ); + Ok(()) +} + +#[tokio::test] +async fn rejection_does_not_echo_unstructured_server_text() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": "do not log secret subject token", + }))) + .expect(1) + .mount(&server) + .await; + let client = test_client(&server).await?; + + assert_eq!( + client + .resolve() + .await + .expect_err("exchange should fail") + .to_string(), + "workload identity token exchange was rejected with HTTP 401 Unauthorized: token endpoint rejected the request" + ); + Ok(()) +} + +#[tokio::test] +async fn rejects_oversized_token_endpoint_response() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with( + ResponseTemplate::new(200).set_body_bytes(vec![b'x'; MAX_TOKEN_RESPONSE_BYTES + 1]), + ) + .expect(1) + .mount(&server) + .await; + let client = test_client(&server).await?; + + assert!(matches!( + client.resolve().await, + Err(WorkloadIdentityError::ResponseTooLarge) + )); + Ok(()) +} diff --git a/codex-rs/workload-identity/src/config.rs b/codex-rs/workload-identity/src/config.rs index 21ab3b0b69..706a19268e 100644 --- a/codex-rs/workload-identity/src/config.rs +++ b/codex-rs/workload-identity/src/config.rs @@ -8,8 +8,16 @@ use url::Host; use url::Url; const DEFAULT_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +const AWS_CONTAINER_AUTHORIZATION_TOKEN_ENV: &str = "AWS_CONTAINER_AUTHORIZATION_TOKEN"; +const AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE_ENV: &str = "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE"; +const AWS_CONTAINER_CREDENTIALS_FULL_URI_ENV: &str = "AWS_CONTAINER_CREDENTIALS_FULL_URI"; +const AWS_CONTAINER_CREDENTIALS_RELATIVE_URI_ENV: &str = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"; +const AWS_WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +const GCP_METADATA_HOST_ENV: &str = "GCE_METADATA_HOST"; const GITHUB_ACTIONS_REQUEST_URL_ENV: &str = "ACTIONS_ID_TOKEN_REQUEST_URL"; const GITHUB_ACTIONS_REQUEST_TOKEN_ENV: &str = "ACTIONS_ID_TOKEN_REQUEST_TOKEN"; +const SPIFFE_ENDPOINT_SOCKET_ENV: &str = "SPIFFE_ENDPOINT_SOCKET"; /// Configuration for exchanging a workload credential for a Codex access token. #[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)] @@ -139,6 +147,15 @@ impl WorkloadIdentityConfig { if token_url.scheme() != "https" && !loopback_http { return Err(WorkloadIdentityConfigError::UnsupportedTokenUrlScheme); } + if token_url.host().is_none() + || !token_url.username().is_empty() + || token_url.password().is_some() + || token_url.fragment().is_some() + { + return Err(WorkloadIdentityConfigError::InvalidTokenUrl( + "host is required and user info and fragments are not allowed".to_string(), + )); + } match &self.credential_source { CredentialSourceConfig::Environment { variable } @@ -179,7 +196,7 @@ impl WorkloadIdentityConfig { CredentialSourceConfig::Spiffe { spiffe_id: Some(spiffe_id), .. - } if !spiffe_id.starts_with("spiffe://") => { + } if !is_valid_spiffe_id(spiffe_id) => { Err(WorkloadIdentityConfigError::InvalidSpiffeId) } CredentialSourceConfig::Spiffe { .. } => Ok(()), @@ -209,17 +226,58 @@ impl CredentialSourceConfig { pub fn sensitive_environment_variables(&self) -> Vec<&str> { match self { Self::Environment { variable } => vec![variable.as_str()], + Self::Azure { .. } => vec![AZURE_FEDERATED_TOKEN_FILE_ENV], + Self::Gcp { .. } => vec![GCP_METADATA_HOST_ENV], Self::GithubActions {} => vec![ GITHUB_ACTIONS_REQUEST_URL_ENV, GITHUB_ACTIONS_REQUEST_TOKEN_ENV, ], - Self::File { .. } - | Self::Azure { .. } - | Self::Gcp { .. } - | Self::Spiffe { .. } - | Self::Aws { .. } => Vec::new(), + Self::Spiffe { .. } => vec![SPIFFE_ENDPOINT_SOCKET_ENV], + Self::Aws { .. } => vec![ + AWS_WEB_IDENTITY_TOKEN_FILE_ENV, + AWS_CONTAINER_CREDENTIALS_RELATIVE_URI_ENV, + AWS_CONTAINER_CREDENTIALS_FULL_URI_ENV, + AWS_CONTAINER_AUTHORIZATION_TOKEN_ENV, + AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE_ENV, + ], + Self::File { .. } => Vec::new(), } } + + /// Files and Unix sockets that can be used to acquire the workload credential. + pub fn credential_file_paths(&self) -> Vec { + match self { + Self::File { path } => vec![path.clone()], + Self::Azure { token_file } => token_file + .clone() + .or_else(|| std::env::var_os(AZURE_FEDERATED_TOKEN_FILE_ENV).map(PathBuf::from)) + .into_iter() + .collect(), + Self::Spiffe { + endpoint_socket, .. + } => endpoint_socket + .clone() + .or_else(|| std::env::var(SPIFFE_ENDPOINT_SOCKET_ENV).ok()) + .and_then(|endpoint| unix_socket_path(&endpoint)) + .into_iter() + .collect(), + Self::Aws { .. } => [ + AWS_WEB_IDENTITY_TOKEN_FILE_ENV, + AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE_ENV, + ] + .into_iter() + .filter_map(std::env::var_os) + .map(PathBuf::from) + .collect(), + Self::Environment { .. } | Self::Gcp { .. } | Self::GithubActions {} => Vec::new(), + } + } +} + +fn unix_socket_path(endpoint: &str) -> Option { + let endpoint = Url::parse(endpoint).ok()?; + (endpoint.scheme() == "unix" && !endpoint.path().is_empty()) + .then(|| PathBuf::from(endpoint.path())) } fn is_valid_environment_variable_name(variable: &str) -> bool { @@ -230,6 +288,17 @@ fn is_valid_environment_variable_name(variable: &str) -> bool { && characters.all(|character| character == '_' || character.is_ascii_alphanumeric()) } +fn is_valid_spiffe_id(value: &str) -> bool { + Url::parse(value).ok().is_some_and(|url| { + url.scheme() == "spiffe" + && url.host().is_some() + && url.username().is_empty() + && url.password().is_none() + && url.query().is_none() + && url.fragment().is_none() + }) +} + #[cfg(test)] #[path = "config_tests.rs"] mod tests; diff --git a/codex-rs/workload-identity/src/config_tests.rs b/codex-rs/workload-identity/src/config_tests.rs index 8a1c69bf79..4661e5c3ae 100644 --- a/codex-rs/workload-identity/src/config_tests.rs +++ b/codex-rs/workload-identity/src/config_tests.rs @@ -53,6 +53,31 @@ fn allows_loopback_http_token_endpoint_for_local_development() { assert_eq!(config.validate(), Ok(())); } +#[test] +fn rejects_token_endpoint_user_info() { + let mut config = valid_config(); + config.token_url = "https://user:secret@auth.openai.com/oauth/token".to_string(); + + assert!(matches!( + config.validate(), + Err(WorkloadIdentityConfigError::InvalidTokenUrl(_)) + )); +} + +#[test] +fn rejects_malformed_spiffe_id() { + let mut config = valid_config(); + config.credential_source = CredentialSourceConfig::Spiffe { + endpoint_socket: Some("unix:///tmp/spire.sock".to_string()), + spiffe_id: Some("spiffe://example.org/workload#fragment".to_string()), + }; + + assert_eq!( + config.validate(), + Err(WorkloadIdentityConfigError::InvalidSpiffeId) + ); +} + #[test] fn rejects_relative_token_file() { let mut config = valid_config(); @@ -102,7 +127,7 @@ fn every_source_variant_has_a_stable_name() { } #[test] -fn reports_only_secret_bearing_environment_inputs() { +fn reports_credential_bearing_environment_inputs() { assert_eq!( CredentialSourceConfig::Environment { variable: "OPENAI_WIF_TOKEN".to_string(), @@ -110,6 +135,17 @@ fn reports_only_secret_bearing_environment_inputs() { .sensitive_environment_variables(), vec!["OPENAI_WIF_TOKEN"] ); + assert_eq!( + CredentialSourceConfig::Azure { token_file: None }.sensitive_environment_variables(), + vec!["AZURE_FEDERATED_TOKEN_FILE"] + ); + assert_eq!( + CredentialSourceConfig::Gcp { + service_account: None + } + .sensitive_environment_variables(), + vec!["GCE_METADATA_HOST"] + ); assert_eq!( CredentialSourceConfig::GithubActions {}.sensitive_environment_variables(), vec![ @@ -117,6 +153,24 @@ fn reports_only_secret_bearing_environment_inputs() { "ACTIONS_ID_TOKEN_REQUEST_TOKEN" ] ); + assert_eq!( + CredentialSourceConfig::Spiffe { + endpoint_socket: None, + spiffe_id: None, + } + .sensitive_environment_variables(), + vec!["SPIFFE_ENDPOINT_SOCKET"] + ); + assert_eq!( + CredentialSourceConfig::Aws { region: None }.sensitive_environment_variables(), + vec![ + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + ] + ); assert!( CredentialSourceConfig::File { path: PathBuf::from("/var/run/openai/token") @@ -125,3 +179,29 @@ fn reports_only_secret_bearing_environment_inputs() { .is_empty() ); } + +#[test] +fn reports_explicit_credential_file_and_unix_socket_paths() { + assert_eq!( + CredentialSourceConfig::File { + path: PathBuf::from("/var/run/openai/token") + } + .credential_file_paths(), + vec![PathBuf::from("/var/run/openai/token")] + ); + assert_eq!( + CredentialSourceConfig::Azure { + token_file: Some(PathBuf::from("/var/run/azure/token")) + } + .credential_file_paths(), + vec![PathBuf::from("/var/run/azure/token")] + ); + assert_eq!( + CredentialSourceConfig::Spiffe { + endpoint_socket: Some("unix:///var/run/spire/agent.sock".to_string()), + spiffe_id: None, + } + .credential_file_paths(), + vec![PathBuf::from("/var/run/spire/agent.sock")] + ); +} diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs index da74c26892..e48ea6fcda 100644 --- a/codex-rs/workload-identity/src/lib.rs +++ b/codex-rs/workload-identity/src/lib.rs @@ -1,5 +1,6 @@ mod client; mod config; +mod response; mod source; pub use client::WorkloadIdentityAccessToken; @@ -9,6 +10,8 @@ pub use config::CredentialSourceConfig; pub use config::WorkloadIdentityConfig; pub use config::WorkloadIdentityConfigError; pub use config::default_token_url; +pub use response::BoundedResponseBodyError; +pub use response::read_bounded_response_body; pub use source::EnvironmentSubjectTokenSource; pub use source::FileSubjectTokenSource; pub use source::MAX_SUBJECT_TOKEN_BYTES; diff --git a/codex-rs/workload-identity/src/response.rs b/codex-rs/workload-identity/src/response.rs new file mode 100644 index 0000000000..f0c91e9cd8 --- /dev/null +++ b/codex-rs/workload-identity/src/response.rs @@ -0,0 +1,30 @@ +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum BoundedResponseBodyError { + #[error(transparent)] + Request(#[from] reqwest::Error), + #[error("response body exceeds the configured size limit")] + TooLarge, +} + +pub async fn read_bounded_response_body( + mut response: reqwest::Response, + max_bytes: usize, +) -> Result, BoundedResponseBodyError> { + if response + .content_length() + .is_some_and(|length| length > max_bytes as u64) + { + return Err(BoundedResponseBodyError::TooLarge); + } + + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await? { + if chunk.len() > max_bytes.saturating_sub(body.len()) { + return Err(BoundedResponseBodyError::TooLarge); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} diff --git a/codex-rs/workload-identity/src/source.rs b/codex-rs/workload-identity/src/source.rs index a5b103ada0..4ab287a67b 100644 --- a/codex-rs/workload-identity/src/source.rs +++ b/codex-rs/workload-identity/src/source.rs @@ -1,6 +1,7 @@ use std::path::PathBuf; use thiserror::Error; +use tokio::io::AsyncReadExt; pub const JWT_SUBJECT_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:jwt"; pub const MAX_SUBJECT_TOKEN_BYTES: usize = 1024 * 1024; @@ -151,14 +152,21 @@ impl FileSubjectTokenSource { impl SubjectTokenProvider for FileSubjectTokenSource { async fn subject_token(&self) -> Result { - let metadata = - tokio::fs::metadata(&self.path) - .await - .map_err(|error| SubjectTokenError::ReadFile { - provider: self.source, - path: self.path.clone(), - kind: error.kind(), - })?; + let file = tokio::fs::File::open(&self.path).await.map_err(|error| { + SubjectTokenError::ReadFile { + provider: self.source, + path: self.path.clone(), + kind: error.kind(), + } + })?; + let metadata = file + .metadata() + .await + .map_err(|error| SubjectTokenError::ReadFile { + provider: self.source, + path: self.path.clone(), + kind: error.kind(), + })?; if !metadata.is_file() { return Err(SubjectTokenError::NotAFile { provider: self.source, @@ -170,13 +178,25 @@ impl SubjectTokenProvider for FileSubjectTokenSource { provider: self.source, }); } - let value = tokio::fs::read_to_string(&self.path) + let mut bytes = Vec::new(); + file.take(MAX_SUBJECT_TOKEN_BYTES as u64 + 1) + .read_to_end(&mut bytes) .await .map_err(|error| SubjectTokenError::ReadFile { provider: self.source, path: self.path.clone(), kind: error.kind(), })?; + if bytes.len() > MAX_SUBJECT_TOKEN_BYTES { + return Err(SubjectTokenError::TooLarge { + provider: self.source, + }); + } + let value = String::from_utf8(bytes).map_err(|_| SubjectTokenError::ReadFile { + provider: self.source, + path: self.path.clone(), + kind: std::io::ErrorKind::InvalidData, + })?; SubjectToken::jwt(value, self.source) } } diff --git a/codex-rs/workload-identity/src/source_tests.rs b/codex-rs/workload-identity/src/source_tests.rs index 094a6f1c04..63565cc6be 100644 --- a/codex-rs/workload-identity/src/source_tests.rs +++ b/codex-rs/workload-identity/src/source_tests.rs @@ -18,13 +18,12 @@ async fn file_source_rereads_rotated_token() -> anyhow::Result<()> { #[tokio::test] async fn environment_source_captures_startup_value() -> anyhow::Result<()> { const VARIABLE: &str = "CODEX_WIF_SOURCE_CAPTURE_TEST"; - // This test owns a unique process variable and restores it before returning. - unsafe { std::env::set_var(VARIABLE, "first.token.value") }; - let source = EnvironmentSubjectTokenSource::capture(VARIABLE); - unsafe { std::env::set_var(VARIABLE, "second.token.value") }; + let source = EnvironmentSubjectTokenSource { + variable: VARIABLE.to_string(), + captured: CapturedEnvironmentValue::Value("first.token.value".to_string()), + }; assert_eq!(source.subject_token().await?.value(), "first.token.value"); - unsafe { std::env::remove_var(VARIABLE) }; Ok(()) }