address auth rejection review feedback

This commit is contained in:
radai
2026-06-23 13:48:27 -07:00
parent c9df866015
commit 553380a1eb
7 changed files with 350 additions and 68 deletions
+1
View File
@@ -3256,6 +3256,7 @@ dependencies = [
"codex-terminal-detection",
"codex-utils-template",
"core_test_support",
"hmac 0.12.1",
"jsonwebtoken",
"keyring",
"once_cell",
+34 -15
View File
@@ -62,6 +62,7 @@ use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_api::response_create_client_metadata;
use codex_app_server_protocol::AuthMode;
use codex_login::AuthFingerprint;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::RefreshTokenError;
@@ -275,7 +276,7 @@ struct LastResponse {
#[derive(Debug, Default)]
struct WebsocketSession {
connection: Option<ApiWebSocketConnection>,
connection_auth: Option<CodexAuth>,
connection_auth: Option<AuthFingerprint>,
last_request: Option<ResponsesApiRequest>,
last_response_rx: Option<oneshot::Receiver<LastResponse>>,
last_response_from_untraced_warmup: bool,
@@ -1193,7 +1194,10 @@ impl ModelClientSession {
client_setup.api_auth.as_ref(),
PendingUnauthorizedRetry::default(),
);
let connection_auth = client_setup.auth.clone();
let connection_auth = client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth);
let connection = self
.client
.connect_websocket(
@@ -1385,7 +1389,10 @@ impl ModelClientSession {
Arc::clone(&self.client.state.provider),
Some(StreamAuthRecovery {
auth_recovery,
rejected_auth: client_setup.auth,
rejected_auth: client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth),
}),
);
return Ok(stream);
@@ -1400,13 +1407,17 @@ impl ModelClientSession {
response_debug_context.request_id.as_deref(),
/*output_items*/ &[],
);
let rejected_auth = client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth);
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
session_telemetry,
&self.client.state.provider,
client_setup.auth.as_ref(),
rejected_auth.as_ref(),
)
.await?,
);
@@ -1505,7 +1516,10 @@ impl ModelClientSession {
session_telemetry,
api_provider: client_setup.api_provider,
api_auth: client_setup.api_auth,
auth: client_setup.auth.clone(),
auth: client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth),
responses_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint(
@@ -1523,13 +1537,17 @@ impl ModelClientSession {
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
let rejected_auth = client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth);
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
session_telemetry,
&self.client.state.provider,
client_setup.auth.as_ref(),
rejected_auth.as_ref(),
)
.await?,
);
@@ -1586,11 +1604,12 @@ impl ModelClientSession {
);
err
})?;
let rejected_auth = self
.websocket_session
.connection_auth
.clone()
.or_else(|| client_setup.auth.clone());
let rejected_auth = self.websocket_session.connection_auth.clone().or_else(|| {
client_setup
.auth
.as_ref()
.and_then(AuthFingerprint::from_auth)
});
let (stream, last_request_rx) = map_response_stream(
stream_result,
request_session_telemetry,
@@ -2088,7 +2107,7 @@ struct WebsocketConnectParams<'a> {
session_telemetry: &'a SessionTelemetry,
api_provider: codex_api::Provider,
api_auth: SharedAuthProvider,
auth: Option<CodexAuth>,
auth: Option<AuthFingerprint>,
responses_metadata: &'a CodexResponsesMetadata,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
@@ -2096,14 +2115,14 @@ struct WebsocketConnectParams<'a> {
struct StreamAuthRecovery {
auth_recovery: Option<UnauthorizedRecovery>,
rejected_auth: Option<CodexAuth>,
rejected_auth: Option<AuthFingerprint>,
}
async fn force_logout_if_workspace_restricted_unauthorized(
transport: &TransportError,
auth_recovery: &mut Option<UnauthorizedRecovery>,
session_telemetry: &SessionTelemetry,
rejected_auth: Option<&CodexAuth>,
rejected_auth: Option<&AuthFingerprint>,
) -> Option<CodexErr> {
let debug = extract_response_debug_context(transport);
if !is_chatgpt_ip_workspace_restricted_unauthorized(transport, &debug) {
@@ -2167,7 +2186,7 @@ async fn handle_unauthorized(
auth_recovery: &mut Option<UnauthorizedRecovery>,
session_telemetry: &SessionTelemetry,
provider: &SharedModelProvider,
rejected_auth: Option<&CodexAuth>,
rejected_auth: Option<&AuthFingerprint>,
) -> Result<UnauthorizedRecoveryExecution> {
if let Some(err) = force_logout_if_workspace_restricted_unauthorized(
&transport,
+7 -2
View File
@@ -20,6 +20,7 @@ use codex_api::TransportError;
use codex_app_server_protocol::AuthMode;
use codex_config::types::AuthCredentialsStoreMode;
use codex_login::AuthDotJson;
use codex_login::AuthFingerprint;
use codex_login::AuthKeyringBackendKind;
use codex_login::AuthManager;
use codex_login::CodexAuth;
@@ -768,6 +769,8 @@ async fn workspace_restricted_401_uses_request_auth_for_forced_logout() -> anyho
let request_auth = request_auth_manager
.auth_cached()
.expect("request auth should load");
let request_auth_fingerprint =
AuthFingerprint::from_auth(&request_auth).expect("request auth can be fingerprinted");
let auth_manager = AuthManager::from_auth_for_testing_with_home(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
codex_home.path().to_path_buf(),
@@ -779,7 +782,7 @@ async fn workspace_restricted_401_uses_request_auth_for_forced_logout() -> anyho
&mut recovery,
&test_session_telemetry(),
&test_model_provider(),
Some(&request_auth),
Some(&request_auth_fingerprint),
)
.await
.expect_err("matching 401 should force reauth without retry");
@@ -892,7 +895,9 @@ async fn workspace_restricted_stream_401_forces_chatgpt_logout() -> anyhow::Resu
test_model_provider(),
Some(super::StreamAuthRecovery {
auth_recovery: Some(auth_manager.unauthorized_recovery()),
rejected_auth: Some(CodexAuth::create_dummy_chatgpt_auth_for_testing()),
rejected_auth: AuthFingerprint::from_auth(
&CodexAuth::create_dummy_chatgpt_auth_for_testing(),
),
}),
);
+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 }
hmac = { workspace = true }
once_cell = { workspace = true }
os_info = { workspace = true }
rand = { workspace = true }
+49
View File
@@ -1594,6 +1594,55 @@ async fn personal_access_token_does_not_offer_unauthorized_recovery() {
server.verify().await;
}
#[test]
fn auth_fingerprint_debug_redacts_auth_material() {
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let fingerprint = AuthFingerprint::from_auth(&auth).expect("ChatGPT auth can be fingerprinted");
let debug = format!("{fingerprint:?}");
assert!(debug.contains("AuthFingerprint"));
assert!(debug.contains("fingerprint_prefix"));
assert!(!debug.contains("Access Token"));
assert!(!debug.contains("test"));
assert!(!debug.contains("account_id"));
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn forced_server_auth_rejection_does_not_reload_access_token_env() {
let codex_home = tempdir().unwrap();
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/user-auth-credential/whoami"))
.and(header("authorization", "Bearer at-env-test"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(personal_access_token_whoami(WORKSPACE_ID_ALLOWED)),
)
.expect(0)
.mount(&server)
.await;
let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri());
let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, "at-env-test");
let rejected_auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let rejected_auth_fingerprint =
AuthFingerprint::from_auth(&rejected_auth).expect("ChatGPT auth can be fingerprinted");
let manager = AuthManager::from_auth_for_testing_with_home(
rejected_auth.clone(),
codex_home.path().to_path_buf(),
);
let changed = manager
.force_logout_due_to_server_auth_rejection(Some(&rejected_auth_fingerprint))
.await
.expect("forced logout should succeed");
assert!(changed);
assert!(manager.auth_cached().is_none());
server.verify().await;
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn load_auth_keeps_codex_api_key_env_precedence() {
+257 -51
View File
@@ -1,4 +1,7 @@
use chrono::Utc;
use hmac::Hmac;
use hmac::Mac;
use rand::RngCore;
use reqwest::StatusCode;
use serde::Deserialize;
use serde::Serialize;
@@ -6,6 +9,7 @@ use serde::Serialize;
use serial_test::serial;
use std::env;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::future::Future;
use std::path::Path;
use std::path::PathBuf;
@@ -62,8 +66,11 @@ use codex_protocol::auth::RefreshTokenFailedError;
use codex_protocol::auth::RefreshTokenFailedReason;
use codex_protocol::protocol::SessionSource;
use serde_json::Value;
use sha2::Sha256;
use thiserror::Error;
type HmacSha256 = Hmac<Sha256>;
/// Authentication mechanism used by the current user.
#[derive(Debug, Clone)]
pub enum CodexAuth {
@@ -84,6 +91,174 @@ pub enum AgentIdentityAuthPolicy {
ChatGptAuth,
}
/// A more secure way to match auth credentials without storing or printing raw key material.
///
/// The fingerprint is keyed with random in-memory material. Debug output can safely include the
/// auth kind and a short fingerprint prefix, while matching can still recompute the keyed
/// fingerprint for candidate stored auth records.
#[derive(Clone)]
pub struct AuthFingerprint {
kind: ApiAuthMode,
fingerprint: [u8; 32],
key: [u8; 32],
}
impl AuthFingerprint {
pub fn from_auth(auth: &CodexAuth) -> Option<Self> {
let kind = auth.api_auth_mode();
let mut key = [0_u8; 32];
rand::rng().fill_bytes(&mut key);
let fingerprint = match auth {
CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => {
let auth_dot_json = auth.get_current_auth_json()?;
fingerprint_chatgpt_auth_dot_json(kind, &auth_dot_json, &key)?
}
CodexAuth::AgentIdentity(auth) => fingerprint_agent_identity_record(
kind,
serde_json::to_vec(auth.record()).ok()?.as_slice(),
&key,
),
CodexAuth::PersonalAccessToken(auth) => {
fingerprint_single_secret(kind, auth.access_token(), &key)
}
CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => return None,
};
Some(Self {
kind,
fingerprint,
key,
})
}
fn matches_auth_dot_json(&self, auth_dot_json: &AuthDotJson) -> bool {
if auth_dot_json.resolved_mode() != self.kind {
return false;
}
let Some(candidate) = fingerprint_auth_dot_json(self.kind, auth_dot_json, &self.key) else {
return false;
};
candidate == self.fingerprint
}
fn is_external_chatgpt_tokens(&self) -> bool {
self.kind == ApiAuthMode::ChatgptAuthTokens
}
}
impl Debug for AuthFingerprint {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthFingerprint")
.field("kind", &self.kind)
.field("fingerprint_prefix", &FingerprintPrefix(&self.fingerprint))
.finish()
}
}
struct FingerprintPrefix<'a>(&'a [u8; 32]);
impl Debug for FingerprintPrefix<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
for byte in self.0.iter().take(4) {
write!(f, "{byte:02x}")?;
}
Ok(())
}
}
fn fingerprint_auth_dot_json(
kind: ApiAuthMode,
auth_dot_json: &AuthDotJson,
key: &[u8; 32],
) -> Option<[u8; 32]> {
match kind {
ApiAuthMode::Chatgpt | ApiAuthMode::ChatgptAuthTokens => {
fingerprint_chatgpt_auth_dot_json(kind, auth_dot_json, key)
}
ApiAuthMode::AgentIdentity => {
let record = auth_dot_json
.agent_identity
.as_deref()
.and_then(|jwt| AgentIdentityAuthRecord::from_agent_identity_jwt(jwt).ok())?;
Some(fingerprint_agent_identity_record(
kind,
serde_json::to_vec(&record).ok()?.as_slice(),
key,
))
}
ApiAuthMode::PersonalAccessToken => auth_dot_json
.personal_access_token
.as_deref()
.map(|access_token| fingerprint_single_secret(kind, access_token, key)),
ApiAuthMode::ApiKey | ApiAuthMode::BedrockApiKey => None,
}
}
fn fingerprint_chatgpt_auth_dot_json(
kind: ApiAuthMode,
auth_dot_json: &AuthDotJson,
key: &[u8; 32],
) -> Option<[u8; 32]> {
let tokens = auth_dot_json.tokens.as_ref()?;
let mut mac = auth_fingerprint_mac(kind, key);
update_mac_field(&mut mac, tokens.access_token.as_bytes());
update_mac_field(&mut mac, tokens.refresh_token.as_bytes());
update_mac_optional_field(
&mut mac,
auth_dot_json_chatgpt_account_id(auth_dot_json).map(str::as_bytes),
);
Some(mac.finalize().into_bytes().into())
}
fn fingerprint_agent_identity_record(
kind: ApiAuthMode,
serialized_record: &[u8],
key: &[u8; 32],
) -> [u8; 32] {
let mut mac = auth_fingerprint_mac(kind, key);
update_mac_field(&mut mac, serialized_record);
mac.finalize().into_bytes().into()
}
fn fingerprint_single_secret(kind: ApiAuthMode, secret: &str, key: &[u8; 32]) -> [u8; 32] {
let mut mac = auth_fingerprint_mac(kind, key);
update_mac_field(&mut mac, secret.as_bytes());
mac.finalize().into_bytes().into()
}
fn auth_fingerprint_mac(kind: ApiAuthMode, key: &[u8; 32]) -> HmacSha256 {
let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any size");
update_mac_field(&mut mac, auth_kind_label(kind));
mac
}
fn auth_kind_label(kind: ApiAuthMode) -> &'static [u8] {
match kind {
ApiAuthMode::ApiKey => b"ApiKey",
ApiAuthMode::Chatgpt => b"Chatgpt",
ApiAuthMode::ChatgptAuthTokens => b"ChatgptAuthTokens",
ApiAuthMode::AgentIdentity => b"AgentIdentity",
ApiAuthMode::PersonalAccessToken => b"PersonalAccessToken",
ApiAuthMode::BedrockApiKey => b"BedrockApiKey",
}
}
fn update_mac_optional_field(mac: &mut HmacSha256, value: Option<&[u8]>) {
match value {
Some(value) => {
mac.update(&[1]);
update_mac_field(mac, value);
}
None => mac.update(&[0]),
}
}
fn update_mac_field(mac: &mut HmacSha256, value: &[u8]) {
mac.update(&(value.len() as u64).to_be_bytes());
mac.update(value);
}
impl PartialEq for CodexAuth {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
@@ -1153,13 +1328,13 @@ fn logout_store_matching_rejected_auth(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
rejected_auth: &CodexAuth,
rejected_auth: &AuthFingerprint,
) -> std::io::Result<bool> {
logout_store_if(
codex_home,
auth_credentials_store_mode,
keyring_backend_kind,
|auth_dot_json| auth_dot_json_matches_rejected_auth(auth_dot_json, rejected_auth),
|auth_dot_json| rejected_auth.matches_auth_dot_json(auth_dot_json),
)
}
@@ -1212,49 +1387,6 @@ fn logout_store_if(
Ok(removed)
}
fn auth_dot_json_matches_rejected_auth(
auth_dot_json: &AuthDotJson,
rejected_auth: &CodexAuth,
) -> bool {
if auth_dot_json.resolved_mode() != rejected_auth.api_auth_mode() {
return false;
}
match rejected_auth {
CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => {
auth_dot_json_matches_rejected_chatgpt_auth(auth_dot_json, rejected_auth)
}
CodexAuth::AgentIdentity(rejected) => auth_dot_json
.agent_identity
.as_deref()
.and_then(|jwt| AgentIdentityAuthRecord::from_agent_identity_jwt(jwt).ok())
.is_some_and(|record| &record == rejected.record()),
CodexAuth::PersonalAccessToken(rejected) => {
auth_dot_json.personal_access_token.as_deref() == Some(rejected.access_token())
}
CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => false,
}
}
fn auth_dot_json_matches_rejected_chatgpt_auth(
auth_dot_json: &AuthDotJson,
rejected_auth: &CodexAuth,
) -> bool {
let Some(rejected_auth_dot_json) = rejected_auth.get_current_auth_json() else {
return false;
};
let (Some(tokens), Some(rejected_tokens)) =
(&auth_dot_json.tokens, &rejected_auth_dot_json.tokens)
else {
return false;
};
tokens.access_token == rejected_tokens.access_token
&& tokens.refresh_token == rejected_tokens.refresh_token
&& auth_dot_json_chatgpt_account_id(auth_dot_json)
== auth_dot_json_chatgpt_account_id(&rejected_auth_dot_json)
}
fn auth_dot_json_chatgpt_account_id(auth_dot_json: &AuthDotJson) -> Option<&str> {
auth_dot_json.tokens.as_ref().and_then(|tokens| {
tokens
@@ -1278,6 +1410,55 @@ async fn load_auth(
keyring_backend_kind: AuthKeyringBackendKind,
agent_identity_authapi_base_url: Option<&str>,
auth_route_config: Option<&AuthRouteConfig>,
) -> std::io::Result<Option<CodexAuth>> {
load_auth_impl(
codex_home,
enable_codex_api_key_env,
/*enable_codex_access_token_env*/ true,
auth_credentials_store_mode,
forced_chatgpt_workspace_id,
chatgpt_base_url,
keyring_backend_kind,
agent_identity_authapi_base_url,
auth_route_config,
)
.await
}
async fn load_auth_without_codex_access_token_env(
codex_home: &Path,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
forced_chatgpt_workspace_id: Option<&[String]>,
chatgpt_base_url: Option<&str>,
keyring_backend_kind: AuthKeyringBackendKind,
agent_identity_authapi_base_url: Option<&str>,
auth_route_config: Option<&AuthRouteConfig>,
) -> std::io::Result<Option<CodexAuth>> {
load_auth_impl(
codex_home,
enable_codex_api_key_env,
/*enable_codex_access_token_env*/ false,
auth_credentials_store_mode,
forced_chatgpt_workspace_id,
chatgpt_base_url,
keyring_backend_kind,
agent_identity_authapi_base_url,
auth_route_config,
)
.await
}
async fn load_auth_impl(
codex_home: &Path,
enable_codex_api_key_env: bool,
enable_codex_access_token_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
forced_chatgpt_workspace_id: Option<&[String]>,
chatgpt_base_url: Option<&str>,
keyring_backend_kind: AuthKeyringBackendKind,
agent_identity_authapi_base_url: Option<&str>,
auth_route_config: Option<&AuthRouteConfig>,
) -> std::io::Result<Option<CodexAuth>> {
// API key via env var takes precedence over any other auth method.
if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() {
@@ -1308,7 +1489,8 @@ async fn load_auth(
return Ok(Some(auth));
}
if let Some(access_token) = read_codex_access_token_from_env() {
if enable_codex_access_token_env && let Some(access_token) = read_codex_access_token_from_env()
{
return match classify_codex_access_token(&access_token) {
CodexAccessToken::PersonalAccessToken(access_token) => {
let auth = PersonalAccessTokenAuth::load(access_token, auth_route_config).await?;
@@ -1788,10 +1970,14 @@ impl UnauthorizedRecovery {
pub async fn force_logout_due_to_server_auth_rejection(
&mut self,
rejected_auth: Option<&CodexAuth>,
rejected_auth: Option<&AuthFingerprint>,
) -> std::io::Result<bool> {
self.step = UnauthorizedRecoveryStep::Done;
let rejected_auth = rejected_auth.or(self.rejected_auth.as_ref());
let fallback_rejected_auth = self
.rejected_auth
.as_ref()
.and_then(AuthFingerprint::from_auth);
let rejected_auth = rejected_auth.or(fallback_rejected_auth.as_ref());
self.manager
.force_logout_due_to_server_auth_rejection(rejected_auth)
.await
@@ -2271,6 +2457,23 @@ impl AuthManager {
.flatten()
}
async fn load_auth_from_storage_without_codex_access_token_env(&self) -> Option<CodexAuth> {
let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id();
load_auth_without_codex_access_token_env(
&self.codex_home,
self.enable_codex_api_key_env,
self.auth_credentials_store_mode,
forced_chatgpt_workspace_id.as_deref(),
self.chatgpt_base_url.as_deref(),
self.keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
self.auth_route_config.as_ref(),
)
.await
.ok()
.flatten()
}
fn set_cached_auth(&self, new_auth: Option<CodexAuth>) -> bool {
if let Ok(mut guard) = self.inner.write() {
let previous = guard.auth.as_ref();
@@ -2539,21 +2742,24 @@ impl AuthManager {
pub async fn force_logout_due_to_server_auth_rejection(
&self,
rejected_auth: Option<&CodexAuth>,
rejected_auth: Option<&AuthFingerprint>,
) -> std::io::Result<bool> {
if !self.current_auth_uses_codex_backend() {
return Ok(false);
}
let removal_result = self.logout_stores_matching_rejected_auth(rejected_auth);
let cache_changed = self.set_cached_auth(self.load_auth_from_storage().await);
let cache_changed = self.set_cached_auth(
self.load_auth_from_storage_without_codex_access_token_env()
.await,
);
let removed = removal_result?;
Ok(removed || cache_changed)
}
fn logout_stores_matching_rejected_auth(
&self,
rejected_auth: Option<&CodexAuth>,
rejected_auth: Option<&AuthFingerprint>,
) -> std::io::Result<bool> {
let Some(rejected_auth) = rejected_auth else {
return logout_all_stores(
+1
View File
@@ -21,6 +21,7 @@ pub use server::run_login_server;
pub use auth::AgentIdentityAuthPolicy;
pub use auth::AuthConfig;
pub use auth::AuthDotJson;
pub use auth::AuthFingerprint;
pub use auth::AuthKeyringBackendKind;
pub use auth::AuthManager;
pub use auth::AuthManagerConfig;