mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Integrate workload identity with Codex authentication (#38188)
## What changed - Select workload identity authentication from `OPENAI_FEDERATION_RULE_ID` and `OPENAI_IDENTITY_TOKEN_FILE`, while preserving explicit API key and access token precedence. - Adapt workload identity exchanges to the shared authentication manager and reuse a compatible process-scoped session for token resolution and refresh. - Reject incomplete, conflicting, or unsupported configurations and exchanged tokens that change the authenticated subject. - Classify exchange failures for authentication retry handling and invalidate rejected cached tokens without discarding a newer concurrent exchange. ## Testing - Cover configuration selection and validation, session reuse and conflicts, identity-preserving refreshes, concurrent refreshes, and retry classification. GitOrigin-RevId: 62afb514b4e0a59effb08f464ac929633462a71b
This commit is contained in:
cooper-oai
Generated
+1
@@ -3475,6 +3475,7 @@ dependencies = [
|
||||
"codex-secrets",
|
||||
"codex-terminal-detection",
|
||||
"codex-utils-template",
|
||||
"codex-workload-identity",
|
||||
"core_test_support",
|
||||
"http 1.4.0",
|
||||
"jsonwebtoken",
|
||||
|
||||
@@ -275,6 +275,7 @@ codex-utils-stream-parser = { path = "utils/stream-parser" }
|
||||
codex-utils-string = { path = "utils/string" }
|
||||
codex-utils-template = { path = "utils/template" }
|
||||
codex-v8-poc = { path = "v8-poc" }
|
||||
codex-workload-identity = { path = "workload-identity" }
|
||||
codex-windows-sandbox = { path = "windows-sandbox-rs" }
|
||||
core_test_support = { path = "core/tests/common" }
|
||||
mcp_test_support = { path = "mcp-server/tests/common" }
|
||||
|
||||
@@ -20,6 +20,7 @@ codex-protocol = { workspace = true }
|
||||
codex-secrets = { workspace = true }
|
||||
codex-terminal-detection = { workspace = true }
|
||||
codex-utils-template = { workspace = true }
|
||||
codex-workload-identity = { workspace = true }
|
||||
http = { workspace = true }
|
||||
once_cell = { workspace = true }
|
||||
os_info = { workspace = true }
|
||||
|
||||
@@ -11,6 +11,8 @@ use codex_protocol::protocol::SessionSource;
|
||||
use base64::Engine;
|
||||
use codex_protocol::config_types::ForcedLoginMethod;
|
||||
use codex_protocol::config_types::ModelProviderAuthInfo;
|
||||
use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR;
|
||||
use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde::Serialize;
|
||||
use serde_json::json;
|
||||
@@ -1702,6 +1704,96 @@ fn remove_access_token_env_var() -> EnvVarGuard {
|
||||
EnvVarGuard::remove(CODEX_ACCESS_TOKEN_ENV_VAR)
|
||||
}
|
||||
|
||||
struct TestAuthManagerConfig(AuthConfig);
|
||||
|
||||
impl AuthManagerConfig for TestAuthManagerConfig {
|
||||
fn codex_home(&self) -> PathBuf {
|
||||
self.0.codex_home.clone()
|
||||
}
|
||||
|
||||
fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode {
|
||||
self.0.auth_credentials_store_mode
|
||||
}
|
||||
|
||||
fn auth_keyring_backend_kind(&self) -> AuthKeyringBackendKind {
|
||||
self.0.keyring_backend_kind
|
||||
}
|
||||
|
||||
fn forced_login_method(&self) -> Option<ForcedLoginMethod> {
|
||||
self.0.forced_login_method
|
||||
}
|
||||
|
||||
fn forced_chatgpt_workspace_id(&self) -> Option<Vec<String>> {
|
||||
self.0.forced_chatgpt_workspace_id.clone()
|
||||
}
|
||||
|
||||
fn managed_auth_policy(&self) -> ManagedAuthPolicy {
|
||||
self.0.managed_auth_policy.clone()
|
||||
}
|
||||
|
||||
fn chatgpt_base_url(&self) -> String {
|
||||
self.0
|
||||
.chatgpt_base_url
|
||||
.clone()
|
||||
.expect("test config should include a ChatGPT base URL")
|
||||
}
|
||||
|
||||
fn auth_route_config(&self) -> AuthRouteConfig {
|
||||
self.0.auth_route_config.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn test_auth_manager_config(codex_home: &Path) -> TestAuthManagerConfig {
|
||||
TestAuthManagerConfig(AuthConfig {
|
||||
codex_home: codex_home.to_path_buf(),
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
|
||||
keyring_backend_kind: AuthKeyringBackendKind::Direct,
|
||||
forced_login_method: Some(ForcedLoginMethod::Chatgpt),
|
||||
chatgpt_base_url: Some("https://chatgpt-staging.com/backend-api".to_string()),
|
||||
forced_chatgpt_workspace_id: Some(vec!["forced-workspace".to_string()]),
|
||||
managed_auth_policy: ManagedAuthPolicy {
|
||||
allowed_login_methods: Some(vec![ForcedLoginMethod::Chatgpt]),
|
||||
allowed_chatgpt_workspaces: Some(vec![
|
||||
"forced-workspace".to_string(),
|
||||
"managed-workspace".to_string(),
|
||||
]),
|
||||
},
|
||||
auth_route_config: AuthRouteConfig::from_http_client_factory(HttpClientFactory::new(
|
||||
OutboundProxyPolicy::RespectSystemProxy,
|
||||
)),
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_config_from_preserves_all_fields() {
|
||||
let codex_home = tempdir().expect("tempdir");
|
||||
let config = test_auth_manager_config(codex_home.path());
|
||||
|
||||
assert_eq!(auth_config_from(&config), config.0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn try_shared_from_config_rejects_partial_workload_identity_configuration() {
|
||||
let codex_home = tempdir().expect("tempdir");
|
||||
let config = test_auth_manager_config(codex_home.path());
|
||||
let _access_token_guard = remove_access_token_env_var();
|
||||
let _rule_guard = EnvVarGuard::set(OPENAI_FEDERATION_RULE_ID_ENV_VAR, "rule-one");
|
||||
let _assertion_file_guard = EnvVarGuard::remove(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR);
|
||||
|
||||
let error =
|
||||
AuthManager::try_shared_from_config(&config, /*enable_codex_api_key_env*/ false)
|
||||
.await
|
||||
.expect_err("partial workload identity config should fail closed");
|
||||
|
||||
assert!(matches!(error, RefreshTokenError::Permanent(_)));
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn load_auth_reads_access_token_from_env() {
|
||||
|
||||
@@ -38,6 +38,8 @@ use super::agent_identity::require_agent_identity_authapi_base_url;
|
||||
use super::agent_identity::verified_record_from_jwt;
|
||||
use super::external_bearer::BearerTokenRefresher;
|
||||
use super::revoke::revoke_auth_tokens;
|
||||
use super::workload_identity::WorkloadIdentityExternalAuth;
|
||||
use super::workload_identity::WorkloadIdentitySessionError;
|
||||
use crate::auth::AuthHeaders;
|
||||
pub use crate::auth::agent_identity::AgentIdentityAuth;
|
||||
pub use crate::auth::agent_identity::AgentIdentityAuthError;
|
||||
@@ -2672,20 +2674,7 @@ impl AuthManager {
|
||||
config: &impl AuthManagerConfig,
|
||||
enable_codex_api_key_env: bool,
|
||||
) -> Arc<Self> {
|
||||
Self::shared_from_auth_config(
|
||||
AuthConfig {
|
||||
codex_home: config.codex_home(),
|
||||
auth_credentials_store_mode: config.cli_auth_credentials_store_mode(),
|
||||
keyring_backend_kind: config.auth_keyring_backend_kind(),
|
||||
forced_login_method: config.forced_login_method(),
|
||||
chatgpt_base_url: Some(config.chatgpt_base_url()),
|
||||
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id(),
|
||||
managed_auth_policy: config.managed_auth_policy(),
|
||||
auth_route_config: config.auth_route_config(),
|
||||
},
|
||||
enable_codex_api_key_env,
|
||||
)
|
||||
.await
|
||||
Self::shared_from_auth_config(auth_config_from(config), enable_codex_api_key_env).await
|
||||
}
|
||||
|
||||
/// Builds a shared manager using restrictions resolved before authentication.
|
||||
@@ -2696,6 +2685,33 @@ impl AuthManager {
|
||||
Arc::new(Self::new_from_auth_config(auth_config, enable_codex_api_key_env).await)
|
||||
}
|
||||
|
||||
/// Constructs a manager and activates process-configured workload identity when selected.
|
||||
pub async fn try_shared_from_config(
|
||||
config: &impl AuthManagerConfig,
|
||||
enable_codex_api_key_env: bool,
|
||||
) -> Result<Arc<Self>, RefreshTokenError> {
|
||||
Self::try_shared_from_auth_config(auth_config_from(config), enable_codex_api_key_env).await
|
||||
}
|
||||
|
||||
/// Activates workload identity against an auth config resolved before full runtime config.
|
||||
pub async fn try_shared_from_auth_config(
|
||||
auth_config: AuthConfig,
|
||||
enable_codex_api_key_env: bool,
|
||||
) -> Result<Arc<Self>, RefreshTokenError> {
|
||||
let external_auth = WorkloadIdentityExternalAuth::from_process_config(
|
||||
&auth_config,
|
||||
enable_codex_api_key_env,
|
||||
)
|
||||
.map_err(configured_workload_identity_error)?;
|
||||
let manager = Self::shared_from_auth_config(auth_config, enable_codex_api_key_env).await;
|
||||
if let Some(external_auth) = external_auth {
|
||||
manager
|
||||
.set_configured_external_auth(Arc::new(external_auth))
|
||||
.await?;
|
||||
}
|
||||
Ok(manager)
|
||||
}
|
||||
|
||||
pub fn unauthorized_recovery(self: &Arc<Self>) -> UnauthorizedRecovery {
|
||||
UnauthorizedRecovery::new(Arc::clone(self))
|
||||
}
|
||||
@@ -3003,6 +3019,26 @@ impl AuthManager {
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_config_from(config: &impl AuthManagerConfig) -> AuthConfig {
|
||||
AuthConfig {
|
||||
codex_home: config.codex_home(),
|
||||
auth_credentials_store_mode: config.cli_auth_credentials_store_mode(),
|
||||
keyring_backend_kind: config.auth_keyring_backend_kind(),
|
||||
forced_login_method: config.forced_login_method(),
|
||||
chatgpt_base_url: Some(config.chatgpt_base_url()),
|
||||
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id(),
|
||||
managed_auth_policy: config.managed_auth_policy(),
|
||||
auth_route_config: config.auth_route_config(),
|
||||
}
|
||||
}
|
||||
|
||||
fn configured_workload_identity_error(error: WorkloadIdentitySessionError) -> RefreshTokenError {
|
||||
RefreshTokenError::Permanent(RefreshTokenFailedError::new(
|
||||
RefreshTokenFailedReason::Other,
|
||||
error.to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "auth_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -7,6 +7,7 @@ pub mod error;
|
||||
mod personal_access_token;
|
||||
mod storage;
|
||||
mod util;
|
||||
mod workload_identity;
|
||||
|
||||
mod external_bearer;
|
||||
mod manager;
|
||||
@@ -18,3 +19,4 @@ pub use bedrock_api_key::login_with_bedrock_api_key;
|
||||
pub use error::RefreshTokenFailedError;
|
||||
pub use error::RefreshTokenFailedReason;
|
||||
pub use manager::*;
|
||||
pub use workload_identity::is_workload_identity_selected;
|
||||
|
||||
@@ -0,0 +1,466 @@
|
||||
use std::ffi::OsString;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::Weak;
|
||||
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_workload_identity::WorkloadIdentityConfig;
|
||||
use codex_workload_identity::WorkloadIdentityError;
|
||||
use codex_workload_identity::WorkloadIdentityExchange;
|
||||
use codex_workload_identity::WorkloadIdentityToken;
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
use super::AuthConfig;
|
||||
use super::CodexAuth;
|
||||
use super::ExternalAuth;
|
||||
use super::ExternalAuthFuture;
|
||||
use super::ExternalAuthRefreshContext;
|
||||
use super::RefreshTokenError;
|
||||
use super::RefreshTokenFailedError;
|
||||
use super::RefreshTokenFailedReason;
|
||||
use super::read_codex_access_token_from_env;
|
||||
use super::read_codex_api_key_from_env;
|
||||
use crate::AuthRouteConfig;
|
||||
|
||||
const PROD_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
|
||||
const STAGING_TOKEN_URL: &str = "https://auth.api.openai.org/oauth/token";
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
enum WorkloadIdentityEnvironment {
|
||||
Production,
|
||||
Staging,
|
||||
}
|
||||
|
||||
impl WorkloadIdentityEnvironment {
|
||||
fn token_url(self) -> Result<Url, WorkloadIdentitySessionError> {
|
||||
Url::parse(match self {
|
||||
Self::Production => PROD_TOKEN_URL,
|
||||
Self::Staging => STAGING_TOKEN_URL,
|
||||
})
|
||||
.map_err(|_| invalid_config("workload identity token endpoint is invalid"))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct WorkloadIdentitySessionConfig {
|
||||
assertion_file: PathBuf,
|
||||
environment: WorkloadIdentityEnvironment,
|
||||
federation_rule_id: String,
|
||||
http_client_factory: HttpClientFactory,
|
||||
token_url: Url,
|
||||
}
|
||||
|
||||
impl WorkloadIdentitySessionConfig {
|
||||
fn fingerprint(&self) -> WorkloadIdentityFingerprint {
|
||||
WorkloadIdentityFingerprint {
|
||||
assertion_file: self.assertion_file.clone(),
|
||||
environment: self.environment,
|
||||
federation_rule_id: self.federation_rule_id.trim().to_string(),
|
||||
token_url: self.token_url.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn into_exchange(self) -> Result<WorkloadIdentityExchange, WorkloadIdentityError> {
|
||||
WorkloadIdentityExchange::new(
|
||||
WorkloadIdentityConfig::new(self.federation_rule_id, self.assertion_file)?,
|
||||
self.token_url,
|
||||
self.http_client_factory,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
struct WorkloadIdentityFingerprint {
|
||||
assertion_file: PathBuf,
|
||||
environment: WorkloadIdentityEnvironment,
|
||||
federation_rule_id: String,
|
||||
token_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(super) enum WorkloadIdentitySessionError {
|
||||
#[error(transparent)]
|
||||
Exchange(#[from] WorkloadIdentityError),
|
||||
#[error("a different workload identity configuration is already active in this process")]
|
||||
ConflictingConfiguration,
|
||||
#[error("the workload identity process-session registry is unavailable")]
|
||||
RegistryUnavailable,
|
||||
#[error("{0}")]
|
||||
InvalidConfiguration(String),
|
||||
}
|
||||
|
||||
pub fn is_workload_identity_selected(
|
||||
auth_config: &AuthConfig,
|
||||
enable_codex_api_key_env: bool,
|
||||
) -> bool {
|
||||
!has_explicit_process_auth(auth_config, enable_codex_api_key_env)
|
||||
&& ProcessEnvironment::read().has_marker()
|
||||
}
|
||||
|
||||
fn resolve_config(
|
||||
chatgpt_base_url: &str,
|
||||
environment: ProcessEnvironment,
|
||||
has_explicit_process_auth: bool,
|
||||
chatgpt_login_allowed: bool,
|
||||
auth_route_config: AuthRouteConfig,
|
||||
) -> Result<Option<WorkloadIdentitySessionConfig>, WorkloadIdentitySessionError> {
|
||||
if has_explicit_process_auth || !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,
|
||||
OPENAI_FEDERATION_RULE_ID_ENV_VAR,
|
||||
)?;
|
||||
let assertion_file = required_path(
|
||||
environment.identity_token_file,
|
||||
OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR,
|
||||
)?;
|
||||
let auth_environment = classify_auth_environment(chatgpt_base_url)?;
|
||||
|
||||
Ok(Some(WorkloadIdentitySessionConfig {
|
||||
assertion_file,
|
||||
environment: auth_environment,
|
||||
federation_rule_id,
|
||||
http_client_factory: auth_route_config.http_client_factory().clone(),
|
||||
token_url: auth_environment.token_url()?,
|
||||
}))
|
||||
}
|
||||
|
||||
fn has_explicit_process_auth(auth_config: &AuthConfig, enable_codex_api_key_env: bool) -> bool {
|
||||
(enable_codex_api_key_env
|
||||
&& auth_config.is_login_method_allowed(ForcedLoginMethod::Api)
|
||||
&& read_codex_api_key_from_env().is_some())
|
||||
|| (auth_config.is_login_method_allowed(ForcedLoginMethod::Chatgpt)
|
||||
&& read_codex_access_token_from_env().is_some())
|
||||
}
|
||||
|
||||
fn required_path(
|
||||
value: Option<OsString>,
|
||||
variable: &'static str,
|
||||
) -> Result<PathBuf, WorkloadIdentitySessionError> {
|
||||
let path = PathBuf::from(
|
||||
value.ok_or_else(|| invalid_config(format!("workload identity requires {variable}")))?,
|
||||
);
|
||||
if !path.is_absolute() {
|
||||
return Err(invalid_config(format!(
|
||||
"{variable} must be an absolute path"
|
||||
)));
|
||||
}
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
fn required_unicode(
|
||||
value: Option<OsString>,
|
||||
variable: &'static str,
|
||||
) -> Result<String, WorkloadIdentitySessionError> {
|
||||
let value = value
|
||||
.ok_or_else(|| invalid_config(format!("workload identity requires {variable}")))?
|
||||
.into_string()
|
||||
.map_err(|_| invalid_config(format!("workload identity variable {variable} is invalid")))?;
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
return Err(invalid_config(format!(
|
||||
"workload identity variable {variable} is invalid"
|
||||
)));
|
||||
}
|
||||
Ok(value.to_string())
|
||||
}
|
||||
|
||||
fn classify_auth_environment(
|
||||
base_url: &str,
|
||||
) -> Result<WorkloadIdentityEnvironment, WorkloadIdentitySessionError> {
|
||||
match base_url.trim().trim_end_matches('/') {
|
||||
"https://chatgpt.com"
|
||||
| "https://chatgpt.com/backend-api"
|
||||
| "https://chatgpt.com/codex"
|
||||
| "https://chatgpt.com/backend-api/codex"
|
||||
| "https://chat.openai.com"
|
||||
| "https://chat.openai.com/backend-api"
|
||||
| "https://chat.openai.com/codex"
|
||||
| "https://chat.openai.com/backend-api/codex" => {
|
||||
Ok(WorkloadIdentityEnvironment::Production)
|
||||
}
|
||||
"https://chatgpt-staging.com"
|
||||
| "https://chatgpt-staging.com/backend-api"
|
||||
| "https://chatgpt-staging.com/codex"
|
||||
| "https://chatgpt-staging.com/backend-api/codex" => {
|
||||
Ok(WorkloadIdentityEnvironment::Staging)
|
||||
}
|
||||
_ => Err(invalid_config(
|
||||
"workload identity auth supports only trusted production and staging app routing",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_config(message: impl Into<String>) -> WorkloadIdentitySessionError {
|
||||
WorkloadIdentitySessionError::InvalidConfiguration(message.into())
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ProcessEnvironment {
|
||||
federation_rule_id: Option<OsString>,
|
||||
identity_token_file: Option<OsString>,
|
||||
}
|
||||
|
||||
impl ProcessEnvironment {
|
||||
fn read() -> Self {
|
||||
Self {
|
||||
federation_rule_id: std::env::var_os(OPENAI_FEDERATION_RULE_ID_ENV_VAR),
|
||||
identity_token_file: std::env::var_os(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR),
|
||||
}
|
||||
}
|
||||
|
||||
fn has_marker(&self) -> bool {
|
||||
self.federation_rule_id.is_some() || self.identity_token_file.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
struct WorkloadIdentitySubject {
|
||||
account_id: String,
|
||||
account_user_id: String,
|
||||
user_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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct WorkloadIdentitySession {
|
||||
exchange: WorkloadIdentityExchange,
|
||||
subject: Mutex<Option<WorkloadIdentitySubject>>,
|
||||
}
|
||||
|
||||
impl WorkloadIdentitySession {
|
||||
fn new(config: WorkloadIdentitySessionConfig) -> Result<Self, WorkloadIdentitySessionError> {
|
||||
Ok(Self {
|
||||
exchange: config.into_exchange()?,
|
||||
subject: Mutex::new(None),
|
||||
})
|
||||
}
|
||||
|
||||
fn accept_subject(
|
||||
&self,
|
||||
token: &WorkloadIdentityToken,
|
||||
previous_account_id: Option<&str>,
|
||||
) -> Result<(), WorkloadIdentityError> {
|
||||
if previous_account_id.is_some_and(|account_id| account_id != token.chatgpt_account_id) {
|
||||
return Err(WorkloadIdentityError::InvalidExchangeResponse);
|
||||
}
|
||||
let subject = WorkloadIdentitySubject::from(token);
|
||||
let mut current = self
|
||||
.subject
|
||||
.lock()
|
||||
.map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)?;
|
||||
match current.as_ref() {
|
||||
Some(current) if current != &subject => {
|
||||
Err(WorkloadIdentityError::InvalidExchangeResponse)
|
||||
}
|
||||
Some(_) => Ok(()),
|
||||
None => {
|
||||
*current = Some(subject);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct WorkloadIdentitySessionRegistry {
|
||||
entry: Mutex<Option<WorkloadIdentitySessionEntry>>,
|
||||
}
|
||||
|
||||
struct WorkloadIdentitySessionEntry {
|
||||
fingerprint: WorkloadIdentityFingerprint,
|
||||
session: Weak<WorkloadIdentitySession>,
|
||||
}
|
||||
|
||||
impl WorkloadIdentitySessionRegistry {
|
||||
fn has_active_session(&self) -> Result<bool, WorkloadIdentitySessionError> {
|
||||
let entry = self
|
||||
.entry
|
||||
.lock()
|
||||
.map_err(|_| WorkloadIdentitySessionError::RegistryUnavailable)?;
|
||||
Ok(entry
|
||||
.as_ref()
|
||||
.and_then(|active| active.session.upgrade())
|
||||
.is_some())
|
||||
}
|
||||
|
||||
fn session(
|
||||
&self,
|
||||
config: WorkloadIdentitySessionConfig,
|
||||
) -> Result<Arc<WorkloadIdentitySession>, WorkloadIdentitySessionError> {
|
||||
let fingerprint = config.fingerprint();
|
||||
let mut entry = self
|
||||
.entry
|
||||
.lock()
|
||||
.map_err(|_| WorkloadIdentitySessionError::RegistryUnavailable)?;
|
||||
if let Some(active) = entry.as_ref()
|
||||
&& let Some(session) = active.session.upgrade()
|
||||
{
|
||||
if active.fingerprint == fingerprint {
|
||||
return Ok(session);
|
||||
}
|
||||
return Err(WorkloadIdentitySessionError::ConflictingConfiguration);
|
||||
}
|
||||
|
||||
let session = Arc::new(WorkloadIdentitySession::new(config)?);
|
||||
*entry = Some(WorkloadIdentitySessionEntry {
|
||||
fingerprint,
|
||||
session: Arc::downgrade(&session),
|
||||
});
|
||||
Ok(session)
|
||||
}
|
||||
}
|
||||
|
||||
fn process_registry() -> &'static WorkloadIdentitySessionRegistry {
|
||||
static REGISTRY: OnceLock<WorkloadIdentitySessionRegistry> = OnceLock::new();
|
||||
REGISTRY.get_or_init(WorkloadIdentitySessionRegistry::default)
|
||||
}
|
||||
|
||||
pub(super) struct WorkloadIdentityExternalAuth {
|
||||
observed_token_version: AtomicU64,
|
||||
session: Arc<WorkloadIdentitySession>,
|
||||
}
|
||||
|
||||
impl WorkloadIdentityExternalAuth {
|
||||
pub(super) fn from_process_config(
|
||||
auth_config: &AuthConfig,
|
||||
enable_codex_api_key_env: bool,
|
||||
) -> Result<Option<Self>, WorkloadIdentitySessionError> {
|
||||
let registry = process_registry();
|
||||
let active_session = registry.has_active_session()?;
|
||||
resolve_config(
|
||||
auth_config
|
||||
.chatgpt_base_url
|
||||
.as_deref()
|
||||
.unwrap_or("https://chatgpt.com/backend-api"),
|
||||
ProcessEnvironment::read(),
|
||||
has_explicit_process_auth(auth_config, enable_codex_api_key_env) && !active_session,
|
||||
auth_config.is_login_method_allowed(ForcedLoginMethod::Chatgpt),
|
||||
auth_config.auth_route_config.clone(),
|
||||
)?
|
||||
.map(|config| Self::from_config_with_registry(config, registry))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn from_config_with_registry(
|
||||
config: WorkloadIdentitySessionConfig,
|
||||
registry: &WorkloadIdentitySessionRegistry,
|
||||
) -> Result<Self, WorkloadIdentitySessionError> {
|
||||
Ok(Self {
|
||||
observed_token_version: AtomicU64::new(0),
|
||||
session: registry.session(config)?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn build_validated_auth(
|
||||
&self,
|
||||
token: WorkloadIdentityToken,
|
||||
previous_account_id: Option<&str>,
|
||||
) -> std::io::Result<CodexAuth> {
|
||||
let token_version = token.version();
|
||||
let result = self.validate_auth(&token, previous_account_id);
|
||||
if result.is_err() {
|
||||
self.session
|
||||
.exchange
|
||||
.invalidate_if_current(token_version)
|
||||
.await;
|
||||
}
|
||||
let auth = result?;
|
||||
self.observed_token_version
|
||||
.store(token_version, Ordering::Release);
|
||||
Ok(auth)
|
||||
}
|
||||
|
||||
fn validate_auth(
|
||||
&self,
|
||||
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,
|
||||
));
|
||||
}
|
||||
self.session
|
||||
.accept_subject(token, previous_account_id)
|
||||
.map_err(std::io::Error::other)?;
|
||||
Ok(auth)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExternalAuth for WorkloadIdentityExternalAuth {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async move {
|
||||
let token = self
|
||||
.session
|
||||
.exchange
|
||||
.resolve()
|
||||
.await
|
||||
.map_err(std::io::Error::other)?;
|
||||
self.build_validated_auth(token, /*previous_account_id*/ None)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async move {
|
||||
let observed_version = self.observed_token_version.load(Ordering::Acquire);
|
||||
let token = self
|
||||
.session
|
||||
.exchange
|
||||
.refresh(observed_version)
|
||||
.await
|
||||
.map_err(std::io::Error::other)?;
|
||||
self.build_validated_auth(token, context.previous_account_id.as_deref())
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn classify_error(&self, error: std::io::Error) -> RefreshTokenError {
|
||||
if error
|
||||
.get_ref()
|
||||
.and_then(|source| source.downcast_ref::<WorkloadIdentityError>())
|
||||
.is_some_and(WorkloadIdentityError::is_transient)
|
||||
{
|
||||
return RefreshTokenError::Transient(error);
|
||||
}
|
||||
|
||||
RefreshTokenError::Permanent(RefreshTokenFailedError::new(
|
||||
RefreshTokenFailedReason::Other,
|
||||
error.to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "workload_identity_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,404 @@
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::time::Duration;
|
||||
|
||||
use base64::Engine as _;
|
||||
use codex_http_client::OutboundProxyPolicy;
|
||||
use pretty_assertions::assert_eq;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
|
||||
use super::*;
|
||||
use crate::auth::ExternalAuthRefreshReason;
|
||||
|
||||
fn auth_route_config(policy: OutboundProxyPolicy) -> AuthRouteConfig {
|
||||
AuthRouteConfig::from_http_client_factory(HttpClientFactory::new(policy))
|
||||
}
|
||||
|
||||
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()),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_for_test(
|
||||
environment: ProcessEnvironment,
|
||||
has_explicit_process_auth: bool,
|
||||
chatgpt_login_allowed: bool,
|
||||
chatgpt_base_url: &str,
|
||||
) -> Result<Option<WorkloadIdentitySessionConfig>, WorkloadIdentitySessionError> {
|
||||
resolve_config(
|
||||
chatgpt_base_url,
|
||||
environment,
|
||||
has_explicit_process_auth,
|
||||
chatgpt_login_allowed,
|
||||
auth_route_config(OutboundProxyPolicy::ReqwestDefault),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markers_select_wif_and_partial_configuration_fails_closed() {
|
||||
assert!(
|
||||
resolve_for_test(
|
||||
ProcessEnvironment::default(),
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
"https://chatgpt.com/backend-api",
|
||||
)
|
||||
.expect("no markers")
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
resolve_for_test(
|
||||
ProcessEnvironment {
|
||||
federation_rule_id: Some("partial".into()),
|
||||
..Default::default()
|
||||
},
|
||||
/*has_explicit_process_auth*/ true,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
"https://chatgpt.com/backend-api",
|
||||
)
|
||||
.expect("explicit auth wins")
|
||||
.is_none()
|
||||
);
|
||||
|
||||
for (environment, missing) in [
|
||||
(
|
||||
ProcessEnvironment {
|
||||
federation_rule_id: None,
|
||||
..complete_environment()
|
||||
},
|
||||
OPENAI_FEDERATION_RULE_ID_ENV_VAR,
|
||||
),
|
||||
(
|
||||
ProcessEnvironment {
|
||||
identity_token_file: None,
|
||||
..complete_environment()
|
||||
},
|
||||
OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR,
|
||||
),
|
||||
] {
|
||||
let error = resolve_for_test(
|
||||
environment,
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
"https://chatgpt.com/backend-api",
|
||||
)
|
||||
.expect_err("partial WIF must not fall back");
|
||||
assert!(error.to_string().contains(missing), "{error}");
|
||||
}
|
||||
|
||||
let relative = ProcessEnvironment {
|
||||
identity_token_file: Some("relative.jwt".into()),
|
||||
..complete_environment()
|
||||
};
|
||||
assert!(
|
||||
resolve_for_test(
|
||||
relative,
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
"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(),
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ false,
|
||||
"https://chatgpt.com/backend-api",
|
||||
)
|
||||
.expect_err("ChatGPT-disallowing policy");
|
||||
assert!(policy_error.to_string().contains("login policy"));
|
||||
|
||||
for (chatgpt_base_url, expected_environment, expected_token_url) in [
|
||||
(
|
||||
"https://chatgpt.com/backend-api/",
|
||||
WorkloadIdentityEnvironment::Production,
|
||||
PROD_TOKEN_URL,
|
||||
),
|
||||
(
|
||||
"https://chatgpt-staging.com/backend-api",
|
||||
WorkloadIdentityEnvironment::Staging,
|
||||
STAGING_TOKEN_URL,
|
||||
),
|
||||
] {
|
||||
let config = resolve_for_test(
|
||||
complete_environment(),
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
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(),
|
||||
/*has_explicit_process_auth*/ false,
|
||||
/*chatgpt_login_allowed*/ true,
|
||||
"https://example.invalid/backend-api",
|
||||
)
|
||||
.expect_err("untrusted auth environment");
|
||||
assert!(error.to_string().contains("app routing"));
|
||||
}
|
||||
|
||||
fn session_config(directory: &Path, server: &MockServer) -> WorkloadIdentitySessionConfig {
|
||||
let assertion_file = directory.join("identity-token");
|
||||
std::fs::write(&assertion_file, "assertion-one").expect("write assertion");
|
||||
WorkloadIdentitySessionConfig {
|
||||
assertion_file,
|
||||
environment: WorkloadIdentityEnvironment::Staging,
|
||||
federation_rule_id: "rule-one".to_string(),
|
||||
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
|
||||
token_url: Url::parse(&format!("{}/oauth/token", server.uri())).expect("token URL"),
|
||||
}
|
||||
}
|
||||
|
||||
fn jwt(label: &str, user_id: &str) -> String {
|
||||
let encode = |value: &serde_json::Value| {
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(serde_json::to_vec(value).expect("serialize JWT part"))
|
||||
};
|
||||
let header = encode(&serde_json::json!({"alg": "none", "typ": "JWT"}));
|
||||
let payload = encode(&serde_json::json!({
|
||||
"jti": label,
|
||||
"https://api.openai.com/auth": {
|
||||
"chatgpt_account_id": "account-one",
|
||||
"chatgpt_plan_type": "enterprise",
|
||||
"chatgpt_user_id": user_id,
|
||||
"user_id": user_id
|
||||
}
|
||||
}));
|
||||
format!("{header}.{payload}.sig")
|
||||
}
|
||||
|
||||
fn success_response(label: &str, account_user_id: &str, user_id: &str) -> ResponseTemplate {
|
||||
ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": jwt(label, user_id),
|
||||
"token_type": "Bearer",
|
||||
"issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
|
||||
"expires_in": 600,
|
||||
"scope": "model.request",
|
||||
"chatgpt_account_id": "account-one",
|
||||
"chatgpt_account_user_id": account_user_id,
|
||||
"chatgpt_plan_type": "enterprise",
|
||||
"user_id": user_id
|
||||
}))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compatible_adapters_share_exchange() {
|
||||
let temp_dir = tempfile::tempdir().expect("tempdir");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(success_response(
|
||||
"access-one",
|
||||
"account-user-one",
|
||||
"user-one",
|
||||
))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let registry = WorkloadIdentitySessionRegistry::default();
|
||||
let first_config = session_config(temp_dir.path(), &server);
|
||||
let second_config = first_config.clone();
|
||||
let first = WorkloadIdentityExternalAuth::from_config_with_registry(first_config, ®istry)
|
||||
.expect("first adapter");
|
||||
let second = WorkloadIdentityExternalAuth::from_config_with_registry(second_config, ®istry)
|
||||
.expect("second adapter");
|
||||
|
||||
assert!(Arc::ptr_eq(&first.session, &second.session));
|
||||
assert_eq!(
|
||||
first
|
||||
.resolve()
|
||||
.await
|
||||
.expect("first auth")
|
||||
.get_token()
|
||||
.expect("first token"),
|
||||
second
|
||||
.resolve()
|
||||
.await
|
||||
.expect("second auth")
|
||||
.get_token()
|
||||
.expect("second token")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn incompatible_process_session_settings_are_rejected() {
|
||||
let temp_dir = tempfile::tempdir().expect("tempdir");
|
||||
let server = MockServer::start().await;
|
||||
let registry = WorkloadIdentitySessionRegistry::default();
|
||||
let base = session_config(temp_dir.path(), &server);
|
||||
let _active = WorkloadIdentityExternalAuth::from_config_with_registry(base.clone(), ®istry)
|
||||
.expect("active adapter");
|
||||
|
||||
let mut different_rule = base.clone();
|
||||
different_rule.federation_rule_id = "rule-two".to_string();
|
||||
let mut different_file = base.clone();
|
||||
different_file.assertion_file = temp_dir.path().join("identity-token-two");
|
||||
std::fs::write(&different_file.assertion_file, "assertion-two").expect("write assertion");
|
||||
let mut different_environment = base.clone();
|
||||
different_environment.environment = WorkloadIdentityEnvironment::Production;
|
||||
let mut different_route = base;
|
||||
different_route.http_client_factory =
|
||||
HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy);
|
||||
|
||||
let different_route_adapter =
|
||||
WorkloadIdentityExternalAuth::from_config_with_registry(different_route, ®istry)
|
||||
.expect("route changes reuse the process-owned session");
|
||||
assert!(Arc::ptr_eq(
|
||||
&_active.session,
|
||||
&different_route_adapter.session
|
||||
));
|
||||
|
||||
for config in [different_rule, different_file, different_environment] {
|
||||
assert!(matches!(
|
||||
WorkloadIdentityExternalAuth::from_config_with_registry(config, ®istry),
|
||||
Err(WorkloadIdentitySessionError::ConflictingConfiguration)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_preserves_identity_and_invalid_tokens_are_reexchanged() {
|
||||
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 => success_response("access-one", "account-user-one", "user-one"),
|
||||
1 => success_response("access-two", "account-user-two", "user-one"),
|
||||
_ => success_response("access-three", "account-user-one", "user-one"),
|
||||
}
|
||||
})
|
||||
.mount(&server)
|
||||
.await;
|
||||
let registry = WorkloadIdentitySessionRegistry::default();
|
||||
let adapter = WorkloadIdentityExternalAuth::from_config_with_registry(
|
||||
session_config(temp_dir.path(), &server),
|
||||
®istry,
|
||||
)
|
||||
.expect("adapter");
|
||||
adapter.resolve().await.expect("initial auth");
|
||||
|
||||
let error = adapter
|
||||
.refresh(ExternalAuthRefreshContext {
|
||||
reason: ExternalAuthRefreshReason::Unauthorized,
|
||||
previous_account_id: Some("account-one".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect_err("identity change must be rejected");
|
||||
assert!(matches!(
|
||||
adapter.classify_error(error),
|
||||
RefreshTokenError::Permanent(_)
|
||||
));
|
||||
assert_eq!(
|
||||
adapter
|
||||
.resolve()
|
||||
.await
|
||||
.expect("corrected token is re-exchanged")
|
||||
.get_token()
|
||||
.expect("token"),
|
||||
jwt("access-three", "user-one")
|
||||
);
|
||||
assert_eq!(request_count.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_refreshes_share_one_exchange() {
|
||||
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| {
|
||||
if response_count.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
success_response("access-one", "account-user-one", "user-one")
|
||||
} else {
|
||||
success_response("access-two", "account-user-one", "user-one")
|
||||
.set_delay(Duration::from_millis(30))
|
||||
}
|
||||
})
|
||||
.mount(&server)
|
||||
.await;
|
||||
let registry = WorkloadIdentitySessionRegistry::default();
|
||||
let first = Arc::new(
|
||||
WorkloadIdentityExternalAuth::from_config_with_registry(
|
||||
session_config(temp_dir.path(), &server),
|
||||
®istry,
|
||||
)
|
||||
.expect("first adapter"),
|
||||
);
|
||||
let second = Arc::new(
|
||||
WorkloadIdentityExternalAuth::from_config_with_registry(
|
||||
session_config(temp_dir.path(), &server),
|
||||
®istry,
|
||||
)
|
||||
.expect("second adapter"),
|
||||
);
|
||||
first.resolve().await.expect("first resolve");
|
||||
second.resolve().await.expect("second resolve");
|
||||
let context = ExternalAuthRefreshContext {
|
||||
reason: ExternalAuthRefreshReason::Unauthorized,
|
||||
previous_account_id: Some("account-one".to_string()),
|
||||
};
|
||||
|
||||
let (first_refresh, second_refresh) =
|
||||
tokio::join!(first.refresh(context.clone()), second.refresh(context));
|
||||
assert_eq!(
|
||||
first_refresh
|
||||
.expect("first refresh")
|
||||
.get_token()
|
||||
.expect("first token"),
|
||||
second_refresh
|
||||
.expect("second refresh")
|
||||
.get_token()
|
||||
.expect("second token")
|
||||
);
|
||||
assert_eq!(request_count.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exchange_errors_map_to_retry_policy() {
|
||||
let temp_dir = tempfile::tempdir().expect("tempdir");
|
||||
let server = MockServer::start().await;
|
||||
let registry = WorkloadIdentitySessionRegistry::default();
|
||||
let adapter = WorkloadIdentityExternalAuth::from_config_with_registry(
|
||||
session_config(temp_dir.path(), &server),
|
||||
®istry,
|
||||
)
|
||||
.expect("adapter");
|
||||
|
||||
let cases = [
|
||||
(WorkloadIdentityError::ExchangeRejected(400), false),
|
||||
(WorkloadIdentityError::ExchangeRejected(408), true),
|
||||
(
|
||||
WorkloadIdentityError::AssertionFile {
|
||||
path: temp_dir.path().join("missing"),
|
||||
source: Arc::new(std::io::Error::from(std::io::ErrorKind::NotFound)),
|
||||
},
|
||||
true,
|
||||
),
|
||||
];
|
||||
for (error, transient) in cases {
|
||||
let classified = adapter.classify_error(std::io::Error::other(error));
|
||||
assert_eq!(
|
||||
matches!(classified, RefreshTokenError::Transient(_)),
|
||||
transient
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -49,6 +49,7 @@ pub use auth::RefreshTokenError;
|
||||
pub use auth::UnauthorizedRecovery;
|
||||
pub use auth::default_client;
|
||||
pub use auth::enforce_login_restrictions;
|
||||
pub use auth::is_workload_identity_selected;
|
||||
pub use auth::load_auth_dot_json;
|
||||
pub use auth::login_with_access_token;
|
||||
pub use auth::login_with_api_key;
|
||||
|
||||
@@ -100,7 +100,7 @@ impl WorkloadIdentityExchange {
|
||||
let valid_from = Instant::now();
|
||||
let result = match self.exchange_uncached().await {
|
||||
Ok(token) => state.store(token, valid_from, Instant::now()),
|
||||
Err(error) if error.allows_cached_fallback() => {
|
||||
Err(error) if error.is_transient() => {
|
||||
let now = Instant::now();
|
||||
match state.cached.as_mut() {
|
||||
Some(cached) if cached.expires_at > now => {
|
||||
@@ -156,6 +156,15 @@ impl WorkloadIdentityExchange {
|
||||
result
|
||||
}
|
||||
|
||||
/// Drops the cached token only when it is still the token the caller rejected.
|
||||
pub async fn invalidate_if_current(&self, observed_token_version: u64) {
|
||||
let mut state = self.state.lock().await;
|
||||
if state.token_generation == observed_token_version {
|
||||
state.cached = None;
|
||||
state.last_attempt_error = None;
|
||||
}
|
||||
}
|
||||
|
||||
fn complete_attempt(&self, state: &mut CacheState, error: Option<&WorkloadIdentityError>) {
|
||||
state.last_attempt_error = error.cloned();
|
||||
self.completed_attempts.fetch_add(1, Ordering::Release);
|
||||
@@ -207,17 +216,6 @@ impl WorkloadIdentityExchange {
|
||||
}
|
||||
}
|
||||
|
||||
impl WorkloadIdentityError {
|
||||
fn allows_cached_fallback(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::AssertionFile { .. }
|
||||
| Self::ExchangeUnavailable
|
||||
| Self::ExchangeRejected(408 | 429 | 500..=599)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_token_url(url: &Url) -> Result<bool, WorkloadIdentityError> {
|
||||
if !url.username().is_empty()
|
||||
|| url.password().is_some()
|
||||
|
||||
@@ -61,3 +61,27 @@ pub enum WorkloadIdentityError {
|
||||
#[error("the workload identity token exchange returned an invalid response")]
|
||||
InvalidExchangeResponse,
|
||||
}
|
||||
|
||||
impl WorkloadIdentityError {
|
||||
/// Whether retrying the operation may succeed without changing configuration.
|
||||
pub fn is_transient(&self) -> bool {
|
||||
match self {
|
||||
Self::AssertionFile { source, .. } => matches!(
|
||||
source.kind(),
|
||||
std::io::ErrorKind::Interrupted
|
||||
| std::io::ErrorKind::NotFound
|
||||
| std::io::ErrorKind::TimedOut
|
||||
| std::io::ErrorKind::WouldBlock
|
||||
),
|
||||
Self::ExchangeUnavailable | Self::ExchangeRejected(408 | 429 | 500..=599) => true,
|
||||
Self::InvalidFederationRuleId
|
||||
| Self::AssertionFileMustBeAbsolute
|
||||
| Self::InvalidAssertion
|
||||
| Self::AssertionTooLarge
|
||||
| Self::HttpClientConfiguration
|
||||
| Self::InvalidTokenUrl
|
||||
| Self::ExchangeRejected(_)
|
||||
| Self::InvalidExchangeResponse => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user