From 18914db7f1c8a3145747c577abf62ae8f1f0f93d Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Fri, 10 Jul 2026 03:43:46 +0000 Subject: [PATCH] [auth] integrate workload identity with Codex auth [ci changed_files] --- codex-rs/Cargo.lock | 1 + codex-rs/config/src/config_toml.rs | 17 + codex-rs/core/config.schema.json | 36 ++ codex-rs/core/src/config/config_tests.rs | 58 +++ codex-rs/core/src/config/mod.rs | 9 + codex-rs/login/Cargo.toml | 1 + codex-rs/login/src/auth/manager.rs | 44 ++- codex-rs/login/src/auth/mod.rs | 1 + codex-rs/login/src/auth/workload_identity.rs | 215 +++++++++++ .../login/src/auth/workload_identity_tests.rs | 345 ++++++++++++++++++ codex-rs/thread-manager-sample/src/main.rs | 1 + 11 files changed, 726 insertions(+), 2 deletions(-) create mode 100644 codex-rs/login/src/auth/workload_identity.rs create mode 100644 codex-rs/login/src/auth/workload_identity_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0b103302b2..dcdb1f0488 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3349,6 +3349,7 @@ dependencies = [ "codex-secrets", "codex-terminal-detection", "codex-utils-template", + "codex-workload-identity", "core_test_support", "http 1.4.0", "jsonwebtoken", diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 823d59c0bc..38c231071f 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -148,6 +148,19 @@ pub struct OrchestratorFeatureToml { pub enabled: Option, } +/// Non-secret settings for exchanging a workload assertion for ChatGPT auth. +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct WorkloadIdentityToml { + pub federation_rule_id: Option, + pub principal_id: Option, + pub tenant_id: Option, + pub workspace_id: Option, + /// Optional assertion file. If omitted, Codex reads the assertion from + /// `OPENAI_IDENTITY_TOKEN` or `OPENAI_IDENTITY_TOKEN_FILE`. + pub identity_token_file: Option, +} + /// Base config deserialized from ~/.codex/config.toml. #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[schemars(deny_unknown_fields)] @@ -372,6 +385,10 @@ pub struct ConfigToml { /// Base URL for requests to ChatGPT (as opposed to the OpenAI API). pub chatgpt_base_url: Option, + /// Workload identity settings for non-interactive ChatGPT authentication. + #[serde(default)] + pub workload_identity: Option, + /// Optional product SKU forwarded on host-owned Codex Apps MCP requests. pub apps_mcp_product_sku: Option, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 4994d50ba1..24ffd9e5e0 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -4600,6 +4600,33 @@ } ] }, + "WorkloadIdentityToml": { + "additionalProperties": false, + "description": "Non-secret settings for exchanging a workload assertion for ChatGPT auth.", + "properties": { + "federation_rule_id": { + "type": "string" + }, + "identity_token_file": { + "allOf": [ + { + "$ref": "#/definitions/AbsolutePathBuf" + } + ], + "description": "Optional assertion file. If omitted, Codex reads the assertion from `OPENAI_IDENTITY_TOKEN` or `OPENAI_IDENTITY_TOKEN_FILE`." + }, + "principal_id": { + "type": "string" + }, + "tenant_id": { + "type": "string" + }, + "workspace_id": { + "type": "string" + } + }, + "type": "object" + }, "WorkspaceRootsToml": { "type": "object" } @@ -5531,6 +5558,15 @@ ], "default": null, "description": "Windows-specific configuration." + }, + "workload_identity": { + "allOf": [ + { + "$ref": "#/definitions/WorkloadIdentityToml" + } + ], + "default": null, + "description": "Workload identity settings for non-interactive ChatGPT authentication." } }, "title": "ConfigToml", diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 1df2bb3e3f..7bbb6c3cf1 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -6976,6 +6976,64 @@ async fn for_config_writes_selected_user_config_file() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn workload_identity_config_merges_from_profile_v2() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let base_config = codex_home.path().join(CONFIG_TOML_FILE); + let selected_config = codex_home.path().join("work.config.toml"); + tokio::fs::write( + &base_config, + r#" +[workload_identity] +tenant_id = "tenant-one" +workspace_id = "workspace-one" +"#, + ) + .await?; + tokio::fs::write( + &selected_config, + r#" +[workload_identity] +federation_rule_id = "rule-one" +principal_id = "user-one" +identity_token_file = "./identity-token" +"#, + ) + .await?; + + let config = ConfigBuilder::without_managed_config_for_tests() + .codex_home(codex_home.path().to_path_buf()) + .loader_overrides(LoaderOverrides { + user_config_path: Some(selected_config.abs()), + user_config_profile: Some("work".parse().expect("profile-v2 name")), + ..LoaderOverrides::without_managed_config_for_tests() + }) + .build() + .await?; + + let workload_identity = config + .workload_identity + .expect("merged workload identity config"); + assert_eq!( + workload_identity.federation_rule_id.as_deref(), + Some("rule-one") + ); + assert_eq!(workload_identity.principal_id.as_deref(), Some("user-one")); + assert_eq!(workload_identity.tenant_id.as_deref(), Some("tenant-one")); + assert_eq!( + workload_identity.workspace_id.as_deref(), + Some("workspace-one") + ); + assert_eq!( + workload_identity + .identity_token_file + .map(codex_utils_absolute_path::AbsolutePathBuf::into_path_buf), + Some(codex_home.path().join("identity-token")) + ); + + Ok(()) +} + #[test] fn profile_v2_config_path_resolves_validated_names() -> anyhow::Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 5cd6748c05..b54ef9ecea 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -28,6 +28,7 @@ use codex_config::config_toml::ProjectConfig; use codex_config::config_toml::RealtimeAudioConfig; use codex_config::config_toml::RealtimeConfig; use codex_config::config_toml::ThreadStoreToml; +use codex_config::config_toml::WorkloadIdentityToml; use codex_config::config_toml::validate_model_providers; use codex_config::loader::load_config_layers_state; use codex_config::loader::project_trust_key; @@ -965,6 +966,9 @@ pub struct Config { /// Base URL for requests to ChatGPT (as opposed to the OpenAI API). pub chatgpt_base_url: String, + /// Workload identity settings resolved from the active config layers. + pub workload_identity: Option, + /// Whether Codex-owned clients should respect host system proxy settings. pub respect_system_proxy: bool, @@ -1225,6 +1229,10 @@ impl AuthManagerConfig for Config { self.chatgpt_base_url.clone() } + fn workload_identity_config(&self) -> Option { + self.workload_identity.clone() + } + fn auth_route_config(&self) -> Option { Config::auth_route_config(self) } @@ -3909,6 +3917,7 @@ impl Config { chatgpt_base_url: cfg .chatgpt_base_url .unwrap_or("https://chatgpt.com/backend-api/".to_string()), + workload_identity: cfg.workload_identity, respect_system_proxy, apps_mcp_product_sku: cfg.apps_mcp_product_sku.clone(), realtime_audio: cfg diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index f50648c3af..43d4638baa 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -20,6 +20,7 @@ codex-protocol = { workspace = true } codex-secrets = { workspace = true } codex-terminal-detection = { workspace = true } codex-utils-template = { workspace = true } +codex-workload-identity = { workspace = true } http = { workspace = true } once_cell = { workspace = true } os_info = { workspace = true } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9509efc348..f1caf39434 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -56,6 +56,7 @@ use crate::outbound_proxy::AuthRouteConfig; use crate::token_data::TokenData; use crate::token_data::parse_chatgpt_jwt_claims; use crate::token_data::parse_jwt_expiration; +use codex_config::config_toml::WorkloadIdentityToml; use codex_config::types::AuthCredentialsStoreMode; use codex_http_client::HttpClient; use codex_protocol::account::PlanType as AccountPlanType; @@ -66,6 +67,8 @@ use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; +use super::workload_identity::WorkloadIdentityExternalAuth; + /// Authentication mechanism used by the current user. #[derive(Debug, Clone)] pub enum CodexAuth { @@ -1803,6 +1806,11 @@ pub trait AuthManagerConfig { /// Returns the ChatGPT backend base URL used for first-party backend authorization. fn chatgpt_base_url(&self) -> String; + /// Returns non-secret workload identity settings from the active config layers. + fn workload_identity_config(&self) -> Option { + None + } + /// Returns route-selection settings for auth-owned clients. fn auth_route_config(&self) -> Option; } @@ -2250,6 +2258,26 @@ impl AuthManager { self.commit_external_auth(auth) } + async fn activate_resolving_external_auth(&self, external_auth: Arc) { + self.set_cached_auth(/*new_auth*/ None); + { + let Ok(mut configured_external_auth) = self.external_auth.write() else { + tracing::error!("failed to activate external auth: external auth lock is poisoned"); + return; + }; + *configured_external_auth = Some(Arc::clone(&external_auth)); + } + + match self.resolve_external_auth(&external_auth).await { + Ok(auth) => { + if let Err(error) = self.commit_external_auth(auth) { + tracing::error!("failed to activate external auth: {error}"); + } + } + Err(error) => tracing::error!("failed to activate external auth: {error}"), + } + } + pub fn clear_external_auth(&self) { if let Ok(mut external_auth) = self.external_auth.write() && external_auth.take().is_some() @@ -2316,7 +2344,7 @@ impl AuthManager { config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { - Self::shared( + let manager = Self::shared( config.codex_home(), enable_codex_api_key_env, config.cli_auth_credentials_store_mode(), @@ -2325,7 +2353,19 @@ impl AuthManager { config.auth_keyring_backend_kind(), config.auth_route_config(), ) - .await + .await; + let explicit_process_auth = (enable_codex_api_key_env + && read_codex_api_key_from_env().is_some()) + || read_codex_access_token_from_env().is_some(); + if !explicit_process_auth + && let Some(workload_identity) = + WorkloadIdentityExternalAuth::from_config(config.workload_identity_config()) + { + manager + .activate_resolving_external_auth(Arc::new(workload_identity)) + .await; + } + manager } pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 16477d7827..56c8a4b3f5 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -7,6 +7,7 @@ pub mod error; mod personal_access_token; mod storage; mod util; +mod workload_identity; mod external_bearer; mod manager; diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs new file mode 100644 index 0000000000..6ee7d969c5 --- /dev/null +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -0,0 +1,215 @@ +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; + +use codex_config::config_toml::WorkloadIdentityToml; +use codex_workload_identity::FEDERATION_RULE_ID_ENV_VAR; +use codex_workload_identity::IDENTITY_TOKEN_ENV_VAR; +use codex_workload_identity::IDENTITY_TOKEN_FILE_ENV_VAR; +use codex_workload_identity::PRINCIPAL_ID_ENV_VAR; +use codex_workload_identity::TENANT_ID_ENV_VAR; +use codex_workload_identity::WORKSPACE_ID_ENV_VAR; +use codex_workload_identity::WorkloadIdentityAssertionSource; +use codex_workload_identity::WorkloadIdentityConfig; +use codex_workload_identity::WorkloadIdentityExchange; +use codex_workload_identity::WorkloadIdentityTarget; +use codex_workload_identity::WorkloadIdentityToken; + +use super::CodexAuth; +use super::ExternalAuth; +use super::ExternalAuthFuture; +use super::ExternalAuthRefreshContext; + +pub(super) struct WorkloadIdentityExternalAuth { + exchange: Result, String>, +} + +impl WorkloadIdentityExternalAuth { + pub(super) fn from_config(config: Option) -> Option { + Self::from_config_and_environment(config, ProcessEnvironment::read()) + } + + fn from_config_and_environment( + config: Option, + environment: ProcessEnvironment, + ) -> Option { + if config.is_none() && !environment.is_configured() { + return None; + } + let exchange = resolve_config(config.unwrap_or_default(), environment) + .and_then(|config| { + WorkloadIdentityExchange::new(config).map_err(|error| error.to_string()) + }) + .map(Arc::new); + Some(Self { exchange }) + } + + fn exchange(&self) -> std::io::Result<&Arc> { + self.exchange + .as_ref() + .map_err(|error| std::io::Error::other(error.clone())) + } + + async fn resolve_auth(&self) -> std::io::Result { + let token = self + .exchange()? + .resolve() + .await + .map_err(std::io::Error::other)?; + codex_auth(token) + } + + async fn refresh_auth( + &self, + context: ExternalAuthRefreshContext, + ) -> std::io::Result { + let token = self + .exchange()? + .refresh() + .await + .map_err(std::io::Error::other)?; + if context + .previous_account_id + .as_deref() + .is_some_and(|account_id| account_id != token.chatgpt_account_id) + { + return Err(std::io::Error::other( + "workload identity refresh changed the ChatGPT workspace", + )); + } + codex_auth(token) + } +} + +impl ExternalAuth for WorkloadIdentityExternalAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(self.resolve_auth()) + } + + fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(self.refresh_auth(context)) + } +} + +fn codex_auth(token: WorkloadIdentityToken) -> std::io::Result { + CodexAuth::from_external_chatgpt_tokens( + &token.access_token, + &token.chatgpt_account_id, + token.chatgpt_plan_type.as_deref(), + ) +} + +#[derive(Default)] +struct ProcessEnvironment { + federation_rule_id: Option, + identity_token: Option, + identity_token_file: Option, + principal_id: Option, + tenant_id: Option, + workspace_id: Option, +} + +impl ProcessEnvironment { + fn read() -> Self { + Self { + federation_rule_id: std::env::var_os(FEDERATION_RULE_ID_ENV_VAR), + identity_token: std::env::var_os(IDENTITY_TOKEN_ENV_VAR), + identity_token_file: std::env::var_os(IDENTITY_TOKEN_FILE_ENV_VAR), + principal_id: std::env::var_os(PRINCIPAL_ID_ENV_VAR), + tenant_id: std::env::var_os(TENANT_ID_ENV_VAR), + workspace_id: std::env::var_os(WORKSPACE_ID_ENV_VAR), + } + } + + fn is_configured(&self) -> bool { + self.federation_rule_id.is_some() + || self.identity_token.is_some() + || self.identity_token_file.is_some() + || self.principal_id.is_some() + || self.tenant_id.is_some() + || self.workspace_id.is_some() + } +} + +fn resolve_config( + config: WorkloadIdentityToml, + environment: ProcessEnvironment, +) -> Result { + let source = match config.identity_token_file { + Some(path) => WorkloadIdentityAssertionSource::File(path.into_path_buf()), + None => match (environment.identity_token, environment.identity_token_file) { + (Some(token), None) => WorkloadIdentityAssertionSource::Environment(required_unicode( + Some(token), + IDENTITY_TOKEN_ENV_VAR, + )?), + (None, Some(path)) => WorkloadIdentityAssertionSource::File(PathBuf::from( + required_unicode(Some(path), IDENTITY_TOKEN_FILE_ENV_VAR)?, + )), + (Some(_), Some(_)) => { + return Err(format!( + "set exactly one of {IDENTITY_TOKEN_ENV_VAR} or {IDENTITY_TOKEN_FILE_ENV_VAR}" + )); + } + (None, None) => { + return Err(format!( + "set one of {IDENTITY_TOKEN_ENV_VAR} or {IDENTITY_TOKEN_FILE_ENV_VAR}" + )); + } + }, + }; + let target = WorkloadIdentityTarget { + federation_rule_id: configured_value( + config.federation_rule_id, + environment.federation_rule_id, + "federation_rule_id", + FEDERATION_RULE_ID_ENV_VAR, + )?, + principal_id: configured_value( + config.principal_id, + environment.principal_id, + "principal_id", + PRINCIPAL_ID_ENV_VAR, + )?, + tenant_id: configured_value( + config.tenant_id, + environment.tenant_id, + "tenant_id", + TENANT_ID_ENV_VAR, + )?, + workspace_id: configured_value( + config.workspace_id, + environment.workspace_id, + "workspace_id", + WORKSPACE_ID_ENV_VAR, + )?, + }; + WorkloadIdentityConfig::new(target, source).map_err(|error| error.to_string()) +} + +fn configured_value( + configured: Option, + environment: Option, + field: &'static str, + variable: &'static str, +) -> Result { + match configured { + Some(value) if !value.trim().is_empty() => Ok(value), + Some(_) => Err(format!("workload identity field {field} must not be empty")), + None => required_unicode(environment, variable), + } +} + +fn required_unicode(value: Option, variable: &'static str) -> Result { + let value = value.ok_or_else(|| format!("workload identity requires {variable}"))?; + let value = value + .into_string() + .map_err(|_| format!("workload identity variable {variable} is invalid"))?; + if value.trim().is_empty() { + return Err(format!("workload identity variable {variable} is invalid")); + } + Ok(value.trim().to_string()) +} + +#[cfg(test)] +#[path = "workload_identity_tests.rs"] +mod tests; diff --git a/codex-rs/login/src/auth/workload_identity_tests.rs b/codex-rs/login/src/auth/workload_identity_tests.rs new file mode 100644 index 0000000000..3818611ffe --- /dev/null +++ b/codex-rs/login/src/auth/workload_identity_tests.rs @@ -0,0 +1,345 @@ +use std::collections::HashMap; +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_config::config_toml::WorkloadIdentityToml; +use codex_config::types::AuthCredentialsStoreMode; +use codex_protocol::auth::AuthMode; +use pretty_assertions::assert_eq; +use serial_test::serial; +use tempfile::TempDir; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; +use crate::auth::AuthKeyringBackendKind; +use crate::auth::AuthManager; +use crate::auth::AuthManagerConfig; +use crate::auth::CODEX_ACCESS_TOKEN_ENV_VAR; +use crate::auth::CODEX_API_KEY_ENV_VAR; +use crate::auth::login_with_api_key; +use crate::outbound_proxy::AuthRouteConfig; + +const TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_WIF_TOKEN_URL_OVERRIDE"; +const WORKSPACE_ID: &str = "workspace-one"; + +struct TestConfig { + codex_home: PathBuf, + workload_identity: Option, +} + +impl AuthManagerConfig for TestConfig { + fn codex_home(&self) -> PathBuf { + self.codex_home.clone() + } + + fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode { + AuthCredentialsStoreMode::File + } + + fn auth_keyring_backend_kind(&self) -> AuthKeyringBackendKind { + AuthKeyringBackendKind::default() + } + + fn forced_chatgpt_workspace_id(&self) -> Option> { + Some(vec![WORKSPACE_ID.to_string()]) + } + + fn chatgpt_base_url(&self) -> String { + "https://chatgpt.com/backend-api/".to_string() + } + + fn workload_identity_config(&self) -> Option { + self.workload_identity.clone() + } + + fn auth_route_config(&self) -> Option { + None + } +} + +struct EnvVarGuard { + name: &'static str, + previous: Option, +} + +impl EnvVarGuard { + fn set(name: &'static str, value: &str) -> Self { + let guard = Self::remove(name); + // SAFETY: these tests are serialized with every auth test that mutates + // the process environment. + unsafe { std::env::set_var(name, value) }; + guard + } + + fn remove(name: &'static str) -> Self { + let previous = std::env::var_os(name); + // SAFETY: these tests are serialized with every auth test that mutates + // the process environment. + unsafe { std::env::remove_var(name) }; + Self { name, previous } + } +} + +impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: these tests are serialized with every auth test that mutates + // the process environment. + unsafe { + match &self.previous { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } +} + +fn test_config(codex_home: &TempDir, token_file: Option) -> TestConfig { + TestConfig { + codex_home: codex_home.path().to_path_buf(), + workload_identity: Some(WorkloadIdentityToml { + federation_rule_id: Some("rule-one".to_string()), + principal_id: Some("user-one".to_string()), + tenant_id: Some("tenant-one".to_string()), + workspace_id: Some(WORKSPACE_ID.to_string()), + identity_token_file: token_file.map(|path| { + codex_config::AbsolutePathBuf::from_absolute_path(path) + .expect("absolute token file") + }), + }), + } +} + +fn fake_access_token(marker: &str) -> String { + let payload = serde_json::json!({ + "marker": marker, + "https://api.openai.com/auth": { + "chatgpt_user_id": "user-one" + } + }); + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).expect("serialize claims")); + format!("e30.{payload}.c2ln") +} + +fn success(access_token: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": 600, + "chatgpt_account_id": WORKSPACE_ID, + "chatgpt_plan_type": "enterprise", + "user_id": "user-one" + })) +} + +fn remove_explicit_auth() -> [EnvVarGuard; 2] { + [ + EnvVarGuard::remove(CODEX_API_KEY_ENV_VAR), + EnvVarGuard::remove(CODEX_ACCESS_TOKEN_ENV_VAR), + ] +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn manager_uses_wif_over_persisted_auth_and_rereads_the_token_file() { + let codex_home = TempDir::new().expect("tempdir"); + login_with_api_key( + codex_home.path(), + "sk-persisted", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("seed persisted auth"); + let auth_path = codex_home.path().join("auth.json"); + let persisted_before = std::fs::read(&auth_path).expect("read persisted auth"); + let token_file = codex_home.path().join("identity-token"); + tokio::fs::write(&token_file, "assertion-one\n") + .await + .expect("write assertion"); + + let server = MockServer::start().await; + let calls = Arc::new(AtomicUsize::new(0)); + let access_tokens = [fake_access_token("one"), fake_access_token("two")]; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with({ + let calls = Arc::clone(&calls); + move |_request: &wiremock::Request| { + let index = calls.fetch_add(1, Ordering::SeqCst).min(1); + success(&access_tokens[index]) + } + }) + .mount(&server) + .await; + let _explicit_auth = remove_explicit_auth(); + let _endpoint = EnvVarGuard::set( + TOKEN_URL_OVERRIDE_ENV_VAR, + &format!("{}/oauth/token", server.uri()), + ); + let _conflicting_environment = [ + EnvVarGuard::set(FEDERATION_RULE_ID_ENV_VAR, "wrong-rule"), + EnvVarGuard::set(PRINCIPAL_ID_ENV_VAR, "wrong-principal"), + EnvVarGuard::set(TENANT_ID_ENV_VAR, "wrong-tenant"), + EnvVarGuard::set(WORKSPACE_ID_ENV_VAR, "wrong-workspace"), + EnvVarGuard::set(IDENTITY_TOKEN_ENV_VAR, "wrong-assertion"), + ]; + let config = test_config(&codex_home, Some(token_file.clone())); + + let manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + let initial = manager.auth_cached().expect("workload auth"); + assert_eq!(initial.api_auth_mode(), AuthMode::ChatgptAuthTokens); + assert_eq!( + initial.get_token().expect("access token"), + fake_access_token("one") + ); + assert_eq!(initial.get_account_id().as_deref(), Some(WORKSPACE_ID)); + assert_eq!( + std::fs::read(&auth_path).expect("read persisted auth"), + persisted_before + ); + + tokio::fs::write(&token_file, "assertion-two\n") + .await + .expect("rotate assertion"); + manager + .refresh_token_from_authority() + .await + .expect("refresh workload auth"); + assert_eq!( + manager + .auth_cached() + .expect("refreshed auth") + .get_token() + .expect("access token"), + fake_access_token("two") + ); + + let requests = server.received_requests().await.expect("received requests"); + let forms = requests + .iter() + .map(|request| { + url::form_urlencoded::parse(&request.body) + .into_owned() + .collect::>() + }) + .collect::>(); + assert_eq!(forms[0]["assertion"], "assertion-one"); + assert_eq!(forms[1]["assertion"], "assertion-two"); + assert_eq!(forms[0]["federation_rule_id"], "rule-one"); + assert_eq!(forms[0]["principal_id"], "user-one"); + assert_eq!(forms[0]["tenant_id"], "tenant-one"); + assert_eq!(forms[0]["workspace_id"], WORKSPACE_ID); + + let second_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + assert!(second_manager.has_external_auth()); + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn manager_accepts_environment_only_configuration() { + let codex_home = TempDir::new().expect("tempdir"); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(success(&fake_access_token("environment"))) + .expect(1) + .mount(&server) + .await; + + let _explicit_auth = remove_explicit_auth(); + let _endpoint = EnvVarGuard::set( + TOKEN_URL_OVERRIDE_ENV_VAR, + &format!("{}/oauth/token", server.uri()), + ); + let _workload_environment = [ + EnvVarGuard::set(FEDERATION_RULE_ID_ENV_VAR, "rule-environment"), + EnvVarGuard::set(PRINCIPAL_ID_ENV_VAR, "user-one"), + EnvVarGuard::set(TENANT_ID_ENV_VAR, "tenant-environment"), + EnvVarGuard::set(WORKSPACE_ID_ENV_VAR, WORKSPACE_ID), + EnvVarGuard::set(IDENTITY_TOKEN_ENV_VAR, "assertion-environment"), + EnvVarGuard::remove(IDENTITY_TOKEN_FILE_ENV_VAR), + ]; + let config = TestConfig { + codex_home: codex_home.path().to_path_buf(), + workload_identity: None, + }; + + let manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + + assert_eq!( + manager + .auth_cached() + .expect("workload auth") + .get_token() + .expect("access token"), + fake_access_token("environment") + ); + let request = server + .received_requests() + .await + .expect("received requests") + .pop() + .expect("token exchange"); + let form = url::form_urlencoded::parse(&request.body) + .into_owned() + .collect::>(); + assert_eq!(form["assertion"], "assertion-environment"); + assert_eq!(form["federation_rule_id"], "rule-environment"); + assert_eq!(form["principal_id"], "user-one"); + assert_eq!(form["tenant_id"], "tenant-environment"); + assert_eq!(form["workspace_id"], WORKSPACE_ID); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn explicit_api_key_precedes_workload_identity() { + let codex_home = TempDir::new().expect("tempdir"); + let token_file = codex_home.path().join("identity-token"); + std::fs::write(&token_file, "assertion-one").expect("write assertion"); + let _access_token = EnvVarGuard::remove(CODEX_ACCESS_TOKEN_ENV_VAR); + let _api_key = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-explicit"); + let config = test_config(&codex_home, Some(token_file)); + + let manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + + assert_eq!( + manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-explicit".to_string()) + ); + assert!(!manager.has_external_auth()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalid_workload_identity_fails_closed_instead_of_using_persisted_auth() { + let codex_home = TempDir::new().expect("tempdir"); + login_with_api_key( + codex_home.path(), + "sk-persisted", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("seed persisted auth"); + let _explicit_auth = remove_explicit_auth(); + let _identity_token = EnvVarGuard::remove(IDENTITY_TOKEN_ENV_VAR); + let _identity_token_file = EnvVarGuard::remove(IDENTITY_TOKEN_FILE_ENV_VAR); + let config = test_config(&codex_home, /*token_file*/ None); + + let manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + + assert!(manager.has_external_auth()); + assert!(manager.auth_cached().is_none()); + assert!(manager.auth().await.is_none()); +} diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index c35daaa497..a78c3ce6af 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -264,6 +264,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R model_catalog: None, model_verbosity: None, chatgpt_base_url: "https://chatgpt.com/backend-api/".to_string(), + workload_identity: None, respect_system_proxy: false, apps_mcp_product_sku: None, realtime_audio: RealtimeAudioConfig::default(),