diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6da768843504..3226b4d70071 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3579,6 +3579,7 @@ dependencies = [ "codex-utils-pty", "futures", "keyring", + "libc", "memchr", "oauth2", "pretty_assertions", @@ -3597,6 +3598,7 @@ dependencies = [ "urlencoding", "webbrowser", "which 8.0.0", + "windows-sys 0.52.0", "wiremock", ] diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 1006ad74e80a..3f177975ae52 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -74,11 +74,20 @@ wiremock = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] keyring = { workspace = true, features = ["linux-native-async-persistent"] } +[target.'cfg(unix)'.dependencies] +libc = { workspace = true } + [target.'cfg(target_os = "macos")'.dependencies] keyring = { workspace = true, features = ["apple-native"] } [target.'cfg(target_os = "windows")'.dependencies] keyring = { workspace = true, features = ["windows-native"] } +windows-sys = { version = "0.52", features = [ + "Win32_Foundation", + "Win32_Security", + "Win32_System_SystemInformation", + "Win32_System_Threading", +] } [target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies] keyring = { workspace = true, features = ["sync-secret-service"] } diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e1ee18c75324..67d459ea5df1 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -17,10 +17,10 @@ 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; -pub(crate) use oauth::load_oauth_tokens; pub use oauth::save_oauth_tokens; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c348460795de..f2dbc6376cde 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; @@ -35,9 +36,12 @@ use sha2::Digest; use sha2::Sha256; 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::sync::OnceLock; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -45,13 +49,24 @@ 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; +use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; +const FALLBACK_LOCK_PREFIX: &str = "codex-mcp-oauth-fallback"; +const FILE_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh-file"; +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 { @@ -94,6 +109,57 @@ pub(crate) fn load_oauth_tokens( } } +pub(crate) async fn load_oauth_tokens_for_runtime( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + fallback_lock_mode: FallbackLockMode, +) -> Result { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => { + match load_oauth_tokens_from_keyring(&keyring_store, server_name, url) { + Ok(Some(tokens)) => Ok(RuntimeOAuthTokenLoad::Loaded(Box::new(tokens))), + Ok(None) => { + load_oauth_tokens_from_file_for_runtime(server_name, url, fallback_lock_mode) + .await + } + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + load_oauth_tokens_from_file_for_runtime(server_name, url, fallback_lock_mode) + .await + .with_context(|| { + format!("failed to read OAuth tokens from keyring: {error}") + }) + } + } + } + OAuthCredentialsStoreMode::File => { + load_oauth_tokens_from_file_for_runtime(server_name, url, fallback_lock_mode).await + } + OAuthCredentialsStoreMode::Keyring => { + match load_oauth_tokens_from_keyring(&keyring_store, server_name, url) + .with_context(|| "failed to read OAuth tokens from keyring".to_string())? + { + Some(tokens) => Ok(RuntimeOAuthTokenLoad::Loaded(Box::new(tokens))), + None => Ok(RuntimeOAuthTokenLoad::Missing), + } + } + } +} + +#[derive(Clone, Copy)] +pub(crate) enum FallbackLockMode { + DeferIfContended, + Wait, +} + +pub(crate) enum RuntimeOAuthTokenLoad { + Loaded(Box), + Missing, + FallbackLockContended, +} + pub(crate) fn has_oauth_tokens( server_name: &str, url: &str, @@ -182,16 +248,23 @@ fn save_oauth_tokens_with_keyring( server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { + let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?; + if let Err(error) = delete_oauth_tokens_from_file(&key) { + warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); + } + Ok(()) +} + +fn save_oauth_tokens_to_keyring( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; let key = compute_store_key(server_name, &tokens.url)?; match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { - Ok(()) => { - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) - } + Ok(()) => Ok(key), Err(error) => { let message = format!( "failed to write OAuth tokens to keyring: {}", @@ -219,6 +292,46 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file( } } +async fn save_oauth_tokens_for_runtime( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => { + match save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await { + Ok(()) => Ok(()), + Err(error) => { + let message = error.to_string(); + warn!("falling back to file storage for OAuth tokens: {message}"); + save_oauth_tokens_to_file_async(tokens) + .await + .with_context(|| { + format!("failed to write OAuth tokens to keyring: {message}") + }) + } + } + } + OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file_async(tokens).await, + OAuthCredentialsStoreMode::Keyring => { + save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await + } + } +} + +async fn save_oauth_tokens_with_keyring_async( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?; + if let Err(error) = delete_oauth_tokens_from_file_async(&key).await { + warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); + } + Ok(()) +} + pub fn delete_oauth_tokens( server_name: &str, url: &str, @@ -235,7 +348,29 @@ fn delete_oauth_tokens_from_keyring_and_file( url: &str, ) -> Result { let key = compute_store_key(server_name, url)?; - let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key); + let keyring_removed = delete_oauth_tokens_from_keyring(keyring_store, store_mode, &key)?; + let file_removed = delete_oauth_tokens_from_file(&key)?; + Ok(keyring_removed || file_removed) +} + +async fn delete_oauth_tokens_for_runtime( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result { + let keyring_store = DefaultKeyringStore; + let key = compute_store_key(server_name, url)?; + let keyring_removed = delete_oauth_tokens_from_keyring(&keyring_store, store_mode, &key)?; + let file_removed = delete_oauth_tokens_from_file_async(&key).await?; + Ok(keyring_removed || file_removed) +} + +fn delete_oauth_tokens_from_keyring( + keyring_store: &K, + store_mode: OAuthCredentialsStoreMode, + key: &str, +) -> Result { + let keyring_result = keyring_store.delete(KEYRING_SERVICE, key); let keyring_removed = match keyring_result { Ok(removed) => removed, Err(error) => { @@ -250,9 +385,7 @@ fn delete_oauth_tokens_from_keyring_and_file( } } }; - - let file_removed = delete_oauth_tokens_from_file(&key)?; - Ok(keyring_removed || file_removed) + Ok(keyring_removed) } #[derive(Clone)] @@ -264,8 +397,25 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, + credential_store: InMemoryCredentialStore, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, + refresh_attempt: Arc>, + refresh_persistence_failure: Mutex>, +} + +struct RefreshPersistenceFailure { + message: String, +} + +static FAILED_REFRESH_LOCKS: OnceLock>> = OnceLock::new(); + +fn retain_failed_refresh_locks(refresh_locks: Vec) { + FAILED_REFRESH_LOCKS + .get_or_init(std::sync::Mutex::default) + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(refresh_locks); } impl OAuthPersistor { @@ -273,6 +423,7 @@ impl OAuthPersistor { server_name: String, url: String, authorization_manager: Arc>, + credential_store: InMemoryCredentialStore, store_mode: OAuthCredentialsStoreMode, initial_credentials: Option, ) -> Self { @@ -281,19 +432,29 @@ impl OAuthPersistor { server_name, url, authorization_manager, + credential_store, store_mode, last_credentials: Mutex::new(initial_credentials), + refresh_attempt: Arc::new(Mutex::new(())), + refresh_persistence_failure: Mutex::new(None), }), } } /// Persists the latest stored credentials if they have changed. /// Deletes the credentials if they are no longer present. + pub(crate) async fn persist_if_needed(&self) -> Result<()> { + let persistor = self.clone(); + tokio::spawn(async move { persistor.persist_if_needed_inner().await }) + .await + .context("OAuth token persistence task failed")? + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] - pub(crate) async fn persist_if_needed(&self) -> Result<()> { + async fn persist_if_needed_inner(&self) -> Result<()> { let (client_id, maybe_credentials) = { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; @@ -321,18 +482,24 @@ impl OAuthPersistor { expires_at, }; if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?; + save_oauth_tokens_for_runtime( + &self.inner.server_name, + &stored, + self.inner.store_mode, + ) + .await?; *last_credentials = Some(stored); } } None => { let mut last_serialized = self.inner.last_credentials.lock().await; if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( + && let Err(error) = delete_oauth_tokens_for_runtime( &self.inner.server_name, &self.inner.url, self.inner.store_mode, ) + .await { warn!( "failed to remove OAuth tokens for server {}: {error}", @@ -350,27 +517,451 @@ impl OAuthPersistor { reason = "AuthorizationManager async access must be serialized through its mutex" )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let expires_at = { + let refresh_attempt_guard = Arc::clone(&self.inner.refresh_attempt).lock_owned().await; + let previous_failure = self + .inner + .refresh_persistence_failure + .lock() + .await + .as_ref() + .map(|failure| failure.message.clone()); + if let Some(message) = previous_failure { + return Err(anyhow!(message)); + } + + let initial_expires_at = { let guard = self.inner.last_credentials.lock().await; guard.as_ref().and_then(|tokens| tokens.expires_at) }; - if !token_needs_refresh(expires_at) { + if !token_needs_refresh(initial_expires_at) { return Ok(()); } - { - let manager = self.inner.authorization_manager.clone(); + // A different process may have rotated the one-time refresh token. + // Reload its durable result while all refreshers share the same lock. + 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) + }; + if credentials_changed { + let token_response = stored.token_response.0.clone(); + let granted_scopes = token_response + .scopes() + .map(|scopes| { + scopes + .iter() + .map(|scope| scope.as_ref().to_string()) + .collect() + }) + .unwrap_or_default(); + let token_received_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + self.inner + .credential_store + .save(StoredCredentials::new( + stored.client_id.clone(), + Some(token_response), + granted_scopes, + Some(token_received_at), + )) + .await + .context("failed to reload persisted OAuth credentials")?; + *self.inner.last_credentials.lock().await = Some(stored); + } + } + + let reloaded_expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + if !token_needs_refresh(reloaded_expires_at) { + return Ok(()); + } + + let persistor = self.clone(); + tokio::spawn(async move { + let _refresh_attempt_guard = refresh_attempt_guard; + let manager = persistor.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 {}", + persistor.inner.server_name + ) + }); + } + } + drop(guard); + + if let Err(error) = persistor.persist_if_needed_inner().await { + let message = format!( + "OAuth refresh succeeded but failed to persist rotated credentials for server {}: {error:#}", + persistor.inner.server_name + ); + warn!("{message}"); + retain_failed_refresh_locks(refresh_locks); + *persistor.inner.refresh_persistence_failure.lock().await = + Some(RefreshPersistenceFailure { + message: message.clone(), + }); + return Err(anyhow!(message)); + } + + Ok(()) + }) + .await + .context("OAuth refresh task failed")? + } +} + +async fn lock_and_load_oauth_tokens( + server_name: &str, + url: &str, + 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(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( + server_name, + url, + )?) + .await?; + let tokens = load_oauth_tokens_from_file_async(server_name, url).await?; + 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( + server_name, + url, + )?) + .await?; + let tokens = load_oauth_tokens_from_file_async(server_name, url) + .await + .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(server_name, url)?) + .await?; + let tokens = load_oauth_tokens_from_file_async(server_name, url).await?; + Ok((vec![lock], tokens)) + } + OAuthCredentialsStoreMode::Keyring => { + 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 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)?; + external_oauth_lock_path(FILE_OAUTH_REFRESH_LOCK_PREFIX, &lock_id) +} + +fn keyring_oauth_refresh_lock_path(server_name: &str, url: &str) -> Result { + let lock_id = oauth_refresh_lock_id(server_name, url)?; + external_oauth_lock_path(KEYRING_OAUTH_REFRESH_LOCK_PREFIX, &lock_id) +} + +fn external_oauth_lock_path(prefix: &str, lock_id: &str) -> Result { + let user_namespace = os_user_namespace()?; + Ok(os_shared_temp_dir()?.join(format!("{prefix}-{user_namespace}-{lock_id}.lock"))) +} + +#[cfg(unix)] +fn os_shared_temp_dir() -> Result { + Ok(PathBuf::from("/tmp")) +} + +#[cfg(unix)] +fn os_user_namespace() -> Result { + // SAFETY: getuid has no preconditions and returns the real UID of this process. + Ok(format!("uid-{}", unsafe { libc::getuid() })) +} + +#[cfg(windows)] +fn os_shared_temp_dir() -> Result { + use std::ffi::OsString; + use std::io; + use std::os::windows::ffi::OsStringExt; + use windows_sys::Win32::System::SystemInformation::GetSystemWindowsDirectoryW; + + let mut buffer = vec![0_u16; 32_768]; + let buffer_len = + u32::try_from(buffer.len()).context("Windows path buffer length did not fit u32")?; + // SAFETY: buffer is writable for the length passed to GetSystemWindowsDirectoryW. + let length = unsafe { GetSystemWindowsDirectoryW(buffer.as_mut_ptr(), buffer_len) }; + if length == 0 { + return Err(io::Error::last_os_error()) + .context("failed to resolve the system Windows directory"); + } + let length = usize::try_from(length).context("Windows directory length did not fit usize")?; + if length >= buffer.len() { + return Err(anyhow!( + "system Windows directory exceeded the fixed path buffer" + )); + } + + Ok(PathBuf::from(OsString::from_wide(&buffer[..length])).join("Temp")) +} + +#[cfg(windows)] +fn os_user_namespace() -> Result { + use std::io; + use windows_sys::Win32::Foundation::CloseHandle; + use windows_sys::Win32::Foundation::HANDLE; + use windows_sys::Win32::Security::GetLengthSid; + use windows_sys::Win32::Security::GetTokenInformation; + use windows_sys::Win32::Security::TOKEN_QUERY; + use windows_sys::Win32::Security::TOKEN_USER; + use windows_sys::Win32::Security::TokenUser; + use windows_sys::Win32::System::Threading::GetCurrentProcess; + use windows_sys::Win32::System::Threading::OpenProcessToken; + + struct OwnedHandle(HANDLE); + + impl Drop for OwnedHandle { + fn drop(&mut self) { + if self.0 != 0 { + // SAFETY: the handle was returned by OpenProcessToken and is owned here. + unsafe { + CloseHandle(self.0); + } + } + } + } + + let mut token = 0; + // SAFETY: token points to writable storage for the returned process-token handle. + if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token) } == 0 { + return Err(io::Error::last_os_error()).context("failed to open the current process token"); + } + let token = OwnedHandle(token); + + let mut required = 0; + // SAFETY: the null-buffer call obtains the required TOKEN_USER buffer length. + unsafe { + GetTokenInformation(token.0, TokenUser, std::ptr::null_mut(), 0, &mut required); + } + if required == 0 { + return Err(io::Error::last_os_error()) + .context("failed to size the current process user SID"); + } + + let mut buffer = vec![0_u8; required as usize]; + // SAFETY: buffer is writable for required bytes and receives a TOKEN_USER value. + if unsafe { + GetTokenInformation( + token.0, + TokenUser, + buffer.as_mut_ptr().cast(), + required, + &mut required, + ) + } == 0 + { + return Err(io::Error::last_os_error()) + .context("failed to read the current process user SID"); + } + + // SAFETY: GetTokenInformation initialized the buffer with TOKEN_USER. The + // unaligned read avoids assuming Vec has TOKEN_USER alignment. + let token_user = unsafe { std::ptr::read_unaligned(buffer.as_ptr().cast::()) }; + // SAFETY: token_user.User.Sid points into buffer and remains valid here. + let sid_length = unsafe { GetLengthSid(token_user.User.Sid) }; + if sid_length == 0 { + return Err(io::Error::last_os_error()).context("failed to size the current user SID"); + } + // SAFETY: GetLengthSid returned the byte length of the valid SID in buffer. + let sid = unsafe { + std::slice::from_raw_parts(token_user.User.Sid.cast::(), sid_length as usize) + }; + Ok(format!("sid-{}", sha_256_bytes_prefix(sid))) +} + +#[cfg(not(any(unix, windows)))] +fn os_shared_temp_dir() -> Result { + Err(anyhow!( + "keyring OAuth refresh locking is unsupported on this platform" + )) +} + +#[cfg(not(any(unix, windows)))] +fn os_user_namespace() -> Result { + Err(anyhow!( + "keyring OAuth refresh locking is unsupported on this platform" + )) +} + +async fn acquire_oauth_refresh_lock(path: &Path) -> Result { + let file = open_oauth_lock_file(path)?; + 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()) + }); + } + } + } +} + +fn acquire_fallback_write_lock() -> Result { + let path = fallback_lock_path()?; + let file = open_oauth_lock_file(&path)?; + file.lock().with_context(|| { + format!( + "failed to acquire OAuth fallback write lock at {}", + path.display() + ) + })?; + Ok(file) +} + +async fn acquire_fallback_write_lock_async() -> Result { + let path = fallback_lock_path()?; + let file = open_oauth_lock_file(&path)?; + 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 fallback write lock at {}", + path.display() + ) + }); + } + } + } +} + +fn acquire_fallback_read_lock() -> Result { + let path = fallback_lock_path()?; + 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) +} + +async fn acquire_fallback_read_lock_async() -> Result { + let path = fallback_lock_path()?; + let file = open_oauth_lock_file(&path)?; + loop { + match file.try_lock_shared() { + 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 fallback read lock at {}", + path.display() + ) + }); + } } + } +} + +fn fallback_lock_path() -> Result { + let codex_home = find_codex_home()?; + let codex_home_id = sha_256_bytes_prefix(codex_home.as_os_str().to_string_lossy().as_bytes()); + external_oauth_lock_path(FALLBACK_LOCK_PREFIX, &codex_home_id) +} - self.persist_if_needed().await +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, } } @@ -394,7 +985,57 @@ 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) +} + +async fn load_oauth_tokens_from_file_async( + server_name: &str, + url: &str, +) -> Result> { + let _read_lock = acquire_fallback_read_lock_async().await?; + load_oauth_tokens_from_file_unlocked(server_name, url) +} + +async fn load_oauth_tokens_from_file_for_runtime( + server_name: &str, + url: &str, + fallback_lock_mode: FallbackLockMode, +) -> Result { + let read_lock = match fallback_lock_mode { + FallbackLockMode::DeferIfContended => { + let path = fallback_lock_path()?; + let file = open_oauth_lock_file(&path)?; + match file.try_lock_shared() { + Ok(()) => file, + Err(fs::TryLockError::WouldBlock) => { + return Ok(RuntimeOAuthTokenLoad::FallbackLockContended); + } + Err(error) => { + return Err(error).with_context(|| { + format!( + "failed to acquire OAuth fallback read lock at {}", + path.display() + ) + }); + } + } + } + FallbackLockMode::Wait => acquire_fallback_read_lock_async().await?, + }; + let tokens = load_oauth_tokens_from_file_unlocked(server_name, url)?; + drop(read_lock); + Ok(match tokens { + Some(tokens) => RuntimeOAuthTokenLoad::Loaded(Box::new(tokens)), + None => RuntimeOAuthTokenLoad::Missing, + }) +} + +fn load_oauth_tokens_from_file_unlocked( + server_name: &str, + url: &str, +) -> Result> { + let Some(store) = read_fallback_file_unlocked()? else { return Ok(None); }; @@ -437,8 +1078,18 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result Result<()> { + let _write_lock = acquire_fallback_write_lock()?; + save_oauth_tokens_to_file_unlocked(tokens) +} + +async fn save_oauth_tokens_to_file_async(tokens: &StoredOAuthTokens) -> Result<()> { + let _write_lock = acquire_fallback_write_lock_async().await?; + save_oauth_tokens_to_file_unlocked(tokens) +} + +fn save_oauth_tokens_to_file_unlocked(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 @@ -462,11 +1113,21 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> { }; store.insert(key, entry); - write_fallback_file(&store) + write_fallback_file_unlocked(&store) } fn delete_oauth_tokens_from_file(key: &str) -> Result { - let mut store = match read_fallback_file()? { + let _write_lock = acquire_fallback_write_lock()?; + delete_oauth_tokens_from_file_unlocked(key) +} + +async fn delete_oauth_tokens_from_file_async(key: &str) -> Result { + let _write_lock = acquire_fallback_write_lock_async().await?; + delete_oauth_tokens_from_file_unlocked(key) +} + +fn delete_oauth_tokens_from_file_unlocked(key: &str) -> Result { + let mut store = match read_fallback_file_unlocked()? { Some(store) => store, None => return Ok(false), }; @@ -474,7 +1135,7 @@ fn delete_oauth_tokens_from_file(key: &str) -> Result { let removed = store.remove(key).is_some(); if removed { - write_fallback_file(&store)?; + write_fallback_file_unlocked(&store)?; } Ok(removed) @@ -537,7 +1198,7 @@ fn fallback_file_path() -> Result { Ok(find_codex_home()?.join(FALLBACK_FILENAME).to_path_buf()) } -fn read_fallback_file() -> Result> { +fn read_fallback_file_unlocked() -> Result> { let path = fallback_file_path()?; let contents = match fs::read_to_string(&path) { Ok(contents) => contents, @@ -559,7 +1220,7 @@ fn read_fallback_file() -> Result> { } } -fn write_fallback_file(store: &FallbackFile) -> Result<()> { +fn write_fallback_file_unlocked(store: &FallbackFile) -> Result<()> { let path = fallback_file_path()?; if store.is_empty() { @@ -589,12 +1250,16 @@ fn write_fallback_file(store: &FallbackFile) -> Result<()> { fn sha_256_prefix(value: &Value) -> Result { let serialized = serde_json::to_string(&value).context("failed to serialize MCP OAuth key payload")?; + Ok(sha_256_bytes_prefix(serialized.as_bytes())) +} + +fn sha_256_bytes_prefix(bytes: &[u8]) -> String { let mut hasher = Sha256::new(); - hasher.update(serialized.as_bytes()); + hasher.update(bytes); let digest = hasher.finalize(); let hex = format!("{digest:x}"); let truncated = &hex[..16]; - Ok(truncated.to_string()) + truncated.to_string() } #[cfg(test)] @@ -607,10 +1272,14 @@ mod tests { use std::sync::MutexGuard; use std::sync::OnceLock; use std::sync::PoisonError; + use std::sync::mpsc; + use std::thread; use tempfile::tempdir; use codex_keyring_store::tests::MockKeyringStore; + static CODEX_HOME_LOCK: OnceLock> = OnceLock::new(); + struct TempCodexHome { _guard: MutexGuard<'static, ()>, _dir: tempfile::TempDir, @@ -618,8 +1287,7 @@ mod tests { impl TempCodexHome { fn new() -> Self { - static LOCK: OnceLock> = OnceLock::new(); - let guard = LOCK + let guard = CODEX_HOME_LOCK .get_or_init(Mutex::default) .lock() .unwrap_or_else(PoisonError::into_inner); @@ -632,6 +1300,28 @@ mod tests { _dir: dir, } } + + fn unusable() -> Self { + let guard = CODEX_HOME_LOCK + .get_or_init(Mutex::default) + .lock() + .unwrap_or_else(PoisonError::into_inner); + let dir = tempdir().expect("create CODEX_HOME parent temp dir"); + let file = dir.path().join("not-a-directory"); + fs::write(&file, "not a directory").expect("create unusable CODEX_HOME path"); + unsafe { + std::env::set_var("CODEX_HOME", file); + } + Self { + _guard: guard, + _dir: dir, + } + } + + #[cfg(unix)] + fn path(&self) -> &Path { + self._dir.path() + } } impl Drop for TempCodexHome { @@ -642,6 +1332,256 @@ mod tests { } } + #[test] + 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)?; + let file_path = file_oauth_refresh_lock_path(&tokens.server_name, &tokens.url)?; + let shared_temp_dir = os_shared_temp_dir()?; + let user_namespace = os_user_namespace()?; + assert_eq!(keyring_path.parent(), Some(shared_temp_dir.as_path())); + assert_eq!(file_path.parent(), Some(shared_temp_dir.as_path())); + assert!( + keyring_path + .file_name() + .is_some_and(|name| name.to_string_lossy().contains(&user_namespace)) + ); + assert!( + file_path + .file_name() + .is_some_and(|name| name.to_string_lossy().contains(&user_namespace)) + ); + #[cfg(unix)] + assert_eq!(shared_temp_dir, PathBuf::from("/tmp")); + 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_path, + file_oauth_refresh_lock_path(&tokens.server_name, other_url)? + ); + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn keyring_refresh_locking_works_without_codex_home_for_auto_and_keyring() -> Result<()> { + let _env = TempCodexHome::unusable(); + 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 held_lock = acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path( + &tokens.server_name, + &tokens.url, + )?) + .await?; + let blocked = tokio::time::timeout( + Duration::from_millis(50), + lock_and_load_oauth_tokens_with_keyring( + &store, + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::Keyring, + ), + ) + .await; + assert!( + blocked.is_err(), + "same keyring credential should wait for its refresh lock" + ); + drop(held_lock); + + for store_mode in [ + OAuthCredentialsStoreMode::Keyring, + OAuthCredentialsStoreMode::Auto, + ] { + let (locks, loaded) = lock_and_load_oauth_tokens_with_keyring( + &store, + &tokens.server_name, + &tokens.url, + store_mode, + ) + .await?; + assert_eq!(locks.len(), 1); + assert_tokens_match_without_expiry( + &loaded.expect("keyring credentials should load"), + &tokens, + ); + } + Ok(()) + } + + #[tokio::test(flavor = "current_thread")] + async fn distinct_keyring_credentials_do_not_block_each_other() -> Result<()> { + let _env = TempCodexHome::unusable(); + let store = MockKeyringStore::default(); + let first = sample_tokens(); + let mut second = sample_tokens(); + second.url = "https://other.example.test/mcp".to_string(); + for tokens in [&first, &second] { + let key = compute_store_key(&tokens.server_name, &tokens.url)?; + store.save(KEYRING_SERVICE, &key, &serde_json::to_string(tokens)?)?; + } + + let held_lock = acquire_oauth_refresh_lock(&keyring_oauth_refresh_lock_path( + &first.server_name, + &first.url, + )?) + .await?; + for store_mode in [ + OAuthCredentialsStoreMode::Keyring, + OAuthCredentialsStoreMode::Auto, + ] { + let result = tokio::time::timeout( + Duration::from_millis(500), + lock_and_load_oauth_tokens_with_keyring( + &store, + &second.server_name, + &second.url, + store_mode, + ), + ) + .await; + assert!( + result.is_ok(), + "a distinct keyring credential should not wait for another lock" + ); + let (locks, loaded) = result.expect("timeout checked above")?; + assert_eq!(locks.len(), 1); + assert_tokens_match_without_expiry( + &loaded.expect("second keyring credential should load"), + &second, + ); + } + drop(held_lock); + 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(()) + } + + #[cfg(unix)] + #[test] + fn file_and_auto_reads_work_with_read_only_codex_home() -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + let env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); + save_oauth_tokens_to_file(&tokens)?; + + let credentials_path = fallback_file_path()?; + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + fs::set_permissions(env.path(), fs::Permissions::from_mode(0o500))?; + + let result = (|| -> Result<(StoredOAuthTokens, StoredOAuthTokens)> { + let file_tokens = load_oauth_tokens_from_file(&tokens.server_name, &tokens.url)? + .expect("file credentials should load"); + let auto_tokens = load_oauth_tokens_from_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens.url, + )? + .expect("auto credentials should load from fallback"); + Ok((file_tokens, auto_tokens)) + })(); + + fs::set_permissions(env.path(), fs::Permissions::from_mode(0o700))?; + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + + let (file_tokens, auto_tokens) = result?; + assert_tokens_match_without_expiry(&file_tokens, &tokens); + assert_tokens_match_without_expiry(&auto_tokens, &tokens); + assert!(!env.path().join(".credentials.json.lock").exists()); + assert_eq!( + fallback_lock_path()?.parent(), + Some(os_shared_temp_dir()?.as_path()) + ); + 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(); @@ -737,7 +1677,8 @@ 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 _read_lock = super::acquire_fallback_read_lock()?; + 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); @@ -815,6 +1756,79 @@ mod tests { Ok(()) } + #[test] + fn concurrent_save_and_delete_preserve_other_fallback_entries() -> Result<()> { + let _env = TempCodexHome::new(); + let target = sample_tokens(); + let target_key = super::compute_store_key(&target.server_name, &target.url)?; + let mut unrelated = sample_tokens(); + unrelated.url = "https://unrelated.example.test".to_string(); + let mut new_tokens = sample_tokens(); + new_tokens.url = "https://new.example.test".to_string(); + let expected_new_tokens = new_tokens.clone(); + + super::save_oauth_tokens_to_file(&target)?; + super::save_oauth_tokens_to_file(&unrelated)?; + + let held_write_lock = super::acquire_fallback_write_lock()?; + let (started_tx, started_rx) = mpsc::channel(); + let (delete_done_tx, delete_done_rx) = mpsc::channel(); + let (save_done_tx, save_done_rx) = mpsc::channel(); + + let delete_started_tx = started_tx.clone(); + let delete_thread = thread::spawn(move || { + let _ = delete_started_tx.send(()); + let result = super::delete_oauth_tokens_from_file(&target_key); + let _ = delete_done_tx.send(()); + result + }); + let save_thread = thread::spawn(move || { + let _ = started_tx.send(()); + let result = super::save_oauth_tokens_to_file(&new_tokens); + let _ = save_done_tx.send(()); + result + }); + + for _ in 0..2 { + started_rx + .recv_timeout(Duration::from_secs(1)) + .context("timed out waiting for fallback mutation thread")?; + } + let save_completed_while_write_locked = + save_done_rx.recv_timeout(Duration::from_secs(1)).is_ok(); + let delete_completed_while_write_locked = delete_done_rx.try_recv().is_ok(); + drop(held_write_lock); + + let removed = delete_thread + .join() + .map_err(|_| anyhow!("fallback delete thread panicked"))??; + save_thread + .join() + .map_err(|_| anyhow!("fallback save thread panicked"))??; + + assert!( + !save_completed_while_write_locked, + "fallback save escaped the competing exclusive lock" + ); + assert!( + !delete_completed_while_write_locked, + "fallback deletion escaped the competing exclusive lock" + ); + assert!(removed); + assert!(super::load_oauth_tokens_from_file(&target.server_name, &target.url)?.is_none()); + let loaded_unrelated = + super::load_oauth_tokens_from_file(&unrelated.server_name, &unrelated.url)? + .context("unrelated fallback entry was lost")?; + let loaded_new = super::load_oauth_tokens_from_file( + &expected_new_tokens.server_name, + &expected_new_tokens.url, + )? + .context("concurrently saved fallback entry was lost")?; + assert_tokens_match_without_expiry(&loaded_unrelated, &unrelated); + assert_tokens_match_without_expiry(&loaded_new, &expected_new_tokens); + Ok(()) + } + #[test] fn refresh_expires_in_from_timestamp_restores_future_durations() { let mut tokens = sample_tokens(); diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c3d..46e32fec7383 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -47,6 +47,7 @@ use rmcp::service::{self}; use rmcp::transport::StreamableHttpClientTransport; use rmcp::transport::auth::AuthClient; use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::InMemoryCredentialStore; use rmcp::transport::auth::OAuthState; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; use rmcp::transport::streamable_http_client::StreamableHttpError; @@ -63,9 +64,11 @@ use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::in_process_transport::InProcessTransportFactory; -use crate::load_oauth_tokens; +use crate::oauth::FallbackLockMode; use crate::oauth::OAuthPersistor; +use crate::oauth::RuntimeOAuthTokenLoad; use crate::oauth::StoredOAuthTokens; +use crate::oauth::load_oauth_tokens_for_runtime; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -75,6 +78,9 @@ use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; enum PendingTransport { + DeferredStreamableHttp { + recipe: TransportRecipe, + }, InProcess { transport: tokio::io::DuplexStream, }, @@ -219,6 +225,8 @@ where enum ClientOperationError { #[error(transparent)] Service(#[from] rmcp::service::ServiceError), + #[error(transparent)] + OAuth(anyhow::Error), #[error("timed out awaiting {label} after {duration:?}")] Timeout { label: String, duration: Duration }, } @@ -320,7 +328,8 @@ impl RmcpClient { .map_err(io::Error::other)?; let stdio_process = match &transport { PendingTransport::Stdio { transport } => Some(transport.process_handle()), - PendingTransport::InProcess { .. } + PendingTransport::DeferredStreamableHttp { .. } + | PendingTransport::InProcess { .. } | PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } => None, }; @@ -425,12 +434,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -439,14 +442,12 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -455,7 +456,6 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -480,7 +480,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -500,14 +499,12 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -516,14 +513,12 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -532,14 +527,12 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -550,7 +543,6 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -597,7 +589,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -606,7 +597,6 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -627,7 +617,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -636,7 +625,6 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -650,7 +638,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } @@ -690,26 +677,29 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); - } + async fn create_pending_transport( + transport_recipe: &TransportRecipe, + ) -> Result { + Self::create_pending_transport_with_fallback_lock_mode( + transport_recipe, + FallbackLockMode::DeferIfContended, + ) + .await } - async fn refresh_oauth_if_needed(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); - } + async fn create_pending_transport_now( + transport_recipe: &TransportRecipe, + ) -> Result { + Self::create_pending_transport_with_fallback_lock_mode( + transport_recipe, + FallbackLockMode::Wait, + ) + .await } - async fn create_pending_transport( + async fn create_pending_transport_with_fallback_lock_mode( transport_recipe: &TransportRecipe, + fallback_lock_mode: FallbackLockMode, ) -> Result { match transport_recipe { TransportRecipe::InProcess { factory } => { @@ -737,8 +727,21 @@ impl RmcpClient { && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(server_name, url, *store_mode) { - Ok(tokens) => tokens, + match load_oauth_tokens_for_runtime( + server_name, + url, + *store_mode, + fallback_lock_mode, + ) + .await + { + Ok(RuntimeOAuthTokenLoad::Loaded(tokens)) => Some(*tokens), + Ok(RuntimeOAuthTokenLoad::Missing) => None, + Ok(RuntimeOAuthTokenLoad::FallbackLockContended) => { + return Ok(PendingTransport::DeferredStreamableHttp { + recipe: transport_recipe.clone(), + }); + } Err(err) => { warn!("failed to read tokens for server `{server_name}`: {err}"); None @@ -823,39 +826,60 @@ 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, - } => ( - service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), - ), - }; + let connect = async move { + let pending_transport = match pending_transport { + PendingTransport::DeferredStreamableHttp { recipe } => { + Self::create_pending_transport_now(&recipe).await? + } + pending_transport => pending_transport, + }; + let (transport, oauth_persistor) = match pending_transport { + PendingTransport::DeferredStreamableHttp { .. } => { + unreachable!("deferred streamable HTTP transport resolved above") + } + 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), + ), + }; - let service = match timeout { - Some(duration) => time::timeout(duration, transport) - .await - .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))? - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, - None => transport + 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}"))?, + .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?; + if let Some(runtime) = oauth_persistor.clone() { + Self::persist_oauth_tokens_in_background( + runtime, + "failed to persist OAuth tokens after initialize", + ); + } + Ok((Arc::new(service), oauth_persistor)) }; - Ok((Arc::new(service), oauth_persistor)) + match timeout { + Some(duration) => time::timeout(duration, connect) + .await + .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))?, + None => connect.await, + } } async fn run_service_operation( @@ -869,8 +893,10 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; + let oauth_persistor = self.oauth_persistor().await; match Self::run_service_operation_once( Arc::clone(&service), + oauth_persistor, label, timeout, self.elicitation_pause_state.clone(), @@ -882,8 +908,10 @@ impl RmcpClient { Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; + let oauth_persistor = self.oauth_persistor().await; Self::run_service_operation_once( recovered_service, + oauth_persistor, label, timeout, self.elicitation_pause_state.clone(), @@ -898,6 +926,7 @@ impl RmcpClient { async fn run_service_operation_once( service: Arc>, + oauth_persistor: Option, label: &str, timeout: Option, pause_state: ElicitationPauseState, @@ -907,20 +936,43 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { + let refresh_and_operation = async move { + if let Some(runtime) = &oauth_persistor { + runtime + .refresh_if_needed() + .await + .map_err(ClientOperationError::OAuth)?; + } + let result = operation(service) + .await + .map_err(ClientOperationError::from)?; + if let Some(runtime) = oauth_persistor { + Self::persist_oauth_tokens_in_background(runtime, "failed to persist OAuth tokens"); + } + Ok(result) + }; + match timeout { Some(duration) => { - active_time_timeout(duration, pause_state.subscribe(), operation(service)) + active_time_timeout(duration, pause_state.subscribe(), refresh_and_operation) .await .map_err(|_| ClientOperationError::Timeout { label: label.to_string(), duration, })? - .map_err(ClientOperationError::from) } - None => operation(service).await.map_err(ClientOperationError::from), + None => refresh_and_operation.await, } } + fn persist_oauth_tokens_in_background(runtime: OAuthPersistor, error_context: &'static str) { + tokio::spawn(async move { + if let Err(error) = runtime.persist_if_needed().await { + warn!("{error_context}: {error}"); + } + }); + } + fn is_session_expired_404(error: &ClientOperationError) -> bool { let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = error @@ -992,12 +1044,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } @@ -1023,6 +1069,15 @@ async fn create_oauth_transport_and_runtime( // reqwest metadata client here. let mut oauth_state = OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?; + // Keep a handle so Codex can replace stale RMCP credentials after another + // process completes a refresh. + let credential_store = InMemoryCredentialStore::new(); + match &mut oauth_state { + OAuthState::Unauthorized(manager) => { + manager.set_credential_store(credential_store.clone()); + } + _ => return Err(anyhow!("unexpected OAuth state during client setup")), + } oauth_state .set_credentials( @@ -1054,6 +1109,7 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, + credential_store, credentials_store, Some(initial_tokens), ); @@ -1063,12 +1119,125 @@ async fn create_oauth_transport_and_runtime( #[cfg(test)] mod tests { + #[cfg(unix)] + use std::ffi::OsString; + #[cfg(unix)] + use std::fs::OpenOptions; + #[cfg(unix)] + use std::path::Path; + #[cfg(unix)] + use std::path::PathBuf; + #[cfg(unix)] + use std::sync::Arc; use std::time::Duration; - + #[cfg(unix)] + use std::time::SystemTime; + #[cfg(unix)] + use std::time::UNIX_EPOCH; + + #[cfg(unix)] + use anyhow::Context; + #[cfg(unix)] + use codex_exec_server::Environment; + #[cfg(unix)] + use oauth2::AccessToken; + #[cfg(unix)] + use oauth2::RefreshToken; + #[cfg(unix)] + use oauth2::basic::BasicTokenType; use pretty_assertions::assert_eq; + #[cfg(unix)] + use rmcp::model::ClientCapabilities; + #[cfg(unix)] + use rmcp::model::Implementation; + #[cfg(unix)] + use rmcp::model::ProtocolVersion; + #[cfg(unix)] + use rmcp::transport::auth::CredentialStore; + #[cfg(unix)] + use rmcp::transport::auth::OAuthTokenResponse; + #[cfg(unix)] + use rmcp::transport::auth::StoredCredentials; + #[cfg(unix)] + use rmcp::transport::auth::VendorExtraTokenFields; + #[cfg(unix)] + use serde_json::Value; + #[cfg(unix)] + use serde_json::json; + #[cfg(unix)] + use sha2::Digest; + #[cfg(unix)] + use sha2::Sha256; + #[cfg(unix)] + use tempfile::TempDir; use tokio::time; + #[cfg(unix)] + use wiremock::Mock; + #[cfg(unix)] + use wiremock::MockServer; + #[cfg(unix)] + use wiremock::Request; + #[cfg(unix)] + use wiremock::ResponseTemplate; + #[cfg(unix)] + use wiremock::matchers::method; + #[cfg(unix)] + use wiremock::matchers::path; use super::*; + #[cfg(unix)] + use crate::oauth::WrappedOAuthTokenResponse; + #[cfg(unix)] + use crate::oauth::load_oauth_tokens; + #[cfg(unix)] + use crate::oauth::save_oauth_tokens; + + #[cfg(unix)] + const OPERATION_TEST_SERVER_NAME: &str = "test-completed-operation-persistence"; + #[cfg(unix)] + const OPERATION_TEST_CLIENT_ID: &str = "test-operation-client-id"; + #[cfg(unix)] + const OPERATION_TEST_OLD_ACCESS_TOKEN: &str = "old-operation-access-token"; + #[cfg(unix)] + const OPERATION_TEST_OLD_REFRESH_TOKEN: &str = "old-operation-refresh-token"; + #[cfg(unix)] + const OPERATION_TEST_NEW_ACCESS_TOKEN: &str = "new-operation-access-token"; + #[cfg(unix)] + const OPERATION_TEST_NEW_REFRESH_TOKEN: &str = "new-operation-refresh-token"; + + #[cfg(unix)] + struct CodexHomeGuard { + original: Option, + } + + #[cfg(unix)] + impl CodexHomeGuard { + fn set(path: &Path) -> Self { + let original = std::env::var_os("CODEX_HOME"); + // SAFETY: this test is serialized and restores the process environment on drop. + unsafe { + std::env::set_var("CODEX_HOME", path); + } + Self { original } + } + } + + #[cfg(unix)] + impl Drop for CodexHomeGuard { + fn drop(&mut self) { + if let Some(original) = &self.original { + // SAFETY: this test is serialized and restores the prior value. + unsafe { + std::env::set_var("CODEX_HOME", original); + } + } else { + // SAFETY: this test is serialized and removes only the value it set. + unsafe { + std::env::remove_var("CODEX_HOME"); + } + } + } + } #[tokio::test] async fn active_time_timeout_pauses_while_elicitation_is_pending() { @@ -1088,4 +1257,278 @@ mod tests { assert_eq!(Ok("done"), result); } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn completed_operation_returns_while_oauth_persistence_is_contended() -> Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400).set_body_string("invalid JSON-RPC request"); + }; + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "serverInfo": { + "name": "completed-operation-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let _codex_home_guard = CodexHomeGuard::set(codex_home.path()); + let server_url = format!("{}/mcp", server.uri()); + let old_response = operation_test_token_response( + OPERATION_TEST_OLD_ACCESS_TOKEN, + OPERATION_TEST_OLD_REFRESH_TOKEN, + ); + let initial_tokens = StoredOAuthTokens { + server_name: OPERATION_TEST_SERVER_NAME.to_string(), + url: server_url.clone(), + client_id: OPERATION_TEST_CLIENT_ID.to_string(), + token_response: WrappedOAuthTokenResponse(old_response.clone()), + expires_at: Some(unix_millis().saturating_add(3_600_000)), + }; + save_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &initial_tokens, + OAuthCredentialsStoreMode::File, + )?; + + let client = RmcpClient::new_streamable_http_client( + OPERATION_TEST_SERVER_NAME, + &server_url, + Some("test-bearer".to_string()), + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(2)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + let service = client.service().await?; + + let credential_store = InMemoryCredentialStore::new(); + let mut oauth_state = + OAuthState::new(server_url.clone(), Some(reqwest::Client::new())).await?; + match &mut oauth_state { + OAuthState::Unauthorized(manager) => { + manager.set_credential_store(credential_store.clone()); + } + _ => { + return Err(anyhow!( + "unexpected OAuth state during operation test setup" + )); + } + } + oauth_state + .set_credentials(OPERATION_TEST_CLIENT_ID, old_response) + .await?; + let authorization_manager = match oauth_state { + OAuthState::Authorized(manager) => manager, + _ => return Err(anyhow!("OAuth operation test setup did not authorize")), + }; + let oauth_persistor = OAuthPersistor::new( + OPERATION_TEST_SERVER_NAME.to_string(), + server_url.clone(), + Arc::new(Mutex::new(authorization_manager)), + credential_store.clone(), + OAuthCredentialsStoreMode::File, + Some(initial_tokens.clone()), + ); + + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(operation_test_fallback_lock_path(codex_home.path())?)?; + fallback_lock.lock_shared()?; + + let new_response = operation_test_token_response( + OPERATION_TEST_NEW_ACCESS_TOKEN, + OPERATION_TEST_NEW_REFRESH_TOKEN, + ); + let operation_store = credential_store.clone(); + let operation_response = new_response.clone(); + let operation = move |_service| { + let operation_store = operation_store.clone(); + let operation_response = operation_response.clone(); + async move { + if let Err(error) = operation_store + .save(StoredCredentials::new( + OPERATION_TEST_CLIENT_ID.to_string(), + Some(operation_response), + Vec::new(), + Some(unix_seconds()), + )) + .await + { + panic!("failed to simulate completed OAuth side effect: {error}"); + } + Ok::<_, rmcp::service::ServiceError>("side effect completed".to_string()) + } + .boxed() + }; + + let result = RmcpClient::run_service_operation_once( + service, + Some(oauth_persistor), + "tools/call", + Some(Duration::from_millis(100)), + ElicitationPauseState::new(), + &operation, + ) + .await?; + assert_eq!(result, "side effect completed"); + time::sleep(Duration::from_millis(200)).await; + let persisted_while_locked = load_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::File, + )? + .context("missing persisted OAuth credentials while fallback lock is held")?; + assert_eq!( + ( + persisted_while_locked + .token_response + .0 + .access_token() + .secret() + .as_str(), + persisted_while_locked + .token_response + .0 + .refresh_token() + .map(|token| token.secret().as_str()), + ), + ( + OPERATION_TEST_OLD_ACCESS_TOKEN, + Some(OPERATION_TEST_OLD_REFRESH_TOKEN), + ) + ); + + drop(fallback_lock); + let persisted = time::timeout(Duration::from_secs(2), async { + loop { + let loaded = load_oauth_tokens( + OPERATION_TEST_SERVER_NAME, + &server_url, + OAuthCredentialsStoreMode::File, + )?; + if loaded.as_ref().is_some_and(|tokens| { + tokens.token_response.0.access_token().secret() + == OPERATION_TEST_NEW_ACCESS_TOKEN + }) { + return Ok::<_, anyhow::Error>(loaded); + } + time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for post-operation OAuth persistence")??; + assert_eq!( + persisted.as_ref().map(|tokens| tokens + .token_response + .0 + .access_token() + .secret() + .as_str()), + Some(OPERATION_TEST_NEW_ACCESS_TOKEN) + ); + + server.verify().await; + Ok(()) + } + + #[cfg(unix)] + fn operation_test_token_response( + access_token: &str, + refresh_token: &str, + ) -> OAuthTokenResponse { + 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_expires_in(Some(&Duration::from_secs(7200))); + response + } + + #[cfg(unix)] + fn operation_test_fallback_lock_path(codex_home: &Path) -> Result { + let canonical_home = codex_home.canonicalize()?; + let mut hasher = Sha256::new(); + hasher.update(canonical_home.as_os_str().to_string_lossy().as_bytes()); + let digest = format!("{:x}", hasher.finalize()); + // SAFETY: getuid has no preconditions and returns the real UID. + let user_namespace = format!("uid-{}", unsafe { libc::getuid() }); + Ok(PathBuf::from("/tmp").join(format!( + "codex-mcp-oauth-fallback-{user_namespace}-{}.lock", + &digest[..16] + ))) + } + + #[cfg(unix)] + fn unix_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 + } + + #[cfg(unix)] + fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + } } 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 new file mode 100644 index 000000000000..c17e5bb191c3 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs @@ -0,0 +1,2007 @@ +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; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +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; +use tokio::time::sleep; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::Times; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use streamable_http_test_support::initialize_client; + +const SERVER_NAME: &str = "test-streamable-http-oauth-refresh-race"; +const CLIENT_ID: &str = "test-client-id"; +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 FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-one-time-refresh-token"; +const FALLBACK_LOCK_PREFIX: &str = "codex-mcp-oauth-fallback"; +const FILE_OAUTH_REFRESH_LOCK_PREFIX: &str = "codex-mcp-oauth-refresh-file"; +const KEYRING_SERVICE: &str = "Codex MCP Credentials"; +const CHILD_WAIT_TIMEOUT: Duration = Duration::from_secs(10); +const OPERATION_TIMEOUT: Duration = Duration::from_millis(100); + +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RACE_SERVER_URL"; +const READY_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_READY_PATH"; +const GO_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_GO_PATH"; +const RESULT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_RESULT_PATH"; +const ATTEMPT_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_ATTEMPT_PATH"; +const ACCESS_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_ACCESS_TOKEN"; +const REFRESH_TOKEN_ENV: &str = "MCP_TEST_OAUTH_RACE_REFRESH_TOKEN"; +const EXPIRES_AT_ENV: &str = "MCP_TEST_OAUTH_RACE_EXPIRES_AT"; +const STORE_MODE_ENV: &str = "MCP_TEST_OAUTH_RACE_STORE_MODE"; +const RELEASE_PATH_ENV: &str = "MCP_TEST_OAUTH_RACE_RELEASE_PATH"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn concurrent_processes_coordinate_one_time_refresh_and_reload_rotated_tokens() +-> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { + 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", + "error_description": "refresh token already used", + })) + } + }) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 4).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 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"), + ]; + let result_paths = [ + control_dir.path().join("client-a.result"), + control_dir.path().join("client-b.result"), + ]; + let mut children = Vec::new(); + 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, + result_path, + )?); + } + + wait_for_paths(&ready_paths).await?; + 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)?; + assert_eq!(outcomes, vec!["success".to_string(), "success".to_string()]); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + let requests = server + .received_requests() + .await + .context("wiremock request recording disabled")?; + let refresh_bodies = request_bodies(&requests, "/oauth/token"); + assert_eq!(refresh_bodies.len(), 1); + assert!( + refresh_bodies + .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 { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + + server.verify().await; + 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 = file_oauth_refresh_lock_path(&codex_home, &server_url)?; + 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 fallback_read_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&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 (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_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")?; + wait_for_paths(std::slice::from_ref(&result_path)).await?; + assert_eq!( + fs::read_to_string(&result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + + drop(fallback_lock); + wait_for_children(vec![child]).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn fallback_write_lock_wait_uses_startup_timeout_budget() -> 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(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .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 (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + 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_paths(std::slice::from_ref(&result_path)).await?; + assert_eq!( + fs::read_to_string(&result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + + drop(fallback_lock); + wait_for_children(vec![child]).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_credentials_persist_after_startup_timeout_before_next_refresh() +-> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .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, + })) + } else { + ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })) + } + }) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + let first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_client_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + assert_eq!( + fs::read_to_string(&first_result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + assert!(!second_result_path.exists()); + + drop(fallback_lock); + wait_for_children(vec![second_child]).await?; + assert_eq!(fs::read_to_string(&second_result_path)?, "success"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn provider_refresh_persists_after_startup_timeout_before_next_refresh() -> anyhow::Result<()> +{ + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + responder_attempts.fetch_add(1, Ordering::SeqCst); + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(6)) + .set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + })) + }) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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 first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_client_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + assert_eq!( + fs::read_to_string(&first_result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + wait_for_persisted_token_snapshot( + &codex_home, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + }, + ) + .await?; + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + wait_for_children(vec![second_child]).await?; + assert_eq!(fs::read_to_string(second_result_path)?, "success"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn canceled_startup_persistence_failure_retains_lock_until_process_exit() -> anyhow::Result<()> +{ + use std::os::unix::fs::PermissionsExt; + + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match responder_attempts.fetch_add(1, Ordering::SeqCst) { + 0 => ResponseTemplate::new(200) + .set_delay(Duration::from_secs(6)) + .set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })), + } + }) + .expect(2) + .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 credentials_path = codex_home.path().join(".credentials.json"); + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + + let first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_client_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + assert_eq!( + fs::read_to_string(&first_result_path)?, + "error:timed out handshaking with MCP server after 5s" + ); + sleep(Duration::from_secs(2)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_attempt_path = control_dir.path().join("client-b.attempt"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child_with_attempt( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_attempt_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + assert!(!second_result_path.exists()); + + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + wait_for_children(vec![second_child]).await?; + assert_eq!( + fs::read_to_string(&second_result_path)?, + "error:Auth required: OAuth refresh token was rejected; reauthentication required" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> 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(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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_operation_timeout_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + sleep(Duration::from_secs(2)).await; + + let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock()?; + + fs::write(&go_path, "go")?; + wait_for_children(vec![child]).await?; + assert_eq!( + fs::read_to_string(result_path)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + + drop(fallback_lock); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn fallback_write_lock_wait_uses_operation_timeout_budget() -> 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(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 31, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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_operation_timeout_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + sleep(Duration::from_secs(2)).await; + + let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + fs::write(&go_path, "go")?; + wait_for_children(vec![child]).await?; + assert_eq!( + fs::read_to_string(result_path)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + + drop(fallback_lock); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_credentials_persist_after_operation_timeout_before_next_refresh() +-> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match 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": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + 1 => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })), + } + }) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 4..=5).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 first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_operation_timeout_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + sleep(Duration::from_secs(2)).await; + + let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?; + let fallback_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path)?; + fallback_lock.lock_shared()?; + + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + assert_eq!( + fs::read_to_string(&first_result_path)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + assert!(!second_result_path.exists()); + + drop(fallback_lock); + wait_for_children(vec![second_child]).await?; + assert_eq!(fs::read_to_string(&second_result_path)?, "success"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: FINAL_ACCESS_TOKEN.to_string(), + refresh_token: Some(FINAL_REFRESH_TOKEN.to_string()), + } + ); + + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match 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": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), + } + }) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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 release_path = control_dir.path().join("client.release"); + let child = spawn_operation_timeout_child_held_open( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + &release_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + sleep(Duration::from_secs(2)).await; + fs::write(&go_path, "go")?; + wait_for_paths(std::slice::from_ref(&result_path)).await?; + + assert_eq!( + fs::read_to_string(&result_path)?, + format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}") + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + wait_for_persisted_token_snapshot( + &codex_home, + PersistedTokenSnapshot { + access_token: FINAL_ACCESS_TOKEN.to_string(), + refresh_token: Some(FINAL_REFRESH_TOKEN.to_string()), + }, + ) + .await?; + + fs::write(release_path, "release")?; + wait_for_children(vec![child]).await?; + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rotated_token_storage_failure_blocks_another_process_refresh() -> anyhow::Result<()> { + use std::os::unix::fs::PermissionsExt; + + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match 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": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + 1 => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(400).set_body_json(json!({ + "error": "invalid_grant", + "error_description": "refresh token already used", + })), + } + }) + .expect(3) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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 first_ready_path = control_dir.path().join("client-a.ready"); + let first_go_path = control_dir.path().join("client-a.go"); + let first_result_path = control_dir.path().join("client-a.result"); + let first_release_path = control_dir.path().join("client-a.release"); + let first_child = spawn_operation_timeout_child_held_open( + &codex_home, + &server_url, + &first_ready_path, + &first_go_path, + &first_result_path, + &first_release_path, + )?; + wait_for_paths(std::slice::from_ref(&first_ready_path)).await?; + + let credentials_path = codex_home.path().join(".credentials.json"); + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + sleep(Duration::from_secs(2)).await; + fs::write(&first_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&first_result_path)).await?; + let first_outcome = fs::read_to_string(&first_result_path)?; + assert!( + first_outcome.starts_with(&format!( + "error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:" + )), + "unexpected first-client outcome: {first_outcome}" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + let second_ready_path = control_dir.path().join("client-b.ready"); + let second_go_path = control_dir.path().join("client-b.go"); + let second_attempt_path = control_dir.path().join("client-b.attempt"); + let second_result_path = control_dir.path().join("client-b.result"); + let second_child = spawn_client_child_with_attempt( + &codex_home, + &server_url, + &second_ready_path, + &second_go_path, + &second_attempt_path, + &second_result_path, + )?; + wait_for_paths(std::slice::from_ref(&second_ready_path)).await?; + fs::write(&second_go_path, "go")?; + wait_for_paths(std::slice::from_ref(&second_attempt_path)).await?; + sleep(Duration::from_millis(250)).await; + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + assert!(!second_result_path.exists()); + + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + fs::write(first_release_path, "release")?; + wait_for_children(vec![first_child]).await?; + wait_for_children(vec![second_child]).await?; + assert_eq!( + fs::read_to_string(&second_result_path)?, + "error:Auth required: OAuth refresh token was rejected; reauthentication required" + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 3); + + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn concurrent_waiter_observes_recorded_persistence_failure() -> anyhow::Result<()> { + use std::os::unix::fs::PermissionsExt; + + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 1).await; + let refresh_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&refresh_attempts); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(move |_request: &Request| { + match 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": 31, + "refresh_token": NEW_REFRESH_TOKEN, + })), + _ => ResponseTemplate::new(200).set_body_json(json!({ + "access_token": FINAL_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": FINAL_REFRESH_TOKEN, + })), + } + }) + .expect(2) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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_concurrent_waiter_child( + &codex_home, + &server_url, + &ready_path, + &go_path, + &result_path, + )?; + wait_for_paths(std::slice::from_ref(&ready_path)).await?; + + let credentials_path = codex_home.path().join(".credentials.json"); + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o400))?; + sleep(Duration::from_secs(2)).await; + fs::write(&go_path, "go")?; + wait_for_children(vec![child]).await?; + fs::set_permissions(&credentials_path, fs::Permissions::from_mode(0o600))?; + + let outcomes = fs::read_to_string(result_path)? + .lines() + .map(str::to_string) + .collect::>(); + assert_eq!(outcomes.len(), 2); + assert_eq!(outcomes[0], outcomes[1]); + assert!( + outcomes[0].starts_with(&format!( + "error:OAuth refresh succeeded but failed to persist rotated credentials for server {SERVER_NAME}:" + )), + "unexpected concurrent waiter outcome: {}", + outcomes[0] + ); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); + + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn distinct_file_credentials_do_not_block_each_other() -> 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(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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 unrelated_lock_path = + file_oauth_refresh_lock_path(&codex_home, "https://unrelated.example.test/mcp")?; + let unrelated_lock = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(unrelated_lock_path)?; + unrelated_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")?; + wait_for_children(vec![child]).await?; + + assert_eq!(fs::read_to_string(result_path)?, "success"); + drop(unrelated_lock); + server.verify().await; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn expired_file_credentials_refresh_with_read_only_codex_home() -> anyhow::Result<()> { + assert_expired_credentials_refresh_with_read_only_codex_home(OAuthCredentialsStoreMode::File) + .await +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn expired_auto_credentials_refresh_with_read_only_codex_home() -> anyhow::Result<()> { + assert_expired_credentials_refresh_with_read_only_codex_home(OAuthCredentialsStoreMode::Auto) + .await +} + +#[cfg(unix)] +async fn assert_expired_credentials_refresh_with_read_only_codex_home( + store_mode: OAuthCredentialsStoreMode, +) -> anyhow::Result<()> { + use std::os::unix::fs::PermissionsExt; + + 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(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 2).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?; + fs::set_permissions(codex_home.path(), fs::Permissions::from_mode(0o500))?; + + let result_path = control_dir.path().join("client.result"); + let child = spawn_read_only_refresh_child(&codex_home, &server_url, &result_path, store_mode); + let child_result = match child { + Ok(child) => wait_for_children(vec![child]).await, + Err(error) => Err(error), + }; + fs::set_permissions(codex_home.path(), fs::Permissions::from_mode(0o700))?; + child_result?; + + assert_eq!(fs::read_to_string(result_path)?, "success"); + assert!(!legacy_file_oauth_refresh_lock_path(&codex_home, &server_url)?.exists()); + assert!(file_oauth_refresh_lock_path(&codex_home, &server_url)?.exists()); + + let requests = server + .received_requests() + .await + .context("wiremock request recording disabled")?; + let authorization_headers = requests + .iter() + .filter(|request| request.method.as_str() == "POST" && request.url.path() == "/mcp") + .map(|request| -> anyhow::Result { + let authorization = request + .headers + .get("authorization") + .context("missing authorization header")?; + Ok(authorization + .to_str() + .context("authorization header is not ASCII")? + .to_string()) + }) + .collect::>>()?; + assert_eq!( + authorization_headers, + vec![format!("Bearer {NEW_ACCESS_TOKEN}"); 2] + ); + + if store_mode == OAuthCredentialsStoreMode::File { + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + } + + server.verify().await; + 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<()> { + let server = MockServer::start().await; + mount_oauth_metadata(&server, /*expected_requests*/ 2).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": NEW_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": NEW_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + mount_mcp_server(&server, /*expected_requests*/ 4).await; + + let codex_home = TempDir::new()?; + let control_dir = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let future_expiry = unix_millis().saturating_add(3_600_000); + seed_tokens(&codex_home, &server_url, future_expiry).await?; + + let stale_ready = control_dir.path().join("stale.ready"); + let stale_go = control_dir.path().join("stale.go"); + let stale_result = control_dir.path().join("stale.result"); + let stale_child = spawn_client_child( + &codex_home, + &server_url, + &stale_ready, + &stale_go, + &stale_result, + )?; + wait_for_paths(std::slice::from_ref(&stale_ready)).await?; + + // The held process has already copied the old credentials into RMCP memory. + // Make the durable copy expired so a second process rotates it. + seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?; + let rotator_ready = control_dir.path().join("rotator.ready"); + let rotator_go = control_dir.path().join("rotator.go"); + let rotator_result = control_dir.path().join("rotator.result"); + let rotator_child = spawn_client_child( + &codex_home, + &server_url, + &rotator_ready, + &rotator_go, + &rotator_result, + )?; + wait_for_paths(std::slice::from_ref(&rotator_ready)).await?; + fs::write(&rotator_go, "go")?; + wait_for_children(vec![rotator_child]).await?; + assert_eq!(fs::read_to_string(&rotator_result)?, "success"); + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + + fs::write(&stale_go, "go")?; + wait_for_children(vec![stale_child]).await?; + assert_eq!(fs::read_to_string(&stale_result)?, "success"); + + let requests = server + .received_requests() + .await + .context("wiremock request recording disabled")?; + let mut 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::>(); + authorization_headers.sort(); + assert_eq!( + authorization_headers, + vec![ + format!("Bearer {NEW_ACCESS_TOKEN}"), + format!("Bearer {NEW_ACCESS_TOKEN}"), + format!("Bearer {OLD_ACCESS_TOKEN}"), + format!("Bearer {OLD_ACCESS_TOKEN}"), + ] + ); + + // initialize() calls Codex's persistor. Because RMCP's stale credentials + // did not change locally, Codex suppresses the write and preserves the + // newer durable credentials. + assert_eq!( + persisted_token_snapshot(&codex_home)?, + PersistedTokenSnapshot { + access_token: NEW_ACCESS_TOKEN.to_string(), + refresh_token: Some(NEW_REFRESH_TOKEN.to_string()), + } + ); + + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_race_seed_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let access_token = std::env::var(ACCESS_TOKEN_ENV)?; + let refresh_token = std::env::var(REFRESH_TOKEN_ENV)?; + let expires_at = std::env::var(EXPIRES_AT_ENV)?.parse()?; + let mut response = OAuthTokenResponse::new( + AccessToken::new(access_token), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(refresh_token))); + response.set_expires_in(Some(&Duration::from_secs(3600))); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url, + client_id: CLIENT_ID.to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(expires_at), + }, + OAuthCredentialsStoreMode::File, + )?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_race_client_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let ready_path = PathBuf::from(std::env::var(READY_PATH_ENV)?); + let go_path = PathBuf::from(std::env::var(GO_PATH_ENV)?); + let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?); + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + fs::write(ready_path, "ready")?; + wait_for_paths(std::slice::from_ref(&go_path)).await?; + if let Ok(attempt_path) = std::env::var(ATTEMPT_PATH_ENV) { + fs::write(attempt_path, "attempt")?; + } + let outcome = match initialize_client(&client).await { + Ok(()) => "success".to_string(), + Err(error) => format!("error:{error:#}"), + }; + fs::write(result_path, outcome)?; + if let Ok(release_path) = std::env::var(RELEASE_PATH_ENV) { + wait_for_paths(&[PathBuf::from(release_path)]).await?; + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_operation_timeout_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let ready_path = PathBuf::from(std::env::var(READY_PATH_ENV)?); + let go_path = PathBuf::from(std::env::var(GO_PATH_ENV)?); + let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?); + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + + fs::write(ready_path, "ready")?; + wait_for_paths(std::slice::from_ref(&go_path)).await?; + let outcome = match client + .list_tools(/*params*/ None, Some(OPERATION_TIMEOUT)) + .await + { + Ok(_) => "success".to_string(), + Err(error) => format!("error:{error:#}"), + }; + fs::write(result_path, outcome)?; + if let Ok(release_path) = std::env::var(RELEASE_PATH_ENV) { + wait_for_paths(&[PathBuf::from(release_path)]).await?; + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_concurrent_waiter_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let ready_path = PathBuf::from(std::env::var(READY_PATH_ENV)?); + let go_path = PathBuf::from(std::env::var(GO_PATH_ENV)?); + let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?); + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + + fs::write(ready_path, "ready")?; + wait_for_paths(std::slice::from_ref(&go_path)).await?; + let outcomes = timeout(Duration::from_secs(3), async { + tokio::join!( + client.list_tools(/*params*/ None, /*timeout*/ None), + client.list_tools(/*params*/ None, /*timeout*/ None), + ) + }) + .await + .map(|(first, second)| { + [first, second].map(|outcome| match outcome { + Ok(_) => "success".to_string(), + Err(error) => format!("error:{error:#}"), + }) + }) + .unwrap_or_else(|_| { + [ + "error:concurrent refresh waiters timed out".to_string(), + "error:concurrent refresh waiters timed out".to_string(), + ] + }); + fs::write(result_path, outcomes.join("\n"))?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by the OAuth refresh race integration tests"] +async fn oauth_refresh_read_only_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let result_path = PathBuf::from(std::env::var(RESULT_PATH_ENV)?); + let store_mode = match std::env::var(STORE_MODE_ENV)?.as_str() { + "file" => OAuthCredentialsStoreMode::File, + "auto" => { + keyring::set_default_credential_builder(keyring::mock::default_credential_builder()); + OAuthCredentialsStoreMode::Auto + } + value => anyhow::bail!("unsupported OAuth credential store mode: {value}"), + }; + let outcome: anyhow::Result<()> = async { + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + store_mode, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await + } + .await; + let outcome = match outcome { + Ok(()) => "success".to_string(), + Err(error) => format!("error:{error:#}"), + }; + fs::write(result_path, outcome)?; + Ok(()) +} + +async fn mount_oauth_metadata(server: &MockServer, expected_requests: u64) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .expect(expected_requests) + .mount(server) + .await; +} + +async fn mount_mcp_server(server: &MockServer, expected_requests: impl Into) { + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let Ok(body) = request.body_json::() else { + return ResponseTemplate::new(400).set_body_string("invalid JSON-RPC request"); + }; + match body.get("method").and_then(Value::as_str) { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-refresh-race-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(expected_requests) + .mount(server) + .await; +} + +async fn seed_tokens( + codex_home: &TempDir, + server_url: &str, + expires_at: u64, +) -> anyhow::Result<()> { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_race_seed_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(ACCESS_TOKEN_ENV, OLD_ACCESS_TOKEN) + .env(REFRESH_TOKEN_ENV, OLD_REFRESH_TOKEN) + .env(EXPIRES_AT_ENV, expires_at.to_string()) + .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(()) +} + +fn spawn_client_child( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_race_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +fn spawn_client_child_held_open( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, + release_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_race_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .env(RELEASE_PATH_ENV, release_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +#[cfg(unix)] +fn spawn_client_child_with_attempt( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + attempt_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_race_client_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(ATTEMPT_PATH_ENV, attempt_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +fn spawn_operation_timeout_child( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_operation_timeout_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +fn spawn_operation_timeout_child_held_open( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, + release_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_operation_timeout_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .env(RELEASE_PATH_ENV, release_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +#[cfg(unix)] +fn spawn_concurrent_waiter_child( + codex_home: &TempDir, + server_url: &str, + ready_path: &Path, + go_path: &Path, + result_path: &Path, +) -> anyhow::Result { + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_concurrent_waiter_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(READY_PATH_ENV, ready_path) + .env(GO_PATH_ENV, go_path) + .env(RESULT_PATH_ENV, result_path) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +#[cfg(unix)] +fn spawn_read_only_refresh_child( + codex_home: &TempDir, + server_url: &str, + result_path: &Path, + store_mode: OAuthCredentialsStoreMode, +) -> anyhow::Result { + let store_mode = match store_mode { + OAuthCredentialsStoreMode::File => "file", + OAuthCredentialsStoreMode::Auto => "auto", + OAuthCredentialsStoreMode::Keyring => { + anyhow::bail!("read-only refresh child does not support keyring-only mode") + } + }; + let mut command = Command::new(std::env::current_exe()?); + command + .args([ + "oauth_refresh_read_only_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, server_url) + .env(RESULT_PATH_ENV, result_path) + .env(STORE_MODE_ENV, store_mode) + .kill_on_drop(true); + Ok(command.spawn()?) +} + +async fn wait_for_children(children: Vec) -> anyhow::Result<()> { + for mut child in children { + 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(()) +} + +async fn wait_for_paths(paths: &[PathBuf]) -> anyhow::Result<()> { + timeout(Duration::from_secs(10), async { + loop { + if paths.iter().all(|path| path.exists()) { + return; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for OAuth race process barrier") +} + +async fn wait_for_persisted_token_snapshot( + codex_home: &TempDir, + expected: PersistedTokenSnapshot, +) -> anyhow::Result<()> { + timeout(Duration::from_secs(10), async { + loop { + if persisted_token_snapshot(codex_home).ok().as_ref() == Some(&expected) { + return; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("timed out waiting for rotated OAuth credentials to persist") +} + +fn read_outcomes(paths: &[PathBuf]) -> anyhow::Result> { + paths + .iter() + .map(fs::read_to_string) + .collect::, _>>() + .map_err(Into::into) +} + +fn request_bodies(requests: &[Request], request_path: &str) -> Vec { + requests + .iter() + .filter(|request| request.method.as_str() == "POST" && request.url.path() == request_path) + .map(|request| String::from_utf8_lossy(&request.body).to_string()) + .collect() +} + +fn file_oauth_refresh_lock_path(codex_home: &TempDir, server_url: &str) -> anyhow::Result { + let (_, user_namespace) = fallback_lock_path_and_user_namespace(codex_home)?; + let lock_id = oauth_refresh_lock_id(server_url)?; + Ok(shared_lock_root()?.join(format!( + "{FILE_OAUTH_REFRESH_LOCK_PREFIX}-{user_namespace}-{lock_id}.lock" + ))) +} + +#[cfg(unix)] +fn legacy_file_oauth_refresh_lock_path( + codex_home: &TempDir, + server_url: &str, +) -> anyhow::Result { + let lock_id = oauth_refresh_lock_id(server_url)?; + Ok(codex_home + .path() + .join(format!(".credentials.{lock_id}.refresh.lock"))) +} + +fn oauth_refresh_lock_id(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)?); + sha_256_prefix(&Value::String(format!("{KEYRING_SERVICE}:{account}"))) +} + +fn fallback_lock_path_and_user_namespace( + codex_home: &TempDir, +) -> anyhow::Result<(PathBuf, String)> { + let canonical_home = codex_home.path().canonicalize()?; + let canonical_home = canonical_home.as_os_str().to_string_lossy().into_owned(); + #[cfg(windows)] + let canonical_home = normalize_windows_device_path(canonical_home); + let codex_home_id = sha_256_bytes_prefix(canonical_home.as_bytes()); + let prefix = format!("{FALLBACK_LOCK_PREFIX}-"); + let suffix = format!("-{codex_home_id}.lock"); + for entry in fs::read_dir(shared_lock_root()?)? { + let entry = entry?; + let file_name = entry.file_name().to_string_lossy().into_owned(); + if let Some(remainder) = file_name.strip_prefix(&prefix) + && let Some(user_namespace) = remainder.strip_suffix(&suffix) + { + return Ok((entry.path(), user_namespace.to_string())); + } + } + anyhow::bail!( + "missing OAuth fallback lock for CODEX_HOME {}", + codex_home.path().display() + ) +} + +#[cfg(unix)] +fn shared_lock_root() -> anyhow::Result { + Ok(PathBuf::from("/tmp")) +} + +#[cfg(windows)] +fn shared_lock_root() -> anyhow::Result { + use std::ffi::OsString; + use std::io; + use std::os::windows::ffi::OsStringExt; + use windows_sys::Win32::System::SystemInformation::GetSystemWindowsDirectoryW; + + let mut buffer = vec![0_u16; 32_768]; + let buffer_len = + u32::try_from(buffer.len()).context("Windows path buffer length did not fit u32")?; + // SAFETY: buffer is writable for the length passed to GetSystemWindowsDirectoryW. + let length = unsafe { GetSystemWindowsDirectoryW(buffer.as_mut_ptr(), buffer_len) }; + if length == 0 { + return Err(io::Error::last_os_error()) + .context("failed to resolve the system Windows directory"); + } + let length = usize::try_from(length).context("Windows directory length did not fit usize")?; + anyhow::ensure!( + length < buffer.len(), + "system Windows directory exceeded the fixed path buffer" + ); + Ok(PathBuf::from(OsString::from_wide(&buffer[..length])).join("Temp")) +} + +#[cfg(windows)] +fn normalize_windows_device_path(path: String) -> String { + if let Some(unc) = path.strip_prefix(r"\\?\UNC\") { + return format!(r"\\{unc}"); + } + if let Some(unc) = path.strip_prefix(r"\\.\UNC\") { + return format!(r"\\{unc}"); + } + if let Some(path) = path + .strip_prefix(r"\\?\") + .or_else(|| path.strip_prefix(r"\\.\")) + { + return path.to_string(); + } + path +} + +#[cfg(not(any(unix, windows)))] +fn shared_lock_root() -> anyhow::Result { + anyhow::bail!("OAuth refresh race tests do not support this platform") +} + +fn sha_256_prefix(value: &Value) -> anyhow::Result { + let serialized = serde_json::to_string(value)?; + Ok(sha_256_bytes_prefix(serialized.as_bytes())) +} + +fn sha_256_bytes_prefix(bytes: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(bytes); + let digest = hasher.finalize(); + format!("{digest:x}")[..16].to_string() +} + +#[derive(Debug, Deserialize)] +struct PersistedTokenEntry { + access_token: String, + refresh_token: Option, +} + +#[derive(Debug, PartialEq)] +struct PersistedTokenSnapshot { + access_token: String, + refresh_token: Option, +} + +fn persisted_token_snapshot(codex_home: &TempDir) -> anyhow::Result { + let contents = fs::read_to_string(codex_home.path().join(".credentials.json"))?; + let entries: BTreeMap = serde_json::from_str(&contents)?; + let entry = entries + .values() + .next() + .context("missing persisted OAuth token entry")?; + anyhow::ensure!(entries.len() == 1, "expected one persisted OAuth entry"); + Ok(PersistedTokenSnapshot { + access_token: entry.access_token.clone(), + refresh_token: entry.refresh_token.clone(), + }) +} + +fn unix_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +}