mirror of
https://github.com/openai/codex.git
synced 2026-08-23 13:09:46 +00:00
Serialize shared MCP OAuth credential stores (#30292)
[Codex Thread 019edd6d-6f14-74e2-853c-345d1803d4a6](https://codex-thread-link.openai.chatgpt-team.site/thread/019edd6d-6f14-74e2-853c-345d1803d4a6) ## Stack Review and merge in order. Every layer is independently correct and documents its safe stopping point. 1. [openai/codex#30292](https://github.com/openai/codex/pull/30292) — aggregate File/Secrets store locking 2. [openai/codex#30293](https://github.com/openai/codex/pull/30293) — resolve and lifecycle-pin the exact OAuth store 3. [openai/codex#30416](https://github.com/openai/codex/pull/30416) — serialized authoritative refresh transaction 4. [openai/codex#30294](https://github.com/openai/codex/pull/30294) — Codex-owned transport refresh and one-shot 401 recovery 5. [openai/codex#30295](https://github.com/openai/codex/pull/30295) — login/logout transaction serialization 6. [openai/codex#30296](https://github.com/openai/codex/pull/30296) — diagnostic-only Auto store drift reporting **This PR is layer 1.** ## Why MCP OAuth credentials stored in File or Secrets share one aggregate map. Concurrent read-modify-write operations for different MCP servers can both read the same snapshot and let the later write discard the earlier update. That is a correctness problem independent of refresh-token rotation. ## What this PR does - Adds a bounded cross-process lock around aggregate File and Secrets loads, saves, and deletes. - Distinguishes aggregate-lock failures from Secrets backend unavailability, so Auto can fall back only for the latter and cannot bypass serialization by reading or writing File. - Keeps Direct keyring operations outside this lock because they are already per credential. - Releases the Secrets aggregate lock before legacy File cleanup so cross-store cleanup cannot create nested aggregate-lock ordering. - Tests actual contention by waiting for an observed `WouldBlock`, rather than assuming a sleeping worker reached the lock. - Tests load and save with only the Secrets lock path broken while fallback File remains readable and writable. ## Decisions and non-goals - This lock protects aggregate-store read-modify-write integrity only. It does not choose a credential authority or serialize an OAuth refresh transaction. - The lock is scoped to the active `CODEX_HOME`, matching the aggregate files it protects. - Lock waits are bounded, and coordination failures are surfaced rather than treated as evidence that Secrets is unavailable. ## Safe stopping point This PR can merge alone. It prevents lost updates and partial aggregate reads. Auto can still resolve again during a client lifecycle until layer 2, and concurrent refreshes remain possible until layer 3. ## Validation - `just test -p codex-rmcp-client` (96 passed; expected environment skips) - Focused aggregate File/Secrets lock contention and Auto fallback tests
This commit is contained in:
@@ -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(())
|
||||
|
||||
181
codex-rs/rmcp-client/src/oauth/store_lock.rs
Normal file
181
codex-rs/rmcp-client/src/oauth/store_lock.rs
Normal file
@@ -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<Self> {
|
||||
// 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<Self> {
|
||||
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;
|
||||
46
codex-rs/rmcp-client/src/oauth/test_support.rs
Normal file
46
codex-rs/rmcp-client/src/oauth/test_support.rs
Normal file
@@ -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<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,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn path(&self) -> &Path {
|
||||
self._dir.path()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TempCodexHome {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
std::env::remove_var("CODEX_HOME");
|
||||
}
|
||||
}
|
||||
}
|
||||
480
codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs
Normal file
480
codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs
Normal file
@@ -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::<OAuthStoreLockFailure>(),
|
||||
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::<OAuthStoreLockFailure>().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::<OAuthStoreLockFailure>().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<tracing::level_filters::LevelFilter> {
|
||||
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<T>(
|
||||
codex_home: &std::path::Path,
|
||||
store: OAuthStore,
|
||||
operation: impl FnOnce() -> Result<T> + Send + 'static,
|
||||
) -> Result<T>
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user