Reload remote executor credentials safely across token rotation (#58251)

## Why

ChatGPT credentials can rotate while a remote executor is idle. Registry requests need to follow stored rotations without continuing after logout or switching the executor's identity. Concurrent refreshes must also avoid overwriting newer credentials.

## What changed

- Reload and validate ChatGPT credentials before registry requests, pinning the credential mode, account, user, and workspace policy. Reuse cached metadata for unchanged personal access tokens.
- Retry registration when temporary credential resolution failures prevent a request from being sent. Bound credential resolution and block reuse of a token generation whose refresh outcome is unknown.
- Restrict supported authentication modes before executor startup loads credentials or cloud configuration.
- Serialize persistent credential writes across cooperating processes and compare stored credentials before committing refreshed tokens.
- Replace `auth.json` atomically so overlapping readers see a complete document, preserving symlink targets and protecting file permissions, including Windows access controls.

## Testing

Add coverage for credential rotation, logout, identity changes, cancelled refresh callers, temporary metadata failures, registration retries, concurrent credential writes, and atomic file replacement.

GitOrigin-RevId: 6b210944bf8eba2f45bc27113f901d2c3cb06b48
This commit is contained in:
tgrue-openai
2026-10-30 05:30:20 +00:00
committed by copyberry
parent c15d0ce323
commit 582a8812e0
25 changed files with 2178 additions and 103 deletions
@@ -1,10 +1,17 @@
//! CLI-only glue between Codex authentication and generic AWS signing.
//! Request-time authentication for remote executors, including AWS signing.
//!
//! Keeping this adapter here leaves `codex-api` and `codex-aws-auth` independent.
//! ChatGPT credentials may rotate while an executor is idle. Reload them for
//! registry requests while retaining the identity that registered the executor.
//! Failed credential reads block headers until revalidation, without counting as logout.
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_api::AuthError;
use codex_api::AuthHeadersFuture;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
use codex_aws_auth::AwsAuthConfig;
@@ -14,7 +21,221 @@ use codex_aws_auth::AwsRequestToSign;
use codex_http_client::Request;
use codex_http_client::RequestBody;
use codex_http_client::RequestCompression;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use http::HeaderMap;
use tokio::sync::Semaphore;
/// Follows stored credential rotation without allowing an executor's owner to change.
pub(super) fn chatgpt_auth_provider(
auth_manager: Arc<AuthManager>,
expected_auth: &CodexAuth,
) -> Result<SharedAuthProvider, AuthError> {
auth_manager
.restrict_auth_modes(vec![expected_auth.auth_mode()])
.map_err(|_| AuthError::Build("executor credential policy is unavailable".to_string()))?;
let account_id = expected_auth.get_account_id();
let user_id = expected_auth.get_chatgpt_user_id();
let identity_complete = account_id.as_ref().is_some_and(|id| !id.trim().is_empty())
&& user_id.as_ref().is_some_and(|id| !id.trim().is_empty());
let unidentified_headers = if identity_complete {
None
} else {
Some(codex_model_provider::auth_provider_from_auth(expected_auth).to_auth_headers()?)
};
Ok(Arc::new(ReloadingChatgptAuthProvider {
managed: codex_model_provider::auth_provider_from_auth_manager(
Arc::clone(&auth_manager),
expected_auth,
),
auth_manager,
account_id,
user_id,
workspace_account: expected_auth.is_workspace_account(),
fedramp_account: expected_auth.is_fedramp_account(),
unidentified_headers,
resolution_lock: Arc::new(Semaphore::new(/*permits*/ 1)),
uncertain_refresh: Arc::new(Mutex::new(None)),
headers_valid: Arc::new(AtomicBool::new(true)),
}))
}
#[derive(Clone)]
struct ReloadingChatgptAuthProvider {
managed: SharedAuthProvider,
auth_manager: Arc<AuthManager>,
account_id: Option<String>,
user_id: Option<String>,
workspace_account: bool,
fedramp_account: bool,
unidentified_headers: Option<HeaderMap>,
resolution_lock: Arc<Semaphore>,
uncertain_refresh: Arc<Mutex<Option<codex_login::TokenData>>>,
headers_valid: Arc<AtomicBool>,
}
impl ReloadingChatgptAuthProvider {
fn checked_headers(&self, auth: Option<CodexAuth>) -> Result<HeaderMap, AuthError> {
let auth = auth
.filter(|auth| {
auth.is_chatgpt_auth()
&& auth.get_account_id() == self.account_id
&& auth.get_chatgpt_user_id() == self.user_id
&& auth.is_workspace_account() == self.workspace_account
&& auth.is_fedramp_account() == self.fedramp_account
})
.ok_or_else(|| AuthError::Build("executor authentication changed".to_string()))?;
let token = auth.get_token().map_err(|_| {
AuthError::Build("executor credential has no usable bearer".to_string())
})?;
if token.trim().is_empty() {
return Err(AuthError::Build(
"executor credential has an empty bearer".to_string(),
));
}
let uncertain = self.uncertain_refresh.lock().map_err(|_| {
AuthError::Transient("executor authentication is unavailable".to_string())
})?;
if uncertain
.as_ref()
.is_some_and(|previous| auth.get_token_data().as_ref().ok() == Some(previous))
{
return Err(AuthError::Transient(
"executor refresh outcome is unknown; replace or renew stored credentials"
.to_string(),
));
}
let headers = codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()?;
if self
.unidentified_headers
.as_ref()
.is_some_and(|expected| expected != &headers)
{
return Err(AuthError::Build(
"executor credential rotation requires a complete account and user identity"
.to_string(),
));
}
Ok(headers)
}
}
impl AuthProvider for ReloadingChatgptAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError> {
if !self.headers_valid.load(Ordering::Acquire) {
return Err(AuthError::Transient(
"executor credentials have not been successfully revalidated".to_string(),
));
}
headers.extend(self.checked_headers(self.auth_manager.auth_cached())?);
Ok(())
}
fn apply_auth_headers(&self, mut headers: HeaderMap) -> AuthHeadersFuture<'_> {
Box::pin(async move {
headers.extend(self.resolve(/*rejected_headers*/ None).await?);
Ok(headers)
})
}
fn refresh_after_unauthorized(
&self,
rejected_headers: HeaderMap,
) -> codex_api::AuthRefreshFuture<'_> {
Box::pin(async move {
self.resolve(Some(rejected_headers.clone()))
.await
.is_ok_and(|headers| {
headers
.iter()
.any(|(name, value)| rejected_headers.get(name) != Some(value))
})
})
}
}
impl ReloadingChatgptAuthProvider {
async fn resolve(&self, rejected_headers: Option<HeaderMap>) -> Result<HeaderMap, AuthError> {
let guard = Arc::clone(&self.resolution_lock)
.acquire_owned()
.await
.map_err(|_| {
AuthError::Transient("executor authentication is unavailable".to_string())
})?;
let provider = self.clone();
// A failed read leaves AuthManager's last known owner intact. Keep those
// cached credentials unavailable to synchronous callers until revalidated.
provider.headers_valid.store(false, Ordering::Release);
// Outlive a cancelled handshake, but never hold the permit indefinitely.
// A timed-out OAuth exchange may already have rotated the token: quarantine
// that generation until another login/refresh supplies different credentials.
tokio::spawn(async move {
let _guard = guard;
let mut attempted_refresh = None;
let result = tokio::time::timeout(Duration::from_secs(30), async {
provider
.auth_manager
.reload_with_cached_auth()
.await
.map_err(|_| {
AuthError::Transient(
"executor credential storage or hydration is unavailable".to_string(),
)
})?;
let auth = provider.auth_manager.auth_cached();
let current = provider.checked_headers(auth.clone())?;
if provider.unidentified_headers.is_some() {
return Ok(current);
}
attempted_refresh = auth.as_ref().and_then(|auth| auth.get_token_data().ok());
if let Some(rejected) = rejected_headers {
if !current
.iter()
.any(|(name, value)| rejected.get(name) != Some(value))
{
provider.managed.refresh_after_unauthorized(rejected).await;
}
} else {
provider.auth_manager.auth().await;
}
// A later storage timeout is not an interrupted OAuth exchange.
attempted_refresh = None;
// auth() deliberately falls back to a snapshot on refresh errors.
// Re-read the selected source before dispatch: logout, replacement,
// or failed hydration must never authorize that fallback snapshot.
provider
.auth_manager
.reload_with_cached_auth()
.await
.map_err(|_| {
AuthError::Transient(
"executor credential storage or hydration is unavailable".to_string(),
)
})?;
provider.checked_headers(provider.auth_manager.auth_cached())
})
.await;
let result = match result {
Ok(result) => result,
Err(_) => {
if let Some(attempted) = attempted_refresh
&& let Ok(mut uncertain) = provider.uncertain_refresh.lock()
{
*uncertain = Some(attempted);
}
Err(AuthError::Transient(
"executor credential resolution timed out; retry after repairing credentials"
.to_string(),
))
}
};
provider.headers_valid.store(result.is_ok(), Ordering::Release);
result
})
.await
.map_err(|_| AuthError::Transient("executor authentication task failed".to_string()))?
}
}
/// Creates a SigV4 provider, preferring an explicit profile over the default credential chain.
pub(super) async fn aws_sigv4_auth_provider(