openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-09 20:50:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bind API workload identity source reuse to session claims (#60433)

Browse Source 
## 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
This commit is contained in:
cooper-oai
2026-11-09 19:26:55 +00:00
committed by copyberry
parent 2728f99126
commit 3ff5d6223b
4 changed files with 199 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/login/src/auth/change_state_tests.rs
+66
View File
@@ -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() {
codex-rs/login/src/auth/source.rs
+32 -25
View File
@@ -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::<OAuthSource>(&previous_token),
decode_jwt_payload::<OAuthSource>(&current_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(&current_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<OAuthSource> {
if token.split('.').count() != 3 {
return None;
}
let source = decode_jwt_payload::<OAuthSource>(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<String>,
wif_principal_kind: Option<String>,
wif_product: Option<String>,
wif_runtime_authorization_version: Option<u32>,
wif_provider_revision: Option<String>,
wif_mapping_revision: Option<String>,
}
codex-rs/login/src/auth/workload_identity.rs
+11 -4
View File
@@ -435,10 +435,17 @@ impl WorkloadIdentityExternalAuth {
previous_account_id: Option<&str>,
) -> std::io::Result<CodexAuth> {
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,
codex-rs/login/src/auth/workload_identity_tests.rs
+90 -10
View File
@@ -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),
&registry,
)
.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")
};