diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index d0005d20a6d0..46e6279a161e 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,12 @@ //! //! 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 store_lock; + +#[cfg(test)] +#[path = "oauth/test_support.rs"] +mod test_support; + use anyhow::Context; use anyhow::Error; use anyhow::Result; @@ -49,6 +55,10 @@ use std::time::SystemTime; use std::time::UNIX_EPOCH; use tracing::warn; +use self::store_lock::OAuthStore; +use self::store_lock::OAuthStoreLock; +use self::store_lock::OAuthStoreLockFailure; + use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use rmcp::transport::auth::AuthorizationManager; @@ -173,6 +183,11 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file 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) @@ -220,6 +235,7 @@ fn load_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -308,12 +324,39 @@ 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(()) +} + +/// Writes one credential to Secrets. The caller must hold the Secrets aggregate-store lock. +fn save_oauth_tokens_to_secrets_keyring_with_lock_held( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + serialized: &str, +) -> Result<()> { let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -323,14 +366,8 @@ fn save_oauth_tokens_to_secrets_keyring( ); let secret_name = compute_secret_name(server_name, &tokens.url)?; manager - .set(&SecretScope::Global, &secret_name, &serialized) - .context("failed to write OAuth tokens to encrypted storage")?; - - 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(()) + .set(&SecretScope::Global, &secret_name, serialized) + .context("failed to write OAuth tokens to encrypted storage") } fn save_oauth_tokens_with_keyring_with_fallback_to_file( @@ -341,6 +378,10 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file Result<()> { match save_oauth_tokens_with_keyring(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 + // behind a stale Secrets entry. + Err(error) if error.downcast_ref::().is_some() => Err(error), Err(error) => { let message = error.to_string(); warn!("falling back to file storage for OAuth tokens: {message}"); @@ -430,6 +471,7 @@ fn delete_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -592,7 +634,8 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { - let Some(store) = read_fallback_file()? else { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -634,9 +677,17 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + save_oauth_tokens_to_file_with_lock_held(tokens) +} + +/// Updates the fallback File. The caller must hold the File aggregate-store lock. +fn save_oauth_tokens_to_file_with_lock_held(tokens: &StoredOAuthTokens) -> Result<()> { let key = compute_store_key(&tokens.server_name, &tokens.url)?; - let mut store = read_fallback_file()?.unwrap_or_default(); + let mut store = read_fallback_file_unlocked()?.unwrap_or_default(); let token_response = &tokens.token_response.0; let expires_at = tokens @@ -664,7 +715,8 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> { } fn delete_oauth_tokens_from_file(key: &str) -> Result { - let mut store = match read_fallback_file()? { + let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let mut store = match read_fallback_file_unlocked()? { Some(store) => store, None => return Ok(false), }; @@ -750,7 +802,7 @@ fn fallback_file_path() -> Result { Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf()) } -fn read_fallback_file() -> Result> { +fn read_fallback_file_unlocked() -> Result> { let path = fallback_file_path()?; let contents = match fs::read_to_string(&path) { Ok(contents) => contents, @@ -814,52 +866,13 @@ fn sha_256_prefix(value: &Value) -> Result { mod tests { use super::*; use anyhow::Result; + use codex_keyring_store::tests::MockKeyringStore; use codex_secrets::compute_keyring_account; use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::sync::Arc; - use std::sync::Mutex; - use std::sync::MutexGuard; - use std::sync::OnceLock; - use std::sync::PoisonError; - use tempfile::tempdir; - - use codex_keyring_store::tests::MockKeyringStore; - - struct TempCodexHome { - _guard: MutexGuard<'static, ()>, - _dir: tempfile::TempDir, - } - - impl TempCodexHome { - fn new() -> Self { - static LOCK: OnceLock> = OnceLock::new(); - let guard = LOCK - .get_or_init(Mutex::default) - .lock() - .unwrap_or_else(PoisonError::into_inner); - let dir = tempdir().expect("create CODEX_HOME temp dir"); - unsafe { - std::env::set_var("CODEX_HOME", dir.path()); - } - Self { - _guard: guard, - _dir: dir, - } - } - fn path(&self) -> &std::path::Path { - self._dir.path() - } - } - - impl Drop for TempCodexHome { - fn drop(&mut self) { - unsafe { - std::env::remove_var("CODEX_HOME"); - } - } - } + use super::test_support::TempCodexHome; #[test] fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { @@ -964,7 +977,7 @@ mod tests { let fallback_path = super::fallback_file_path()?; assert!(fallback_path.exists(), "fallback file should be created"); - let saved = super::read_fallback_file()?.expect("fallback file should load"); + let saved = super::read_fallback_file_unlocked()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; let entry = saved.get(&key).expect("entry for key"); assert_eq!(entry.server_name, tokens.server_name); @@ -1076,7 +1089,7 @@ mod tests { &tokens, )?; - let saved = super::read_fallback_file()?.expect("fallback file should load"); + let saved = super::read_fallback_file_unlocked()?.expect("fallback file should load"); let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; assert!(saved.contains_key(&key)); Ok(()) diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs new file mode 100644 index 000000000000..5f83ebb8a63f --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -0,0 +1,181 @@ +//! Cross-process serialization for MCP OAuth stores shared by multiple credentials. +//! +//! File and Secrets each keep credentials for multiple MCP servers in one aggregate document. +//! Their lock therefore protects the complete read-modify-write operation. Direct keyring entries +//! are already stored independently per credential and do not use this lock. + +use std::fs; +use std::fs::File; +use std::fs::OpenOptions; +use std::io; +use std::path::Path; +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"; +const STORE_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60); +const STORE_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(50); +// Tests listen for this event so they prove a contender reached the real WouldBlock branch. +const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention"; + +#[derive(Clone, Copy, Debug)] +pub(super) enum OAuthStore { + File, + Secrets, +} + +impl OAuthStore { + fn lock_filename(self) -> &'static str { + match self { + Self::File => "file-store.lock", + Self::Secrets => "secrets-store.lock", + } + } +} + +impl std::fmt::Display for OAuthStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::File => f.write_str("fallback file"), + Self::Secrets => f.write_str("encrypted secrets"), + } + } +} + +/// Serializes one complete operation on an aggregate OAuth credential store. +pub(super) struct OAuthStoreLock { + _file: File, +} + +impl OAuthStoreLock { + 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 + // those stores already cannot operate. A future provider-backed credential store should + // provide its own matching lock authority instead of using this local path. + let codex_home = find_codex_home() + .map_err(|source| OAuthStoreLockFailure::CodexHome { store, source })?; + Self::acquire_in(&codex_home, store, STORE_LOCK_ACQUIRE_TIMEOUT) + } + + pub(super) fn acquire_in( + codex_home: &Path, + store: OAuthStore, + acquire_timeout: Duration, + ) -> 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 { + store, + path: parent.to_path_buf(), + source, + })?; + } + + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&path) + .map_err(|source| OAuthStoreLockFailure::Open { + store, + path: path.clone(), + source, + })?; + let started = Instant::now(); + let mut reported_contention = false; + + loop { + match file.try_lock() { + Ok(()) => return Ok(Self { _file: file }), + Err(std::fs::TryLockError::WouldBlock) if started.elapsed() >= acquire_timeout => { + return Err(OAuthStoreLockFailure::Timeout { + store, + path, + acquire_timeout, + } + .into()); + } + Err(std::fs::TryLockError::WouldBlock) => { + if !reported_contention { + tracing::debug!( + target: LOCK_CONTENTION_EVENT_TARGET, + store = %store, + lock_path = %path.display(), + "waiting for another process to finish updating MCP OAuth store state" + ); + reported_contention = true; + } + std::thread::sleep(STORE_LOCK_RETRY_SLEEP.min(acquire_timeout)); + } + Err(error) => { + return Err(OAuthStoreLockFailure::Lock { + 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. +#[derive(Debug, thiserror::Error)] +pub(super) enum OAuthStoreLockFailure { + #[error("failed to resolve CODEX_HOME for MCP OAuth {store} aggregate-store lock")] + CodexHome { + store: OAuthStore, + #[source] + source: io::Error, + }, + #[error("failed to create MCP OAuth {store} aggregate-store lock directory {}", path.display())] + CreateDir { + store: OAuthStore, + path: PathBuf, + #[source] + source: io::Error, + }, + #[error("failed to open MCP OAuth {store} aggregate-store lock {}", path.display())] + Open { + store: OAuthStore, + path: PathBuf, + #[source] + source: io::Error, + }, + #[error( + "timed out after {acquire_timeout:?} waiting for MCP OAuth {store} aggregate-store lock {}", + path.display() + )] + Timeout { + store: OAuthStore, + path: PathBuf, + acquire_timeout: Duration, + }, + #[error("failed to lock MCP OAuth {store} aggregate-store lock {}", path.display())] + Lock { + store: OAuthStore, + path: PathBuf, + #[source] + source: io::Error, + }, +} + +fn oauth_store_lock_path(codex_home: &Path, store: OAuthStore) -> PathBuf { + codex_home.join(OAUTH_LOCK_DIR).join(store.lock_filename()) +} + +#[cfg(test)] +#[path = "tests/store_lock_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/test_support.rs b/codex-rs/rmcp-client/src/oauth/test_support.rs new file mode 100644 index 000000000000..6ca2f510f5fa --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/test_support.rs @@ -0,0 +1,46 @@ +use std::path::Path; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::OnceLock; +use std::sync::PoisonError; + +use tempfile::tempdir; + +/// Serializes tests that mutate process-wide CODEX_HOME. +/// +/// Keep OAuth tests on this one guard instead of defining per-module helpers; otherwise +/// concurrently running test modules can point File/Secrets storage at different homes. +pub(super) struct TempCodexHome { + _guard: MutexGuard<'static, ()>, + _dir: tempfile::TempDir, +} + +impl TempCodexHome { + pub(super) fn new() -> Self { + static LOCK: OnceLock> = OnceLock::new(); + let guard = LOCK + .get_or_init(Mutex::default) + .lock() + .unwrap_or_else(PoisonError::into_inner); + let dir = tempdir().expect("create CODEX_HOME temp dir"); + unsafe { + std::env::set_var("CODEX_HOME", dir.path()); + } + Self { + _guard: guard, + _dir: dir, + } + } + + pub(super) fn path(&self) -> &Path { + self._dir.path() + } +} + +impl Drop for TempCodexHome { + fn drop(&mut self) { + unsafe { + std::env::remove_var("CODEX_HOME"); + } + } +} 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 new file mode 100644 index 000000000000..aafc362c204b --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -0,0 +1,480 @@ +use std::process::Command; +use std::sync::mpsc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_keyring_store::tests::MockKeyringStore; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::Scope; +use oauth2::TokenResponse; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use tracing::Event; +use tracing::Id; +use tracing::Metadata; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Record; +use tracing::subscriber::Interest; + +use super::OAuthStore; +use super::OAuthStoreLock; +use super::OAuthStoreLockFailure; +use crate::oauth::StoredOAuthTokens; +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::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; + +const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention"; + +fn assert_tokens_match_without_expiry(actual: &StoredOAuthTokens, expected: &StoredOAuthTokens) { + assert_eq!(actual.server_name, expected.server_name); + assert_eq!(actual.url, expected.url); + assert_eq!(actual.client_id, expected.client_id); + assert_eq!(actual.expires_at, expected.expires_at); + assert_token_response_match_without_expiry(&actual.token_response, &expected.token_response); +} + +fn assert_token_response_match_without_expiry( + actual: &WrappedOAuthTokenResponse, + expected: &WrappedOAuthTokenResponse, +) { + let actual_response = &actual.0; + let expected_response = &expected.0; + + assert_eq!( + actual_response.access_token().secret(), + expected_response.access_token().secret() + ); + assert_eq!(actual_response.token_type(), expected_response.token_type()); + assert_eq!( + actual_response.refresh_token().map(RefreshToken::secret), + expected_response.refresh_token().map(RefreshToken::secret), + ); + assert_eq!(actual_response.scopes(), expected_response.scopes()); + assert_eq!( + actual_response.extra_fields().0, + expected_response.extra_fields().0 + ); + assert_eq!( + actual_response.expires_in().is_some(), + expected_response.expires_in().is_some() + ); +} + +fn sample_tokens() -> StoredOAuthTokens { + let mut response = OAuthTokenResponse::new( + AccessToken::new("access-token".to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new("refresh-token".to_string()))); + response.set_scopes(Some(vec![ + Scope::new("scope-a".to_string()), + Scope::new("scope-b".to_string()), + ])); + let expires_in = Duration::from_secs(3600); + response.set_expires_in(Some(&expires_in)); + let expires_at = crate::oauth::compute_expires_at_millis(&response); + + StoredOAuthTokens { + server_name: "test-server".to_string(), + url: "https://example.test".to_string(), + client_id: "client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at, + } +} + +const LOCK_HOLDER_CHILD_TEST: &str = + "oauth::store_lock::tests::store_lock_is_released_when_holder_process_exits_child"; +const LOCK_HOLDER_READY_PATH_ENV: &str = "CODEX_OAUTH_STORE_LOCK_CHILD_READY_PATH"; + +#[test] +fn store_lock_is_released_when_holder_process_exits() -> Result<()> { + let env = TempCodexHome::new(); + let ready_file = env.path().join("lock-holder-ready"); + let mut child = Command::new(std::env::current_exe()?) + .arg("--exact") + .arg(LOCK_HOLDER_CHILD_TEST) + .arg("--ignored") + .env("CODEX_HOME", env.path()) + .env(LOCK_HOLDER_READY_PATH_ENV, &ready_file) + .spawn() + .context("spawn OAuth store lock holder test process")?; + + let test_result = (|| -> Result<()> { + let started = Instant::now(); + while !ready_file.exists() { + if started.elapsed() > Duration::from_secs(/*secs*/ 5) { + anyhow::bail!("timed out waiting for child process to acquire OAuth store lock"); + } + std::thread::sleep(Duration::from_millis(/*millis*/ 20)); + } + + let error = match OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::File, + Duration::from_millis(/*millis*/ 100), + ) { + Ok(_) => { + anyhow::bail!("live holder process should keep the OAuth store lock unavailable") + } + Err(error) => error, + }; + assert!(matches!( + error.downcast_ref::(), + Some(OAuthStoreLockFailure::Timeout { .. }) + )); + + child + .kill() + .context("kill OAuth store lock holder process")?; + let status = child + .wait() + .context("wait for killed OAuth store lock holder process")?; + assert!(!status.success()); + let _lock = OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::File, + Duration::from_secs(/*secs*/ 1), + )?; + Ok(()) + })(); + + if let Ok(None) = child.try_wait() { + let _ = child.kill(); + let _ = child.wait(); + } + + test_result +} + +#[test] +#[ignore = "child process for store_lock_is_released_when_holder_process_exits"] +fn store_lock_is_released_when_holder_process_exits_child() -> Result<()> { + let ready_file = match std::env::var_os(LOCK_HOLDER_READY_PATH_ENV) { + Some(path) => std::path::PathBuf::from(path), + None => return Ok(()), + }; + let _lock = OAuthStoreLock::acquire(OAuthStore::File)?; + std::fs::write(ready_file, b"ready")?; + loop { + std::thread::sleep(Duration::from_secs(/*secs*/ 60)); + } +} + +#[test] +fn auto_save_secrets_lock_failure_does_not_fall_back_to_file() -> Result<()> { + let env = TempCodexHome::new(); + let lock_dir = env.path().join("mcp-oauth-locks"); + std::fs::create_dir_all(&lock_dir)?; + // Break only the Secrets lock path. The distinct File lock remains usable, so Auto would + // successfully write fallback credentials if it mistook coordination failure for backend + // unavailability. + std::fs::create_dir(lock_dir.join("secrets-store.lock"))?; + let keyring_store = MockKeyringStore::default(); + let tokens = sample_tokens(); + + let error = save_oauth_tokens_with_keyring_with_fallback_to_file( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens, + ) + .expect_err("aggregate-store lock failure must abort Auto persistence"); + + assert!(error.downcast_ref::().is_some()); + assert!(!fallback_file_path()?.exists()); + save_oauth_tokens_to_file(&tokens)?; + let loaded = load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("fallback File should remain independently writable"); + assert_tokens_match_without_expiry(&loaded, &tokens); + Ok(()) +} + +#[test] +fn auto_load_secrets_lock_failure_does_not_fall_back_to_file() -> Result<()> { + let env = TempCodexHome::new(); + let keyring_store = MockKeyringStore::default(); + let tokens = sample_tokens(); + save_oauth_tokens_to_file(&tokens)?; + + 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( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens.url, + ) + .expect_err("aggregate-store lock failure must abort Auto resolution"); + + assert!(error.downcast_ref::().is_some()); + let loaded = load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("fallback File should remain independently readable"); + assert_tokens_match_without_expiry(&loaded, &tokens); + Ok(()) +} + +struct LockContentionSubscriber { + contended_tx: mpsc::Sender<()>, +} + +impl Subscriber for LockContentionSubscriber { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target() == STORE_LOCK_CONTENTION_EVENT_TARGET + } + + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if self.enabled(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn max_level_hint(&self) -> Option { + Some(tracing::level_filters::LevelFilter::DEBUG) + } + + fn new_span(&self, _span: &Attributes<'_>) -> Id { + Id::from_u64(/*u*/ 1) + } + + fn record(&self, _span: &Id, _values: &Record<'_>) {} + + fn record_follows_from(&self, _span: &Id, _follows_from: &Id) {} + + fn event(&self, event: &Event<'_>) { + if self.enabled(event.metadata()) { + self.contended_tx + .send(()) + .expect("signal actual OAuth store lock contention"); + } + } + + fn enter(&self, _span: &Id) {} + + fn exit(&self, _span: &Id) {} +} + +fn complete_after_store_lock_contention( + codex_home: &std::path::Path, + store: OAuthStore, + operation: impl FnOnce() -> Result + Send + 'static, +) -> Result +where + T: Send + 'static, +{ + let held_lock = + OAuthStoreLock::acquire_in(codex_home, store, Duration::from_millis(/*millis*/ 100))?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let worker = std::thread::spawn(move || { + tracing::subscriber::with_default(LockContentionSubscriber { contended_tx }, || { + result_tx + .send(operation()) + .expect("send contending OAuth store operation result"); + }); + }); + + // This event is emitted only after `try_lock()` returns WouldBlock, so the test fails if the + // operation stops acquiring the aggregate-store lock. + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + drop(held_lock); + let result = result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; + worker + .join() + .expect("contending OAuth store worker should finish"); + Ok(result) +} + +#[test] +fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let env = TempCodexHome::new(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::File, + Duration::from_millis(/*millis*/ 100), + )?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + tracing::subscriber::with_default(LockContentionSubscriber { contended_tx }, || { + result_tx + .send(save_oauth_tokens_to_file(&second_for_writer)) + .expect("send writer result"); + }); + }); + + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + save_oauth_tokens_to_file_with_lock_held(&first)?; + drop(held_lock); + result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; + writer.join().expect("file store writer should finish"); + + let loaded_first = load_oauth_tokens_from_file(&first.server_name, &first.url)? + .expect("first server tokens should remain stored"); + let loaded_second = load_oauth_tokens_from_file(&second.server_name, &second.url)? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) +} + +#[test] +fn file_store_load_and_delete_observe_aggregate_lock() -> Result<()> { + let env = TempCodexHome::new(); + let tokens = sample_tokens(); + save_oauth_tokens_to_file(&tokens)?; + + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let loaded = complete_after_store_lock_contention(env.path(), OAuthStore::File, move || { + load_oauth_tokens_from_file(&server_name, &url) + })? + .expect("file credentials should remain readable after contention"); + assert_tokens_match_without_expiry(&loaded, &tokens); + + let key = crate::oauth::compute_store_key(&tokens.server_name, &tokens.url)?; + let removed = complete_after_store_lock_contention(env.path(), OAuthStore::File, move || { + crate::oauth::delete_oauth_tokens_from_file(&key) + })?; + assert!(removed); + assert!(load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)?.is_none()); + Ok(()) +} + +#[test] +fn secrets_store_lock_preserves_updates_for_different_servers() -> Result<()> { + let env = TempCodexHome::new(); + let keyring_store = MockKeyringStore::default(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.server_name = "second-server".to_string(); + second.url = "https://second.example.test".to_string(); + + let held_lock = OAuthStoreLock::acquire_in( + env.path(), + OAuthStore::Secrets, + Duration::from_millis(/*millis*/ 100), + )?; + let (contended_tx, contended_rx) = mpsc::channel(); + let (result_tx, result_rx) = mpsc::channel(); + let store_for_writer = keyring_store.clone(); + let second_for_writer = second.clone(); + let writer = std::thread::spawn(move || { + tracing::subscriber::with_default(LockContentionSubscriber { contended_tx }, || { + result_tx + .send(save_oauth_tokens_with_keyring( + &store_for_writer, + AuthKeyringBackendKind::Secrets, + &second_for_writer.server_name, + &second_for_writer, + )) + .expect("send writer result"); + }); + }); + + contended_rx.recv_timeout(Duration::from_secs(/*secs*/ 1))?; + let first_serialized = serde_json::to_string(&first)?; + save_oauth_tokens_to_secrets_keyring_with_lock_held( + &keyring_store, + &first.server_name, + &first, + &first_serialized, + )?; + drop(held_lock); + result_rx.recv_timeout(Duration::from_secs(/*secs*/ 10))??; + writer.join().expect("secrets store writer should finish"); + + let loaded_first = load_oauth_tokens_from_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &first.server_name, + &first.url, + )? + .expect("first server tokens should remain stored"); + let loaded_second = load_oauth_tokens_from_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &second.server_name, + &second.url, + )? + .expect("second server tokens should be stored"); + assert_tokens_match_without_expiry(&loaded_first, &first); + assert_tokens_match_without_expiry(&loaded_second, &second); + Ok(()) +} + +#[test] +fn secrets_store_load_and_delete_observe_aggregate_lock() -> Result<()> { + let env = TempCodexHome::new(); + let keyring_store = MockKeyringStore::default(); + let tokens = sample_tokens(); + save_oauth_tokens_with_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens, + )?; + + let store_for_load = keyring_store.clone(); + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let loaded = + complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { + 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); + + let store_for_delete = keyring_store.clone(); + let server_name = tokens.server_name.clone(); + let url = tokens.url.clone(); + let removed = + complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { + crate::oauth::delete_oauth_tokens_from_secrets_keyring( + &store_for_delete, + &server_name, + &url, + ) + })?; + assert!(removed); + assert!( + load_oauth_tokens_from_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens.url, + )? + .is_none() + ); + Ok(()) +}