diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index d567c52fec..29890ff09a 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -93,7 +93,7 @@ keyring = { workspace = true, features = ["sync-secret-service"] } # This test is compiled through `#[path]` inside the inline `oauth::tests` module. Cargo-shear # cannot resolve that nested module path and otherwise reports the linked file as unlinked. [package.metadata.cargo-shear] -ignored-paths = ["src/oauth/tests/persistor_tests.rs"] +ignored-paths = ["src/oauth/tests/persistor_tests.rs", "src/oauth/tests/credential_store_tests.rs"] [lib] doctest = false diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 2880d8cba1..3b28dc2273 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,8 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +#[cfg(test)] +mod credential_store; mod ema_identity; mod issuer_binding; mod refresh_lock; @@ -72,6 +74,8 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; +#[cfg(test)] +pub(crate) use self::credential_store::OAuthCredentialStore; pub(crate) use self::ema_identity::stored_oidc_identity; pub(crate) use self::issuer_binding::validate_authorization_server_endpoints; pub(crate) use self::issuer_binding::validate_refresh_token_issuer; @@ -1121,6 +1125,8 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::sync::Arc; + #[path = "credential_store_tests.rs"] + mod credential_store_tests; #[path = "persistor_tests.rs"] mod persistor_tests; diff --git a/codex-rs/rmcp-client/src/oauth/credential_store.rs b/codex-rs/rmcp-client/src/oauth/credential_store.rs new file mode 100644 index 0000000000..01c0114332 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/credential_store.rs @@ -0,0 +1,261 @@ +//! Adapts Codex's pinned credential storage to RMCP-owned OAuth refreshes. +//! +//! Attach this store only after `AuthorizationManager::initialize_from_store` completes +//! against an in-memory store. Initialization may save tokenless client credentials, which +//! this refresh adapter does not support. +//! +//! RMCP must hold the credential lock across its authoritative reread, token exchange, and +//! save. Saves and clears require an active guard; no store operation reacquires the lock. +//! The runtime snapshot advances only for credentials compatible with this connection, +//! so replacement logins still cause a rebuild. +//! Synchronous store operations run off the Tokio workers so caller deadlines remain pollable. +//! Blocking mutations retain the caller's transaction guard even if their await is cancelled. + +use std::sync::Arc; +use std::sync::Weak; + +use anyhow::Context; +use anyhow::Result; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use futures::future::BoxFuture; +use oauth2::Scope; +use oauth2::TokenResponse; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::CredentialStore; +use rmcp::transport::auth::StoredCredentials; +use tokio::sync::Mutex; +use tracing::warn; + +use super::RefreshCredentialLock; +use super::ResolvedOAuthCredentialStore; +use super::StoredOAuthTokens; +use super::WrappedOAuthTokenResponse; +use super::normalized_oauth_credentials; +use super::refresh_expires_in_from_timestamp; + +#[derive(Clone)] +pub(crate) struct OAuthCredentialStore { + inner: Arc>, +} + +struct OAuthCredentialStoreInner { + server_name: String, + url: String, + client_id: String, + issuer: Option, + store: ResolvedOAuthCredentialStore, + keyring: K, + last_credentials: Mutex>, + refresh_guard: Mutex>, +} + +impl OAuthCredentialStore { + pub(crate) fn new( + tokens: StoredOAuthTokens, + store: ResolvedOAuthCredentialStore, + keyring: K, + ) -> Self { + Self { + inner: Arc::new(OAuthCredentialStoreInner { + server_name: tokens.server_name.clone(), + url: tokens.url.clone(), + client_id: tokens.client_id.clone(), + issuer: tokens.bound_issuer().map(str::to_owned), + store, + keyring, + last_credentials: Mutex::new(Some(tokens)), + refresh_guard: Mutex::new(Weak::new()), + }), + } + } + + pub(crate) async fn stored_credentials(&self) -> Option { + normalized_oauth_credentials(self.inner.last_credentials.lock().await.as_ref()) + } + + pub(crate) async fn acquire_transaction_guard(&self) -> Result> { + let guard = Arc::new( + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?, + ); + *self.inner.refresh_guard.lock().await = Arc::downgrade(&guard); + Ok(guard) + } + + fn validate_connection(&self, tokens: &StoredOAuthTokens) -> Result<(), AuthError> { + if tokens.client_id != self.inner.client_id + || tokens.bound_issuer() != self.inner.issuer.as_deref() + || tokens.has_refresh_token() && tokens.bound_issuer().is_none() + { + warn!( + server_name = %self.inner.server_name, + "stored OAuth credentials no longer match this connection's client and issuer; authorization must be rebuilt" + ); + return Err(AuthError::AuthorizationRequired); + } + if !tokens.has_refresh_token() && !tokens.access_token_is_usable_without_refresh() { + return Err(AuthError::TokenExpired); + } + Ok(()) + } +} + +// Spell out RMCP's object-safe boxed-future interface without introducing async-trait here. +impl CredentialStore for OAuthCredentialStore { + fn load<'life0, 'async_trait>( + &'life0 self, + ) -> BoxFuture<'async_trait, Result, AuthError>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let inner = Arc::clone(&self.inner); + let Some(tokens) = tokio::task::spawn_blocking(move || { + inner + .store + .load(&inner.keyring, &inner.server_name, &inner.url) + }) + .await + .context("OAuth credential load task failed") + .map_err(credential_store_error)? + .map_err(credential_store_error)? + else { + // Do not acknowledge removal as a live connection update: reconciliation + // must see that this connection still belongs to the previous credentials. + return Ok(None); + }; + self.validate_connection(&tokens)?; + let credentials = rmcp_credentials(&tokens); + *self.inner.last_credentials.lock().await = Some(tokens); + Ok(Some(credentials)) + }) + } + + fn save<'life0, 'async_trait>( + &'life0 self, + credentials: StoredCredentials, + ) -> BoxFuture<'async_trait, Result<(), AuthError>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let mut token_response = credentials + .token_response + .ok_or(AuthError::AuthorizationRequired)?; + // Codex stores granted scopes in the token response rather than a separate field. + if token_response.scopes().is_none() && !credentials.granted_scopes.is_empty() { + token_response.set_scopes(Some( + credentials + .granted_scopes + .into_iter() + .map(Scope::new) + .collect(), + )); + } + // The SDK's receipt time, rather than the time this save completes, owns expiry. + if credentials.token_received_at.is_none() { + token_response.set_expires_in(None); + } + let expires_at = credentials.token_received_at.and_then(|received_at| { + token_response.expires_in().map(|expires_in| { + received_at + .saturating_mul(1000) + .saturating_add(u64::try_from(expires_in.as_millis()).unwrap_or(u64::MAX)) + }) + }); + let tokens = StoredOAuthTokens { + server_name: self.inner.server_name.clone(), + url: self.inner.url.clone(), + issuer: credentials.issuer, + client_id: credentials.client_id, + token_response: WrappedOAuthTokenResponse(token_response), + expires_at, + }; + self.validate_connection(&tokens)?; + let inner = Arc::clone(&self.inner); + let guard = self + .inner + .refresh_guard + .lock() + .await + .upgrade() + .context("OAuth credential mutation requires an active refresh guard") + .map_err(credential_store_error)?; + let tokens = tokio::task::spawn_blocking(move || -> Result { + let _guard = guard; + inner + .store + .save(&inner.keyring, &inner.server_name, &tokens)?; + Ok(tokens) + }) + .await + .context("OAuth credential save task failed") + .map_err(credential_store_error)? + .map_err(credential_store_error)?; + *self.inner.last_credentials.lock().await = Some(tokens); + Ok(()) + }) + } + + fn clear<'life0, 'async_trait>(&'life0 self) -> BoxFuture<'async_trait, Result<(), AuthError>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let inner = Arc::clone(&self.inner); + let guard = self + .inner + .refresh_guard + .lock() + .await + .upgrade() + .context("OAuth credential mutation requires an active refresh guard") + .map_err(credential_store_error)?; + tokio::task::spawn_blocking(move || -> Result<()> { + let _guard = guard; + inner + .store + .delete(&inner.keyring, &inner.server_name, &inner.url)?; + Ok(()) + }) + .await + .context("OAuth credential removal task failed") + .map_err(credential_store_error)? + .map_err(credential_store_error) + }) + } +} + +fn credential_store_error(error: anyhow::Error) -> AuthError { + AuthError::InternalError(format!("{error:#}")) +} + +fn rmcp_credentials(tokens: &StoredOAuthTokens) -> StoredCredentials { + let mut tokens = tokens.clone(); + refresh_expires_in_from_timestamp(&mut tokens); + let token_received_at = tokens.expires_at.map(|expires_at| { + let remaining = tokens.token_response.0.expires_in().unwrap_or_default(); + // Reconstruct the receipt time from the authority's deadline; a second clock read + // could otherwise extend the grant if this task paused between the two reads. + expires_at.saturating_sub(u64::try_from(remaining.as_millis()).unwrap_or(u64::MAX)) / 1000 + }); + if token_received_at.is_none() { + tokens.token_response.0.set_expires_in(None); + } + let token_response = tokens.token_response.0; + let granted_scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); + StoredCredentials::new( + tokens.client_id, + Some(token_response), + granted_scopes, + token_received_at, + ) + .with_issuer(tokens.issuer) +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs new file mode 100644 index 0000000000..bd5d2659cd --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs @@ -0,0 +1,260 @@ +//! Verifies guarded mutations, pinned storage, token mapping, and saved runtime snapshots. + +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use keyring::Error as KeyringError; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::TokenResponse; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::CredentialStore; + +use super::MockKeyringStore; +use super::TempCodexHome; +use super::assert_tokens_match_without_expiry; +use super::sample_tokens; +use crate::oauth::OAuthCredentialStore; +use crate::oauth::RefreshCredentialLock; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::compute_store_key; +use crate::oauth::load_oauth_tokens_from_file; +use crate::oauth::normalized_oauth_credentials; +use crate::oauth::save_oauth_tokens_to_file; +use crate::oauth::store_lock::OAuthStore; +use crate::oauth::store_lock::OAuthStoreLock; + +#[tokio::test(flavor = "current_thread")] +async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> { + for clear in [false, true] { + let _env = TempCodexHome::new(); + let initial = sample_tokens(); + save_oauth_tokens_to_file(&initial)?; + let store = OAuthCredentialStore::new( + initial.clone(), + ResolvedOAuthCredentialStore::File, + MockKeyringStore::default(), + ); + let guard = store.acquire_transaction_guard().await?; + let credentials = store.load().await?.unwrap(); + if !clear { + store.clear().await?; + } + drop(guard); + let result = if clear { + store.clear().await + } else { + store.save(credentials.clone()).await + }; + assert!( + matches!(result, Err(AuthError::InternalError(error)) if error.contains("active refresh guard")) + ); + assert_eq!( + load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none(), + !clear, + ); + let aggregate_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; + let guard = store.acquire_transaction_guard().await?; + let mut mutation = if clear { + store.clear() + } else { + store.save(credentials) + }; + assert!(futures::poll!(&mut mutation).is_pending()); + drop(mutation); + drop(guard); + + // The aggregate lock gates I/O, but only the retained refresh guard blocks this probe. + let mut contender = Box::pin(RefreshCredentialLock::acquire_for_server( + &initial.server_name, + &initial.url, + )); + assert!(futures::poll!(&mut contender).is_pending()); + drop(aggregate_lock); + let _guard = tokio::time::timeout(Duration::from_secs(/*secs*/ 5), contender).await??; + let durable = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?; + assert_eq!(durable.is_none(), clear); + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn save_publishes_only_persisted_credentials() -> Result<()> { + for fail_save in [false, true] { + let _env = TempCodexHome::new(); + let initial = sample_tokens(); + let keyring = MockKeyringStore::default(); + let authority = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct); + authority.save(&keyring, &initial.server_name, &initial)?; + let mut fallback = initial.clone(); + fallback + .token_response + .0 + .set_access_token(AccessToken::new("fallback-token".into())); + save_oauth_tokens_to_file(&fallback)?; + let store = OAuthCredentialStore::new(initial.clone(), authority, keyring.clone()); + let _guard = store.acquire_transaction_guard().await?; + let mut credentials = store.load().await?.expect("stored credentials"); + let token_response = credentials + .token_response + .as_mut() + .expect("stored token response"); + token_response.set_access_token(AccessToken::new("rotated-access".into())); + token_response.set_refresh_token(Some(RefreshToken::new("rotated-refresh".into()))); + // The adapter must encode the separate authoritative grant in the stored response. + token_response.set_scopes(/*scopes*/ None); + if fail_save { + let key = compute_store_key(&initial.server_name, &initial.url)?; + keyring.set_error(&key, KeyringError::Invalid("test".into(), "save".into())); + } + let result = store.save(credentials).await; + let durable = authority + .load(&keyring, &initial.server_name, &initial.url)? + .expect("durable keyring credentials"); + let mut expected = initial.clone(); + if fail_save { + assert!(matches!(result, Err(AuthError::InternalError(_)))); + } else { + result?; + expected + .token_response + .0 + .set_access_token(AccessToken::new("rotated-access".into())); + expected + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new("rotated-refresh".into()))); + expected.expires_at = durable.expires_at; + } + assert_tokens_match_without_expiry(&durable, &expected); + assert_eq!( + store.stored_credentials().await, + normalized_oauth_credentials(Some(&expected)) + ); + assert_tokens_match_without_expiry( + &load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.unwrap(), + &fallback, + ); + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn pinned_read_failure_does_not_adopt_fallback_credentials() -> Result<()> { + let _env = TempCodexHome::new(); + let initial = sample_tokens(); + save_oauth_tokens_to_file(&initial)?; + let keyring = MockKeyringStore::default(); + let key = compute_store_key(&initial.server_name, &initial.url)?; + keyring.set_error(&key, KeyringError::Invalid("test".into(), "load".into())); + let store = OAuthCredentialStore::new( + initial.clone(), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + keyring, + ); + assert!(matches!( + store.load().await, + Err(AuthError::InternalError(_)) + )); + assert_eq!( + store.stored_credentials().await, + normalized_oauth_credentials(Some(&initial)) + ); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn replacement_or_removal_does_not_acknowledge_a_new_runtime_snapshot() -> Result<()> { + let _env = TempCodexHome::new(); + let initial = sample_tokens(); + save_oauth_tokens_to_file(&initial)?; + let store = OAuthCredentialStore::new( + initial.clone(), + ResolvedOAuthCredentialStore::File, + MockKeyringStore::default(), + ); + let original_snapshot = store.stored_credentials().await; + for (client_id, issuer) in [ + ("replacement-client", initial.issuer.clone()), + ( + initial.client_id.as_str(), + Some("https://replacement.example.test".into()), + ), + (initial.client_id.as_str(), None), + ] { + let mut replacement = initial.clone(); + replacement.client_id = client_id.into(); + replacement.issuer = issuer; + save_oauth_tokens_to_file(&replacement)?; + assert!(matches!( + store.load().await, + Err(AuthError::AuthorizationRequired) + )); + assert_eq!(store.stored_credentials().await, original_snapshot); + } + let _guard = store.acquire_transaction_guard().await?; + store.clear().await?; + assert!(store.load().await?.is_none()); + assert_eq!(store.stored_credentials().await, original_snapshot); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()> { + let _env = TempCodexHome::new(); + let initial = sample_tokens(); + save_oauth_tokens_to_file(&initial)?; + let store = OAuthCredentialStore::new( + initial.clone(), + ResolvedOAuthCredentialStore::File, + MockKeyringStore::default(), + ); + let future_received_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); + let future_deadline = (future_received_at + 120) * 1000; + for (received_at, expected_deadline) in [ + (Some(future_received_at), Some(future_deadline)), + (None, None), + ] { + let _guard = store.acquire_transaction_guard().await?; + let mut credentials = store.load().await?.unwrap(); + credentials.token_received_at = received_at; + credentials + .token_response + .as_mut() + .unwrap() + .set_expires_in(Some(&Duration::from_secs(/*secs*/ 120))); + store.save(credentials).await?; + let durable = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.unwrap(); + let mut expected = initial.clone(); + expected.expires_at = expected_deadline; + expected + .token_response + .0 + .set_expires_in(expected_deadline.map(|_| Duration::ZERO).as_ref()); + assert_tokens_match_without_expiry(&durable, &expected); + let reloaded = store.load().await?.unwrap(); + let expires_in = reloaded.token_response.unwrap().expires_in(); + if let Some(expected_deadline) = expected_deadline { + let expires_in = expires_in.expect("future token expiry"); + assert!(!expires_in.is_zero()); + let reconstructed_deadline = + Duration::from_secs(reloaded.token_received_at.expect("token receipt time")) + + expires_in; + let expected_deadline = Duration::from_millis(expected_deadline); + // Receipt times use whole seconds, so reconstruction may lose less than a second. + assert!(reconstructed_deadline <= expected_deadline); + assert!(expected_deadline - reconstructed_deadline < Duration::from_secs(/*secs*/ 1)); + } else { + assert_eq!(expires_in, None); + } + assert_eq!( + store.stored_credentials().await, + normalized_oauth_credentials(Some(&expected)) + ); + } + Ok(()) +}