openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Prototype Azure workload identity auth [ci changed_files]

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