diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 913108b8fe..6af8d95f2f 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -282,7 +282,7 @@ fn load_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> std::result::Result, OAuthKeyringLoadError> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_read(OAuthStore::Secrets)?; 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(), @@ -386,7 +386,7 @@ fn save_oauth_tokens_to_secrets_keyring( tokens: &StoredOAuthTokens, ) -> Result<()> { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?; save_oauth_tokens_to_secrets_keyring_with_lock_held( keyring_store, server_name, @@ -541,7 +541,7 @@ fn delete_oauth_tokens_from_secrets_keyring( server_name: &str, url: &str, ) -> Result { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::Secrets)?; let codex_home = find_codex_home()?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), @@ -681,7 +681,7 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_read(OAuthStore::File)?; let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -739,7 +739,7 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; save_oauth_tokens_to_file_with_lock_held(tokens) } @@ -779,7 +779,7 @@ fn save_oauth_tokens_to_file_with_lock_held(tokens: &StoredOAuthTokens) -> Resul } fn delete_oauth_tokens_from_file(key: &str) -> Result { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _store_lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; let mut store = match read_fallback_file_unlocked()? { Some(store) => store, None => return Ok(false), diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index b053f4e5ae..0a6f3c6494 100644 --- a/codex-rs/rmcp-client/src/oauth/store_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -1,8 +1,9 @@ //! 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. +//! Writers hold an exclusive lock across the complete read-modify-write operation. Readers share +//! the same lock so concurrent MCP startup and status checks do not serialize behind one another. +//! Direct keyring entries are already stored independently per credential and do not use this lock. use std::fs; use std::fs::File; @@ -50,8 +51,25 @@ pub(super) struct OAuthStoreLock { _file: File, } +#[derive(Clone, Copy, Debug)] +enum OAuthStoreLockMode { + Shared, + Exclusive, +} + impl OAuthStoreLock { - pub(super) fn acquire(store: OAuthStore) -> Result { + pub(super) fn acquire_for_write(store: OAuthStore) -> Result { + Self::acquire_with_mode(store, OAuthStoreLockMode::Exclusive) + } + + pub(super) fn acquire_for_read(store: OAuthStore) -> Result { + Self::acquire_with_mode(store, OAuthStoreLockMode::Shared) + } + + fn acquire_with_mode( + store: OAuthStore, + mode: OAuthStoreLockMode, + ) -> 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 @@ -59,13 +77,14 @@ impl OAuthStoreLock { // 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) + Self::acquire_in_with_mode(&codex_home, store, STORE_LOCK_ACQUIRE_TIMEOUT, mode) } - pub(super) fn acquire_in( + fn acquire_in_with_mode( codex_home: &Path, store: OAuthStore, acquire_timeout: Duration, + mode: OAuthStoreLockMode, ) -> Result { let path = oauth_store_lock_path(codex_home, store); if let Some(parent) = path.parent() { @@ -91,7 +110,11 @@ impl OAuthStoreLock { let mut reported_contention = false; loop { - match file.try_lock() { + let result = match mode { + OAuthStoreLockMode::Shared => file.try_lock_shared(), + OAuthStoreLockMode::Exclusive => file.try_lock(), + }; + match result { Ok(()) => return Ok(Self { _file: file }), Err(std::fs::TryLockError::WouldBlock) if started.elapsed() >= acquire_timeout => { return Err(OAuthStoreLockFailure::Timeout { 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 3278365e01..037f05cff4 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 @@ -27,6 +27,7 @@ use tracing::subscriber::Interest; use super::OAuthStore; use super::OAuthStoreLock; use super::OAuthStoreLockFailure; +use super::OAuthStoreLockMode; use crate::oauth::StoredOAuthTokens; use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::fallback_file_path; @@ -173,10 +174,11 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { std::thread::sleep(Duration::from_millis(/*millis*/ 20)); } - let error = match OAuthStoreLock::acquire_in( + let error = match OAuthStoreLock::acquire_in_with_mode( env.path(), OAuthStore::File, Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, ) { Ok(_) => { anyhow::bail!("live holder process should keep the OAuth store lock unavailable") @@ -192,10 +194,11 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { .wait() .context("wait for killed OAuth store lock holder process")?; assert!(!status.success()); - let _lock = OAuthStoreLock::acquire_in( + let _lock = OAuthStoreLock::acquire_in_with_mode( env.path(), OAuthStore::File, Duration::from_secs(/*secs*/ 1), + OAuthStoreLockMode::Exclusive, )?; Ok(()) })(); @@ -215,7 +218,7 @@ fn store_lock_is_released_when_holder_process_exits_child() -> Result<()> { Some(path) => std::path::PathBuf::from(path), None => return Ok(()), }; - let _lock = OAuthStoreLock::acquire(OAuthStore::File)?; + let _lock = OAuthStoreLock::acquire_for_write(OAuthStore::File)?; std::fs::write(ready_file, b"ready")?; loop { std::thread::sleep(Duration::from_secs(/*secs*/ 60)); @@ -328,8 +331,12 @@ where T: Send + 'static, { std::thread::scope(|scope| { - let held_lock = - OAuthStoreLock::acquire_in(codex_home, store, Duration::from_millis(/*millis*/ 100))?; + let held_lock = OAuthStoreLock::acquire_in_with_mode( + codex_home, + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + )?; let (contended_tx, contended_rx) = mpsc::channel(); let worker = scope.spawn(move || { tracing::subscriber::with_default(LockContentionSubscriber { contended_tx }, operation) @@ -348,6 +355,90 @@ where }) } +#[test] +fn aggregate_store_readers_share_access_while_writers_remain_exclusive() -> Result<()> { + let env = TempCodexHome::new(); + + for store in [OAuthStore::File, OAuthStore::Secrets] { + let readers = (0..2) + .map(|_| { + OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + ) + }) + .collect::, _>>()?; + + let writer_error = match OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + ) { + Ok(_) => anyhow::bail!("an active OAuth store reader must exclude writers"), + Err(error) => error, + }; + assert!(matches!( + writer_error, + OAuthStoreLockFailure::Timeout { .. } + )); + drop(readers); + let _writer = OAuthStoreLock::acquire_in_with_mode( + env.path(), + store, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Exclusive, + )?; + } + + Ok(()) +} + +#[test] +fn aggregate_store_credential_loads_can_share_an_existing_reader() -> Result<()> { + let env = TempCodexHome::new(); + let keyring_store = MockKeyringStore::default(); + let tokens = sample_tokens(); + save_oauth_tokens_to_file(&tokens)?; + save_oauth_tokens_with_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens, + )?; + + let file_reader = OAuthStoreLock::acquire_in_with_mode( + env.path(), + OAuthStore::File, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + )?; + let file_tokens = load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("a File credential read should coexist with another reader"); + assert_tokens_match_without_expiry(&file_tokens, &tokens); + drop(file_reader); + + let secrets_reader = OAuthStoreLock::acquire_in_with_mode( + env.path(), + OAuthStore::Secrets, + Duration::from_millis(/*millis*/ 100), + OAuthStoreLockMode::Shared, + )?; + let secrets_tokens = load_oauth_tokens_from_keyring( + &keyring_store, + AuthKeyringBackendKind::Secrets, + &tokens.server_name, + &tokens.url, + )? + .expect("a Secrets credential read should coexist with another reader"); + assert_tokens_match_without_expiry(&secrets_tokens, &tokens); + drop(secrets_reader); + + Ok(()) +} + #[test] fn file_store_lock_preserves_updates_for_different_servers() -> Result<()> { let env = TempCodexHome::new(); diff --git a/codex-rs/secrets/src/local.rs b/codex-rs/secrets/src/local.rs index 366be386b9..e76e769dd3 100644 --- a/codex-rs/secrets/src/local.rs +++ b/codex-rs/secrets/src/local.rs @@ -4,6 +4,8 @@ use std::io::Write; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; use std::sync::atomic::Ordering; use std::sync::atomic::compiler_fence; use std::time::SystemTime; @@ -24,6 +26,8 @@ use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; use tracing::warn; use super::SecretListEntry; @@ -37,6 +41,7 @@ const SECRETS_VERSION: u8 = 1; const LOCAL_SECRETS_FILENAME: &str = "local.age"; const CODEX_AUTH_SECRETS_FILENAME: &str = "codex_auth.age"; const MCP_OAUTH_SECRETS_FILENAME: &str = "mcp_oauth.age"; +static MCP_OAUTH_CACHE: Mutex> = Mutex::new(None); /// Selects the local encrypted file used by a `LocalSecretsBackend`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -56,6 +61,13 @@ struct SecretsFile { secrets: BTreeMap, } +struct CachedMcpSecrets { + path: PathBuf, + ciphertext_hash: [u8; 32], + passphrase_hash: [u8; 32], + file: Arc, +} + impl SecretsFile { fn new_empty() -> Self { Self { @@ -157,6 +169,23 @@ impl LocalSecretsBackend { let ciphertext = fs::read(&path) .with_context(|| format!("failed to read secrets file at {}", path.display()))?; let passphrase = self.load_or_create_passphrase()?; + let cache = (self.namespace == LocalSecretsNamespace::McpOAuth).then(|| { + let ciphertext_hash: [u8; 32] = Sha256::digest(&ciphertext).into(); + let passphrase_hash: [u8; 32] = + Sha256::digest(passphrase.expose_secret().as_bytes()).into(); + let cache = MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner); + (cache, ciphertext_hash, passphrase_hash) + }); + if let Some((cache, ciphertext_hash, passphrase_hash)) = cache.as_ref() + && let Some(cached) = cache.as_ref() + && cached.path == path + && cached.ciphertext_hash == *ciphertext_hash + && cached.passphrase_hash == *passphrase_hash + { + return Ok(cached.file.as_ref().clone()); + } let plaintext = decrypt_with_passphrase(&ciphertext, &passphrase)?; let mut parsed: SecretsFile = serde_json::from_slice(&plaintext).with_context(|| { format!( @@ -173,6 +202,14 @@ impl LocalSecretsBackend { parsed.version, SECRETS_VERSION ); + if let Some((mut cache, ciphertext_hash, passphrase_hash)) = cache { + *cache = Some(CachedMcpSecrets { + path, + ciphertext_hash, + passphrase_hash, + file: Arc::new(parsed.clone()), + }); + } Ok(parsed) } @@ -186,6 +223,14 @@ impl LocalSecretsBackend { let ciphertext = encrypt_with_passphrase(&plaintext, &passphrase)?; let path = self.secrets_path(); write_file_atomically(&path, &ciphertext)?; + if self.namespace == LocalSecretsNamespace::McpOAuth { + let mut cache = MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner); + if cache.as_ref().is_some_and(|cached| cached.path == path) { + *cache = None; + } + } Ok(()) } @@ -376,6 +421,8 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; + static MCP_OAUTH_CACHE_TEST_LOCK: Mutex<()> = Mutex::new(()); + #[test] fn load_file_rejects_newer_schema_versions() -> Result<()> { let codex_home = tempfile::tempdir().expect("tempdir"); @@ -446,11 +493,21 @@ mod tests { .collect(); assert_eq!(filenames, vec![LOCAL_SECRETS_FILENAME.to_string()]); assert_eq!(backend.get(&scope, &name)?, Some("two".to_string())); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != backend.secrets_path()) + ); Ok(()) } #[test] fn local_namespaces_write_separate_files() -> Result<()> { + let _cache_lock = MCP_OAUTH_CACHE_TEST_LOCK + .lock() + .unwrap_or_else(PoisonError::into_inner); let codex_home = tempfile::tempdir().expect("tempdir"); let keyring = Arc::new(MockKeyringStore::default()); let codex_auth_backend = LocalSecretsBackend::new_with_namespace( @@ -473,6 +530,13 @@ mod tests { codex_auth_backend.get(&scope, &name)?, Some("codex-auth-value".to_string()) ); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != codex_auth_backend.secrets_path()) + ); assert_eq!( mcp_backend.get(&scope, &name)?, Some("mcp-value".to_string()) @@ -494,4 +558,78 @@ mod tests { assert!(!codex_home.path().join("secrets").join("local.age").exists()); Ok(()) } + + #[test] + fn mcp_oauth_cache_reuses_plaintext_and_invalidates_when_ciphertext_changes() -> Result<()> { + let _cache_lock = MCP_OAUTH_CACHE_TEST_LOCK + .lock() + .unwrap_or_else(PoisonError::into_inner); + let codex_home = tempfile::tempdir().expect("tempdir"); + let keyring = Arc::new(MockKeyringStore::default()); + let first = LocalSecretsBackend::new_with_namespace( + codex_home.path().to_path_buf(), + keyring.clone(), + LocalSecretsNamespace::McpOAuth, + ); + let second = LocalSecretsBackend::new_with_namespace( + codex_home.path().to_path_buf(), + keyring, + LocalSecretsNamespace::McpOAuth, + ); + let scope = SecretScope::Global; + let name = SecretName::new("TEST_SECRET")?; + let cached_file = || { + Arc::clone( + &MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .expect("MCP OAuth credentials should be cached") + .file, + ) + }; + + first.set(&scope, &name, "one")?; + let (first_cached, second_cached) = std::thread::scope(|threads| { + let first_reader = threads.spawn(|| { + assert_eq!(first.get(&scope, &name)?, Some("one".to_string())); + Ok::<_, anyhow::Error>(cached_file()) + }); + let second_reader = threads.spawn(|| { + assert_eq!(second.get(&scope, &name)?, Some("one".to_string())); + Ok::<_, anyhow::Error>(cached_file()) + }); + Ok::<_, anyhow::Error>(( + first_reader.join().expect("first credential reader")?, + second_reader.join().expect("second credential reader")?, + )) + })?; + assert!(Arc::ptr_eq(&first_cached, &second_cached)); + + assert_eq!(second.get(&scope, &name)?, Some("one".to_string())); + assert!(Arc::ptr_eq(&first_cached, &cached_file())); + + first.set(&scope, &name, "two")?; + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != first.secrets_path()) + ); + assert_eq!(second.get(&scope, &name)?, Some("two".to_string())); + assert!(!Arc::ptr_eq(&first_cached, &cached_file())); + + assert!(first.delete(&scope, &name)?); + assert!( + MCP_OAUTH_CACHE + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_none_or(|cached| cached.path != first.secrets_path()) + ); + assert_eq!(second.get(&scope, &name)?, None); + + Ok(()) + } }