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

Support API workload identity authentication (#59275)

Browse Source 
## Why

Workload identity authentication assumed every token exchange resolved to ChatGPT, preventing API service accounts from using the same flow.

## What changed

- Accept `wif_product: "api"` responses with organization, project, and service account metadata. Preserve ChatGPT compatibility when `wif_product` is omitted.
- Resolve API workload identity tokens as API key authentication and apply login restrictions to the resolved product. Read the actual login status during TUI startup instead of assuming ChatGPT.
- Preserve auth ownership across token renewal for the same API identity, and reject product or identity changes during refresh.
- Redact API keys from debug output.

## Testing

Add coverage for both exchange products, malformed metadata, assertion rotation and unauthorized recovery, identity changes, login and workspace restrictions, and CLI login status under both forced login methods.

GitOrigin-RevId: f4fc4da74b86e4a9ae47b28221d3e065500f5b6e
This commit is contained in:
cooper-oai
2026-11-04 00:52:58 +00:00
committed by copyberry
parent c141c31323
commit f39e45e4a6
11 changed files with 615 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/cli/tests/login.rs
+17 -8
View File
@@ -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(())
}
codex-rs/login/src/auth/change_state.rs
+5 -2
View File
@@ -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;
codex-rs/login/src/auth/change_state_tests.rs
+47
View File
@@ -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() {
codex-rs/login/src/auth/manager.rs
+25 -3
View File
@@ -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<ApiTokenMetadata>,
}
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<TokenData, std::io::Error> {
let auth_dot_json: Option<AuthDotJson> = 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()
}
codex-rs/login/src/auth/source.rs
+5
View File
@@ -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;
codex-rs/login/src/auth/workload_identity.rs
+58 -33
View File
@@ -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<Option<WorkloadIdentitySessionConfig>, 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<CodexAuth> {
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)
}
codex-rs/login/src/auth/workload_identity_tests.rs
+252 -38
View File
@@ -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, &registry)
.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<_, _>>(),
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<Option<WorkloadIdentitySessionConfig>, 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),
&registry,
)
.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),
&registry,
)
.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");
codex-rs/tui/src/lib.rs
+15 -20
View File
@@ -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);
codex-rs/workload-identity/src/exchange.rs
+64 -20
View File
@@ -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<String>,
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<String>,
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,
})
}
codex-rs/workload-identity/src/lib.rs
+4 -1
View File
@@ -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,
codex-rs/workload-identity/src/workload_identity_tests.rs
+123 -43
View File
@@ -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<_>>(),
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<_>>(),
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]