From 3ff5d6223b4f659805c48b8fe7bd5a4c7475a9d1 Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Mon, 9 Nov 2026 19:10:39 +0000 Subject: [PATCH] Bind API workload identity source reuse to session claims (#60433) ## Why Matching API workload identity metadata alone can preserve a credential source across session-family or provider and mapping revision changes. ## What changed - Compare API workload identity tokens using OAuth source claims so ordinary token refreshes preserve the source while session or identity changes retire it. - Include `wif_runtime_authorization_version`, `wif_provider_revision`, and `wif_mapping_revision` in source comparisons. - Reject API workload identity exchange tokens without parseable source claims and nonempty `iss` and `session_id` values. Evict rejected credentials so a subsequent resolve exchanges them again. ## Testing Add regression coverage for refresh continuity, session-family and revision changes, and rejection and re-exchange of tokens with missing, empty, or whitespace-only `session_id` values. GitOrigin-RevId: 67189c27efae8a233250a320e6d4fac71d865e90 --- codex-rs/login/src/auth/change_state_tests.rs | 66 ++++++++++++ codex-rs/login/src/auth/source.rs | 57 +++++----- codex-rs/login/src/auth/workload_identity.rs | 15 ++- .../login/src/auth/workload_identity_tests.rs | 100 ++++++++++++++++-- 4 files changed, 199 insertions(+), 39 deletions(-) diff --git a/codex-rs/login/src/auth/change_state_tests.rs b/codex-rs/login/src/auth/change_state_tests.rs index 873b8e6fe0..0b2f2c89ba 100644 --- a/codex-rs/login/src/auth/change_state_tests.rs +++ b/codex-rs/login/src/auth/change_state_tests.rs @@ -115,6 +115,72 @@ fn api_workload_identity_changes_retire_previous_owner() { } } +#[test] +fn platform_wif_refresh_preserves_source_but_family_changes_retire_it() { + let claims = json!({ + "iss": "https://auth.openai.com", + "sub": "service-account", + "session_id": "family-a", + "iat": 100, + "exp": 200, + "jti": "token-a", + "https://api.openai.com/auth": { + "user_id": "service-account", + "wif_product": "api", + "wif_identity_provider_id": "provider-a", + "wif_federation_rule_id": "mapping-a", + "wif_runtime_authorization_version": 1, + "wif_provider_revision": "provider-revision-a", + "wif_mapping_revision": "mapping-revision-a" + } + }); + let auth = |claims: &serde_json::Value| { + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + CodexAuth::ApiKey(ApiKeyAuth { + api_key: format!("header.{payload}.signature"), + workload_identity: Some(ApiTokenMetadata { + api_organization_id: "organization".into(), + api_project_id: "project".into(), + user_id: "service-account".into(), + }), + }) + }; + let mut refreshed = claims.clone(); + refreshed["iat"] = json!(200); + refreshed["exp"] = json!(300); + refreshed["jti"] = json!("token-b"); + let mut new_family = refreshed.clone(); + new_family["session_id"] = json!("family-b"); + let mut new_provider = refreshed.clone(); + new_provider["https://api.openai.com/auth"]["wif_provider_revision"] = + json!("provider-revision-b"); + let mut new_mapping = refreshed.clone(); + new_mapping["https://api.openai.com/auth"]["wif_mapping_revision"] = + json!("mapping-revision-b"); + let mut missing_family = refreshed.clone(); + missing_family.as_object_mut().unwrap().remove("session_id"); + let mut empty_family = refreshed.clone(); + empty_family["session_id"] = json!(""); + for (replacement, same_source) in [ + (refreshed, true), + (new_family, false), + (new_provider, false), + (new_mapping, false), + (missing_family, false), + (empty_family, false), + ] { + let manager = AuthManager::from_auth_for_testing(auth(&claims)); + let before = *manager.auth_change_state_receiver().borrow(); + manager.set_cached_auth(Some(auth(&replacement))); + let after = *manager.auth_change_state_receiver().borrow(); + assert_eq!( + after.source_generation == before.source_generation, + same_source + ); + assert_eq!(after.generation, before.generation + 1); + } +} + #[tokio::test] #[serial_test::serial(codex_auth_env)] async fn auth_reload_and_refresh_read_failure_preserves_network_permit() { diff --git a/codex-rs/login/src/auth/source.rs b/codex-rs/login/src/auth/source.rs index b50fdbfd70..2233a9c28d 100644 --- a/codex-rs/login/src/auth/source.rs +++ b/codex-rs/login/src/auth/source.rs @@ -19,8 +19,8 @@ pub(super) fn same_source(previous: Option<&CodexAuth>, current: Option<&CodexAu } 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; + if previous_subject != current_subject { + return false; } if let (CodexAuth::AgentIdentity(previous), CodexAuth::AgentIdentity(current)) = (previous, current) @@ -34,30 +34,34 @@ pub(super) fn same_source(previous: Option<&CodexAuth>, current: Option<&CodexAu if previous_token == current_token { return true; } - match previous { - CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { - if previous_token.split('.').count() != 3 || current_token.split('.').count() != 3 { - return false; - } - let (Ok(previous), Ok(current)) = ( - decode_jwt_payload::(&previous_token), - decode_jwt_payload::(¤t_token), - ) else { - // Sessionless credentials must refresh during bootstrap before a - // source-bound runtime can be registered. Never guess continuity. - return false; - }; - !previous.iss.trim().is_empty() - && !previous.session_id.trim().is_empty() - && previous == current - } - CodexAuth::ApiKey(_) - | CodexAuth::PersonalAccessToken(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::BedrockApiKey(_) - | CodexAuth::BedrockAccessKeys(_) => false, + let platform_wif = previous_subject.is_some(); + if !platform_wif + && !matches!( + previous, + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) + ) + { + return false; } + match ( + parse_oauth_source(&previous_token), + parse_oauth_source(¤t_token), + ) { + (Some(previous), Some(current)) => previous == current, + _ => false, + } +} + +pub(super) fn has_session_claim(token: &str) -> bool { + parse_oauth_source(token).is_some() +} + +fn parse_oauth_source(token: &str) -> Option { + if token.split('.').count() != 3 { + return None; + } + let source = decode_jwt_payload::(token).ok()?; + (!source.iss.trim().is_empty() && !source.session_id.trim().is_empty()).then_some(source) } #[derive(Deserialize, PartialEq, Eq)] @@ -84,4 +88,7 @@ struct OAuthIdentity { wif_external_issuer: Option, wif_principal_kind: Option, wif_product: Option, + wif_runtime_authorization_version: Option, + wif_provider_revision: Option, + wif_mapping_revision: Option, } diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs index 6ed55bf158..c33009cce6 100644 --- a/codex-rs/login/src/auth/workload_identity.rs +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -435,10 +435,17 @@ impl WorkloadIdentityExternalAuth { previous_account_id: Option<&str>, ) -> std::io::Result { let auth = match &token.metadata { - WorkloadIdentityMetadata::Api(metadata) => CodexAuth::ApiKey(ApiKeyAuth { - api_key: token.access_token.clone(), - workload_identity: Some(metadata.clone()), - }), + WorkloadIdentityMetadata::Api(metadata) => { + if !super::source::has_session_claim(&token.access_token) { + return Err(std::io::Error::other( + WorkloadIdentityError::InvalidExchangeResponse, + )); + } + 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, diff --git a/codex-rs/login/src/auth/workload_identity_tests.rs b/codex-rs/login/src/auth/workload_identity_tests.rs index 750031f2b8..fe635dd589 100644 --- a/codex-rs/login/src/auth/workload_identity_tests.rs +++ b/codex-rs/login/src/auth/workload_identity_tests.rs @@ -42,6 +42,27 @@ fn api_response(access_token: &str) -> serde_json::Value { }) } +fn api_access_token(assertion: &str) -> String { + let claims = serde_json::json!({ + "iss": "https://auth.api.openai.org", + "sub": "service-account-one", + "session_id": "family-one", + "jti": assertion, + "https://api.openai.com/auth": { + "user_id": "service-account-one", + "wif_product": "api", + "wif_identity_provider_id": "provider-one", + "wif_federation_rule_id": "rule-one", + "wif_runtime_authorization_version": 1, + "wif_provider_revision": "provider-revision-one", + "wif_mapping_revision": "mapping-revision-one" + } + }); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&claims).expect("serialize API WIF claims")); + format!("e30.{payload}.sig") +} + #[tokio::test] async fn api_platform_auth_refreshes_with_a_rotated_assertion() { let temp_dir = tempfile::tempdir().expect("tempdir"); @@ -54,8 +75,7 @@ async fn api_platform_auth_refreshes_with_a_rotated_assertion() { .find(|(name, _)| name == "assertion") .expect("assertion") .1; - ResponseTemplate::new(200) - .set_body_json(api_response(&format!("access-for-{assertion}"))) + ResponseTemplate::new(200).set_body_json(api_response(&api_access_token(&assertion))) }) .expect(2) .mount(&server) @@ -79,7 +99,7 @@ async fn api_platform_auth_refreshes_with_a_rotated_assertion() { .expect("initial auth") .get_token() .expect("token"), - "access-for-assertion-one" + api_access_token("assertion-one") ); let before = *manager.auth_change_state_receiver().borrow(); std::fs::write(&assertion_file, "assertion-two").expect("rotate assertion"); @@ -96,7 +116,7 @@ async fn api_platform_auth_refreshes_with_a_rotated_assertion() { .expect("renewed auth") .get_token() .expect("token"), - "access-for-assertion-two" + api_access_token("assertion-two") ); assert_eq!( *manager.auth_change_state_receiver().borrow(), @@ -127,6 +147,64 @@ async fn api_platform_auth_refreshes_with_a_rotated_assertion() { } } +#[tokio::test] +async fn api_platform_auth_requires_a_runtime_family() { + for family in [None, Some(""), Some(" ")] { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let mut claims: serde_json::Value = + crate::token_data::decode_jwt_payload(&api_access_token("unsupported")) + .expect("valid API WIF claims"); + if let Some(family) = family { + claims["session_id"] = family.into(); + } else { + claims + .as_object_mut() + .expect("claims object") + .remove("session_id"); + } + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&claims).expect("serialize API WIF claims")); + let invalid = format!("e30.{payload}.sig"); + let server = MockServer::start().await; + let count = AtomicUsize::new(0); + Mock::given(method("POST")) + .respond_with(move |_request: &wiremock::Request| { + let token = if count.fetch_add(1, Ordering::SeqCst) == 0 { + invalid.clone() + } else { + api_access_token("supported") + }; + ResponseTemplate::new(200).set_body_json(api_response(&token)) + }) + .expect(2) + .mount(&server) + .await; + let registry = WorkloadIdentitySessionRegistry::default(); + let adapter = WorkloadIdentityExternalAuth::from_config_with_registry( + session_config(temp_dir.path(), &server), + ®istry, + ) + .expect("adapter"); + let error = adapter + .resolve() + .await + .expect_err("Platform WIF requires a nonempty runtime family"); + assert!(matches!( + adapter.classify_error(error), + RefreshTokenError::Permanent(_) + )); + assert_eq!( + adapter + .resolve() + .await + .expect("unsupported credentials are evicted and re-exchanged") + .get_token() + .expect("token"), + api_access_token("supported") + ); + } +} + fn resolve_for_test( environment: ProcessEnvironment, chatgpt_base_url: &str, @@ -436,16 +514,18 @@ async fn api_refresh_preserves_product_and_identity() { 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")), + 0 => ResponseTemplate::new(200) + .set_body_json(api_response(&api_access_token("access-one"))), 1 => match changed_field { Some(field) => { - let mut body = api_response("access-two"); + let mut body = api_response(&api_access_token("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")), + _ => ResponseTemplate::new(200) + .set_body_json(api_response(&api_access_token("access-three"))), } }) .mount(&server) @@ -458,7 +538,7 @@ async fn api_refresh_preserves_product_and_identity() { .expect("adapter"); assert_eq!( adapter.resolve().await.expect("initial API auth"), - CodexAuth::from_api_key("access-one") + CodexAuth::from_api_key(&api_access_token("access-one")) ); let error = adapter .refresh(ExternalAuthRefreshContext { @@ -476,7 +556,7 @@ async fn api_refresh_preserves_product_and_identity() { .resolve() .await .expect("corrected token is re-exchanged"), - CodexAuth::from_api_key("access-three") + CodexAuth::from_api_key(&api_access_token("access-three")) ); assert_eq!(request_count.load(Ordering::SeqCst), 3); } @@ -512,7 +592,7 @@ async fn resolved_product_obeys_login_and_workspace_restrictions() { 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")) + ResponseTemplate::new(200).set_body_json(api_response(&api_access_token("api-access"))) } else { success_response("chatgpt-access", "account-user-one", "user-one") };