From 96c8be200cf1e96f75449b4751a5fab6756a16aa Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Wed, 12 Aug 2026 13:22:52 +0000 Subject: [PATCH] Integrate workload identity with Codex authentication (#38188) ## What changed - Select workload identity authentication from `OPENAI_FEDERATION_RULE_ID` and `OPENAI_IDENTITY_TOKEN_FILE`, while preserving explicit API key and access token precedence. - Adapt workload identity exchanges to the shared authentication manager and reuse a compatible process-scoped session for token resolution and refresh. - Reject incomplete, conflicting, or unsupported configurations and exchanged tokens that change the authenticated subject. - Classify exchange failures for authentication retry handling and invalidate rejected cached tokens without discarding a newer concurrent exchange. ## Testing - Cover configuration selection and validation, session reuse and conflicts, identity-preserving refreshes, concurrent refreshes, and retry classification. GitOrigin-RevId: 62afb514b4e0a59effb08f464ac929633462a71b --- codex-rs/Cargo.lock | 1 + codex-rs/Cargo.toml | 1 + codex-rs/login/Cargo.toml | 1 + codex-rs/login/src/auth/auth_tests.rs | 92 ++++ codex-rs/login/src/auth/manager.rs | 64 ++- codex-rs/login/src/auth/mod.rs | 2 + codex-rs/login/src/auth/workload_identity.rs | 466 ++++++++++++++++++ .../login/src/auth/workload_identity_tests.rs | 404 +++++++++++++++ codex-rs/login/src/lib.rs | 1 + codex-rs/workload-identity/src/exchange.rs | 22 +- codex-rs/workload-identity/src/lib.rs | 24 + 11 files changed, 1052 insertions(+), 26 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 50d942875b..fdaeaf51a1 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3475,6 +3475,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/Cargo.toml b/codex-rs/Cargo.toml index 34eb438d9e..c5bbcb5920 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -275,6 +275,7 @@ codex-utils-stream-parser = { path = "utils/stream-parser" } codex-utils-string = { path = "utils/string" } codex-utils-template = { path = "utils/template" } codex-v8-poc = { path = "v8-poc" } +codex-workload-identity = { path = "workload-identity" } codex-windows-sandbox = { path = "windows-sandbox-rs" } core_test_support = { path = "core/tests/common" } mcp_test_support = { path = "mcp-server/tests/common" } diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index 424e7383b3..2971dec6c1 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/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 9aa566f90e..eb9c266acb 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -11,6 +11,8 @@ use codex_protocol::protocol::SessionSource; use base64::Engine; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; +use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR; +use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR; use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; @@ -1702,6 +1704,96 @@ fn remove_access_token_env_var() -> EnvVarGuard { EnvVarGuard::remove(CODEX_ACCESS_TOKEN_ENV_VAR) } +struct TestAuthManagerConfig(AuthConfig); + +impl AuthManagerConfig for TestAuthManagerConfig { + fn codex_home(&self) -> PathBuf { + self.0.codex_home.clone() + } + + fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode { + self.0.auth_credentials_store_mode + } + + fn auth_keyring_backend_kind(&self) -> AuthKeyringBackendKind { + self.0.keyring_backend_kind + } + + fn forced_login_method(&self) -> Option { + self.0.forced_login_method + } + + fn forced_chatgpt_workspace_id(&self) -> Option> { + self.0.forced_chatgpt_workspace_id.clone() + } + + fn managed_auth_policy(&self) -> ManagedAuthPolicy { + self.0.managed_auth_policy.clone() + } + + fn chatgpt_base_url(&self) -> String { + self.0 + .chatgpt_base_url + .clone() + .expect("test config should include a ChatGPT base URL") + } + + fn auth_route_config(&self) -> AuthRouteConfig { + self.0.auth_route_config.clone() + } +} + +fn test_auth_manager_config(codex_home: &Path) -> TestAuthManagerConfig { + TestAuthManagerConfig(AuthConfig { + codex_home: codex_home.to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + keyring_backend_kind: AuthKeyringBackendKind::Direct, + forced_login_method: Some(ForcedLoginMethod::Chatgpt), + chatgpt_base_url: Some("https://chatgpt-staging.com/backend-api".to_string()), + forced_chatgpt_workspace_id: Some(vec!["forced-workspace".to_string()]), + managed_auth_policy: ManagedAuthPolicy { + allowed_login_methods: Some(vec![ForcedLoginMethod::Chatgpt]), + allowed_chatgpt_workspaces: Some(vec![ + "forced-workspace".to_string(), + "managed-workspace".to_string(), + ]), + }, + auth_route_config: AuthRouteConfig::from_http_client_factory(HttpClientFactory::new( + OutboundProxyPolicy::RespectSystemProxy, + )), + }) +} + +#[test] +fn auth_config_from_preserves_all_fields() { + let codex_home = tempdir().expect("tempdir"); + let config = test_auth_manager_config(codex_home.path()); + + assert_eq!(auth_config_from(&config), config.0); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn try_shared_from_config_rejects_partial_workload_identity_configuration() { + let codex_home = tempdir().expect("tempdir"); + let config = test_auth_manager_config(codex_home.path()); + let _access_token_guard = remove_access_token_env_var(); + let _rule_guard = EnvVarGuard::set(OPENAI_FEDERATION_RULE_ID_ENV_VAR, "rule-one"); + let _assertion_file_guard = EnvVarGuard::remove(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR); + + let error = + AuthManager::try_shared_from_config(&config, /*enable_codex_api_key_env*/ false) + .await + .expect_err("partial workload identity config should fail closed"); + + assert!(matches!(error, RefreshTokenError::Permanent(_))); + assert!( + error + .to_string() + .contains(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR) + ); +} + #[tokio::test] #[serial(codex_auth_env)] async fn load_auth_reads_access_token_from_env() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index a7c022656b..96d3a39439 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -38,6 +38,8 @@ use super::agent_identity::require_agent_identity_authapi_base_url; use super::agent_identity::verified_record_from_jwt; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; +use super::workload_identity::WorkloadIdentityExternalAuth; +use super::workload_identity::WorkloadIdentitySessionError; use crate::auth::AuthHeaders; pub use crate::auth::agent_identity::AgentIdentityAuth; pub use crate::auth::agent_identity::AgentIdentityAuthError; @@ -2672,20 +2674,7 @@ impl AuthManager { config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { - Self::shared_from_auth_config( - AuthConfig { - codex_home: config.codex_home(), - auth_credentials_store_mode: config.cli_auth_credentials_store_mode(), - keyring_backend_kind: config.auth_keyring_backend_kind(), - forced_login_method: config.forced_login_method(), - chatgpt_base_url: Some(config.chatgpt_base_url()), - forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id(), - managed_auth_policy: config.managed_auth_policy(), - auth_route_config: config.auth_route_config(), - }, - enable_codex_api_key_env, - ) - .await + Self::shared_from_auth_config(auth_config_from(config), enable_codex_api_key_env).await } /// Builds a shared manager using restrictions resolved before authentication. @@ -2696,6 +2685,33 @@ impl AuthManager { Arc::new(Self::new_from_auth_config(auth_config, enable_codex_api_key_env).await) } + /// Constructs a manager and activates process-configured workload identity when selected. + pub async fn try_shared_from_config( + config: &impl AuthManagerConfig, + enable_codex_api_key_env: bool, + ) -> Result, RefreshTokenError> { + Self::try_shared_from_auth_config(auth_config_from(config), enable_codex_api_key_env).await + } + + /// Activates workload identity against an auth config resolved before full runtime config. + pub async fn try_shared_from_auth_config( + auth_config: AuthConfig, + enable_codex_api_key_env: bool, + ) -> Result, RefreshTokenError> { + let external_auth = WorkloadIdentityExternalAuth::from_process_config( + &auth_config, + enable_codex_api_key_env, + ) + .map_err(configured_workload_identity_error)?; + let manager = Self::shared_from_auth_config(auth_config, enable_codex_api_key_env).await; + if let Some(external_auth) = external_auth { + manager + .set_configured_external_auth(Arc::new(external_auth)) + .await?; + } + Ok(manager) + } + pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { UnauthorizedRecovery::new(Arc::clone(self)) } @@ -3003,6 +3019,26 @@ impl AuthManager { } } +fn auth_config_from(config: &impl AuthManagerConfig) -> AuthConfig { + AuthConfig { + codex_home: config.codex_home(), + auth_credentials_store_mode: config.cli_auth_credentials_store_mode(), + keyring_backend_kind: config.auth_keyring_backend_kind(), + forced_login_method: config.forced_login_method(), + chatgpt_base_url: Some(config.chatgpt_base_url()), + forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id(), + managed_auth_policy: config.managed_auth_policy(), + auth_route_config: config.auth_route_config(), + } +} + +fn configured_workload_identity_error(error: WorkloadIdentitySessionError) -> RefreshTokenError { + RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + error.to_string(), + )) +} + #[cfg(test)] #[path = "auth_tests.rs"] mod tests; diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 16477d7827..8c76713b8f 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; @@ -18,3 +19,4 @@ pub use bedrock_api_key::login_with_bedrock_api_key; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; +pub use workload_identity::is_workload_identity_selected; 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..13cc1920ad --- /dev/null +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -0,0 +1,466 @@ +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::Weak; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_http_client::HttpClientFactory; +use codex_protocol::config_types::ForcedLoginMethod; +use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR; +use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR; +use codex_workload_identity::WorkloadIdentityConfig; +use codex_workload_identity::WorkloadIdentityError; +use codex_workload_identity::WorkloadIdentityExchange; +use codex_workload_identity::WorkloadIdentityToken; +use thiserror::Error; +use url::Url; + +use super::AuthConfig; +use super::CodexAuth; +use super::ExternalAuth; +use super::ExternalAuthFuture; +use super::ExternalAuthRefreshContext; +use super::RefreshTokenError; +use super::RefreshTokenFailedError; +use super::RefreshTokenFailedReason; +use super::read_codex_access_token_from_env; +use super::read_codex_api_key_from_env; +use crate::AuthRouteConfig; + +const PROD_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +const STAGING_TOKEN_URL: &str = "https://auth.api.openai.org/oauth/token"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum WorkloadIdentityEnvironment { + Production, + Staging, +} + +impl WorkloadIdentityEnvironment { + fn token_url(self) -> Result { + Url::parse(match self { + Self::Production => PROD_TOKEN_URL, + Self::Staging => STAGING_TOKEN_URL, + }) + .map_err(|_| invalid_config("workload identity token endpoint is invalid")) + } +} + +#[derive(Clone, Debug)] +struct WorkloadIdentitySessionConfig { + assertion_file: PathBuf, + environment: WorkloadIdentityEnvironment, + federation_rule_id: String, + http_client_factory: HttpClientFactory, + token_url: Url, +} + +impl WorkloadIdentitySessionConfig { + fn fingerprint(&self) -> WorkloadIdentityFingerprint { + WorkloadIdentityFingerprint { + assertion_file: self.assertion_file.clone(), + environment: self.environment, + federation_rule_id: self.federation_rule_id.trim().to_string(), + token_url: self.token_url.to_string(), + } + } + + fn into_exchange(self) -> Result { + WorkloadIdentityExchange::new( + WorkloadIdentityConfig::new(self.federation_rule_id, self.assertion_file)?, + self.token_url, + self.http_client_factory, + ) + } +} + +#[derive(Clone, PartialEq, Eq)] +struct WorkloadIdentityFingerprint { + assertion_file: PathBuf, + environment: WorkloadIdentityEnvironment, + federation_rule_id: String, + token_url: String, +} + +#[derive(Debug, Error)] +pub(super) enum WorkloadIdentitySessionError { + #[error(transparent)] + Exchange(#[from] WorkloadIdentityError), + #[error("a different workload identity configuration is already active in this process")] + ConflictingConfiguration, + #[error("the workload identity process-session registry is unavailable")] + RegistryUnavailable, + #[error("{0}")] + InvalidConfiguration(String), +} + +pub fn is_workload_identity_selected( + auth_config: &AuthConfig, + enable_codex_api_key_env: bool, +) -> bool { + !has_explicit_process_auth(auth_config, enable_codex_api_key_env) + && ProcessEnvironment::read().has_marker() +} + +fn resolve_config( + chatgpt_base_url: &str, + environment: ProcessEnvironment, + has_explicit_process_auth: bool, + chatgpt_login_allowed: bool, + auth_route_config: AuthRouteConfig, +) -> Result, WorkloadIdentitySessionError> { + if has_explicit_process_auth || !environment.has_marker() { + return Ok(None); + } + if !chatgpt_login_allowed { + return Err(invalid_config( + "workload identity requires a login policy that permits ChatGPT authentication", + )); + } + + let federation_rule_id = required_unicode( + environment.federation_rule_id, + OPENAI_FEDERATION_RULE_ID_ENV_VAR, + )?; + let assertion_file = required_path( + environment.identity_token_file, + OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR, + )?; + let auth_environment = classify_auth_environment(chatgpt_base_url)?; + + Ok(Some(WorkloadIdentitySessionConfig { + assertion_file, + environment: auth_environment, + federation_rule_id, + http_client_factory: auth_route_config.http_client_factory().clone(), + token_url: auth_environment.token_url()?, + })) +} + +fn has_explicit_process_auth(auth_config: &AuthConfig, enable_codex_api_key_env: bool) -> bool { + (enable_codex_api_key_env + && auth_config.is_login_method_allowed(ForcedLoginMethod::Api) + && read_codex_api_key_from_env().is_some()) + || (auth_config.is_login_method_allowed(ForcedLoginMethod::Chatgpt) + && read_codex_access_token_from_env().is_some()) +} + +fn required_path( + value: Option, + variable: &'static str, +) -> Result { + let path = PathBuf::from( + value.ok_or_else(|| invalid_config(format!("workload identity requires {variable}")))?, + ); + if !path.is_absolute() { + return Err(invalid_config(format!( + "{variable} must be an absolute path" + ))); + } + Ok(path) +} + +fn required_unicode( + value: Option, + variable: &'static str, +) -> Result { + let value = value + .ok_or_else(|| invalid_config(format!("workload identity requires {variable}")))? + .into_string() + .map_err(|_| invalid_config(format!("workload identity variable {variable} is invalid")))?; + let value = value.trim(); + if value.is_empty() { + return Err(invalid_config(format!( + "workload identity variable {variable} is invalid" + ))); + } + Ok(value.to_string()) +} + +fn classify_auth_environment( + base_url: &str, +) -> Result { + match base_url.trim().trim_end_matches('/') { + "https://chatgpt.com" + | "https://chatgpt.com/backend-api" + | "https://chatgpt.com/codex" + | "https://chatgpt.com/backend-api/codex" + | "https://chat.openai.com" + | "https://chat.openai.com/backend-api" + | "https://chat.openai.com/codex" + | "https://chat.openai.com/backend-api/codex" => { + Ok(WorkloadIdentityEnvironment::Production) + } + "https://chatgpt-staging.com" + | "https://chatgpt-staging.com/backend-api" + | "https://chatgpt-staging.com/codex" + | "https://chatgpt-staging.com/backend-api/codex" => { + Ok(WorkloadIdentityEnvironment::Staging) + } + _ => Err(invalid_config( + "workload identity auth supports only trusted production and staging app routing", + )), + } +} + +fn invalid_config(message: impl Into) -> WorkloadIdentitySessionError { + WorkloadIdentitySessionError::InvalidConfiguration(message.into()) +} + +#[derive(Default)] +struct ProcessEnvironment { + federation_rule_id: Option, + identity_token_file: Option, +} + +impl ProcessEnvironment { + fn read() -> Self { + Self { + federation_rule_id: std::env::var_os(OPENAI_FEDERATION_RULE_ID_ENV_VAR), + identity_token_file: std::env::var_os(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR), + } + } + + fn has_marker(&self) -> bool { + self.federation_rule_id.is_some() || self.identity_token_file.is_some() + } +} + +#[derive(Clone, PartialEq, Eq)] +struct WorkloadIdentitySubject { + account_id: String, + account_user_id: String, + user_id: String, +} + +impl From<&WorkloadIdentityToken> for WorkloadIdentitySubject { + fn from(token: &WorkloadIdentityToken) -> Self { + Self { + account_id: token.chatgpt_account_id.clone(), + account_user_id: token.chatgpt_account_user_id.clone(), + user_id: token.user_id.clone(), + } + } +} + +struct WorkloadIdentitySession { + exchange: WorkloadIdentityExchange, + subject: Mutex>, +} + +impl WorkloadIdentitySession { + fn new(config: WorkloadIdentitySessionConfig) -> Result { + Ok(Self { + exchange: config.into_exchange()?, + subject: Mutex::new(None), + }) + } + + fn accept_subject( + &self, + token: &WorkloadIdentityToken, + previous_account_id: Option<&str>, + ) -> Result<(), WorkloadIdentityError> { + if previous_account_id.is_some_and(|account_id| account_id != token.chatgpt_account_id) { + return Err(WorkloadIdentityError::InvalidExchangeResponse); + } + let subject = WorkloadIdentitySubject::from(token); + let mut current = self + .subject + .lock() + .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?; + match current.as_ref() { + Some(current) if current != &subject => { + Err(WorkloadIdentityError::InvalidExchangeResponse) + } + Some(_) => Ok(()), + None => { + *current = Some(subject); + Ok(()) + } + } + } +} + +#[derive(Default)] +struct WorkloadIdentitySessionRegistry { + entry: Mutex>, +} + +struct WorkloadIdentitySessionEntry { + fingerprint: WorkloadIdentityFingerprint, + session: Weak, +} + +impl WorkloadIdentitySessionRegistry { + fn has_active_session(&self) -> Result { + let entry = self + .entry + .lock() + .map_err(|_| WorkloadIdentitySessionError::RegistryUnavailable)?; + Ok(entry + .as_ref() + .and_then(|active| active.session.upgrade()) + .is_some()) + } + + fn session( + &self, + config: WorkloadIdentitySessionConfig, + ) -> Result, WorkloadIdentitySessionError> { + let fingerprint = config.fingerprint(); + let mut entry = self + .entry + .lock() + .map_err(|_| WorkloadIdentitySessionError::RegistryUnavailable)?; + if let Some(active) = entry.as_ref() + && let Some(session) = active.session.upgrade() + { + if active.fingerprint == fingerprint { + return Ok(session); + } + return Err(WorkloadIdentitySessionError::ConflictingConfiguration); + } + + let session = Arc::new(WorkloadIdentitySession::new(config)?); + *entry = Some(WorkloadIdentitySessionEntry { + fingerprint, + session: Arc::downgrade(&session), + }); + Ok(session) + } +} + +fn process_registry() -> &'static WorkloadIdentitySessionRegistry { + static REGISTRY: OnceLock = OnceLock::new(); + REGISTRY.get_or_init(WorkloadIdentitySessionRegistry::default) +} + +pub(super) struct WorkloadIdentityExternalAuth { + observed_token_version: AtomicU64, + session: Arc, +} + +impl WorkloadIdentityExternalAuth { + pub(super) fn from_process_config( + auth_config: &AuthConfig, + enable_codex_api_key_env: bool, + ) -> Result, WorkloadIdentitySessionError> { + let registry = process_registry(); + let active_session = registry.has_active_session()?; + resolve_config( + auth_config + .chatgpt_base_url + .as_deref() + .unwrap_or("https://chatgpt.com/backend-api"), + ProcessEnvironment::read(), + has_explicit_process_auth(auth_config, enable_codex_api_key_env) && !active_session, + auth_config.is_login_method_allowed(ForcedLoginMethod::Chatgpt), + auth_config.auth_route_config.clone(), + )? + .map(|config| Self::from_config_with_registry(config, registry)) + .transpose() + } + + fn from_config_with_registry( + config: WorkloadIdentitySessionConfig, + registry: &WorkloadIdentitySessionRegistry, + ) -> Result { + Ok(Self { + observed_token_version: AtomicU64::new(0), + session: registry.session(config)?, + }) + } + + async fn build_validated_auth( + &self, + token: WorkloadIdentityToken, + previous_account_id: Option<&str>, + ) -> std::io::Result { + let token_version = token.version(); + let result = self.validate_auth(&token, previous_account_id); + if result.is_err() { + self.session + .exchange + .invalidate_if_current(token_version) + .await; + } + let auth = result?; + self.observed_token_version + .store(token_version, Ordering::Release); + Ok(auth) + } + + fn validate_auth( + &self, + token: &WorkloadIdentityToken, + previous_account_id: Option<&str>, + ) -> std::io::Result { + let auth = CodexAuth::from_external_chatgpt_tokens( + &token.access_token, + &token.chatgpt_account_id, + token.chatgpt_plan_type.as_deref(), + ) + .map_err(|_| std::io::Error::other(WorkloadIdentityError::InvalidExchangeResponse))?; + if auth.get_chatgpt_user_id().as_deref() != Some(token.user_id.as_str()) { + return Err(std::io::Error::other( + WorkloadIdentityError::InvalidExchangeResponse, + )); + } + self.session + .accept_subject(token, previous_account_id) + .map_err(std::io::Error::other)?; + Ok(auth) + } +} + +impl ExternalAuth for WorkloadIdentityExternalAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async move { + let token = self + .session + .exchange + .resolve() + .await + .map_err(std::io::Error::other)?; + self.build_validated_auth(token, /*previous_account_id*/ None) + .await + }) + } + + fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async move { + let observed_version = self.observed_token_version.load(Ordering::Acquire); + let token = self + .session + .exchange + .refresh(observed_version) + .await + .map_err(std::io::Error::other)?; + self.build_validated_auth(token, context.previous_account_id.as_deref()) + .await + }) + } + + fn classify_error(&self, error: std::io::Error) -> RefreshTokenError { + if error + .get_ref() + .and_then(|source| source.downcast_ref::()) + .is_some_and(WorkloadIdentityError::is_transient) + { + return RefreshTokenError::Transient(error); + } + + RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + error.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..24d0040576 --- /dev/null +++ b/codex-rs/login/src/auth/workload_identity_tests.rs @@ -0,0 +1,404 @@ +use std::path::Path; +use std::sync::atomic::AtomicUsize; +use std::time::Duration; + +use base64::Engine as _; +use codex_http_client::OutboundProxyPolicy; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; + +use super::*; +use crate::auth::ExternalAuthRefreshReason; + +fn auth_route_config(policy: OutboundProxyPolicy) -> AuthRouteConfig { + AuthRouteConfig::from_http_client_factory(HttpClientFactory::new(policy)) +} + +fn complete_environment() -> ProcessEnvironment { + ProcessEnvironment { + federation_rule_id: Some("rule-one".into()), + identity_token_file: Some(std::env::temp_dir().join("identity-token").into_os_string()), + } +} + +fn resolve_for_test( + environment: ProcessEnvironment, + has_explicit_process_auth: bool, + chatgpt_login_allowed: bool, + chatgpt_base_url: &str, +) -> Result, WorkloadIdentitySessionError> { + resolve_config( + chatgpt_base_url, + environment, + has_explicit_process_auth, + chatgpt_login_allowed, + auth_route_config(OutboundProxyPolicy::ReqwestDefault), + ) +} + +#[test] +fn markers_select_wif_and_partial_configuration_fails_closed() { + assert!( + resolve_for_test( + ProcessEnvironment::default(), + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ true, + "https://chatgpt.com/backend-api", + ) + .expect("no markers") + .is_none() + ); + assert!( + resolve_for_test( + ProcessEnvironment { + federation_rule_id: Some("partial".into()), + ..Default::default() + }, + /*has_explicit_process_auth*/ true, + /*chatgpt_login_allowed*/ true, + "https://chatgpt.com/backend-api", + ) + .expect("explicit auth wins") + .is_none() + ); + + for (environment, missing) in [ + ( + ProcessEnvironment { + federation_rule_id: None, + ..complete_environment() + }, + OPENAI_FEDERATION_RULE_ID_ENV_VAR, + ), + ( + ProcessEnvironment { + identity_token_file: None, + ..complete_environment() + }, + OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR, + ), + ] { + let error = resolve_for_test( + environment, + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ true, + "https://chatgpt.com/backend-api", + ) + .expect_err("partial WIF must not fall back"); + assert!(error.to_string().contains(missing), "{error}"); + } + + let relative = ProcessEnvironment { + identity_token_file: Some("relative.jwt".into()), + ..complete_environment() + }; + assert!( + resolve_for_test( + relative, + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ true, + "https://chatgpt.com/backend-api", + ) + .expect_err("relative assertion path") + .to_string() + .contains("absolute path") + ); +} + +#[test] +fn auth_policy_and_app_environment_are_enforced() { + let policy_error = resolve_for_test( + complete_environment(), + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ false, + "https://chatgpt.com/backend-api", + ) + .expect_err("ChatGPT-disallowing policy"); + assert!(policy_error.to_string().contains("login policy")); + + for (chatgpt_base_url, expected_environment, expected_token_url) in [ + ( + "https://chatgpt.com/backend-api/", + WorkloadIdentityEnvironment::Production, + PROD_TOKEN_URL, + ), + ( + "https://chatgpt-staging.com/backend-api", + WorkloadIdentityEnvironment::Staging, + STAGING_TOKEN_URL, + ), + ] { + let config = resolve_for_test( + complete_environment(), + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ true, + chatgpt_base_url, + ) + .expect("trusted app routing") + .expect("WIF selected"); + assert_eq!(config.environment, expected_environment); + assert_eq!(config.token_url.as_str(), expected_token_url); + } + + let error = resolve_for_test( + complete_environment(), + /*has_explicit_process_auth*/ false, + /*chatgpt_login_allowed*/ true, + "https://example.invalid/backend-api", + ) + .expect_err("untrusted auth environment"); + assert!(error.to_string().contains("app routing")); +} + +fn session_config(directory: &Path, server: &MockServer) -> WorkloadIdentitySessionConfig { + let assertion_file = directory.join("identity-token"); + std::fs::write(&assertion_file, "assertion-one").expect("write assertion"); + WorkloadIdentitySessionConfig { + assertion_file, + environment: WorkloadIdentityEnvironment::Staging, + federation_rule_id: "rule-one".to_string(), + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + token_url: Url::parse(&format!("{}/oauth/token", server.uri())).expect("token URL"), + } +} + +fn jwt(label: &str, user_id: &str) -> String { + let encode = |value: &serde_json::Value| { + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(value).expect("serialize JWT part")) + }; + let header = encode(&serde_json::json!({"alg": "none", "typ": "JWT"})); + let payload = encode(&serde_json::json!({ + "jti": label, + "https://api.openai.com/auth": { + "chatgpt_account_id": "account-one", + "chatgpt_plan_type": "enterprise", + "chatgpt_user_id": user_id, + "user_id": user_id + } + })); + format!("{header}.{payload}.sig") +} + +fn success_response(label: &str, account_user_id: &str, user_id: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": jwt(label, user_id), + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "expires_in": 600, + "scope": "model.request", + "chatgpt_account_id": "account-one", + "chatgpt_account_user_id": account_user_id, + "chatgpt_plan_type": "enterprise", + "user_id": user_id + })) +} + +#[tokio::test] +async fn compatible_adapters_share_exchange() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(success_response( + "access-one", + "account-user-one", + "user-one", + )) + .expect(1) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let first_config = session_config(temp_dir.path(), &server); + let second_config = first_config.clone(); + let first = WorkloadIdentityExternalAuth::from_config_with_registry(first_config, ®istry) + .expect("first adapter"); + let second = WorkloadIdentityExternalAuth::from_config_with_registry(second_config, ®istry) + .expect("second adapter"); + + assert!(Arc::ptr_eq(&first.session, &second.session)); + assert_eq!( + first + .resolve() + .await + .expect("first auth") + .get_token() + .expect("first token"), + second + .resolve() + .await + .expect("second auth") + .get_token() + .expect("second token") + ); +} + +#[tokio::test] +async fn incompatible_process_session_settings_are_rejected() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + let registry = WorkloadIdentitySessionRegistry::default(); + let base = session_config(temp_dir.path(), &server); + let _active = WorkloadIdentityExternalAuth::from_config_with_registry(base.clone(), ®istry) + .expect("active adapter"); + + let mut different_rule = base.clone(); + different_rule.federation_rule_id = "rule-two".to_string(); + let mut different_file = base.clone(); + different_file.assertion_file = temp_dir.path().join("identity-token-two"); + std::fs::write(&different_file.assertion_file, "assertion-two").expect("write assertion"); + let mut different_environment = base.clone(); + different_environment.environment = WorkloadIdentityEnvironment::Production; + let mut different_route = base; + different_route.http_client_factory = + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy); + + let different_route_adapter = + WorkloadIdentityExternalAuth::from_config_with_registry(different_route, ®istry) + .expect("route changes reuse the process-owned session"); + assert!(Arc::ptr_eq( + &_active.session, + &different_route_adapter.session + )); + + for config in [different_rule, different_file, different_environment] { + assert!(matches!( + WorkloadIdentityExternalAuth::from_config_with_registry(config, ®istry), + Err(WorkloadIdentitySessionError::ConflictingConfiguration) + )); + } +} + +#[tokio::test] +async fn refresh_preserves_identity_and_invalid_tokens_are_reexchanged() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + let request_count = Arc::new(AtomicUsize::new(0)); + let response_count = Arc::clone(&request_count); + Mock::given(method("POST")) + .respond_with(move |_request: &wiremock::Request| { + match response_count.fetch_add(1, Ordering::SeqCst) { + 0 => success_response("access-one", "account-user-one", "user-one"), + 1 => success_response("access-two", "account-user-two", "user-one"), + _ => success_response("access-three", "account-user-one", "user-one"), + } + }) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("adapter"); + adapter.resolve().await.expect("initial auth"); + + let error = adapter + .refresh(ExternalAuthRefreshContext { + reason: ExternalAuthRefreshReason::Unauthorized, + previous_account_id: Some("account-one".to_string()), + }) + .await + .expect_err("identity change must be rejected"); + assert!(matches!( + adapter.classify_error(error), + RefreshTokenError::Permanent(_) + )); + assert_eq!( + adapter + .resolve() + .await + .expect("corrected token is re-exchanged") + .get_token() + .expect("token"), + jwt("access-three", "user-one") + ); + assert_eq!(request_count.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn concurrent_refreshes_share_one_exchange() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + let request_count = Arc::new(AtomicUsize::new(0)); + let response_count = Arc::clone(&request_count); + Mock::given(method("POST")) + .respond_with(move |_request: &wiremock::Request| { + if response_count.fetch_add(1, Ordering::SeqCst) == 0 { + success_response("access-one", "account-user-one", "user-one") + } else { + success_response("access-two", "account-user-one", "user-one") + .set_delay(Duration::from_millis(30)) + } + }) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let first = Arc::new( + WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("first adapter"), + ); + let second = Arc::new( + WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("second adapter"), + ); + first.resolve().await.expect("first resolve"); + second.resolve().await.expect("second resolve"); + let context = ExternalAuthRefreshContext { + reason: ExternalAuthRefreshReason::Unauthorized, + previous_account_id: Some("account-one".to_string()), + }; + + let (first_refresh, second_refresh) = + tokio::join!(first.refresh(context.clone()), second.refresh(context)); + assert_eq!( + first_refresh + .expect("first refresh") + .get_token() + .expect("first token"), + second_refresh + .expect("second refresh") + .get_token() + .expect("second token") + ); + assert_eq!(request_count.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn exchange_errors_map_to_retry_policy() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("adapter"); + + let cases = [ + (WorkloadIdentityError::ExchangeRejected(400), false), + (WorkloadIdentityError::ExchangeRejected(408), true), + ( + WorkloadIdentityError::AssertionFile { + path: temp_dir.path().join("missing"), + source: Arc::new(std::io::Error::from(std::io::ErrorKind::NotFound)), + }, + true, + ), + ]; + for (error, transient) in cases { + let classified = adapter.classify_error(std::io::Error::other(error)); + assert_eq!( + matches!(classified, RefreshTokenError::Transient(_)), + transient + ); + } +} diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 1576afe548..91573159d5 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -49,6 +49,7 @@ pub use auth::RefreshTokenError; pub use auth::UnauthorizedRecovery; pub use auth::default_client; pub use auth::enforce_login_restrictions; +pub use auth::is_workload_identity_selected; pub use auth::load_auth_dot_json; pub use auth::login_with_access_token; pub use auth::login_with_api_key; diff --git a/codex-rs/workload-identity/src/exchange.rs b/codex-rs/workload-identity/src/exchange.rs index 7bbd73014c..0427eed221 100644 --- a/codex-rs/workload-identity/src/exchange.rs +++ b/codex-rs/workload-identity/src/exchange.rs @@ -100,7 +100,7 @@ impl WorkloadIdentityExchange { let valid_from = Instant::now(); let result = match self.exchange_uncached().await { Ok(token) => state.store(token, valid_from, Instant::now()), - Err(error) if error.allows_cached_fallback() => { + Err(error) if error.is_transient() => { let now = Instant::now(); match state.cached.as_mut() { Some(cached) if cached.expires_at > now => { @@ -156,6 +156,15 @@ impl WorkloadIdentityExchange { result } + /// Drops the cached token only when it is still the token the caller rejected. + pub async fn invalidate_if_current(&self, observed_token_version: u64) { + let mut state = self.state.lock().await; + if state.token_generation == observed_token_version { + state.cached = None; + state.last_attempt_error = None; + } + } + fn complete_attempt(&self, state: &mut CacheState, error: Option<&WorkloadIdentityError>) { state.last_attempt_error = error.cloned(); self.completed_attempts.fetch_add(1, Ordering::Release); @@ -207,17 +216,6 @@ impl WorkloadIdentityExchange { } } -impl WorkloadIdentityError { - fn allows_cached_fallback(&self) -> bool { - matches!( - self, - Self::AssertionFile { .. } - | Self::ExchangeUnavailable - | Self::ExchangeRejected(408 | 429 | 500..=599) - ) - } -} - fn validate_token_url(url: &Url) -> Result { if !url.username().is_empty() || url.password().is_some() diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs index 570ffec28c..a75c9decdb 100644 --- a/codex-rs/workload-identity/src/lib.rs +++ b/codex-rs/workload-identity/src/lib.rs @@ -61,3 +61,27 @@ pub enum WorkloadIdentityError { #[error("the workload identity token exchange returned an invalid response")] InvalidExchangeResponse, } + +impl WorkloadIdentityError { + /// Whether retrying the operation may succeed without changing configuration. + pub fn is_transient(&self) -> bool { + match self { + Self::AssertionFile { source, .. } => matches!( + source.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::NotFound + | std::io::ErrorKind::TimedOut + | std::io::ErrorKind::WouldBlock + ), + Self::ExchangeUnavailable | Self::ExchangeRejected(408 | 429 | 500..=599) => true, + Self::InvalidFederationRuleId + | Self::AssertionFileMustBeAbsolute + | Self::InvalidAssertion + | Self::AssertionTooLarge + | Self::HttpClientConfiguration + | Self::InvalidTokenUrl + | Self::ExchangeRejected(_) + | Self::InvalidExchangeResponse => false, + } + } +}