mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user