diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 2040955c1c..9eae95f008 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -69,7 +69,6 @@ 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::load_oauth_tokens_from_store; pub(crate) use self::resolved_store::resolve_oauth_tokens; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; @@ -168,10 +167,11 @@ fn load_oauth_tokens_from_keyring( 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) @@ -201,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, @@ -225,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, @@ -334,7 +345,12 @@ fn save_oauth_tokens_with_keyring_and_cleanup_file { let mut last_credentials = self.inner.last_credentials.lock().await; if last_credentials.take().is_some() - && let Err(error) = delete_from_resolved_store( + && let Err(error) = self.inner.credential_store.delete( &DefaultKeyringStore, &self.inner.server_name, &self.inner.url, - self.inner.credential_store, ) { warn!( @@ -582,44 +601,6 @@ impl OAuthPersistor { } } -fn save_to_resolved_store( - keyring_store: &K, - inner: &OAuthPersistorInner, - tokens: &StoredOAuthTokens, -) -> Result<()> { - match inner.credential_store { - ResolvedOAuthCredentialStore::File => save_oauth_tokens_to_file(tokens), - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - save_oauth_tokens_with_keyring( - keyring_store, - keyring_backend_kind, - &inner.server_name, - tokens, - ) - } - } -} - -fn delete_from_resolved_store( - keyring_store: &K, - server_name: &str, - url: &str, - credential_store: ResolvedOAuthCredentialStore, -) -> Result { - match credential_store { - ResolvedOAuthCredentialStore::File => { - let key = compute_store_key(server_name, url)?; - delete_oauth_tokens_from_file(&key) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) => { - delete_oauth_tokens_from_direct_keyring(keyring_store, server_name, url) - } - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Secrets) => { - delete_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url) - } - } -} - const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; @@ -974,13 +955,9 @@ mod tests { )?; assert_eq!(fs::read(fallback_path)?, fallback_before); - let loaded = super::load_oauth_tokens_from_store( - &store, - &keyring_tokens.server_name, - &keyring_tokens.url, - ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), - )? - .expect("tokens should load from the selected keyring store"); + 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 index e551c4ed45..e77b92b0ee 100644 --- a/codex-rs/rmcp-client/src/oauth/resolved_store.rs +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -7,10 +7,16 @@ 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::store_lock::OAuthStoreLockFailure; +use super::save_oauth_tokens_to_file; +use super::save_oauth_tokens_with_keyring; /// Concrete credential store resolved for one MCP OAuth client lifecycle. /// @@ -24,6 +30,72 @@ pub(crate) enum ResolvedOAuthCredentialStore { Keyring(AuthKeyringBackendKind), } +impl ResolvedOAuthCredentialStore { + /// Loads credentials only from this already-resolved authority. + /// + /// Unlike `resolve_oauth_tokens`, 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, @@ -67,7 +139,7 @@ pub(crate) fn resolve_oauth_tokens( // 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(error) if error.downcast_ref::().is_some() => Err(error), + 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) @@ -93,33 +165,11 @@ pub(crate) fn resolve_oauth_tokens( server_name, url, ) - .with_context(|| "failed to read OAuth tokens from keyring".to_string())? + .map_err(anyhow::Error::from) + .context("failed to read OAuth tokens from keyring")? .map(|tokens| ResolvedOAuthTokens { tokens, store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), })), } } - -pub(crate) fn load_oauth_tokens_from_store( - keyring_store: &K, - server_name: &str, - url: &str, - store: ResolvedOAuthCredentialStore, -) -> Result> { - match store { - ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) - .context("failed to reread OAuth tokens from resolved file storage"), - ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { - load_oauth_tokens_from_keyring( - keyring_store, - keyring_backend_kind, - server_name, - url, - ) - .context( - "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", - ) - } - } -} diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index 5f83ebb8a6..b053f4e5ae 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 45278e2d12..0fe6d40874 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 @@ -137,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() @@ -447,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 4a5d62e0d2..afd1e3e24d 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -70,7 +70,6 @@ use crate::oauth::OAuthPersistor; use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; -use crate::oauth::load_oauth_tokens_from_store; use crate::oauth::resolve_oauth_tokens; use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::stdio_server_launcher::StdioServerCommand; @@ -806,7 +805,10 @@ impl RmcpClient { && !default_headers.contains_key(AUTHORIZATION) { if let Some(store) = resolved_store.get().copied() { - load_oauth_tokens_from_store(&DefaultKeyringStore, server_name, url, store)? + // 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( 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 0000000000..0fd7f0aa05 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs @@ -0,0 +1,219 @@ +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_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +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(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 client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + Environment::default_for_tests().get_http_client(), + /*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); + + let error = call_echo_tool(&client, "recovery-must-not-fallback") + .await + .expect_err("recovery must surface failure from the pinned keyring store"); + let error_chain = format!("{error:#}"); + assert!( + error_chain.contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected recovery error: {error_chain}" + ); + 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, + } +}