Serialize shared MCP OAuth credential stores

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<K: KeyringStore + Clone
     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::<OAuthStoreLockFailure>().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<K: KeyringStore + Clone + 'static>(
     server_name: &str,
     url: &str,
 ) -> Result<Option<StoredOAuthTokens>> {
+    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<K: KeyringStore>(
     }
 }
 
+/// 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<K: KeyringStore + Clone + 'static>(
     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<K: KeyringStore + Clone + 'static>(
+    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<K: KeyringStore + Clone + 'static>(
     );
     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<K: KeyringStore + Clone + 'static>(
@@ -341,6 +378,10 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file<K: KeyringStore + Clone
 ) -> 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::<OAuthStoreLockFailure>().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<K: KeyringStore + Clone + 'static>(
     server_name: &str,
     url: &str,
 ) -> Result<bool> {
+    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<Option<StoredOAuthTokens>> {
-    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<Option<St
     Ok(None)
 }
 
+/// Saves one credential while holding the File aggregate-store lock across the full
+/// read-modify-write operation.
 fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> 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<bool> {
-    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<PathBuf> {
     Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf())
 }
 
-fn read_fallback_file() -> Result<Option<FallbackFile>> {
+fn read_fallback_file_unlocked() -> Result<Option<FallbackFile>> {
     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<String> {
 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<Mutex<()>> = 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(())