From 0d5e92d6fad337acab101df3efc2a4701a295668 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:15:59 +0000 Subject: [PATCH] fix(rmcp): bound OAuth refresh coordination --- codex-rs/rmcp-client/src/oauth.rs | 123 ++++++++++++++---- codex-rs/rmcp-client/src/rmcp_client.rs | 63 ++++----- .../streamable_http_oauth_refresh_race.rs | 100 ++++++++++++-- 3 files changed, 224 insertions(+), 62 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index b5930656fd..2401888328 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -37,6 +37,7 @@ use std::collections::BTreeMap; use std::fs; use std::fs::OpenOptions; use std::io::ErrorKind; +use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -55,7 +56,8 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; -const OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock"; +const FILE_OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock"; +const KEYRING_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh"; const REFRESH_SKEW_MILLIS: u64 = 30_000; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -369,13 +371,14 @@ impl OAuthPersistor { // A different process may have rotated the one-time refresh token. // Reload its durable result while all refreshers share the same lock. - let _refresh_lock = acquire_oauth_refresh_lock().await?; - - if let Some(stored) = load_oauth_tokens( + let (_refresh_locks, stored) = lock_and_load_oauth_tokens( &self.inner.server_name, &self.inner.url, self.inner.store_mode, - )? { + ) + .await?; + + if let Some(stored) = stored { let credentials_changed = { let last_credentials = self.inner.last_credentials.lock().await; last_credentials.as_ref() != Some(&stored) @@ -432,27 +435,96 @@ impl OAuthPersistor { } } -async fn acquire_oauth_refresh_lock() -> Result { - let path = find_codex_home()?.join(OAUTH_REFRESH_LOCK_FILENAME); +async fn lock_and_load_oauth_tokens( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result<(Vec, Option)> { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => { + let keyring_lock = + acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path()?).await?; + match load_oauth_tokens_from_keyring(&keyring_store, server_name, url) { + Ok(Some(tokens)) => Ok((vec![keyring_lock], Some(tokens))), + Ok(None) => { + let file_lock = + acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?; + let tokens = load_oauth_tokens_from_file(server_name, url)?; + Ok((vec![keyring_lock, file_lock], tokens)) + } + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + let file_lock = + acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?; + let tokens = + load_oauth_tokens_from_file(server_name, url).with_context(|| { + format!("failed to read OAuth tokens from keyring: {error}") + })?; + Ok((vec![keyring_lock, file_lock], tokens)) + } + } + } + OAuthCredentialsStoreMode::File => { + let lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?; + let tokens = load_oauth_tokens_from_file(server_name, url)?; + Ok((vec![lock], tokens)) + } + OAuthCredentialsStoreMode::Keyring => { + let lock = acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path()?).await?; + let tokens = load_oauth_tokens_from_keyring(&keyring_store, server_name, url) + .with_context(|| "failed to read OAuth tokens from keyring".to_string())?; + Ok((vec![lock], tokens)) + } + } +} + +fn file_oauth_refresh_lock_path() -> Result { + Ok(find_codex_home()? + .join(FILE_OAUTH_REFRESH_LOCK_FILENAME) + .to_path_buf()) +} + +fn keyring_oauth_refresh_lock_path() -> Result { + let user_namespace = ["HOME", "USERPROFILE", "USER", "USERNAME"] + .into_iter() + .find_map(std::env::var_os) + .unwrap_or_default(); + let lock_id = sha_256_prefix(&Value::String( + user_namespace.to_string_lossy().into_owned(), + ))?; + Ok(std::env::temp_dir().join(format!( + "{KEYRING_OAUTH_REFRESH_LOCK_PREFIX}-{lock_id}.lock" + ))) +} + +async fn acquire_oauth_refresh_lock(path: &Path) -> Result { if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; } - let file = OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(&path) + let mut options = OpenOptions::new(); + options.read(true).write(true).create(true).truncate(false); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + let file = options + .open(path) .with_context(|| format!("failed to open OAuth refresh lock at {}", path.display()))?; - - tokio::task::spawn_blocking(move || { - file.lock().with_context(|| { - format!("failed to acquire OAuth refresh lock at {}", path.display()) - })?; - Ok(file) - }) - .await - .context("failed to join OAuth refresh lock task")? + loop { + match file.try_lock() { + Ok(()) => return Ok(file), + Err(fs::TryLockError::WouldBlock) => { + tokio::time::sleep(Duration::from_millis(10)).await; + } + Err(error) => { + return Err(error).with_context(|| { + format!("failed to acquire OAuth refresh lock at {}", path.display()) + }); + } + } + } } const FALLBACK_FILENAME: &str = ".credentials.json"; @@ -723,6 +795,13 @@ mod tests { } } + #[test] + fn keyring_refresh_lock_does_not_use_codex_home() -> Result<()> { + let lock_path = keyring_oauth_refresh_lock_path()?; + assert_eq!(lock_path.parent(), Some(std::env::temp_dir().as_path())); + Ok(()) + } + #[test] fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 50d9fab8bd..2396b21248 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -824,42 +824,45 @@ impl RmcpClient { Arc>, Option, )> { - let (transport, oauth_persistor) = match pending_transport { - PendingTransport::InProcess { transport } => ( - service::serve_client(client_service, transport).boxed(), - None, - ), - PendingTransport::Stdio { transport } => ( - service::serve_client(client_service, transport).boxed(), - None, - ), - PendingTransport::StreamableHttp { transport } => ( - service::serve_client(client_service, transport).boxed(), - None, - ), - PendingTransport::StreamableHttpWithOAuth { - transport, - oauth_persistor, - } => { - oauth_persistor.refresh_if_needed().await?; - ( + let connect = async move { + let (transport, oauth_persistor) = match pending_transport { + PendingTransport::InProcess { transport } => ( + service::serve_client(client_service, transport).boxed(), + None, + ), + PendingTransport::Stdio { transport } => ( + service::serve_client(client_service, transport).boxed(), + None, + ), + PendingTransport::StreamableHttp { transport } => ( + service::serve_client(client_service, transport).boxed(), + None, + ), + PendingTransport::StreamableHttpWithOAuth { + transport, + oauth_persistor, + } => ( service::serve_client(client_service, transport).boxed(), Some(oauth_persistor), - ) + ), + }; + + if let Some(runtime) = &oauth_persistor { + runtime.refresh_if_needed().await?; } + + let service = transport + .await + .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?; + Ok((Arc::new(service), oauth_persistor)) }; - let service = match timeout { - Some(duration) => time::timeout(duration, transport) + match timeout { + Some(duration) => time::timeout(duration, connect) .await - .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))? - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, - None => transport - .await - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, - }; - - Ok((Arc::new(service), oauth_persistor)) + .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))?, + None => connect.await, + } } async fn run_service_operation( diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs index c8e64e24ce..8706facb47 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs @@ -2,6 +2,7 @@ mod streamable_http_test_support; use std::collections::BTreeMap; use std::fs; +use std::fs::OpenOptions; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; @@ -47,6 +48,7 @@ const OLD_ACCESS_TOKEN: &str = "old-access-token"; const OLD_REFRESH_TOKEN: &str = "old-one-time-refresh-token"; const NEW_ACCESS_TOKEN: &str = "new-access-token"; const NEW_REFRESH_TOKEN: &str = "new-one-time-refresh-token"; +const FILE_OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock"; const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL"; const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH"; @@ -68,12 +70,14 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok .and(path("/oauth/token")) .respond_with(move |_request: &Request| { if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { - ResponseTemplate::new(200).set_body_json(json!({ - "access_token": NEW_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": NEW_REFRESH_TOKEN, - })) + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + })) } else { ResponseTemplate::new(400).set_body_json(json!({ "error": "invalid_grant", @@ -91,7 +95,10 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok let server_url = format!("{}/mcp", server.uri()); seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; - let go_path = control_dir.path().join("go"); + let go_paths = [ + control_dir.path().join("client-a.go"), + control_dir.path().join("client-b.go"), + ]; let ready_paths = [ control_dir.path().join("client-a.ready"), control_dir.path().join("client-b.ready"), @@ -101,18 +108,30 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok control_dir.path().join("client-b.result"), ]; let mut children = Vec::new(); - for (ready_path, result_path) in ready_paths.iter().zip(result_paths.iter()) { + for ((ready_path, go_path), result_path) in ready_paths + .iter() + .zip(go_paths.iter()) + .zip(result_paths.iter()) + { children.push(spawn_client_child( &codex_home, &server_url, ready_path, - &go_path, + go_path, result_path, )?); } wait_for_paths(&ready_paths).await?; - fs::write(&go_path, "go")?; + fs::write(&go_paths[0], "go")?; + timeout(Duration::from_secs(5), async { + while refresh_attempts.load(Ordering::SeqCst) == 0 { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for the first process to hold the refresh lock")?; + fs::write(&go_paths[1], "go")?; wait_for_children(children).await?; let outcomes = read_outcomes(&result_paths)?; @@ -130,6 +149,23 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok .iter() .all(|body| body.contains(&format!("refresh_token={OLD_REFRESH_TOKEN}"))) ); + let authorization_headers = requests + .iter() + .filter(|request| request.method.as_str() == "POST" && request.url.path() == "/mcp") + .map(|request| { + request + .headers + .get("authorization") + .expect("authorization header") + .to_str() + .expect("ASCII authorization header") + .to_string() + }) + .collect::>(); + assert_eq!( + authorization_headers, + vec![format!("Bearer {NEW_ACCESS_TOKEN}"); 4] + ); assert_eq!( persisted_token_snapshot(&codex_home)?, PersistedTokenSnapshot { @@ -142,6 +178,50 @@ async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tok Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refresh_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + + let lock_path = codex_home.path().join(FILE_OAUTH_REFRESH_LOCK_FILENAME); + let refresh_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + refresh_lock.lock()?; + + let ready_path = control_dir.path().join("client.ready"); + let go_path = control_dir.path().join("client.go"); + let result_path = control_dir.path().join("client.result"); + let child = spawn_client_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + fs::write(&go_path, "go")?; + timeout(Duration::from_secs(7), wait_for_children(vec![child])) + .await + .context("startup did not time out while waiting for the refresh lock")??; + + assert_eq!( + fs::read_to_string(result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn process_keeps_stale_in_memory_tokens_but_does_not_overwrite_rotated_tokens() -> anyhow::Result<()> {