diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 5303fccd43b..a66ff3f3a71 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -103,7 +103,9 @@ pub enum AuthCredentialsStoreMode { #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "lowercase")] pub enum OAuthCredentialsStoreMode { - /// `Keyring` when available; otherwise, `File`. + /// Prefer `Keyring` and use `File` when keyring storage is unavailable. + /// Once an MCP client loads credentials from one store, that client keeps the resolved store + /// for its lifetime so refreshes cannot switch to a possibly stale credential source. /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. #[default] Auto, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 0a48882b09c..0bc23730af3 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2129,7 +2129,7 @@ "description": "Determine where Codex should store and read MCP credentials.", "oneOf": [ { - "description": "`Keyring` when available; otherwise, `File`. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", + "description": "Prefer `Keyring` and use `File` when keyring storage is unavailable. Once an MCP client loads credentials from one store, that client keeps the resolved store for its lifetime so refreshes cannot switch to a possibly stale credential source. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", "enum": [ "auto" ], diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e57eb3686c4..74dd1697beb 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -26,7 +26,6 @@ pub use in_process_transport::InProcessTransportFactory; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; -pub(crate) use oauth::load_oauth_tokens; pub use oauth::save_oauth_tokens; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 46e6279a161..601d0f45f4e 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,7 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +mod resolved_store; mod store_lock; #[cfg(test)] @@ -66,6 +67,10 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; +pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; +pub(crate) use self::resolved_store::ResolvedOAuthTokens; +pub(crate) use self::resolved_store::resolve_oauth_tokens_from_store_policy; + const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; @@ -100,41 +105,24 @@ pub(crate) enum StoredOAuthTokenStatus { AuthorizationRequired, } -pub(crate) fn load_oauth_tokens( - server_name: &str, - url: &str, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, -) -> Result> { - let keyring_store = DefaultKeyringStore; - match store_mode { - OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( - &keyring_store, - keyring_backend_kind, - server_name, - url, - ), - OAuthCredentialsStoreMode::File => load_oauth_tokens_from_file(server_name, url), - OAuthCredentialsStoreMode::Keyring => { - load_oauth_tokens_from_keyring(&keyring_store, keyring_backend_kind, server_name, url) - .with_context(|| "failed to read OAuth tokens from keyring".to_string()) - } - } -} - pub(crate) fn oauth_token_status( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { - Ok( - match load_oauth_tokens(server_name, url, store_mode, keyring_backend_kind)?.as_ref() { - None => StoredOAuthTokenStatus::Missing, - Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, - Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, - }, - ) + let resolved = resolve_oauth_tokens_from_store_policy( + &DefaultKeyringStore, + server_name, + url, + store_mode, + keyring_backend_kind, + )?; + Ok(match resolved.as_ref().map(|resolved| &resolved.tokens) { + None => StoredOAuthTokenStatus::Missing, + Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, + Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, + }) } fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { @@ -174,37 +162,16 @@ fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { } } -fn load_oauth_tokens_from_keyring_with_fallback_to_file( - keyring_store: &K, - keyring_backend_kind: AuthKeyringBackendKind, - server_name: &str, - url: &str, -) -> Result> { - match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { - Ok(Some(tokens)) => Ok(Some(tokens)), - Ok(None) => load_oauth_tokens_from_file(server_name, url), - // A store lock failure means the configured aggregate authority could be changing, or - // that coordination itself is unavailable. It is not evidence that the keyring backend - // is unavailable, so consulting File here could replay credentials hidden behind a - // newer Secrets entry. This is the load-side counterpart of the save guard below. - Err(error) if error.downcast_ref::().is_some() => Err(error), - Err(error) => { - warn!("failed to read OAuth tokens from keyring: {error}"); - load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")) - } - } -} - fn load_oauth_tokens_from_keyring( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, server_name: &str, url: &str, -) -> Result> { +) -> std::result::Result, OAuthKeyringLoadError> { match keyring_backend_kind { AuthKeyringBackendKind::Direct => { load_oauth_tokens_from_direct_keyring(keyring_store, server_name, url) + .map_err(OAuthKeyringLoadError::Backend) } AuthKeyringBackendKind::Secrets => { load_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url) @@ -234,9 +201,9 @@ fn load_oauth_tokens_from_secrets_keyring( keyring_store: &K, server_name: &str, url: &str, -) -> Result> { +) -> std::result::Result, OAuthKeyringLoadError> { let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; - let codex_home = find_codex_home()?; + let codex_home = find_codex_home().map_err(anyhow::Error::from)?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), SecretsBackendKind::Local, @@ -258,6 +225,17 @@ fn load_oauth_tokens_from_secrets_keyring( } } +/// Classifies keyring load failures that affect Auto fallback policy. +#[derive(Debug, thiserror::Error)] +enum OAuthKeyringLoadError { + /// Store coordination failed, so consulting another authority would be unsafe. + #[error(transparent)] + StoreLock(#[from] OAuthStoreLockFailure), + /// The selected keyring backend itself was unavailable or its data was invalid. + #[error(transparent)] + Backend(#[from] anyhow::Error), +} + pub fn save_oauth_tokens( server_name: &str, tokens: &StoredOAuthTokens, @@ -273,7 +251,7 @@ pub fn save_oauth_tokens( tokens, ), OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file(tokens), - OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring( + OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring_and_cleanup_file( &keyring_store, keyring_backend_kind, server_name, @@ -288,6 +266,8 @@ fn save_oauth_tokens_with_keyring( server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { + // This exact-store writer is used after a client resolves its authority. Only login-time + // policy resolution may clean up or update the non-selected store. match keyring_backend_kind { AuthKeyringBackendKind::Direct => { save_oauth_tokens_to_direct_keyring(keyring_store, server_name, tokens) @@ -307,12 +287,7 @@ fn save_oauth_tokens_to_direct_keyring( let key = compute_store_key(server_name, &tokens.url)?; match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { - Ok(()) => { - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) - } + Ok(()) => Ok(()), Err(error) => { let message = format!( "failed to write OAuth tokens to keyring: {}", @@ -325,29 +300,19 @@ fn save_oauth_tokens_to_direct_keyring( } /// Saves one credential while holding the Secrets aggregate-store lock across the mutation. -/// -/// The Secrets lock is released before fallback File cleanup to preserve aggregate-lock ordering. fn save_oauth_tokens_to_secrets_keyring( keyring_store: &K, server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; - { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; - save_oauth_tokens_to_secrets_keyring_with_lock_held( - keyring_store, - server_name, - tokens, - &serialized, - )?; - } - - let key = compute_store_key(server_name, &tokens.url)?; - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + save_oauth_tokens_to_secrets_keyring_with_lock_held( + keyring_store, + server_name, + tokens, + &serialized, + ) } /// Writes one credential to Secrets. The caller must hold the Secrets aggregate-store lock. @@ -370,13 +335,38 @@ fn save_oauth_tokens_to_secrets_keyring_with_lock_held( + keyring_store: &K, + keyring_backend_kind: AuthKeyringBackendKind, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens)?; + let key = compute_store_key(server_name, &tokens.url)?; + if let Err(error) = delete_oauth_tokens_from_file(&key) { + warn!( + server_name, + keyring_backend = ?keyring_backend_kind, + error = %error, + "failed to remove OAuth tokens from fallback storage" + ); + } + Ok(()) +} + fn save_oauth_tokens_with_keyring_with_fallback_to_file( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { - match save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens) { + match save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store, + keyring_backend_kind, + server_name, + tokens, + ) { Ok(()) => Ok(()), // As on load, a store lock failure is a coordination failure rather than evidence that // the keyring backend is unavailable. Falling back could leave a newer File token hidden @@ -495,8 +485,7 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, last_credentials: Mutex>, } @@ -505,8 +494,7 @@ impl OAuthPersistor { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, initial_credentials: Option, ) -> Self { Self { @@ -514,15 +502,13 @@ impl OAuthPersistor { server_name, url, authorization_manager, - store_mode, - keyring_backend_kind, + credential_store, last_credentials: Mutex::new(initial_credentials), }), } } - /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. + /// Persists RMCP-managed credential changes back to this client's resolved authority. #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" @@ -540,10 +526,12 @@ impl OAuthPersistor { let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); let same_token = last_credentials .as_ref() - .map(|prev| prev.token_response == new_token_response) + .map(|previous| previous.token_response == new_token_response) .unwrap_or(false); let expires_at = if same_token { - last_credentials.as_ref().and_then(|prev| prev.expires_at) + last_credentials + .as_ref() + .and_then(|previous| previous.expires_at) } else { compute_expires_at_millis(&credentials) }; @@ -555,28 +543,27 @@ impl OAuthPersistor { expires_at, }; if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens( + self.inner.credential_store.save( + &DefaultKeyringStore, &self.inner.server_name, &stored, - self.inner.store_mode, - self.inner.keyring_backend_kind, )?; *last_credentials = Some(stored); } } None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( + let mut last_credentials = self.inner.last_credentials.lock().await; + if last_credentials.take().is_some() + && let Err(error) = self.inner.credential_store.delete( + &DefaultKeyringStore, &self.inner.server_name, &self.inner.url, - self.inner.store_mode, - self.inner.keyring_backend_kind, ) { warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name + server_name = %self.inner.server_name, + error = %error, + "failed to remove MCP OAuth credentials from the resolved store" ); } } @@ -875,7 +862,7 @@ mod tests { use super::test_support::TempCodexHome; #[test] - fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { + fn resolve_oauth_tokens_from_store_policy_uses_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); let store = MockKeyringStore::default(); let tokens = sample_tokens(); @@ -884,14 +871,19 @@ mod tests { let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; store.save(KEYRING_SERVICE, &key, &serialized)?; - let loaded = super::load_oauth_tokens_from_keyring( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from keyring"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!( + resolved.store, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) + ); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -904,14 +896,16 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -926,14 +920,45 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); + Ok(()) + } + + #[test] + fn exact_store_operations_do_not_adopt_or_mutate_the_other_store() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let file_tokens = sample_tokens(); + let mut keyring_tokens = file_tokens.clone(); + keyring_tokens + .token_response + .0 + .set_access_token(AccessToken::new("keyring-access-token".to_string())); + + super::save_oauth_tokens_to_file(&file_tokens)?; + let fallback_path = super::fallback_file_path()?; + let fallback_before = fs::read(&fallback_path)?; + super::save_oauth_tokens_with_keyring( + &store, + AuthKeyringBackendKind::Direct, + &keyring_tokens.server_name, + &keyring_tokens, + )?; + + assert_eq!(fs::read(fallback_path)?, fallback_before); + let loaded = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) + .load(&store, &keyring_tokens.server_name, &keyring_tokens.url)? + .expect("tokens should load from the selected keyring store"); + assert_tokens_match_without_expiry(&loaded, &keyring_tokens); Ok(()) } diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs new file mode 100644 index 00000000000..6a688ea5393 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -0,0 +1,176 @@ +//! Resolves the configured MCP OAuth store and pins that concrete source for one client lifecycle. + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_keyring_store::KeyringStore; +use tracing::warn; + +use super::OAuthKeyringLoadError; +use super::StoredOAuthTokens; +use super::compute_store_key; +use super::delete_oauth_tokens_from_direct_keyring; +use super::delete_oauth_tokens_from_file; +use super::delete_oauth_tokens_from_secrets_keyring; +use super::load_oauth_tokens_from_file; +use super::load_oauth_tokens_from_keyring; +use super::save_oauth_tokens_to_file; +use super::save_oauth_tokens_with_keyring; + +/// Concrete credential store resolved for one MCP OAuth client lifecycle. +/// +/// This is intentionally not durable. `Auto` may resolve differently in a later process, but a +/// client that loaded credentials from one store must reread, refresh, persist, and remove only +/// through that store. A mid-lifecycle backend failure is unexpected and must return an error +/// rather than falling back to another possibly stale refresh token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolvedOAuthCredentialStore { + File, + Keyring(AuthKeyringBackendKind), +} + +impl ResolvedOAuthCredentialStore { + /// Loads credentials only from this already-resolved authority. + /// + /// Unlike `resolve_oauth_tokens_from_store_policy`, this never evaluates configured + /// `Auto` fallback policy. + pub(crate) fn load( + self, + keyring_store: &K, + server_name: &str, + url: &str, + ) -> Result> { + match self { + Self::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to reread OAuth tokens from resolved file storage"), + Self::Keyring(keyring_backend_kind) => load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .map_err(anyhow::Error::from) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ), + } + } + + /// Saves credentials only to this already-resolved authority. + pub(crate) fn save( + self, + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + ) -> Result<()> { + match self { + Self::File => save_oauth_tokens_to_file(tokens), + Self::Keyring(keyring_backend_kind) => save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + server_name, + tokens, + ), + } + } + + /// Deletes credentials only from this already-resolved authority. + pub(crate) fn delete( + self, + keyring_store: &K, + server_name: &str, + url: &str, + ) -> Result { + match self { + Self::File => { + let key = compute_store_key(server_name, url)?; + delete_oauth_tokens_from_file(&key) + } + Self::Keyring(AuthKeyringBackendKind::Direct) => { + delete_oauth_tokens_from_direct_keyring(keyring_store, server_name, url) + } + Self::Keyring(AuthKeyringBackendKind::Secrets) => { + delete_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url) + } + } + } +} + +#[derive(Debug)] +pub(crate) struct ResolvedOAuthTokens { + pub(crate) tokens: StoredOAuthTokens, + pub(crate) store: ResolvedOAuthCredentialStore, +} + +pub(crate) fn resolve_oauth_tokens_from_store_policy( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + match store_mode { + OAuthCredentialsStoreMode::Auto => { + // Auto remains keyring-first at lifecycle startup. The returned source is then pinned + // by the client transport recipe and OAuth persistor so retries, recovery, and + // refresh work cannot hot-switch stores. + // TODO(stevenlee): Different processes can still resolve Auto to different stores + // when keyring availability differs. Solving that safely requires durable backend + // selection or reconciliation of legacy entries and is intentionally outside this + // stack. + match load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) { + Ok(Some(tokens)) => Ok(Some(ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + Ok(None) => Ok( + load_oauth_tokens_from_file(server_name, url)?.map(|tokens| { + ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + } + }), + ), + // Auto may fall back when the keyring backend is unavailable, but a Secrets + // aggregate-lock failure means authority may be changing. Consulting File in + // that state could replay credentials hidden behind a newer Secrets entry. + Err(OAuthKeyringLoadError::StoreLock(error)) => Err(error.into()), + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + Ok(load_oauth_tokens_from_file(server_name, url) + .with_context(|| { + format!("failed to read OAuth tokens from keyring: {error}") + })? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + })) + } + } + } + OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( + |tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }, + )), + OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .map_err(anyhow::Error::from) + .context("failed to read OAuth tokens from keyring")? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + } +} diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index 5f83ebb8a63..b053f4e5ae2 100644 --- a/codex-rs/rmcp-client/src/oauth/store_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -13,7 +13,6 @@ use std::path::PathBuf; use std::time::Duration; use std::time::Instant; -use anyhow::Result; use codex_utils_home_dir::find_codex_home; const OAUTH_LOCK_DIR: &str = "mcp-oauth-locks"; @@ -52,7 +51,7 @@ pub(super) struct OAuthStoreLock { } impl OAuthStoreLock { - pub(super) fn acquire(store: OAuthStore) -> Result { + pub(super) fn acquire(store: OAuthStore) -> Result { // This lock intentionally follows the existing local File/Secrets credential-store // authority. Those stores are CODEX_HOME-backed today: if CODEX_HOME is unset they use // the default home (`~/.codex`), and if an embedder has no local home/filesystem authority @@ -67,7 +66,7 @@ impl OAuthStoreLock { codex_home: &Path, store: OAuthStore, acquire_timeout: Duration, - ) -> Result { + ) -> Result { let path = oauth_store_lock_path(codex_home, store); if let Some(parent) = path.parent() { fs::create_dir_all(parent).map_err(|source| OAuthStoreLockFailure::CreateDir { @@ -99,8 +98,7 @@ impl OAuthStoreLock { store, path, acquire_timeout, - } - .into()); + }); } Err(std::fs::TryLockError::WouldBlock) => { if !reported_contention { @@ -119,16 +117,13 @@ impl OAuthStoreLock { store, path, source: io::Error::from(error), - } - .into()); + }); } } } } } -/// Marks aggregate-store coordination failures in an [`anyhow::Error`] chain. -/// /// Auto may fall back when the configured keyring backend is unavailable, but it must surface a /// lock failure. Falling back while another process owns the aggregate-store lock could leave the /// newer credential in File while a stale Secrets entry remains preferred. diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index aafc362c204..9ecbf2fff6f 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -31,13 +31,14 @@ use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::fallback_file_path; use crate::oauth::load_oauth_tokens_from_file; use crate::oauth::load_oauth_tokens_from_keyring; -use crate::oauth::load_oauth_tokens_from_keyring_with_fallback_to_file; +use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth::save_oauth_tokens_to_file; use crate::oauth::save_oauth_tokens_to_file_with_lock_held; use crate::oauth::save_oauth_tokens_to_secrets_keyring_with_lock_held; use crate::oauth::save_oauth_tokens_with_keyring; use crate::oauth::save_oauth_tokens_with_keyring_with_fallback_to_file; use crate::oauth::test_support::TempCodexHome; +use codex_config::types::OAuthCredentialsStoreMode; const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention"; @@ -136,10 +137,7 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { } Err(error) => error, }; - assert!(matches!( - error.downcast_ref::(), - Some(OAuthStoreLockFailure::Timeout { .. }) - )); + assert!(matches!(error, OAuthStoreLockFailure::Timeout { .. })); child .kill() @@ -216,11 +214,12 @@ fn auto_load_secrets_lock_failure_does_not_fall_back_to_file() -> Result<()> { let lock_dir = env.path().join("mcp-oauth-locks"); std::fs::create_dir(lock_dir.join("secrets-store.lock"))?; - let error = load_oauth_tokens_from_keyring_with_fallback_to_file( + let error = resolve_oauth_tokens_from_store_policy( &keyring_store, - AuthKeyringBackendKind::Secrets, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Secrets, ) .expect_err("aggregate-store lock failure must abort Auto resolution"); @@ -445,12 +444,12 @@ fn secrets_store_load_and_delete_observe_aggregate_lock() -> Result<()> { let url = tokens.url.clone(); let loaded = complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { - load_oauth_tokens_from_keyring( + Ok(load_oauth_tokens_from_keyring( &store_for_load, AuthKeyringBackendKind::Secrets, &server_name, &url, - ) + )?) })? .expect("encrypted credentials should remain readable after contention"); assert_tokens_match_without_expiry(&loaded, &tokens); diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index c6527990fac..d6600d8e2a6 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -3,6 +3,7 @@ use std::ffi::OsString; use std::future::Future; use std::io; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -14,6 +15,7 @@ use codex_api::SharedAuthProvider; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerEnvVar; use codex_exec_server::HttpClient; +use codex_keyring_store::DefaultKeyringStore; use futures::FutureExt; use futures::future::BoxFuture; use oauth2::TokenResponse; @@ -64,9 +66,11 @@ use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::in_process_transport::InProcessTransportFactory; -use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; +use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; @@ -126,6 +130,7 @@ enum TransportRecipe { env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, + pinned_credential_store: Arc>, http_client: Arc, auth_provider: Option, }, @@ -402,6 +407,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + pinned_credential_store: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -781,6 +787,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + pinned_credential_store, http_client, auth_provider, } => { @@ -793,28 +800,56 @@ impl RmcpClient { auth_provider.clone() }; - let initial_oauth_tokens = if bearer_token.is_none() + let resolved_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(server_name, url, *store_mode, *keyring_backend_kind) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None + if let Some(store) = pinned_credential_store.get().copied() { + // Rebuilds reread the source selected during first construction. Only the + // initial construction below evaluates configured store policy. + store + .load(&DefaultKeyringStore, server_name, url)? + .map(|tokens| ResolvedOAuthTokens { tokens, store }) + } else { + match resolve_oauth_tokens_from_store_policy( + &DefaultKeyringStore, + server_name, + url, + *store_mode, + *keyring_backend_kind, + ) { + Ok(tokens) => { + if let Some(resolved) = tokens.as_ref() { + // Retries and session recovery rebuild this transport. Pin the + // first concrete source so Auto is not reevaluated mid-client. + pinned_credential_store.set(resolved.store).map_err(|_| { + anyhow!( + "OAuth credential store pinned concurrently for MCP server `{server_name}`" + ) + })?; + } + tokens + } + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None + } } } } else { None }; - if let Some(initial_tokens) = initial_oauth_tokens.clone() { + if let Some(ResolvedOAuthTokens { + tokens: initial_tokens, + store: credential_store, + }) = resolved_oauth_tokens + { match create_oauth_transport_and_runtime( server_name, url, initial_tokens.clone(), - *store_mode, - *keyring_backend_kind, + credential_store, default_headers.clone(), Arc::clone(http_client), ) @@ -1157,8 +1192,7 @@ async fn create_oauth_transport_and_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, - credentials_store: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, ) -> Result<( @@ -1202,8 +1236,7 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, - credentials_store, - keyring_backend_kind, + credential_store, Some(initial_tokens), ); diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs new file mode 100644 index 00000000000..d92df03dd25 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs @@ -0,0 +1,296 @@ +mod streamable_http_test_support; + +use std::any::Any; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use futures::future::BoxFuture; +use keyring::credential::Credential; +use keyring::credential::CredentialApi; +use keyring::credential::CredentialBuilderApi; +use keyring::credential::CredentialPersistence; +use oauth2::AccessToken; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use tempfile::TempDir; +use tokio::process::Command; + +use streamable_http_test_support::arm_session_post_failure; +use streamable_http_test_support::call_echo_tool; +use streamable_http_test_support::expected_echo_result; +use streamable_http_test_support::initialize_client; +use streamable_http_test_support::spawn_streamable_http_server; + +const SERVER_NAME: &str = "test-streamable-http-oauth-store-pinning"; +const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_PINNED_STORE_SERVER_URL"; +const KEYRING_ACCESS_TOKEN: &str = "keyring-access-token"; +const FILE_ACCESS_TOKEN: &str = "stale-file-access-token"; + +#[derive(Clone)] +struct RecordingHttpClient { + inner: Arc, + bearer_tokens: Arc>>, +} + +impl RecordingHttpClient { + fn new(inner: Arc) -> Self { + Self { + inner, + bearer_tokens: Arc::new(Mutex::new(Vec::new())), + } + } + + fn record_bearer_token(&self, params: &HttpRequestParams) { + let Some(header) = params + .headers + .iter() + .find(|header| header.name.eq_ignore_ascii_case("authorization")) + else { + return; + }; + self.bearer_tokens + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(header.value.clone()); + } + + fn bearer_tokens(&self) -> Vec { + self.bearer_tokens + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } +} + +impl HttpClient for RecordingHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + self.record_bearer_token(¶ms); + self.inner.http_request(params) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + self.record_bearer_token(¶ms); + self.inner.http_request_stream(params) + } +} + +#[derive(Debug, Default)] +struct TestKeyringState { + secret: Mutex>>, + fail_reads: AtomicBool, +} + +#[derive(Clone, Debug)] +struct TestCredential { + state: Arc, +} + +impl CredentialApi for TestCredential { + fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> { + *self + .state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some(secret.to_vec()); + Ok(()) + } + + fn get_secret(&self) -> keyring::Result> { + if self.state.fail_reads.load(Ordering::SeqCst) { + return Err(keyring::Error::Invalid( + "simulated keyring read failure".to_string(), + "load".to_string(), + )); + } + + self.state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + .ok_or(keyring::Error::NoEntry) + } + + fn delete_credential(&self) -> keyring::Result<()> { + self.state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .map(|_| ()) + .ok_or(keyring::Error::NoEntry) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[derive(Debug)] +struct TestCredentialBuilder { + state: Arc, +} + +impl CredentialBuilderApi for TestCredentialBuilder { + fn build( + &self, + _target: Option<&str>, + _service: &str, + _user: &str, + ) -> keyring::Result> { + Ok(Box::new(TestCredential { + state: Arc::clone(&self.state), + })) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn persistence(&self) -> CredentialPersistence { + CredentialPersistence::ProcessOnly + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn auto_store_remains_pinned_across_session_recovery() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let codex_home = TempDir::new()?; + + let status = Command::new(std::env::current_exe()?) + .args([ + "auto_store_remains_pinned_across_session_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &base_url) + .status() + .await?; + + assert!( + status.success(), + "OAuth store-pinning child failed: {status}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by auto_store_remains_pinned_across_session_recovery"] +async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Result<()> { + let state = Arc::new(TestKeyringState::default()); + keyring::set_default_credential_builder(Box::new(TestCredentialBuilder { + state: Arc::clone(&state), + })); + + let base_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let server_url = format!("{base_url}/mcp"); + let keyring_tokens = stored_tokens(&server_url, KEYRING_ACCESS_TOKEN); + save_oauth_tokens( + SERVER_NAME, + &keyring_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + let file_tokens = stored_tokens(&server_url, FILE_ACCESS_TOKEN); + save_oauth_tokens( + SERVER_NAME, + &file_tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + let http_client = RecordingHttpClient::new(Environment::default_for_tests().get_http_client()); + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + Arc::new(http_client.clone()), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + assert_eq!( + call_echo_tool(&client, "warmup").await?, + expected_echo_result("warmup") + ); + + arm_session_post_failure( + &base_url, + /*status*/ 404, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + // The selected keyring becomes unavailable only after initial construction. If recovery + // reevaluates Auto, it adopts the stale File token and this operation incorrectly succeeds. + state.fail_reads.store(true, Ordering::SeqCst); + + match call_echo_tool(&client, "recovery-must-not-fallback").await { + Ok(result) => assert_eq!(result, expected_echo_result("recovery-must-not-fallback")), + Err(error) => { + let error_chain = format!("{error:#}"); + assert!( + error_chain.contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected recovery error: {error_chain}" + ); + } + } + + let bearer_tokens = http_client.bearer_tokens(); + assert!( + bearer_tokens + .iter() + .any(|token| token == &format!("Bearer {KEYRING_ACCESS_TOKEN}")), + "expected requests authenticated by the keyring token: {bearer_tokens:?}" + ); + assert!( + bearer_tokens + .iter() + .all(|token| token != &format!("Bearer {FILE_ACCESS_TOKEN}")), + "stale File token must never be sent during recovery: {bearer_tokens:?}" + ); + Ok(()) +} + +fn stored_tokens(server_url: &str, access_token: &str) -> StoredOAuthTokens { + StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(OAuthTokenResponse::new( + AccessToken::new(access_token.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + )), + expires_at: None, + } +}