diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 08d4a21c50..45bd55302b 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -592,11 +592,11 @@ pub fn save_auth( storage.save(auth) } -/// Load CLI auth data using the configured credential store backend. -/// Returns `None` when no credentials are stored. This function is -/// provided only for tests. Production code should not directly load -/// from the auth.json storage. It should use the AuthManager abstraction -/// instead. +/// Load the raw stored auth payload without applying environment overrides. +/// +/// Returns `None` when no credentials are stored. Prefer `AuthManager` for +/// ordinary production reads; this helper is for tests and write-side +/// maintenance that must inspect the exact payload in storage. pub fn load_auth_dot_json( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 07c44983a9..167ccd3f75 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -11,3 +11,5 @@ mod revoke; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; +pub(crate) use revoke::revoke_auth_tokens; +pub(crate) use revoke::should_revoke_auth_tokens; diff --git a/codex-rs/login/src/auth/revoke.rs b/codex-rs/login/src/auth/revoke.rs index 71164523a9..b6a9fef4bb 100644 --- a/codex-rs/login/src/auth/revoke.rs +++ b/codex-rs/login/src/auth/revoke.rs @@ -1,8 +1,9 @@ -//! Best-effort OAuth token revocation used during logout. +//! Best-effort OAuth token revocation for managed auth cleanup. //! -//! Managed ChatGPT auth stores OAuth tokens locally. Logout attempts to revoke the -//! refresh token, falling back to the access token when no refresh token is -//! available, and callers still remove local auth if the revoke request fails. +//! Managed ChatGPT auth stores OAuth tokens locally. Cleanup attempts to revoke +//! the refresh token, falling back to the access token when no refresh token is +//! available, and callers still complete their primary work if the revoke request +//! fails. use serde::Serialize; use std::time::Duration; @@ -51,35 +52,43 @@ struct RevokeTokenRequest<'a> { client_id: Option<&'static str>, } -pub(super) async fn revoke_auth_tokens( +pub(crate) async fn revoke_auth_tokens( auth_dot_json: Option<&AuthDotJson>, ) -> Result<(), std::io::Error> { - let Some(tokens) = auth_dot_json.and_then(managed_chatgpt_tokens) else { + let Some((token, kind)) = auth_dot_json.and_then(revocable_token) else { return Ok(()); }; let client = create_client(); let endpoint = revoke_token_endpoint(); + revoke_oauth_token(&client, endpoint.as_str(), token, kind, REVOKE_HTTP_TIMEOUT).await +} + +pub(crate) fn should_revoke_auth_tokens( + auth_dot_json: Option<&AuthDotJson>, + replacement_auth: &AuthDotJson, +) -> bool { + let Some((token, kind)) = auth_dot_json.and_then(revocable_token) else { + return false; + }; + let Some(replacement_tokens) = managed_chatgpt_tokens(replacement_auth) else { + return true; + }; + + match kind { + RevokeTokenKind::Access => replacement_tokens.access_token != token, + RevokeTokenKind::Refresh => replacement_tokens.refresh_token != token, + } +} + +fn revocable_token(auth_dot_json: &AuthDotJson) -> Option<(&str, RevokeTokenKind)> { + let tokens = managed_chatgpt_tokens(auth_dot_json)?; if !tokens.refresh_token.is_empty() { - revoke_oauth_token( - &client, - endpoint.as_str(), - tokens.refresh_token.as_str(), - RevokeTokenKind::Refresh, - REVOKE_HTTP_TIMEOUT, - ) - .await + Some((tokens.refresh_token.as_str(), RevokeTokenKind::Refresh)) } else if !tokens.access_token.is_empty() { - revoke_oauth_token( - &client, - endpoint.as_str(), - tokens.access_token.as_str(), - RevokeTokenKind::Access, - REVOKE_HTTP_TIMEOUT, - ) - .await + Some((tokens.access_token.as_str(), RevokeTokenKind::Access)) } else { - Ok(()) + None } } diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 9b6f835bb4..ad0d6b5e5e 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -25,7 +25,10 @@ use std::thread; use std::time::Duration; use crate::auth::AuthDotJson; +use crate::auth::load_auth_dot_json; +use crate::auth::revoke_auth_tokens; use crate::auth::save_auth; +use crate::auth::should_revoke_auth_tokens; use crate::default_client::originator; use crate::pkce::PkceCodes; use crate::pkce::generate_pkce; @@ -781,7 +784,8 @@ pub(crate) async fn exchange_code_for_tokens( }) } -/// Persists exchanged credentials using the configured local auth store. +/// Persists exchanged credentials using the configured local auth store, then +/// best-effort revokes any superseded managed ChatGPT tokens. pub(crate) async fn persist_tokens_async( codex_home: &Path, api_key: Option, @@ -792,7 +796,14 @@ pub(crate) async fn persist_tokens_async( ) -> io::Result<()> { // Reuse existing synchronous logic but run it off the async runtime. let codex_home = codex_home.to_path_buf(); - tokio::task::spawn_blocking(move || { + let (previous_auth, auth) = tokio::task::spawn_blocking(move || { + let previous_auth = match load_auth_dot_json(&codex_home, auth_credentials_store_mode) { + Ok(auth) => auth, + Err(err) => { + warn!("failed to load previous auth before saving new login: {err}"); + None + } + }; let mut tokens = TokenData { id_token: parse_chatgpt_jwt_claims(&id_token).map_err(io::Error::other)?, access_token, @@ -812,10 +823,19 @@ pub(crate) async fn persist_tokens_async( last_refresh: Some(Utc::now()), agent_identity: None, }; - save_auth(&codex_home, &auth, auth_credentials_store_mode) + save_auth(&codex_home, &auth, auth_credentials_store_mode)?; + Ok::<_, io::Error>((previous_auth, auth)) }) .await - .map_err(|e| io::Error::other(format!("persist task failed: {e}")))? + .map_err(|e| io::Error::other(format!("persist task failed: {e}")))??; + + if should_revoke_auth_tokens(previous_auth.as_ref(), &auth) + && let Err(err) = revoke_auth_tokens(previous_auth.as_ref()).await + { + warn!("failed to revoke superseded auth tokens after login: {err}"); + } + + Ok(()) } fn compose_success_url( @@ -1132,6 +1152,28 @@ pub(crate) async fn obtain_api_key( } #[cfg(test)] mod tests { + use std::ffi::OsString; + + use anyhow::Context; + use base64::Engine; + use codex_app_server_protocol::AuthMode; + use codex_config::types::AuthCredentialsStoreMode; + use serde_json::Value; + use serde_json::json; + use tempfile::tempdir; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + use crate::auth::AuthDotJson; + use crate::auth::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR; + use crate::auth::load_auth_dot_json; + use crate::auth::save_auth; + use crate::token_data::TokenData; + use crate::token_data::parse_chatgpt_jwt_claims; + use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use super::DEFAULT_ISSUER; @@ -1140,11 +1182,179 @@ mod tests { use super::html_escape; use super::is_missing_codex_entitlement_error; use super::parse_token_endpoint_error; + use super::persist_tokens_async; use super::redact_sensitive_query_value; use super::redact_sensitive_url_parts; use super::render_login_error_page; use super::sanitize_url_for_logging; + #[serial_test::serial(logout_revoke)] + #[tokio::test] + async fn persist_tokens_async_revokes_previous_auth_without_failing_login() -> anyhow::Result<()> + { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "error": { + "message": "revoke failed" + } + }))) + .expect(1) + .mount(&server) + .await; + let _env_guard = EnvGuard::set( + REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/revoke", server.uri()), + ); + + let codex_home = tempdir()?; + save_auth( + codex_home.path(), + &chatgpt_auth("old-access", "old-refresh", "old-account"), + AuthCredentialsStoreMode::File, + )?; + + persist_tokens_async( + codex_home.path(), + /*api_key*/ None, + jwt_for_account("new-account"), + "new-access".to_string(), + "new-refresh".to_string(), + AuthCredentialsStoreMode::File, + ) + .await?; + + let auth = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)? + .context("auth.json should exist after login")?; + assert_eq!( + auth.tokens.context("new tokens should be persisted")?, + TokenData { + id_token: parse_chatgpt_jwt_claims(&jwt_for_account("new-account")) + .expect("new JWT should parse"), + access_token: "new-access".to_string(), + refresh_token: "new-refresh".to_string(), + account_id: Some("new-account".to_string()), + } + ); + + let requests = server + .received_requests() + .await + .context("failed to fetch revoke requests")?; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0] + .body_json::() + .context("revoke request should be JSON")?, + json!({ + "token": "old-refresh", + "token_type_hint": "refresh_token", + "client_id": crate::auth::CLIENT_ID, + }) + ); + server.verify().await; + Ok(()) + } + + #[serial_test::serial(logout_revoke)] + #[tokio::test] + async fn persist_tokens_async_does_not_revoke_reused_refresh_token() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let _env_guard = EnvGuard::set( + REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/revoke", server.uri()), + ); + + let codex_home = tempdir()?; + save_auth( + codex_home.path(), + &chatgpt_auth("old-access", "shared-refresh", "old-account"), + AuthCredentialsStoreMode::File, + )?; + + persist_tokens_async( + codex_home.path(), + /*api_key*/ None, + jwt_for_account("new-account"), + "new-access".to_string(), + "shared-refresh".to_string(), + AuthCredentialsStoreMode::File, + ) + .await?; + + let requests = server + .received_requests() + .await + .context("failed to fetch revoke requests")?; + assert_eq!(requests.len(), 0); + Ok(()) + } + + fn chatgpt_auth(access_token: &str, refresh_token: &str, account_id: &str) -> AuthDotJson { + AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: parse_chatgpt_jwt_claims(&jwt_for_account(account_id)) + .expect("test JWT should parse"), + access_token: access_token.to_string(), + refresh_token: refresh_token.to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: None, + agent_identity: None, + } + } + + fn jwt_for_account(account_id: &str) -> String { + let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); + let header_b64 = encode(br#"{"alg":"none","typ":"JWT"}"#); + let payload_b64 = encode( + serde_json::to_string(&json!({ + "https://api.openai.com/auth": { + "chatgpt_account_id": account_id, + } + })) + .expect("payload should serialize") + .as_bytes(), + ); + let signature_b64 = encode(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") + } + + struct EnvGuard { + key: &'static str, + original: Option, + } + + impl EnvGuard { + fn set(key: &'static str, value: String) -> Self { + let original = std::env::var_os(key); + // SAFETY: this test executes serially with other revoke tests. + unsafe { + std::env::set_var(key, &value); + } + Self { key, original } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: the guard restores the original environment before other revoke tests run. + unsafe { + match &self.original { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } + } + } + #[test] fn parse_token_endpoint_error_prefers_error_description() { let detail = parse_token_endpoint_error(