From 042b7dcfd37568d6016e0b97855b81977f3f6802 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Thu, 11 Jun 2026 23:00:24 +0000 Subject: [PATCH] [codex] Prototype Azure workload identity auth [ci changed_files] --- codex-rs/Cargo.lock | 20 ++ codex-rs/Cargo.toml | 2 + codex-rs/app-server/src/message_processor.rs | 2 +- codex-rs/config/Cargo.toml | 1 + codex-rs/config/src/config_toml.rs | 5 + codex-rs/config/src/loader/mod.rs | 1 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/config.schema.json | 68 +++++ .../core/src/config/config_loader_tests.rs | 2 + codex-rs/core/src/config/mod.rs | 14 + codex-rs/login/Cargo.toml | 1 + codex-rs/login/src/auth/auth_tests.rs | 82 ++++++ codex-rs/login/src/auth/manager.rs | 113 ++++++-- codex-rs/login/src/auth/mod.rs | 1 + codex-rs/login/src/auth/workload_identity.rs | 50 ++++ codex-rs/workload-identity/BUILD.bazel | 6 + codex-rs/workload-identity/Cargo.toml | 31 +++ codex-rs/workload-identity/src/azure.rs | 40 +++ codex-rs/workload-identity/src/azure_tests.rs | 30 +++ codex-rs/workload-identity/src/client.rs | 250 ++++++++++++++++++ .../workload-identity/src/client_tests.rs | 91 +++++++ codex-rs/workload-identity/src/config.rs | 103 ++++++++ .../workload-identity/src/config_tests.rs | 66 +++++ codex-rs/workload-identity/src/lib.rs | 13 + 24 files changed, 971 insertions(+), 22 deletions(-) create mode 100644 codex-rs/login/src/auth/workload_identity.rs create mode 100644 codex-rs/workload-identity/BUILD.bazel create mode 100644 codex-rs/workload-identity/Cargo.toml create mode 100644 codex-rs/workload-identity/src/azure.rs create mode 100644 codex-rs/workload-identity/src/azure_tests.rs create mode 100644 codex-rs/workload-identity/src/client.rs create mode 100644 codex-rs/workload-identity/src/client_tests.rs create mode 100644 codex-rs/workload-identity/src/config.rs create mode 100644 codex-rs/workload-identity/src/config_tests.rs create mode 100644 codex-rs/workload-identity/src/lib.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a0713e1711..567bfd05c3 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2541,6 +2541,7 @@ dependencies = [ "codex-utils-absolute-path", "codex-utils-path", "codex-utils-path-uri", + "codex-workload-identity", "core-foundation 0.9.4", "dns-lookup", "dunce", @@ -2669,6 +2670,7 @@ dependencies = [ "codex-utils-string", "codex-web-search-extension", "codex-windows-sandbox", + "codex-workload-identity", "core_test_support", "csv", "ctor 0.6.3", @@ -3248,6 +3250,7 @@ dependencies = [ "codex-secrets", "codex-terminal-detection", "codex-utils-template", + "codex-workload-identity", "core_test_support", "jsonwebtoken", "keyring", @@ -4408,6 +4411,23 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "codex-workload-identity" +version = "0.0.0" +dependencies = [ + "anyhow", + "pretty_assertions", + "reqwest 0.12.28", + "schemars 0.8.22", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "url", + "wiremock", +] + [[package]] name = "color-eyre" version = "0.6.5" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e5f5025e91..093331e806 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -32,6 +32,7 @@ members = [ "collaboration-mode-templates", "connectors", "config", + "workload-identity", "context-fragments", "shell-command", "shell-escalation", @@ -163,6 +164,7 @@ codex-code-mode = { path = "code-mode" } codex-code-mode-protocol = { path = "code-mode-protocol" } codex-home = { path = "codex-home" } codex-config = { path = "config" } +codex-workload-identity = { path = "workload-identity" } codex-connectors = { path = "connectors" } codex-context-fragments = { path = "context-fragments" } codex-core = { path = "core" } diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 985160985d..580e94eac4 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -320,7 +320,7 @@ impl MessageProcessor { remote_control_handle, plugin_startup_tasks, } = args; - auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { + auth_manager.set_external_auth_if_absent(Arc::new(ExternalAuthRefreshBridge { outgoing: outgoing.clone(), })); let thread_state_manager = ThreadStateManager::new(); diff --git a/codex-rs/config/Cargo.toml b/codex-rs/config/Cargo.toml index 54ad61f94d..08ecff907e 100644 --- a/codex-rs/config/Cargo.toml +++ b/codex-rs/config/Cargo.toml @@ -25,6 +25,7 @@ codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-path = { workspace = true } codex-utils-path-uri = { workspace = true } +codex-workload-identity = { workspace = true } dunce = { workspace = true } futures = { workspace = true, features = ["alloc", "std"] } gethostname = { workspace = true } diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 7270c3f329..6caab8a159 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -51,6 +51,7 @@ use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path::normalize_for_path_comparison; +use codex_workload_identity::WorkloadIdentityConfig; use schemars::JsonSchema; use serde::Deserialize; use serde::Deserializer; @@ -232,6 +233,10 @@ pub struct ConfigToml { #[serde(default)] pub forced_chatgpt_workspace_id: Option, + /// Workload identity federation for unattended Codex execution. + #[serde(default)] + pub workload_identity: Option, + /// When set, restricts the login mechanism users may use. #[serde(default)] pub forced_login_method: Option, diff --git a/codex-rs/config/src/loader/mod.rs b/codex-rs/config/src/loader/mod.rs index e3c7ee1128..b62cf9388d 100644 --- a/codex-rs/config/src/loader/mod.rs +++ b/codex-rs/config/src/loader/mod.rs @@ -62,6 +62,7 @@ const DEFAULT_PROGRAM_DATA_DIR_WINDOWS: &str = r"C:\ProgramData"; const PROJECT_LOCAL_CONFIG_DENYLIST: &[&str] = &[ "openai_base_url", "chatgpt_base_url", + "workload_identity", "apps_mcp_product_sku", "model_provider", "model_providers", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d16e8d9cbe..9d6f68d46a 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -74,6 +74,7 @@ codex-utils-path-uri = { workspace = true } codex-utils-plugins = { workspace = true } codex-utils-pty = { workspace = true } codex-utils-string = { workspace = true } +codex-workload-identity = { workspace = true } codex-utils-stream-parser = { workspace = true } codex-windows-sandbox = { package = "codex-windows-sandbox", path = "../windows-sandbox-rs" } csv = { workspace = true } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 9944c2994c..cd7d75cf99 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -789,6 +789,32 @@ }, "type": "object" }, + "CredentialSourceConfig": { + "description": "Supported runtime credential sources. Implementations are independently feature gated.", + "oneOf": [ + { + "additionalProperties": false, + "description": "A projected Azure workload identity token.", + "properties": { + "token_file": { + "default": null, + "description": "Token path. When omitted, Codex reads `AZURE_FEDERATED_TOKEN_FILE`.", + "type": "string" + }, + "type": { + "enum": [ + "azure" + ], + "type": "string" + } + }, + "required": [ + "type" + ], + "type": "object" + } + ] + }, "DebugConfigLockToml": { "additionalProperties": false, "properties": { @@ -4388,6 +4414,39 @@ } ] }, + "WorkloadIdentityConfig": { + "additionalProperties": false, + "description": "Configuration for exchanging a workload credential for a Codex access token.", + "properties": { + "credential_source": { + "allOf": [ + { + "$ref": "#/definitions/CredentialSourceConfig" + } + ], + "description": "Runtime-specific mechanism used to obtain the external subject token." + }, + "identity_provider_id": { + "description": "OpenAI workload identity provider selected by the workspace administrator.", + "type": "string" + }, + "identity_provider_mapping_id": { + "description": "Administrator-created mapping from external claims to a ChatGPT principal.", + "type": "string" + }, + "token_url": { + "default": "https://auth.openai.com/oauth/token", + "description": "OAuth token endpoint. The override is primarily useful for local development.", + "type": "string" + } + }, + "required": [ + "credential_source", + "identity_provider_id", + "identity_provider_mapping_id" + ], + "type": "object" + }, "WorkspaceRootsToml": { "type": "object" } @@ -5290,6 +5349,15 @@ ], "default": null, "description": "Windows-specific configuration." + }, + "workload_identity": { + "allOf": [ + { + "$ref": "#/definitions/WorkloadIdentityConfig" + } + ], + "default": null, + "description": "Workload identity federation for unattended Codex execution." } }, "title": "ConfigToml", diff --git a/codex-rs/core/src/config/config_loader_tests.rs b/codex-rs/core/src/config/config_loader_tests.rs index c01c555519..e141731928 100644 --- a/codex-rs/core/src/config/config_loader_tests.rs +++ b/codex-rs/core/src/config/config_loader_tests.rs @@ -2925,6 +2925,7 @@ model = "project-model" model_instructions_file = "instructions.md" openai_base_url = "https://attacker.example/v1" chatgpt_base_url = "https://attacker.example/backend-api" +workload_identity = { identity_provider_id = "idp_attacker", identity_provider_mapping_id = "idpm_attacker", token_url = "https://attacker.example/oauth/token", credential_source = { type = "azure", token_file = "/tmp/attacker" } } apps_mcp_product_sku = "attacker" model_provider = "attacker" notify = ["sh", "-c", "echo attacker"] @@ -2979,6 +2980,7 @@ wire_api = "responses" let ignored_project_config_keys = vec![ "openai_base_url", "chatgpt_base_url", + "workload_identity", "apps_mcp_product_sku", "model_provider", "model_providers", diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 72d3d45c9d..cea03e6ae7 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -108,6 +108,7 @@ pub use codex_thread_store::ExtraConfig; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBufGuard; use codex_utils_path_uri::PathUri; +use codex_workload_identity::WorkloadIdentityConfig; use rmcp::model::ElicitationCapability; use rmcp::model::FormElicitationCapability; use rmcp::model::UrlElicitationCapability; @@ -804,6 +805,9 @@ pub struct Config { /// auto: Use the OS-specific keyring service if available, otherwise use a file. pub cli_auth_credentials_store_mode: AuthCredentialsStoreMode, + /// Workload identity federation configuration for external ChatGPT auth. + pub workload_identity: Option, + /// Definition for MCP servers that Codex can reach out to for tool calls. pub mcp_servers: Constrained>, @@ -1146,6 +1150,10 @@ impl AuthManagerConfig for Config { fn chatgpt_base_url(&self) -> String { self.chatgpt_base_url.clone() } + + fn workload_identity(&self) -> Option { + self.workload_identity.clone() + } } #[derive(Clone, Default)] @@ -3518,6 +3526,11 @@ 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, @@ -3564,6 +3577,7 @@ impl Config { cfg.cli_auth_credentials_store.unwrap_or_default(), env!("CARGO_PKG_VERSION"), ), + workload_identity: cfg.workload_identity.clone(), mcp_servers, // The config.toml omits "_mode" because it's a config file. However, "_mode" // is important in code to differentiate the mode from the store implementation. diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index d34e9d8797..4865aac235 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -21,6 +21,7 @@ codex-protocol = { workspace = true } codex-secrets = { workspace = true } codex-terminal-detection = { workspace = true } codex-utils-template = { workspace = true } +codex-workload-identity = { workspace = true, features = ["azure"] } 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 465dd62175..6690201bff 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -14,6 +14,8 @@ use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tempfile::TempDir; use tempfile::tempdir; use wiremock::Mock; @@ -547,6 +549,86 @@ fn external_auth_tokens_without_chatgpt_metadata_cannot_seed_chatgpt_auth() { ); } +struct TestExternalChatgptAuth { + initial: ExternalAuthTokens, + refreshed: ExternalAuthTokens, + refresh_count: AtomicUsize, +} + +#[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())) + } + + async fn refresh( + &self, + _context: ExternalAuthRefreshContext, + ) -> std::io::Result { + self.refresh_count.fetch_add(1, Ordering::Relaxed); + Ok(self.refreshed.clone()) + } +} + +#[tokio::test] +async fn external_chatgpt_auth_resolves_and_refreshes_without_writing_auth_file() { + 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 refreshed_token = format!("{}.refreshed", token.rsplit_once('.').unwrap().0); + let external = Arc::new(TestExternalChatgptAuth { + initial: ExternalAuthTokens::chatgpt( + token.clone(), + WORKSPACE_ID_ALLOWED, + Some("enterprise".to_string()), + ), + refreshed: ExternalAuthTokens::chatgpt( + refreshed_token.clone(), + WORKSPACE_ID_ALLOWED, + 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, + ) + .await; + manager.set_external_auth(external.clone()); + + assert_eq!(manager.auth_mode(), Some(AuthMode::Chatgpt)); + assert_eq!( + manager.get_api_auth_mode(), + Some(AuthMode::ChatgptAuthTokens) + ); + let auth = manager.auth().await.expect("external auth should resolve"); + assert!(manager.is_external_chatgpt_auth_active()); + assert_eq!(auth.get_token().unwrap(), token); + assert_eq!(auth.get_account_id().as_deref(), Some(WORKSPACE_ID_ALLOWED)); + assert!(!get_auth_file(codex_home.path()).exists()); + + manager + .refresh_token() + .await + .expect("external auth refresh"); + assert_eq!(external.refresh_count.load(Ordering::Relaxed), 1); + assert_eq!( + manager.auth_cached().unwrap().get_token().unwrap(), + refreshed_token + ); + assert!(!get_auth_file(codex_home.path()).exists()); +} + #[tokio::test] async fn external_bearer_only_auth_manager_uses_cached_provider_token() { let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap(); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index b7d61153fd..2e1b8eabbf 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -29,6 +29,7 @@ use super::access_token::CodexAccessToken; use super::access_token::classify_codex_access_token; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; +use super::workload_identity::WorkloadIdentityExternalAuth; pub use crate::auth::agent_identity::AgentIdentityAuth; pub use crate::auth::bedrock_api_key::BedrockApiKeyAuth; pub use crate::auth::personal_access_token::PersonalAccessTokenAuth; @@ -49,6 +50,8 @@ use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; +use codex_workload_identity::WorkloadIdentityClient; +use codex_workload_identity::WorkloadIdentityConfig; use serde_json::Value; use thiserror::Error; @@ -1286,10 +1289,8 @@ enum UnauthorizedRecoveryMode { // // For external auth sources, UnauthorizedRecovery retries once. // -// - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed by asking -// the parent app for new tokens through the configured -// `ExternalAuth`, persisting them in the ephemeral auth store, and -// reloading the cached auth snapshot. +// - External ChatGPT auth tokens (`chatgptAuthTokens`) are refreshed through +// the configured `ExternalAuth` and replace the in-memory auth snapshot. // - External bearer auth sources for custom model providers rerun the provider // auth command without touching disk. pub struct UnauthorizedRecovery { @@ -1512,6 +1513,11 @@ pub trait AuthManagerConfig { /// Returns the ChatGPT backend base URL used for first-party backend authorization. fn chatgpt_base_url(&self) -> String; + + /// Returns workload identity configuration for externally managed ChatGPT auth. + fn workload_identity(&self) -> Option { + None + } } impl Debug for AuthManager { @@ -1683,8 +1689,8 @@ impl AuthManager { /// For managed ChatGPT auth that needs a proactive refresh, first performs /// a guarded reload and then refreshes only if the on-disk auth is unchanged. pub async fn auth(&self) -> Option { - if let Some(auth) = self.resolve_external_api_key_auth().await { - return Some(auth); + if self.has_external_auth() { + return self.resolve_external_auth().await; } let auth = self.auth_cached()?; @@ -1831,6 +1837,14 @@ impl AuthManager { } } + pub fn set_external_auth_if_absent(&self, external_auth: Arc) { + if let Ok(mut guard) = self.external_auth.write() + && guard.is_none() + { + *guard = Some(external_auth); + } + } + pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; @@ -1891,7 +1905,7 @@ impl AuthManager { config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { - Arc::new( + let auth_manager = Arc::new( Self::new_with_workspace_restriction( config.codex_home(), enable_codex_api_key_env, @@ -1901,7 +1915,13 @@ impl AuthManager { config.auth_keyring_backend_kind(), ) .await, - ) + ); + if let Some(workload_identity) = config.workload_identity() { + let client = + WorkloadIdentityClient::new(workload_identity, CLIENT_ID, build_reqwest_client()); + auth_manager.set_external_auth(Arc::new(WorkloadIdentityExternalAuth::new(client))); + } + auth_manager } pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { @@ -1925,23 +1945,58 @@ impl AuthManager { self.external_auth_mode() == Some(AuthMode::ApiKey) } - async fn resolve_external_api_key_auth(&self) -> Option { - if !self.has_external_api_key_auth() { - return None; - } - + async fn resolve_external_auth(&self) -> Option { let external_auth = self.external_auth()?; match external_auth.resolve().await { - Ok(Some(tokens)) => Some(CodexAuth::from_api_key(&tokens.access_token)), - Ok(None) => None, + Ok(Some(tokens)) if external_auth.auth_mode() == AuthMode::ApiKey => { + Some(CodexAuth::from_api_key(&tokens.access_token)) + } + Ok(Some(tokens)) => match self.external_chatgpt_auth(&tokens).await { + Ok(auth) => { + self.set_cached_auth(Some(auth.clone())); + Some(auth) + } + Err(error) => { + tracing::error!("Failed to resolve external ChatGPT auth: {error}"); + None + } + }, + Ok(None) => self.auth_cached(), Err(err) => { - tracing::error!("Failed to resolve external API key auth: {err}"); + tracing::error!("Failed to resolve external auth: {err}"); None } } } + async fn external_chatgpt_auth( + &self, + tokens: &ExternalAuthTokens, + ) -> std::io::Result { + let Some(chatgpt_metadata) = tokens.chatgpt_metadata() else { + return Err(std::io::Error::other( + "external auth tokens are missing ChatGPT metadata", + )); + }; + if let Some(expected_workspace_ids) = self.forced_chatgpt_workspace_id().as_deref() + && !expected_workspace_ids.contains(&chatgpt_metadata.account_id) + { + return Err(std::io::Error::other(format!( + "external auth returned workspace {:?}, expected one of {:?}", + chatgpt_metadata.account_id, expected_workspace_ids, + ))); + } + let auth_dot_json = AuthDotJson::from_external_tokens(tokens)?; + CodexAuth::from_auth_dot_json( + &self.codex_home, + auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + self.chatgpt_base_url.as_deref(), + ) + .await + } + /// Attempt to refresh the token by first performing a guarded reload. Auth /// is reloaded from storage only when the account id matches the currently /// cached account id. If the persisted token differs from the cached token, we @@ -1955,6 +2010,15 @@ impl AuthManager { )) })?; let auth_before_reload = self.auth_cached(); + if self.has_external_auth() + && auth_before_reload + .as_ref() + .is_some_and(CodexAuth::is_external_chatgpt_tokens) + { + return self + .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await; + } if auth_before_reload .as_ref() .is_some_and(|auth| auth.is_api_key_auth() || auth.is_personal_access_token_auth()) @@ -2067,15 +2131,23 @@ impl AuthManager { } pub fn get_api_auth_mode(&self) -> Option { - if self.has_external_api_key_auth() { - return Some(ApiAuthMode::ApiKey); + match self.external_auth_mode() { + Some(AuthMode::ApiKey) => return Some(ApiAuthMode::ApiKey), + Some(AuthMode::Chatgpt) => return Some(ApiAuthMode::ChatgptAuthTokens), + Some( + AuthMode::ChatgptAuthTokens + | AuthMode::AgentIdentity + | AuthMode::PersonalAccessToken + | AuthMode::BedrockApiKey, + ) + | None => {} } self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) } pub fn auth_mode(&self) -> Option { - if self.has_external_api_key_auth() { - return Some(AuthMode::ApiKey); + if let Some(mode @ (AuthMode::ApiKey | AuthMode::Chatgpt)) = self.external_auth_mode() { + return Some(mode); } self.auth_cached().as_ref().map(CodexAuth::auth_mode) } @@ -2118,7 +2190,6 @@ impl AuthManager { "external auth is not configured", ))); }; - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); let previous_account_id = self .auth_cached() .as_ref() diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index d87c7cc6ad..8002f1f742 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -10,6 +10,7 @@ mod util; mod external_bearer; mod manager; mod revoke; +mod workload_identity; pub use bedrock_api_key::BedrockApiKeyAuth; pub use bedrock_api_key::login_with_bedrock_api_key; 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..fe0feb8679 --- /dev/null +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -0,0 +1,50 @@ +use async_trait::async_trait; +use codex_app_server_protocol::AuthMode; +use codex_workload_identity::WorkloadIdentityClient; + +use super::ExternalAuth; +use super::ExternalAuthRefreshContext; +use super::ExternalAuthTokens; + +pub(super) struct WorkloadIdentityExternalAuth { + client: WorkloadIdentityClient, +} + +impl WorkloadIdentityExternalAuth { + pub(super) fn new(client: WorkloadIdentityClient) -> Self { + Self { client } + } + + async fn tokens(&self, force_refresh: bool) -> std::io::Result { + let token = if force_refresh { + self.client.refresh().await + } else { + self.client.resolve().await + } + .map_err(std::io::Error::other)?; + + Ok(ExternalAuthTokens::chatgpt( + token.access_token, + token.chatgpt_account_id, + token.chatgpt_plan_type, + )) + } +} + +#[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) + } + + async fn refresh( + &self, + _context: ExternalAuthRefreshContext, + ) -> std::io::Result { + self.tokens(true).await + } +} diff --git a/codex-rs/workload-identity/BUILD.bazel b/codex-rs/workload-identity/BUILD.bazel new file mode 100644 index 0000000000..ac046d02d4 --- /dev/null +++ b/codex-rs/workload-identity/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "workload-identity", + crate_name = "codex_workload_identity", +) diff --git a/codex-rs/workload-identity/Cargo.toml b/codex-rs/workload-identity/Cargo.toml new file mode 100644 index 0000000000..c9ee639aa5 --- /dev/null +++ b/codex-rs/workload-identity/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "codex-workload-identity" +version.workspace = true +edition.workspace = true +license.workspace = true + +[features] +default = [] +azure = [] + +[lints] +workspace = true + +[dependencies] +reqwest = { workspace = true, features = ["json"] } +schemars = { workspace = true } +serde = { workspace = true, features = ["derive"] } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["fs"] } +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 } + +[lib] +doctest = false diff --git a/codex-rs/workload-identity/src/azure.rs b/codex-rs/workload-identity/src/azure.rs new file mode 100644 index 0000000000..7ad25eaffa --- /dev/null +++ b/codex-rs/workload-identity/src/azure.rs @@ -0,0 +1,40 @@ +use std::path::PathBuf; + +use crate::WorkloadIdentityError; + +const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; + +#[derive(Clone, Debug)] +pub(crate) struct AzureSubjectTokenSource { + token_file: Option, +} + +impl AzureSubjectTokenSource { + pub(crate) fn new(token_file: Option) -> Self { + Self { token_file } + } + + pub(crate) async fn subject_token(&self) -> Result { + let token_file = match self.token_file.clone() { + Some(token_file) => token_file, + None => std::env::var_os(AZURE_FEDERATED_TOKEN_FILE_ENV) + .map(PathBuf::from) + .ok_or(WorkloadIdentityError::MissingAzureFederatedTokenFile)?, + }; + let token = tokio::fs::read_to_string(&token_file) + .await + .map_err(|source| WorkloadIdentityError::ReadSubjectToken { + path: token_file, + source, + })?; + let token = token.trim(); + if token.is_empty() { + return Err(WorkloadIdentityError::EmptySubjectToken); + } + Ok(token.to_string()) + } +} + +#[cfg(test)] +#[path = "azure_tests.rs"] +mod tests; diff --git a/codex-rs/workload-identity/src/azure_tests.rs b/codex-rs/workload-identity/src/azure_tests.rs new file mode 100644 index 0000000000..c6c8d75e71 --- /dev/null +++ b/codex-rs/workload-identity/src/azure_tests.rs @@ -0,0 +1,30 @@ +use std::path::PathBuf; + +use pretty_assertions::assert_eq; + +use super::*; + +#[tokio::test] +async fn reads_and_trims_projected_token_file() -> anyhow::Result<()> { + let temp_dir = tempfile::tempdir()?; + let token_file = temp_dir.path().join("azure-token"); + tokio::fs::write(&token_file, " header.payload.signature\n").await?; + let source = AzureSubjectTokenSource::new(Some(token_file)); + + assert_eq!(source.subject_token().await?, "header.payload.signature"); + Ok(()) +} + +#[tokio::test] +async fn rejects_empty_projected_token_file() -> anyhow::Result<()> { + let temp_dir = tempfile::tempdir()?; + let token_file = temp_dir.path().join("azure-token"); + tokio::fs::write(&token_file, " \n").await?; + let source = AzureSubjectTokenSource::new(Some(PathBuf::from(token_file))); + + assert!(matches!( + source.subject_token().await, + Err(WorkloadIdentityError::EmptySubjectToken) + )); + Ok(()) +} diff --git a/codex-rs/workload-identity/src/client.rs b/codex-rs/workload-identity/src/client.rs new file mode 100644 index 0000000000..c7da7d36b9 --- /dev/null +++ b/codex-rs/workload-identity/src/client.rs @@ -0,0 +1,250 @@ +use std::path::PathBuf; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; + +use crate::CredentialSourceConfig; +use crate::WorkloadIdentityConfig; +#[cfg(feature = "azure")] +use crate::azure::AzureSubjectTokenSource; +use reqwest::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; + +const TOKEN_EXCHANGE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:token-exchange"; +const JWT_TOKEN_TYPE: &str = "urn:ietf:params:oauth:token-type:jwt"; +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); + +#[derive(Clone, PartialEq, Eq)] +pub struct WorkloadIdentityAccessToken { + pub access_token: String, + pub chatgpt_account_id: String, + pub chatgpt_plan_type: Option, +} + +impl std::fmt::Debug for WorkloadIdentityAccessToken { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkloadIdentityAccessToken") + .field("access_token", &"[REDACTED]") + .field("chatgpt_account_id", &self.chatgpt_account_id) + .field("chatgpt_plan_type", &self.chatgpt_plan_type) + .finish() + } +} + +#[derive(Debug, Error)] +pub enum WorkloadIdentityError { + #[error("AZURE_FEDERATED_TOKEN_FILE is not set and no Azure token_file was configured")] + MissingAzureFederatedTokenFile, + #[error("failed to read workload identity subject token from {path}: {source}")] + ReadSubjectToken { + path: PathBuf, + #[source] + source: std::io::Error, + }, + #[error("workload identity subject token file is empty")] + EmptySubjectToken, + #[error("the configured credential source is not enabled in this Codex build")] + CredentialSourceNotEnabled, + #[error("workload identity token exchange request failed: {0}")] + Request(#[from] reqwest::Error), + #[error("workload identity token exchange was rejected with HTTP {status}: {message}")] + Rejected { status: StatusCode, message: String }, + #[error("workload identity token exchange returned an empty access token")] + EmptyAccessToken, + #[error("workload identity token exchange returned an empty ChatGPT account ID")] + EmptyAccountId, + #[error("workload identity token exchange returned a token with no usable lifetime")] + InvalidLifetime, + #[error("workload identity token exchange returned an unexpected token type")] + UnexpectedTokenType, +} + +enum SubjectTokenSource { + #[cfg(feature = "azure")] + Azure(AzureSubjectTokenSource), + #[cfg(not(feature = "azure"))] + Disabled, +} + +impl SubjectTokenSource { + fn from_config(config: CredentialSourceConfig) -> Self { + match config { + CredentialSourceConfig::Azure { token_file } => { + #[cfg(feature = "azure")] + { + Self::Azure(AzureSubjectTokenSource::new(token_file)) + } + #[cfg(not(feature = "azure"))] + { + let _ = token_file; + Self::Disabled + } + } + } + } + + async fn subject_token(&self) -> Result { + match self { + #[cfg(feature = "azure")] + Self::Azure(source) => source.subject_token().await, + #[cfg(not(feature = "azure"))] + Self::Disabled => Err(WorkloadIdentityError::CredentialSourceNotEnabled), + } + } +} + +struct CachedAccessToken { + token: WorkloadIdentityAccessToken, + refresh_at: Instant, +} + +pub struct WorkloadIdentityClient { + identity_provider_id: String, + identity_provider_mapping_id: String, + token_url: String, + client_id: String, + source: SubjectTokenSource, + http: reqwest::Client, + cache: Mutex>, +} + +impl WorkloadIdentityClient { + pub fn new( + config: WorkloadIdentityConfig, + client_id: impl Into, + http: reqwest::Client, + ) -> Self { + Self { + identity_provider_id: config.identity_provider_id, + identity_provider_mapping_id: config.identity_provider_mapping_id, + token_url: config.token_url, + client_id: client_id.into(), + source: SubjectTokenSource::from_config(config.credential_source), + http, + cache: Mutex::new(None), + } + } + + pub async fn resolve(&self) -> Result { + { + let cache = self + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cache.as_ref() + && Instant::now() < cached.refresh_at + { + return Ok(cached.token.clone()); + } + } + self.exchange_and_cache().await + } + + pub async fn refresh(&self) -> Result { + self.exchange_and_cache().await + } + + async fn exchange_and_cache( + &self, + ) -> Result { + let subject_token = self.source.subject_token().await?; + let response = self + .http + .post(&self.token_url) + .timeout(TOKEN_EXCHANGE_TIMEOUT) + .json(&TokenExchangeRequest { + grant_type: TOKEN_EXCHANGE_GRANT_TYPE, + requested_token_type: ACCESS_TOKEN_TYPE, + subject_token: &subject_token, + subject_token_type: JWT_TOKEN_TYPE, + identity_provider_id: &self.identity_provider_id, + identity_provider_mapping_id: &self.identity_provider_mapping_id, + client_id: &self.client_id, + }) + .send() + .await?; + let status = response.status(); + if !status.is_success() { + let message = response + .json::() + .await + .ok() + .and_then(|response| response.error) + .unwrap_or_else(|| "token endpoint rejected the request".to_string()); + return Err(WorkloadIdentityError::Rejected { status, message }); + } + let response: TokenExchangeResponse = response.json().await?; + if !response.token_type.eq_ignore_ascii_case("bearer") + || response.issued_token_type != ACCESS_TOKEN_TYPE + { + return Err(WorkloadIdentityError::UnexpectedTokenType); + } + if response.access_token.trim().is_empty() { + return Err(WorkloadIdentityError::EmptyAccessToken); + } + if response.chatgpt_account_id.trim().is_empty() { + return Err(WorkloadIdentityError::EmptyAccountId); + } + if response.expires_in == 0 { + return Err(WorkloadIdentityError::InvalidLifetime); + } + + let token = WorkloadIdentityAccessToken { + access_token: response.access_token, + chatgpt_account_id: response.chatgpt_account_id, + chatgpt_plan_type: response.chatgpt_plan_type, + }; + let lifetime = Duration::from_secs(response.expires_in); + let refresh_lead = MAX_REFRESH_LEAD.min(Duration::from_secs( + response.expires_in.saturating_div(10).max(1), + )); + let refresh_at = Instant::now() + .checked_add(lifetime.saturating_sub(refresh_lead)) + .ok_or(WorkloadIdentityError::InvalidLifetime)?; + *self + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(CachedAccessToken { + token: token.clone(), + refresh_at, + }); + Ok(token) + } +} + +#[derive(Serialize)] +struct TokenExchangeRequest<'a> { + grant_type: &'static str, + requested_token_type: &'static str, + subject_token: &'a str, + subject_token_type: &'static str, + identity_provider_id: &'a str, + identity_provider_mapping_id: &'a str, + client_id: &'a str, +} + +#[derive(Deserialize)] +struct TokenExchangeResponse { + access_token: String, + issued_token_type: String, + token_type: String, + expires_in: u64, + chatgpt_account_id: String, + #[serde(default)] + chatgpt_plan_type: Option, +} + +#[derive(Deserialize)] +struct TokenExchangeErrorResponse { + #[serde(default)] + error: Option, +} + +#[cfg(all(test, feature = "azure"))] +#[path = "client_tests.rs"] +mod tests; diff --git a/codex-rs/workload-identity/src/client_tests.rs b/codex-rs/workload-identity/src/client_tests.rs new file mode 100644 index 0000000000..79210ffdf9 --- /dev/null +++ b/codex-rs/workload-identity/src/client_tests.rs @@ -0,0 +1,91 @@ +use std::path::PathBuf; + +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_json; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; + +async fn test_client(server: &MockServer) -> anyhow::Result { + let temp_dir = tempfile::tempdir()?; + let token_file = temp_dir.keep().join("subject-token"); + tokio::fs::write(&token_file, "azure.subject.token\n").await?; + Ok(WorkloadIdentityClient::new( + WorkloadIdentityConfig { + identity_provider_id: "idp_example".to_string(), + identity_provider_mapping_id: "idpm_example".to_string(), + token_url: format!("{}/oauth/token", server.uri()), + credential_source: CredentialSourceConfig::Azure { + token_file: Some(PathBuf::from(token_file)), + }, + }, + "app_codex_cli", + reqwest::Client::new(), + )) +} + +fn exchange_request() -> serde_json::Value { + json!({ + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "requested_token_type": ACCESS_TOKEN_TYPE, + "subject_token": "azure.subject.token", + "subject_token_type": JWT_TOKEN_TYPE, + "identity_provider_id": "idp_example", + "identity_provider_mapping_id": "idpm_example", + "client_id": "app_codex_cli", + }) +} + +fn exchange_response(access_token: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": access_token, + "issued_token_type": ACCESS_TOKEN_TYPE, + "token_type": "Bearer", + "expires_in": 3600, + "chatgpt_account_id": "workspace_example", + "chatgpt_plan_type": "enterprise", + })) +} + +#[tokio::test] +async fn exchanges_azure_subject_token_and_caches_access_token() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_json(exchange_request())) + .respond_with(exchange_response("first.access.token")) + .expect(1) + .mount(&server) + .await; + let client = test_client(&server).await?; + + let expected = WorkloadIdentityAccessToken { + access_token: "first.access.token".to_string(), + chatgpt_account_id: "workspace_example".to_string(), + chatgpt_plan_type: Some("enterprise".to_string()), + }; + assert_eq!(client.resolve().await?, expected); + assert_eq!(client.resolve().await?, expected); + Ok(()) +} + +#[tokio::test] +async fn forced_refresh_performs_a_new_exchange() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(exchange_response("access.token")) + .expect(2) + .mount(&server) + .await; + let client = test_client(&server).await?; + + client.resolve().await?; + client.refresh().await?; + Ok(()) +} diff --git a/codex-rs/workload-identity/src/config.rs b/codex-rs/workload-identity/src/config.rs new file mode 100644 index 0000000000..0bdad08f17 --- /dev/null +++ b/codex-rs/workload-identity/src/config.rs @@ -0,0 +1,103 @@ +use std::path::PathBuf; + +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; +use url::Host; +use url::Url; + +const DEFAULT_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; + +/// Configuration for exchanging a workload credential for a Codex access token. +#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WorkloadIdentityConfig { + /// OpenAI workload identity provider selected by the workspace administrator. + pub identity_provider_id: String, + + /// Administrator-created mapping from external claims to a ChatGPT principal. + pub identity_provider_mapping_id: String, + + /// OAuth token endpoint. The override is primarily useful for local development. + #[serde(default = "default_token_url")] + pub token_url: String, + + /// Runtime-specific mechanism used to obtain the external subject token. + pub credential_source: CredentialSourceConfig, +} + +/// Supported runtime credential sources. Implementations are independently feature gated. +#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)] +#[serde(deny_unknown_fields, rename_all = "snake_case", tag = "type")] +pub enum CredentialSourceConfig { + /// A projected Azure workload identity token. + Azure { + /// Token path. When omitted, Codex reads `AZURE_FEDERATED_TOKEN_FILE`. + #[serde(default)] + token_file: Option, + }, +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum WorkloadIdentityConfigError { + #[error("workload_identity.{0} must not be empty")] + EmptyField(&'static str), + #[error("workload_identity.token_url is invalid: {0}")] + InvalidTokenUrl(String), + #[error("workload_identity.token_url must use https or loopback http")] + UnsupportedTokenUrlScheme, + #[error("workload_identity.credential_source.token_file must be an absolute path")] + RelativeTokenFile, + #[error("workload_identity.credential_source.token_file must not be empty")] + EmptyTokenFile, +} + +pub fn default_token_url() -> String { + DEFAULT_TOKEN_URL.to_string() +} + +impl WorkloadIdentityConfig { + pub fn validate(&self) -> Result<(), WorkloadIdentityConfigError> { + for (field, value) in [ + ("identity_provider_id", self.identity_provider_id.as_str()), + ( + "identity_provider_mapping_id", + self.identity_provider_mapping_id.as_str(), + ), + ("token_url", self.token_url.as_str()), + ] { + if value.trim().is_empty() { + return Err(WorkloadIdentityConfigError::EmptyField(field)); + } + } + + let token_url = Url::parse(&self.token_url) + .map_err(|error| WorkloadIdentityConfigError::InvalidTokenUrl(error.to_string()))?; + let loopback_http = token_url.scheme() == "http" + && token_url.host().is_some_and(|host| match host { + Host::Domain(domain) => domain.eq_ignore_ascii_case("localhost"), + Host::Ipv4(address) => address.is_loopback(), + Host::Ipv6(address) => address.is_loopback(), + }); + if token_url.scheme() != "https" && !loopback_http { + return Err(WorkloadIdentityConfigError::UnsupportedTokenUrlScheme); + } + + match &self.credential_source { + CredentialSourceConfig::Azure { + token_file: Some(token_file), + } if token_file.as_os_str().is_empty() => { + Err(WorkloadIdentityConfigError::EmptyTokenFile) + } + CredentialSourceConfig::Azure { + token_file: Some(token_file), + } if !token_file.is_absolute() => Err(WorkloadIdentityConfigError::RelativeTokenFile), + CredentialSourceConfig::Azure { .. } => Ok(()), + } + } +} + +#[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 new file mode 100644 index 0000000000..ebbf520860 --- /dev/null +++ b/codex-rs/workload-identity/src/config_tests.rs @@ -0,0 +1,66 @@ +use std::path::PathBuf; + +use pretty_assertions::assert_eq; + +use super::*; + +fn valid_config() -> WorkloadIdentityConfig { + WorkloadIdentityConfig { + identity_provider_id: "idp_example".to_string(), + identity_provider_mapping_id: "idpm_example".to_string(), + token_url: "https://auth.openai.com/oauth/token".to_string(), + credential_source: CredentialSourceConfig::Azure { + token_file: Some(PathBuf::from( + "/var/run/secrets/azure/tokens/azure-identity-token", + )), + }, + } +} + +#[test] +fn validates_complete_configuration() { + assert_eq!(valid_config().validate(), Ok(())); +} + +#[test] +fn rejects_non_http_token_endpoint() { + let mut config = valid_config(); + config.token_url = "file:///tmp/token".to_string(); + + assert_eq!( + config.validate(), + Err(WorkloadIdentityConfigError::UnsupportedTokenUrlScheme) + ); +} + +#[test] +fn rejects_non_loopback_http_token_endpoint() { + let mut config = valid_config(); + config.token_url = "http://attacker.example/oauth/token".to_string(); + + assert_eq!( + config.validate(), + Err(WorkloadIdentityConfigError::UnsupportedTokenUrlScheme) + ); +} + +#[test] +fn allows_loopback_http_token_endpoint_for_local_development() { + let mut config = valid_config(); + config.token_url = "http://127.0.0.1:3007/oauth/token".to_string(); + + assert_eq!(config.validate(), Ok(())); +} + +#[test] +fn rejects_relative_token_file() { + let mut config = valid_config(); + config.credential_source = CredentialSourceConfig::Azure { + token_file: Some(PathBuf::from("azure-token")), + }; + + assert_eq!( + config.validate(), + Err(WorkloadIdentityConfigError::RelativeTokenFile) + ); +} diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs new file mode 100644 index 0000000000..ea69db6958 --- /dev/null +++ b/codex-rs/workload-identity/src/lib.rs @@ -0,0 +1,13 @@ +mod client; +mod config; + +#[cfg(feature = "azure")] +mod azure; + +pub use client::WorkloadIdentityAccessToken; +pub use client::WorkloadIdentityClient; +pub use client::WorkloadIdentityError; +pub use config::CredentialSourceConfig; +pub use config::WorkloadIdentityConfig; +pub use config::WorkloadIdentityConfigError; +pub use config::default_token_url;