diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs index b5e729f4ac..88a1b98d44 100644 --- a/codex-rs/cli/tests/login.rs +++ b/codex-rs/cli/tests/login.rs @@ -155,14 +155,23 @@ fn login_status_validates_configured_workload_identity() -> Result<()> { write_file_auth_config(codex_home.path())?; let missing_assertion = codex_home.path().join("missing-identity-token"); - codex_command(codex_home.path())? - .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) - .env(OPENAI_FEDERATION_RULE_ID_ENV_VAR, "rule-test") - .env(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR, &missing_assertion) - .args(["login", "status"]) - .assert() - .failure() - .stderr(contains("could not read workload identity assertion file")); + for login_method in ["chatgpt", "api"] { + let mut command = codex_command(codex_home.path())?; + command + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .env(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR, &missing_assertion) + .env(OPENAI_FEDERATION_RULE_ID_ENV_VAR, "rule-test") + .args([ + "-c", + &format!("forced_login_method=\"{login_method}\""), + "login", + "status", + ]); + command + .assert() + .failure() + .stderr(contains("could not read workload identity assertion file")); + } Ok(()) } diff --git a/codex-rs/login/src/auth/change_state.rs b/codex-rs/login/src/auth/change_state.rs index cdb1ab3cf0..6fbf5cc717 100644 --- a/codex-rs/login/src/auth/change_state.rs +++ b/codex-rs/login/src/auth/change_state.rs @@ -7,13 +7,13 @@ use super::CodexAuth; pub struct AuthChangeState { /// Advances on cached credential changes and explicit invalidation of pending reads. pub generation: u64, - /// Advances on login, logout, or changes of user, workspace, or auth mode. + /// Advances on login, logout, or changes of user, workspace, service account, or auth mode. /// A verified legacy-to-OAuth credential migration preserves the same owner. /// Credential changes with incomplete owner identity also advance this revision. /// Losing login metadata alone preserves the account owner. pub owner_generation: u64, /// Advances on an underlying actor credential/source change, including a new OAuth login. - /// Ordinary refresh within an identifiable OAuth session preserves this revision. + /// Renewal within an identifiable OAuth session or WIF subject preserves this revision. pub source_generation: u64, /// Preserves strict login ownership for EMA, including loss of login metadata. pub login_generation: u64, @@ -47,6 +47,9 @@ pub(super) fn same_owner(previous: Option<&CodexAuth>, current: Option<&CodexAut if previous.api_auth_mode() != current.api_auth_mode() { return false; } + if let Some(subject) = previous.api_workload_identity() { + return current.api_workload_identity() == Some(subject); + } let (Some(user), Some(workspace)) = (previous.get_chatgpt_user_id(), previous.get_account_id()) else { return false; diff --git a/codex-rs/login/src/auth/change_state_tests.rs b/codex-rs/login/src/auth/change_state_tests.rs index 4a13338454..2dc462752d 100644 --- a/codex-rs/login/src/auth/change_state_tests.rs +++ b/codex-rs/login/src/auth/change_state_tests.rs @@ -68,6 +68,53 @@ fn auth_owner_generation_distinguishes_refreshes_from_identity_changes() { } } +#[test] +fn api_workload_identity_changes_retire_previous_owner() { + let identity = ApiTokenMetadata { + api_organization_id: "organization".into(), + api_project_id: "project".into(), + user_id: "service-account".into(), + }; + let auth = |workload_identity| { + CodexAuth::ApiKey(ApiKeyAuth { + api_key: "same-bearer".into(), + workload_identity, + }) + }; + for replacement in [ + None, + Some(ApiTokenMetadata { + api_organization_id: "another-organization".into(), + ..identity.clone() + }), + Some(ApiTokenMetadata { + api_project_id: "another-project".into(), + ..identity.clone() + }), + Some(ApiTokenMetadata { + user_id: "another-service-account".into(), + ..identity.clone() + }), + ] { + let previous = auth(Some(identity.clone())); + let current = auth(replacement); + for (previous, current) in [(previous.clone(), current.clone()), (current, previous)] { + let manager = AuthManager::from_auth_for_testing(previous); + let before = *manager.auth_change_state_receiver().borrow(); + manager.set_cached_auth(Some(current)); + assert_eq!( + *manager.auth_change_state_receiver().borrow(), + AuthChangeState { + generation: before.generation + 1, + owner_generation: before.owner_generation + 1, + source_generation: before.source_generation + 1, + login_generation: before.login_generation + 1, + } + ); + } + } +} + #[tokio::test] #[serial_test::serial(codex_auth_env)] async fn auth_refresh_read_failure_is_transient_and_preserves_network_permit() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index ced49b3dc6..ed3bcc9dbb 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -41,6 +41,7 @@ use codex_agent_identity::ChatGptEnvironment; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; +use codex_workload_identity::ApiTokenMetadata; use super::access_token::CodexAccessToken; use super::access_token::classify_codex_access_token; @@ -196,9 +197,18 @@ impl PartialEq for CodexAuth { } } -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct ApiKeyAuth { - api_key: String, + pub(super) api_key: String, + pub(super) workload_identity: Option, +} + +impl Debug for ApiKeyAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ApiKeyAuth") + .field("api_key", &"[redacted]") + .finish_non_exhaustive() + } } #[derive(Debug, Clone)] @@ -622,6 +632,14 @@ impl CodexAuth { } } + pub(super) fn api_workload_identity(&self) -> Option<&ApiTokenMetadata> { + if let Self::ApiKey(auth) = self { + auth.workload_identity.as_ref() + } else { + None + } + } + /// Returns `Err` if token-backed ChatGPT auth is unavailable. pub fn get_token_data(&self) -> Result { let auth_dot_json: Option = self.get_current_auth_json(); @@ -916,6 +934,7 @@ impl CodexAuth { pub fn from_api_key(api_key: &str) -> Self { Self::ApiKey(ApiKeyAuth { api_key: api_key.to_owned(), + workload_identity: None, }) } } @@ -2908,7 +2927,10 @@ impl AuthManager { match (a, b) { (None, None) => true, (Some(a), Some(b)) => match (a.api_auth_mode(), b.api_auth_mode()) { - (AuthMode::ApiKey, AuthMode::ApiKey) => a.api_key() == b.api_key(), + (AuthMode::ApiKey, AuthMode::ApiKey) => { + a.api_key() == b.api_key() + && a.api_workload_identity() == b.api_workload_identity() + } (AuthMode::Chatgpt, AuthMode::Chatgpt) => { a.get_current_auth_json() == b.get_current_auth_json() } diff --git a/codex-rs/login/src/auth/source.rs b/codex-rs/login/src/auth/source.rs index afe03ea705..3e25bdb0e1 100644 --- a/codex-rs/login/src/auth/source.rs +++ b/codex-rs/login/src/auth/source.rs @@ -17,6 +17,11 @@ pub(super) fn same_source(previous: Option<&CodexAuth>, current: Option<&CodexAu { return false; } + let previous_subject = previous.api_workload_identity(); + let current_subject = current.api_workload_identity(); + if previous_subject.is_some() || current_subject.is_some() { + return previous_subject == current_subject; + } let (Ok(previous_token), Ok(current_token)) = (previous.get_token(), current.get_token()) else { return false; diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs index 06cf649027..6ed55bf158 100644 --- a/codex-rs/login/src/auth/workload_identity.rs +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -8,17 +8,18 @@ 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_protocol::shell_environment::OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR; use codex_workload_identity::WorkloadIdentityConfig; use codex_workload_identity::WorkloadIdentityError; use codex_workload_identity::WorkloadIdentityExchange; +use codex_workload_identity::WorkloadIdentityMetadata; use codex_workload_identity::WorkloadIdentityToken; use thiserror::Error; use url::Url; +use super::ApiKeyAuth; use super::AuthConfig; use super::CodexAuth; use super::ExternalAuth; @@ -131,17 +132,11 @@ pub fn is_workload_identity_selected() -> bool { fn resolve_config( chatgpt_base_url: &str, environment: ProcessEnvironment, - chatgpt_login_allowed: bool, auth_route_config: AuthRouteConfig, ) -> Result, WorkloadIdentitySessionError> { if !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, @@ -265,18 +260,32 @@ impl ProcessEnvironment { } #[derive(Clone, PartialEq, Eq)] -struct WorkloadIdentitySubject { - account_id: String, - account_user_id: String, - user_id: String, +enum WorkloadIdentitySubject { + Chatgpt { + account_id: String, + account_user_id: String, + user_id: String, + }, + Api { + organization_id: String, + project_id: String, + service_account_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(), +impl From<&WorkloadIdentityMetadata> for WorkloadIdentitySubject { + fn from(metadata: &WorkloadIdentityMetadata) -> Self { + match metadata { + WorkloadIdentityMetadata::Chatgpt(metadata) => Self::Chatgpt { + account_id: metadata.chatgpt_account_id.clone(), + account_user_id: metadata.chatgpt_account_user_id.clone(), + user_id: metadata.user_id.clone(), + }, + WorkloadIdentityMetadata::Api(metadata) => Self::Api { + organization_id: metadata.api_organization_id.clone(), + project_id: metadata.api_project_id.clone(), + service_account_id: metadata.user_id.clone(), + }, } } } @@ -296,13 +305,19 @@ impl WorkloadIdentitySession { fn accept_subject( &self, - token: &WorkloadIdentityToken, + metadata: &WorkloadIdentityMetadata, previous_account_id: Option<&str>, ) -> Result<(), WorkloadIdentityError> { - if previous_account_id.is_some_and(|account_id| account_id != token.chatgpt_account_id) { + let account_id = match metadata { + WorkloadIdentityMetadata::Chatgpt(metadata) => { + Some(metadata.chatgpt_account_id.as_str()) + } + WorkloadIdentityMetadata::Api(_) => None, + }; + if previous_account_id.is_some() && previous_account_id != account_id { return Err(WorkloadIdentityError::InvalidExchangeResponse); } - let subject = WorkloadIdentitySubject::from(token); + let subject = WorkloadIdentitySubject::from(metadata); let mut current = self .subject .lock() @@ -379,7 +394,6 @@ impl WorkloadIdentityExternalAuth { .as_deref() .unwrap_or("https://chatgpt.com/backend-api"), ProcessEnvironment::read(), - auth_config.is_login_method_allowed(ForcedLoginMethod::Chatgpt), auth_config.auth_route_config.clone(), )? .map(|config| Self::from_config_with_registry(config, registry)) @@ -420,19 +434,30 @@ impl WorkloadIdentityExternalAuth { 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, - )); - } + let auth = match &token.metadata { + WorkloadIdentityMetadata::Api(metadata) => CodexAuth::ApiKey(ApiKeyAuth { + api_key: token.access_token.clone(), + workload_identity: Some(metadata.clone()), + }), + WorkloadIdentityMetadata::Chatgpt(metadata) => { + let auth = CodexAuth::from_external_chatgpt_tokens( + &token.access_token, + &metadata.chatgpt_account_id, + metadata.chatgpt_plan_type.as_deref(), + ) + .map_err(|_| { + std::io::Error::other(WorkloadIdentityError::InvalidExchangeResponse) + })?; + if auth.get_chatgpt_user_id().as_deref() != Some(metadata.user_id.as_str()) { + return Err(std::io::Error::other( + WorkloadIdentityError::InvalidExchangeResponse, + )); + } + auth + } + }; self.session - .accept_subject(token, previous_account_id) + .accept_subject(&token.metadata, previous_account_id) .map_err(std::io::Error::other)?; Ok(auth) } diff --git a/codex-rs/login/src/auth/workload_identity_tests.rs b/codex-rs/login/src/auth/workload_identity_tests.rs index e11eca1116..750031f2b8 100644 --- a/codex-rs/login/src/auth/workload_identity_tests.rs +++ b/codex-rs/login/src/auth/workload_identity_tests.rs @@ -3,14 +3,18 @@ use std::sync::atomic::AtomicUsize; use std::time::Duration; use base64::Engine as _; +use codex_config::types::AuthCredentialsStoreMode; use codex_http_client::OutboundProxyPolicy; +use codex_protocol::config_types::ForcedLoginMethod; use pretty_assertions::assert_eq; +use serial_test::serial; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::method; use super::*; +use crate::AuthManager; use crate::auth::ExternalAuthRefreshReason; fn auth_route_config(policy: OutboundProxyPolicy) -> AuthRouteConfig { @@ -21,19 +25,115 @@ 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()), - workload_identity_context: None, + ..ProcessEnvironment::default() + } +} + +fn api_response(access_token: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access_token, + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "expires_in": 600, + "wif_product": "api", + "api_organization_id": "org-one", + "api_project_id": "proj-one", + "user_id": "service-account-one" + }) +} + +#[tokio::test] +async fn api_platform_auth_refreshes_with_a_rotated_assertion() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let assertion_file = temp_dir.path().join("identity-token"); + std::fs::write(&assertion_file, "assertion-one").expect("write assertion"); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(|request: &wiremock::Request| { + let assertion = url::form_urlencoded::parse(&request.body) + .find(|(name, _)| name == "assertion") + .expect("assertion") + .1; + ResponseTemplate::new(200) + .set_body_json(api_response(&format!("access-for-{assertion}"))) + }) + .expect(2) + .mount(&server) + .await; + let config = session_config(temp_dir.path(), &server); + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry(config, ®istry) + .expect("adapter"); + let manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("unused-fallback"), + temp_dir.path().to_path_buf(), + ); + manager + .set_external_auth(Arc::new(adapter)) + .await + .expect("install WIF"); + assert_eq!( + manager + .auth() + .await + .expect("initial auth") + .get_token() + .expect("token"), + "access-for-assertion-one" + ); + let before = *manager.auth_change_state_receiver().borrow(); + std::fs::write(&assertion_file, "assertion-two").expect("rotate assertion"); + let mut recovery = manager.unauthorized_recovery(); + assert!( + recovery.has_next(), + "API WIF supports unauthorized recovery" + ); + recovery.next().await.expect("refresh WIF token"); + assert_eq!( + manager + .auth() + .await + .expect("renewed auth") + .get_token() + .expect("token"), + "access-for-assertion-two" + ); + assert_eq!( + *manager.auth_change_state_receiver().borrow(), + crate::AuthChangeState { + generation: before.generation + 1, + ..before + } + ); + assert!(!temp_dir.path().join("auth.json").exists()); + + let requests = server.received_requests().await.expect("requests"); + assert_eq!(requests.len(), 2); + for (request, assertion) in requests.iter().zip(["assertion-one", "assertion-two"]) { + assert_eq!(request.url.path(), "/oauth/token"); + assert_eq!( + url::form_urlencoded::parse(&request.body) + .into_owned() + .collect::>(), + std::collections::BTreeMap::from([ + ( + "grant_type".into(), + "urn:ietf:params:oauth:grant-type:jwt-bearer".into() + ), + ("assertion".into(), assertion.into()), + ("federation_rule_id".into(), "rule-one".into()), + ]), + ); } } fn resolve_for_test( environment: ProcessEnvironment, - chatgpt_login_allowed: bool, chatgpt_base_url: &str, ) -> Result, WorkloadIdentitySessionError> { resolve_config( chatgpt_base_url, environment, - chatgpt_login_allowed, auth_route_config(OutboundProxyPolicy::ReqwestDefault), ) } @@ -43,7 +143,6 @@ fn markers_select_wif_and_partial_configuration_fails_closed() { assert!( resolve_for_test( ProcessEnvironment::default(), - /*chatgpt_login_allowed*/ true, "https://chatgpt.com/backend-api", ) .expect("no markers") @@ -55,7 +154,6 @@ fn markers_select_wif_and_partial_configuration_fails_closed() { workload_identity_context: Some(r#"{"instance_id":"box-one"}"#.into()), ..ProcessEnvironment::default() }, - /*chatgpt_login_allowed*/ true, "https://chatgpt.com/backend-api", ) .expect("context alone is not a WIF marker") @@ -77,12 +175,8 @@ fn markers_select_wif_and_partial_configuration_fails_closed() { OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR, ), ] { - let error = resolve_for_test( - environment, - /*chatgpt_login_allowed*/ true, - "https://chatgpt.com/backend-api", - ) - .expect_err("partial WIF must not fall back"); + let error = resolve_for_test(environment, "https://chatgpt.com/backend-api") + .expect_err("partial WIF must not fall back"); assert!(error.to_string().contains(missing), "{error}"); } @@ -91,27 +185,15 @@ fn markers_select_wif_and_partial_configuration_fails_closed() { ..complete_environment() }; assert!( - resolve_for_test( - relative, - /*chatgpt_login_allowed*/ true, - "https://chatgpt.com/backend-api", - ) - .expect_err("relative assertion path") - .to_string() - .contains("absolute path") + resolve_for_test(relative, "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(), - /*chatgpt_login_allowed*/ false, - "https://chatgpt.com/backend-api", - ) - .expect_err("ChatGPT-disallowing policy"); - assert!(policy_error.to_string().contains("login policy")); - +fn app_environment_is_enforced() { for (chatgpt_base_url, expected_environment, expected_token_url) in [ ( "https://chatgpt.com/backend-api/", @@ -124,20 +206,15 @@ fn auth_policy_and_app_environment_are_enforced() { STAGING_TOKEN_URL, ), ] { - let config = resolve_for_test( - complete_environment(), - /*chatgpt_login_allowed*/ true, - chatgpt_base_url, - ) - .expect("trusted app routing") - .expect("WIF selected"); + let config = resolve_for_test(complete_environment(), 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(), - /*chatgpt_login_allowed*/ true, "https://example.invalid/backend-api", ) .expect_err("untrusted auth environment"); @@ -152,13 +229,11 @@ fn workload_context_is_preserved_without_logging_its_value() { workload_identity_context: Some(context.into()), ..complete_environment() }, - /*chatgpt_login_allowed*/ true, "https://chatgpt.com/backend-api", ) .expect("valid configuration") .expect("WIF selected"); - assert_eq!(config.workload_identity_context.as_deref(), Some(context)); assert!(!format!("{config:?}").contains(context)); } @@ -346,6 +421,145 @@ async fn refresh_preserves_identity_and_invalid_tokens_are_reexchanged() { assert_eq!(request_count.load(Ordering::SeqCst), 3); } +#[tokio::test] +async fn api_refresh_preserves_product_and_identity() { + for changed_field in [ + Some("api_organization_id"), + Some("api_project_id"), + Some("user_id"), + None, + ] { + 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 => ResponseTemplate::new(200).set_body_json(api_response("access-one")), + 1 => match changed_field { + Some(field) => { + let mut body = api_response("access-two"); + body[field] = "another-identity".into(); + ResponseTemplate::new(200).set_body_json(body) + } + None => success_response("access-two", "account-user-one", "user-one"), + }, + _ => ResponseTemplate::new(200).set_body_json(api_response("access-three")), + } + }) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("adapter"); + assert_eq!( + adapter.resolve().await.expect("initial API auth"), + CodexAuth::from_api_key("access-one") + ); + let error = adapter + .refresh(ExternalAuthRefreshContext { + reason: ExternalAuthRefreshReason::Unauthorized, + previous_account_id: None, + }) + .await + .expect_err("product or identity change must be rejected"); + assert!(matches!( + adapter.classify_error(error), + RefreshTokenError::Permanent(_) + )); + assert_eq!( + adapter + .resolve() + .await + .expect("corrected token is re-exchanged"), + CodexAuth::from_api_key("access-three") + ); + assert_eq!(request_count.load(Ordering::SeqCst), 3); + } +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn resolved_product_obeys_login_and_workspace_restrictions() { + for (product, login_method, workspace, allowed) in [ + ("api", ForcedLoginMethod::Api, None, true), + ("api", ForcedLoginMethod::Chatgpt, None, false), + ("chatgpt", ForcedLoginMethod::Chatgpt, None, true), + ("chatgpt", ForcedLoginMethod::Api, None, false), + ( + "chatgpt", + ForcedLoginMethod::Chatgpt, + Some("account-one"), + true, + ), + ( + "chatgpt", + ForcedLoginMethod::Chatgpt, + Some("another-account"), + false, + ), + ( + "api", + ForcedLoginMethod::Chatgpt, + Some("account-one"), + false, + ), + ] { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let server = MockServer::start().await; + let response = if product == "api" { + ResponseTemplate::new(200).set_body_json(api_response("api-access")) + } else { + success_response("chatgpt-access", "account-user-one", "user-one") + }; + Mock::given(method("POST")) + .respond_with(response) + .expect(1) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("adapter"); + let manager = AuthManager::shared_from_auth_config( + AuthConfig { + codex_home: temp_dir.path().to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::Ephemeral, + keyring_backend_kind: Default::default(), + forced_login_method: Some(login_method), + chatgpt_base_url: None, + forced_chatgpt_workspace_id: workspace.map(|id| vec![id.to_string()]), + managed_auth_policy: Default::default(), + auth_route_config: auth_route_config(OutboundProxyPolicy::ReqwestDefault), + }, + /*enable_codex_api_key_env*/ false, + ) + .await + .expect("valid auth policy"); + let result = manager.set_external_auth(Arc::new(adapter)).await; + assert_eq!( + result.is_ok(), + allowed, + "{product} with {login_method:?} and {workspace:?}" + ); + if allowed { + let auth = manager.auth().await.expect("accepted WIF auth"); + assert_eq!(auth.is_api_key_auth(), product == "api"); + } else { + assert!( + manager.auth().await.is_none(), + "rejected WIF must not install credentials" + ); + } + } +} + #[tokio::test] async fn concurrent_refreshes_share_one_exchange() { let temp_dir = tempfile::tempdir().expect("tempdir"); diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index c25fb9ff2e..362bbd7228 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -1381,28 +1381,23 @@ async fn run_ratatui_app( #[cfg(target_os = "windows")] let mut trust_decision_was_made = false; let startup_model_provider = initial_config.model_provider_id.clone(); - let (login_status, mut startup_account) = if workload_identity_selected { - (LoginStatus::AuthMode(AuthMode::Chatgpt), None) - } else { - let Some(active_app_server) = app_server.as_mut() else { - unreachable!("app server should exist when auth is required"); - }; - let login_status = startup_draft - .run_until(&mut tui, get_login_status(active_app_server)) - .await; - match login_status { - Ok(Ok((login_status, account))) => (login_status, Some(account)), - Ok(Err(err)) => { - shutdown_startup_session(app_server.take(), &mut terminal_restore_guard).await; - return Err(err); - } - Err(err) => { - shutdown_startup_session(app_server.take(), &mut terminal_restore_guard).await; - return Err(err.into()); - } + let Some(active_app_server) = app_server.as_mut() else { + unreachable!("app server should exist when auth is required"); + }; + let login_status = startup_draft + .run_until(&mut tui, get_login_status(active_app_server)) + .await; + let (login_status, mut startup_account) = match login_status { + Ok(Ok((login_status, account))) => (login_status, Some(account)), + Ok(Err(err)) => { + shutdown_startup_session(app_server.take(), &mut terminal_restore_guard).await; + return Err(err); + } + Err(err) => { + shutdown_startup_session(app_server.take(), &mut terminal_restore_guard).await; + return Err(err.into()); } }; - // Workload identity bypasses interactive login; every other provider uses account/read. let requires_openai_auth = startup_account .as_ref() .is_some_and(|account| account.requires_openai_auth); diff --git a/codex-rs/workload-identity/src/exchange.rs b/codex-rs/workload-identity/src/exchange.rs index d91078650c..147dca8761 100644 --- a/codex-rs/workload-identity/src/exchange.rs +++ b/codex-rs/workload-identity/src/exchange.rs @@ -306,13 +306,32 @@ impl CachedToken { #[derive(Clone, PartialEq, Eq)] pub struct WorkloadIdentityToken { pub access_token: String, + pub expires_in: u64, + pub metadata: WorkloadIdentityMetadata, + version: u64, +} + +/// Product and identity resolved by the federation rule's token exchange. +#[derive(Clone, PartialEq, Eq)] +pub enum WorkloadIdentityMetadata { + Chatgpt(ChatgptTokenMetadata), + Api(ApiTokenMetadata), +} + +#[derive(Clone, Deserialize, PartialEq, Eq)] +pub struct ApiTokenMetadata { + pub api_organization_id: String, + pub api_project_id: String, + pub user_id: String, +} + +#[derive(Clone, Deserialize, PartialEq, Eq)] +pub struct ChatgptTokenMetadata { pub chatgpt_account_id: String, pub chatgpt_account_user_id: String, pub chatgpt_plan_type: Option, - pub expires_in: u64, pub scope: String, pub user_id: String, - version: u64, } impl WorkloadIdentityToken { @@ -332,17 +351,25 @@ impl fmt::Debug for WorkloadIdentityToken { } } +#[derive(Default, Deserialize)] +#[serde(rename_all = "snake_case")] +enum WorkloadIdentityProduct { + #[default] + Chatgpt, + Api, +} + #[derive(Deserialize)] struct TokenExchangeResponse { access_token: String, - chatgpt_account_id: String, - chatgpt_account_user_id: String, - chatgpt_plan_type: Option, expires_in: u64, issued_token_type: String, - scope: String, token_type: String, - user_id: String, + // Older ChatGPT token endpoints omit the product marker. + #[serde(default)] + wif_product: WorkloadIdentityProduct, + #[serde(flatten)] + metadata: serde_json::Value, } impl TokenExchangeResponse { @@ -353,25 +380,42 @@ impl TokenExchangeResponse { || !self.token_type.eq_ignore_ascii_case("bearer") || lifetime.is_zero() || lifetime > MAX_ACCESS_TOKEN_LIFETIME - || self.scope.trim().is_empty() - || self.chatgpt_account_id.trim().is_empty() - || self.chatgpt_account_user_id.trim().is_empty() - || self.user_id.trim().is_empty() - || self - .chatgpt_plan_type - .as_deref() - .is_some_and(|plan_type| plan_type.trim().is_empty()) { return Err(WorkloadIdentityError::InvalidExchangeResponse); } + let metadata = match self.wif_product { + WorkloadIdentityProduct::Chatgpt => { + let metadata: ChatgptTokenMetadata = serde_json::from_value(self.metadata) + .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?; + if metadata.scope.trim().is_empty() + || metadata.chatgpt_account_id.trim().is_empty() + || metadata.chatgpt_account_user_id.trim().is_empty() + || metadata.user_id.trim().is_empty() + || metadata + .chatgpt_plan_type + .as_deref() + .is_some_and(|plan_type| plan_type.trim().is_empty()) + { + return Err(WorkloadIdentityError::InvalidExchangeResponse); + } + WorkloadIdentityMetadata::Chatgpt(metadata) + } + WorkloadIdentityProduct::Api => { + let metadata: ApiTokenMetadata = serde_json::from_value(self.metadata) + .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?; + if metadata.api_organization_id.trim().is_empty() + || metadata.api_project_id.trim().is_empty() + || metadata.user_id.trim().is_empty() + { + return Err(WorkloadIdentityError::InvalidExchangeResponse); + } + WorkloadIdentityMetadata::Api(metadata) + } + }; Ok(WorkloadIdentityToken { access_token: self.access_token, - chatgpt_account_id: self.chatgpt_account_id, - chatgpt_account_user_id: self.chatgpt_account_user_id, - chatgpt_plan_type: self.chatgpt_plan_type, expires_in: self.expires_in, - scope: self.scope, - user_id: self.user_id, + metadata, version: 0, }) } diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs index e693772b19..14415b643f 100644 --- a/codex-rs/workload-identity/src/lib.rs +++ b/codex-rs/workload-identity/src/lib.rs @@ -4,11 +4,14 @@ mod exchange; use std::path::PathBuf; use std::sync::Arc; +pub use exchange::ApiTokenMetadata; +pub use exchange::ChatgptTokenMetadata; pub use exchange::WorkloadIdentityExchange; +pub use exchange::WorkloadIdentityMetadata; pub use exchange::WorkloadIdentityToken; use thiserror::Error; -/// The inputs needed to exchange a file-backed assertion for ChatGPT auth. +/// The inputs needed to exchange a file-backed assertion for workload identity auth. #[derive(Clone)] pub struct WorkloadIdentityConfig { pub(crate) assertion_file: PathBuf, diff --git a/codex-rs/workload-identity/src/workload_identity_tests.rs b/codex-rs/workload-identity/src/workload_identity_tests.rs index 59ca8f75a3..f64dbb88b1 100644 --- a/codex-rs/workload-identity/src/workload_identity_tests.rs +++ b/codex-rs/workload-identity/src/workload_identity_tests.rs @@ -20,8 +20,11 @@ use super::ACCESS_TOKEN_TYPE; use super::JWT_BEARER_GRANT_TYPE; use super::WorkloadIdentityExchange; use super::WorkloadIdentityToken; +use crate::ApiTokenMetadata; +use crate::ChatgptTokenMetadata; use crate::WorkloadIdentityConfig; use crate::WorkloadIdentityError; +use crate::WorkloadIdentityMetadata; fn assertion_file(assertion: &str) -> (TempDir, PathBuf) { let temp_dir = TempDir::new().expect("tempdir"); @@ -44,8 +47,8 @@ fn make_exchange(path: PathBuf, server: &MockServer) -> WorkloadIdentityExchange .expect("valid exchange") } -fn success(access_token: &str, expires_in: u64) -> ResponseTemplate { - ResponseTemplate::new(200).set_body_json(serde_json::json!({ +fn success_body(access_token: &str, expires_in: u64) -> serde_json::Value { + serde_json::json!({ "access_token": access_token, "issued_token_type": ACCESS_TOKEN_TYPE, "token_type": "Bearer", @@ -55,51 +58,128 @@ fn success(access_token: &str, expires_in: u64) -> ResponseTemplate { "chatgpt_account_user_id": "membership-one", "user_id": "user-one", "chatgpt_plan_type": "enterprise" - })) + }) +} + +fn success(access_token: &str, expires_in: u64) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(success_body(access_token, expires_in)) +} + +fn api_success_body(access_token: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access_token, + "issued_token_type": ACCESS_TOKEN_TYPE, + "token_type": "Bearer", + "expires_in": 600, + "wif_product": "api", + "api_organization_id": "org-one", + "api_project_id": "proj-one", + "user_id": "service-account-one" + }) } #[tokio::test] -async fn exchange_sends_three_field_contract_and_caches_valid_response() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/oauth/token")) - .and(header("content-type", "application/x-www-form-urlencoded")) - .respond_with(success("sensitive-access-token", /*expires_in*/ 600)) - .mount(&server) - .await; - let (_temp_dir, assertion_path) = assertion_file("assertion-one\n"); - let exchange = make_exchange(assertion_path, &server); - - let expected = WorkloadIdentityToken { - access_token: "sensitive-access-token".to_string(), - chatgpt_account_id: "workspace-one".to_string(), - chatgpt_account_user_id: "membership-one".to_string(), - chatgpt_plan_type: Some("enterprise".to_string()), - expires_in: 600, - scope: "openid profile email chatgpt.workspace.feature.allow-codex-local-access.access" - .to_string(), - user_id: "user-one".to_string(), - version: 1, - }; - assert_eq!(exchange.resolve().await.expect("exchange"), expected); - assert_eq!(exchange.resolve().await.expect("cached token"), expected); - assert!(!format!("{expected:?}").contains("sensitive-access-token")); - - let requests = server.received_requests().await.expect("requests"); - assert_eq!(requests.len(), 1); - assert_eq!( - url::form_urlencoded::parse(&requests[0].body) - .into_owned() - .collect::>(), - vec![ - ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()), - ("assertion".to_string(), "assertion-one".to_string()), +async fn exchange_sends_same_contract_for_both_products_and_caches_valid_response() { + for product in [None, Some("chatgpt"), Some("api")] { + let (mut body, metadata) = if product == Some("api") { ( - "federation_rule_id".to_string(), - "idpm_rule_one".to_string() - ), - ] - ); + api_success_body("sensitive-access-token"), + WorkloadIdentityMetadata::Api(ApiTokenMetadata { + api_organization_id: "org-one".to_string(), + api_project_id: "proj-one".to_string(), + user_id: "service-account-one".to_string(), + }), + ) + } else { + ( + success_body("sensitive-access-token", /*expires_in*/ 600), + WorkloadIdentityMetadata::Chatgpt(ChatgptTokenMetadata { + chatgpt_account_id: "workspace-one".to_string(), + chatgpt_account_user_id: "membership-one".to_string(), + chatgpt_plan_type: Some("enterprise".to_string()), + scope: "openid profile email chatgpt.workspace.feature.allow-codex-local-access.access" + .to_string(), + user_id: "user-one".to_string(), + }), + ) + }; + if let Some(product) = product { + body["wif_product"] = product.into(); + } + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(header("content-type", "application/x-www-form-urlencoded")) + .respond_with(ResponseTemplate::new(200).set_body_json(body)) + .mount(&server) + .await; + let (_temp_dir, assertion_path) = assertion_file("assertion-one\n"); + let exchange = make_exchange(assertion_path, &server); + + let expected = WorkloadIdentityToken { + access_token: "sensitive-access-token".to_string(), + expires_in: 600, + metadata, + version: 1, + }; + assert_eq!(exchange.resolve().await.expect("exchange"), expected); + assert_eq!(exchange.resolve().await.expect("cached token"), expected); + assert!(!format!("{expected:?}").contains("sensitive-access-token")); + + let requests = server.received_requests().await.expect("requests"); + assert_eq!(requests.len(), 1); + assert_eq!( + url::form_urlencoded::parse(&requests[0].body) + .into_owned() + .collect::>(), + vec![ + ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()), + ("assertion".to_string(), "assertion-one".to_string()), + ( + "federation_rule_id".to_string(), + "idpm_rule_one".to_string() + ), + ] + ); + } +} + +#[tokio::test] +async fn exchange_rejects_unknown_products_and_incomplete_api_identity() { + let mut invalid_responses = Vec::new(); + for product in [serde_json::json!("unknown"), serde_json::Value::Null] { + let mut body = success_body("access-one", /*expires_in*/ 600); + body["wif_product"] = product; + invalid_responses.push(body); + } + for field in [ + "wif_product", + "api_organization_id", + "api_project_id", + "user_id", + ] { + let mut body = api_success_body("access-one"); + body.as_object_mut().expect("response object").remove(field); + invalid_responses.push(body); + let mut body = api_success_body("access-one"); + body[field] = " ".into(); + invalid_responses.push(body); + } + let (_temp_dir, assertion_path) = assertion_file("assertion-one"); + for body in invalid_responses { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(body)) + .expect(1) + .mount(&server) + .await; + let exchange = make_exchange(assertion_path.clone(), &server); + assert!(matches!( + exchange.resolve().await, + Err(WorkloadIdentityError::InvalidExchangeResponse) + )); + } } #[tokio::test]