From 0d63af1ab7b5acdb4fc183e0b929bd96f119cb7d Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:45:00 +0000 Subject: [PATCH] fix(rmcp): finalize credential-scoped refresh locking --- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/oauth.rs | 281 +++++++++++++++--- .../streamable_http_oauth_refresh_race.rs | 84 +++++- 3 files changed, 316 insertions(+), 50 deletions(-) diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1ee18c753..f3d2270d03 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -17,6 +17,7 @@ pub use auth_status::discover_streamable_http_oauth; pub use auth_status::supports_oauth_login; pub use codex_protocol::protocol::McpAuthStatus; pub use in_process_transport::InProcessTransportFactory; +pub use oauth::OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 2401888328..4816cca315 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -19,6 +19,7 @@ use anyhow::Context; use anyhow::Error; use anyhow::Result; +use anyhow::anyhow; use codex_config::types::OAuthCredentialsStoreMode; use oauth2::AccessToken; use oauth2::RefreshToken; @@ -47,6 +48,7 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use rmcp::transport::auth::AuthError; use rmcp::transport::auth::AuthorizationManager; use rmcp::transport::auth::CredentialStore; use rmcp::transport::auth::InMemoryCredentialStore; @@ -56,9 +58,14 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; -const FILE_OAUTH_REFRESH_LOCK_FILENAME: &str = ".credentials.json.refresh.lock"; +const FALLBACK_WRITE_LOCK_FILENAME: &str = ".credentials.json.lock"; +const FILE_OAUTH_REFRESH_LOCK_PREFIX: &str = ".credentials"; const KEYRING_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh"; +const MISSING_REFRESH_TOKEN_ERROR: &str = "No refresh token available"; +const OAUTH_SERVER_ERROR_PREFIX: &str = "Server returned error response: "; const REFRESH_SKEW_MILLIS: u64 = 30_000; +pub const OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR: &str = + "OAuth refresh token was rejected; reauthentication required"; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -423,12 +430,22 @@ impl OAuthPersistor { { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", - self.inner.server_name - ) - })?; + match guard.refresh_token().await { + Ok(_) => {} + Err(error) if refresh_requires_reauthentication(&error) => { + return Err(anyhow!( + "Auth required: {OAUTH_REFRESH_REAUTHENTICATION_REQUIRED_ERROR}" + )); + } + Err(error) => { + return Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + }); + } + } } self.persist_if_needed().await @@ -441,22 +458,38 @@ async fn lock_and_load_oauth_tokens( store_mode: OAuthCredentialsStoreMode, ) -> Result<(Vec, Option)> { let keyring_store = DefaultKeyringStore; + lock_and_load_oauth_tokens_with_keyring(&keyring_store, server_name, url, store_mode).await +} + +async fn lock_and_load_oauth_tokens_with_keyring( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result<(Vec, Option)> { 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) { + acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path(server_name, url)?) + .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 file_lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path( + server_name, + url, + )?) + .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 file_lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path( + server_name, + url, + )?) + .await?; let tokens = load_oauth_tokens_from_file(server_name, url).with_context(|| { format!("failed to read OAuth tokens from keyring: {error}") @@ -466,52 +499,45 @@ async fn lock_and_load_oauth_tokens( } } OAuthCredentialsStoreMode::File => { - let lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path()?).await?; + let lock = acquire_oauth_refresh_lock(&file_oauth_refresh_lock_path(server_name, url)?) + .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) + let lock = + acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path(server_name, url)?) + .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 { +fn oauth_refresh_lock_id(server_name: &str, url: &str) -> Result { + let account = compute_store_key(server_name, url)?; + sha_256_prefix(&Value::String(format!("{KEYRING_SERVICE}:{account}"))) +} + +fn file_oauth_refresh_lock_path(server_name: &str, url: &str) -> Result { + let lock_id = oauth_refresh_lock_id(server_name, url)?; Ok(find_codex_home()? - .join(FILE_OAUTH_REFRESH_LOCK_FILENAME) + .join(format!( + "{FILE_OAUTH_REFRESH_LOCK_PREFIX}.{lock_id}.refresh.lock" + )) .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(), - ))?; +fn keyring_oauth_refresh_lock_path(server_name: &str, url: &str) -> Result { + let lock_id = oauth_refresh_lock_id(server_name, url)?; 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 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()))?; + let file = open_oauth_lock_file(path)?; loop { match file.try_lock() { Ok(()) => return Ok(file), @@ -527,6 +553,68 @@ async fn acquire_oauth_refresh_lock(path: &Path) -> Result { } } +fn acquire_fallback_write_lock() -> Result { + let path = find_codex_home()? + .join(FALLBACK_WRITE_LOCK_FILENAME) + .to_path_buf(); + let file = open_oauth_lock_file(&path)?; + file.lock().with_context(|| { + format!( + "failed to acquire OAuth fallback write lock at {}", + path.display() + ) + })?; + Ok(file) +} + +fn acquire_fallback_read_lock() -> Result { + let path = find_codex_home()? + .join(FALLBACK_WRITE_LOCK_FILENAME) + .to_path_buf(); + let file = open_oauth_lock_file(&path)?; + file.lock_shared().with_context(|| { + format!( + "failed to acquire OAuth fallback read lock at {}", + path.display() + ) + })?; + Ok(file) +} + +fn open_oauth_lock_file(path: &Path) -> Result { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + 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); + } + options + .open(path) + .with_context(|| format!("failed to open OAuth lock at {}", path.display())) +} + +fn refresh_requires_reauthentication(error: &AuthError) -> bool { + match error { + AuthError::AuthorizationRequired => true, + AuthError::TokenRefreshFailed(message) => { + if message == MISSING_REFRESH_TOKEN_ERROR { + return true; + } + message + .strip_prefix(OAUTH_SERVER_ERROR_PREFIX) + .is_some_and(|response| { + let error_code_end = response.find([':', ' ']).unwrap_or(response.len()); + &response[..error_code_end] == "invalid_grant" + }) + } + _ => false, + } +} + const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; @@ -547,7 +635,15 @@ struct FallbackTokenEntry { } fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result> { - let Some(store) = read_fallback_file()? else { + let _read_lock = acquire_fallback_read_lock()?; + load_oauth_tokens_from_file_unlocked(server_name, url) +} + +fn load_oauth_tokens_from_file_unlocked( + server_name: &str, + url: &str, +) -> Result> { + let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -590,8 +686,9 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { + let _write_lock = acquire_fallback_write_lock()?; 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 @@ -691,6 +788,11 @@ fn fallback_file_path() -> Result { } fn read_fallback_file() -> Result> { + let _read_lock = acquire_fallback_read_lock()?; + read_fallback_file_unlocked() +} + +fn read_fallback_file_unlocked() -> Result> { let path = fallback_file_path()?; let contents = match fs::read_to_string(&path) { Ok(contents) => contents, @@ -796,12 +898,103 @@ 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())); + fn refresh_lock_identity_is_credential_specific_and_home_independent() -> Result<()> { + let _env = TempCodexHome::new(); + let tokens = sample_tokens(); + let other_url = "https://other.example.test/mcp"; + + let keyring_path = keyring_oauth_refresh_lock_path(&tokens.server_name, &tokens.url)?; + assert_eq!(keyring_path.parent(), Some(std::env::temp_dir().as_path())); + assert_eq!( + keyring_path, + keyring_oauth_refresh_lock_path(&tokens.server_name, &tokens.url)? + ); + assert_ne!( + keyring_path, + keyring_oauth_refresh_lock_path(&tokens.server_name, other_url)? + ); + assert_ne!( + file_oauth_refresh_lock_path(&tokens.server_name, &tokens.url)?, + file_oauth_refresh_lock_path(&tokens.server_name, other_url)? + ); Ok(()) } + #[tokio::test(flavor = "current_thread")] + async fn refresh_lock_loading_covers_keyring_and_auto_branches() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); + let serialized = serde_json::to_string(&tokens)?; + let key = compute_store_key(&tokens.server_name, &tokens.url)?; + store.save(KEYRING_SERVICE, &key, &serialized)?; + + let (locks, loaded) = lock_and_load_oauth_tokens_with_keyring( + &store, + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::Keyring, + ) + .await?; + assert_eq!(locks.len(), 1); + assert_tokens_match_without_expiry( + &loaded.expect("keyring credentials should load"), + &tokens, + ); + drop(locks); + + let (locks, loaded) = lock_and_load_oauth_tokens_with_keyring( + &store, + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::Auto, + ) + .await?; + assert_eq!(locks.len(), 1); + assert_tokens_match_without_expiry( + &loaded.expect("auto mode should prefer keyring credentials"), + &tokens, + ); + drop(locks); + + store.delete(KEYRING_SERVICE, &key)?; + save_oauth_tokens_to_file(&tokens)?; + let (locks, loaded) = lock_and_load_oauth_tokens_with_keyring( + &store, + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::Auto, + ) + .await?; + assert_eq!(locks.len(), 2); + assert_tokens_match_without_expiry( + &loaded.expect("auto mode should load fallback file credentials"), + &tokens, + ); + Ok(()) + } + + #[test] + fn refresh_reauthentication_classification_is_narrow() { + assert!(refresh_requires_reauthentication( + &AuthError::AuthorizationRequired + )); + assert!(refresh_requires_reauthentication( + &AuthError::TokenRefreshFailed(MISSING_REFRESH_TOKEN_ERROR.to_string()) + )); + assert!(refresh_requires_reauthentication( + &AuthError::TokenRefreshFailed( + "Server returned error response: invalid_grant: revoked".to_string() + ) + )); + assert!(!refresh_requires_reauthentication( + &AuthError::TokenRefreshFailed( + "Server returned error response: invalid_request: invalid_grant in description" + .to_string() + ) + )); + } + #[test] fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); 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 8706facb47..c2387ec10f 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 @@ -28,6 +28,8 @@ use rmcp::transport::auth::VendorExtraTokenFields; use serde::Deserialize; use serde_json::Value; use serde_json::json; +use sha2::Digest; +use sha2::Sha256; use tempfile::TempDir; use tokio::process::Child; use tokio::process::Command; @@ -48,7 +50,8 @@ 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 KEYRING_SERVICE: &str = "Codex MCP Credentials"; +const CHILD_WAIT_TIMEOUT: Duration = Duration::from_secs(10); const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL"; const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH"; @@ -188,7 +191,7 @@ async fn refresh_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> { 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 lock_path = file_oauth_refresh_lock_path(&codex_home, &server_url)?; let refresh_lock = OpenOptions::new() .read(true) .write(true) @@ -222,6 +225,49 @@ async fn refresh_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rejected_refresh_requires_login_without_reaching_mcp() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token revoked", + }))) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 0).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 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")?; + wait_for_children(vec![child]).await?; + + assert_eq!( + fs::read_to_string(result_path)?, + "error:Auth required: OAuth refresh token was rejected; reauthentication required" + ); + + 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<()> { @@ -438,7 +484,8 @@ async fn seed_tokens( server_url: &str, expires_at: u64, ) -> anyhow::Result<()> { - let status = Command::new(std::env::current_exe()?) + let mut command = Command::new(std::env::current_exe()?); + command .args([ "oauth_refresh_race_seed_child", "--exact", @@ -450,8 +497,10 @@ async fn seed_tokens( .env(ACCESS_TOKEN_ENV, OLD_ACCESS_TOKEN) .env(REFRESH_TOKEN_ENV, OLD_REFRESH_TOKEN) .env(EXPIRES_AT_ENV, expires_at.to_string()) - .status() - .await?; + .kill_on_drop(true); + let status = timeout(CHILD_WAIT_TIMEOUT, command.status()) + .await + .context("timed out waiting for OAuth token seed child")??; anyhow::ensure!(status.success(), "OAuth token seed child failed: {status}"); Ok(()) } @@ -482,7 +531,9 @@ fn spawn_client_child( async fn wait_for_children(children: Vec) -> anyhow::Result<()> { for mut child in children { - let status = child.wait().await?; + let status = timeout(CHILD_WAIT_TIMEOUT, child.wait()) + .await + .context("timed out waiting for OAuth client child")??; anyhow::ensure!(status.success(), "OAuth client child failed: {status}"); } Ok(()) @@ -517,6 +568,27 @@ fn request_bodies(requests: &[Request], request_path: &str) -> Vec { .collect() } +fn file_oauth_refresh_lock_path(codex_home: &TempDir, server_url: &str) -> anyhow::Result { + let store_key_payload = json!({ + "type": "http", + "url": server_url, + "headers": {}, + }); + let account = format!("{SERVER_NAME}|{}", sha_256_prefix(&store_key_payload)?); + let lock_id = sha_256_prefix(&Value::String(format!("{KEYRING_SERVICE}:{account}")))?; + Ok(codex_home + .path() + .join(format!(".credentials.{lock_id}.refresh.lock"))) +} + +fn sha_256_prefix(value: &Value) -> anyhow::Result { + let serialized = serde_json::to_string(value)?; + let mut hasher = Sha256::new(); + hasher.update(serialized.as_bytes()); + let digest = hasher.finalize(); + Ok(format!("{digest:x}")[..16].to_string()) +} + #[derive(Debug, Deserialize)] struct PersistedTokenEntry { access_token: String,