From ecc78e4cf5607ecf5f080d682eae0ecb650868ae Mon Sep 17 00:00:00 2001 From: Nick Steele Date: Thu, 1 Oct 2026 13:48:35 +0000 Subject: [PATCH] Add renewable EMA HTTP authentication and credential versioning (#49987) ## What changed - Add `EmaAuthenticatedHttpClient` to cache and renew resource bearers, serialize concurrent renewals, and include authentication time in the request deadline. Confine bearers to the configured MCP endpoint and reject redirects. - Return authentication-required errors for `401` responses without replaying MCP operations. Preserve a newer cached bearer when a delayed response rejects an older one. - Add revocable request authority that cancels pending requests and ends response streams. - Pin enterprise credentials to a persistent version advanced on save or deletion. Use short credential locks and invalidate rejected grants only when their record and version are still current, preserving newer logins. - Hide keyring account identifiers from enterprise credential errors and traces while preserving authentication failure classification. ## Testing Add tests for token expiry and renewal, endpoint confinement, redirects, request deadlines, revocation during renewal and streaming, delayed `401` responses, credential replacement and logout races, and credential error privacy. GitOrigin-RevId: ee3e5cf7ecbba3d09e1ab29ed4fbb2a35eacc60d --- codex-rs/exec-server/src/client.rs | 2 + .../src/client/connection_failure.rs | 1 + .../src/client/http_response_body_stream.rs | 22 + codex-rs/rmcp-client/Cargo.toml | 2 +- codex-rs/rmcp-client/src/ema_auth.rs | 12 +- codex-rs/rmcp-client/src/ema_auth_tests.rs | 95 ++++ codex-rs/rmcp-client/src/ema_exchange.rs | 4 + .../rmcp-client/src/ema_exchange_tests.rs | 4 +- codex-rs/rmcp-client/src/ema_http_client.rs | 308 ++++++++++++ .../src/ema_http_client_regression_tests.rs | 387 ++++++++++++++++ .../rmcp-client/src/ema_http_client_tests.rs | 437 ++++++++++++++++++ codex-rs/rmcp-client/src/ema_identity.rs | 32 +- .../rmcp-client/src/ema_identity_tests.rs | 32 +- .../rmcp-client/src/enterprise_oauth_login.rs | 10 +- .../src/enterprise_oauth_login_tests.rs | 65 +++ .../src/enterprise_oauth_logout_tests.rs | 142 +++++- codex-rs/rmcp-client/src/lib.rs | 6 + codex-rs/rmcp-client/src/oauth.rs | 10 +- .../rmcp-client/src/oauth/ema_identity.rs | 205 +++++++- .../src/oauth/ema_identity_tests.rs | 330 ++++++++++++- .../src/oauth/enterprise_generation.rs | 37 +- .../rmcp-client/src/oauth/refresh_lock.rs | 1 + codex-rs/rmcp-client/src/startup_error.rs | 26 ++ 23 files changed, 2109 insertions(+), 61 deletions(-) create mode 100644 codex-rs/rmcp-client/src/ema_http_client.rs create mode 100644 codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs create mode 100644 codex-rs/rmcp-client/src/ema_http_client_tests.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index c372edadc0..0369bb7185 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -688,6 +688,8 @@ pub enum ExecServerError { Json(#[from] serde_json::Error), #[error("HTTP request failed: {0}")] HttpRequest(String), + #[error("authentication required: {0}")] + AuthenticationRequired(String), #[error("exec-server protocol error: {0}")] Protocol(String), #[error( diff --git a/codex-rs/exec-server/src/client/connection_failure.rs b/codex-rs/exec-server/src/client/connection_failure.rs index fb26363337..5265d610c2 100644 --- a/codex-rs/exec-server/src/client/connection_failure.rs +++ b/codex-rs/exec-server/src/client/connection_failure.rs @@ -43,6 +43,7 @@ impl ExecServerError { | Self::Json(_) | Self::HttpRequest(_) | Self::Protocol(_) + | Self::AuthenticationRequired(_) | Self::ProvisioningModeConflict { .. } | Self::Server { .. } | Self::EnvironmentRegistryHttp { .. } diff --git a/codex-rs/exec-server/src/client/http_response_body_stream.rs b/codex-rs/exec-server/src/client/http_response_body_stream.rs index be52e5cb6a..36f22176e1 100644 --- a/codex-rs/exec-server/src/client/http_response_body_stream.rs +++ b/codex-rs/exec-server/src/client/http_response_body_stream.rs @@ -19,6 +19,7 @@ use tokio::runtime::Handle; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TrySendError; +use tokio_util::sync::CancellationToken; use tracing::debug; use crate::client::ExecServerError; @@ -75,6 +76,7 @@ enum HttpResponseBodyStreamInner { /// id until EOF or a terminal error. pub struct HttpResponseBodyStream { inner: HttpResponseBodyStreamInner, + cancellation: Option, } impl HttpResponseBodyStream { @@ -93,6 +95,7 @@ impl HttpResponseBodyStream { inner: HttpResponseBodyStreamInner::Local { body: Box::pin(body), }, + cancellation: None, } } @@ -101,6 +104,7 @@ impl HttpResponseBodyStream { inner: HttpResponseBodyStreamInner::Local { body: Box::pin(response.bytes_stream()), }, + cancellation: None, } } @@ -118,14 +122,32 @@ impl HttpResponseBodyStream { pending_eof: false, closed: false, }, + cancellation: None, } } + /// Ends this stream when its owning request authority is revoked. + pub fn cancel_on(mut self, cancellation: CancellationToken) -> Self { + self.cancellation = Some(cancellation); + self + } + /// Receives the next response-body chunk. /// /// Returns `Ok(None)` at EOF and converts sequence gaps or stream-side /// stream errors into protocol errors. pub async fn recv(&mut self) -> Result>, ExecServerError> { + let Some(cancellation) = self.cancellation.clone() else { + return self.recv_inner().await; + }; + tokio::select! { + biased; + () = cancellation.cancelled() => Ok(None), + result = self.recv_inner() => result, + } + } + + async fn recv_inner(&mut self) -> Result>, ExecServerError> { match &mut self.inner { HttpResponseBodyStreamInner::Local { body } => match body.next().await { Some(chunk) => match chunk { diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index d44f91328d..81758721ca 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -49,6 +49,7 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha2 = { workspace = true } sse-stream = "0.2.5" +tempfile = { workspace = true } thiserror = { workspace = true } tiny_http = { workspace = true } tokio = { workspace = true, features = [ @@ -74,7 +75,6 @@ codex-http-client = { workspace = true } codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serial_test = { workspace = true } -tempfile = { workspace = true } tracing-subscriber = { workspace = true } tracing-test = { workspace = true, features = ["no-env-filter"] } wiremock = { workspace = true } diff --git a/codex-rs/rmcp-client/src/ema_auth.rs b/codex-rs/rmcp-client/src/ema_auth.rs index 5d5f5de64b..02f1697e85 100644 --- a/codex-rs/rmcp-client/src/ema_auth.rs +++ b/codex-rs/rmcp-client/src/ema_auth.rs @@ -14,6 +14,7 @@ use serde::Deserialize; use serde_json::Value; use url::Url; +use crate::EmaCredentialLease; use crate::ema_auth_policy::advertised_capability; use crate::ema_auth_policy::validate_ema_auth_resource; use crate::ema_auth_policy::validate_ema_oauth_endpoint; @@ -67,7 +68,7 @@ impl<'a> From<&EmaAuthTokenExchangeRequest<'a>> for EmaDiscoveryRequest<'a> { /// Discover the resource's authorization server, then perform the ID-JAG flow. pub async fn exchange_ema_auth_token( request: EmaAuthTokenExchangeRequest<'_>, -) -> Result { +) -> Result<(EmaAccessToken, EmaCredentialLease)> { validate_ema_auth_resource(request.server_url, request.resource)?; let client_id = request .mcp_client_id @@ -86,9 +87,9 @@ pub async fn exchange_ema_auth_token( let EmaIdpIdentity { token_endpoint, refresh_token, - credential_lock, + credentials, } = request.idp_identity.await?; - let result = exchange_id_jag(EmaIdJagExchangeRequest { + let token = exchange_id_jag(EmaIdJagExchangeRequest { resource: &resource, scopes: request.scopes, mcp_client_id: client_id, @@ -101,9 +102,8 @@ pub async fn exchange_ema_auth_token( idp_http_client: request.idp_http_client, resource_http_client: request.resource_http_client, }) - .await; - drop(credential_lock); - result + .await?; + Ok((token, credentials)) } async fn discover_ema_authorization_metadata( diff --git a/codex-rs/rmcp-client/src/ema_auth_tests.rs b/codex-rs/rmcp-client/src/ema_auth_tests.rs index d327ae34c2..cf263a3352 100644 --- a/codex-rs/rmcp-client/src/ema_auth_tests.rs +++ b/codex-rs/rmcp-client/src/ema_auth_tests.rs @@ -3,6 +3,9 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_config::types::AuthKeyringBackendKind; use codex_exec_server::RouteAwareHttpClient; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; @@ -15,6 +18,13 @@ use wiremock::matchers::method; use wiremock::matchers::path; use super::*; +use crate::ema_auth_policy::EmaAuthFailure; +use crate::ema_auth_policy::EmaInvalidGrantSource; +use crate::oauth::RefreshCredentialLock; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthCredentialSnapshot; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::test_support::TempCodexHome; fn request<'a>(resource: &'a str, issuer: Option<&'a str>) -> EmaAuthTokenExchangeRequest<'a> { let client: Arc = Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( @@ -375,3 +385,88 @@ fn discovery_rejects_private_or_ambiguous_issuer_destinations() -> Result<()> { )?; Ok(()) } + +#[tokio::test] +async fn rejected_grants_preserve_their_source_and_release_the_credential_lease() -> Result<()> { + let _home = TempCodexHome::new(); + for failure_source in [ + EmaInvalidGrantSource::EnterpriseIdentity, + EmaInvalidGrantSource::ResourceAuthorization, + ] { + let server = MockServer::start().await; + let resource = format!("{}/mcp", server.uri()); + let issuer = format!("{}/as", server.uri()); + let idp = format!("{}/idp", server.uri()); + let assertion = format!( + "{}.{}.{}", + URL_SAFE_NO_PAD.encode(br#"{"alg":"ES256","typ":"oauth-id-jag+jwt"}"#), + URL_SAFE_NO_PAD.encode(serde_json::to_vec(&json!({"iss": idp, "aud": issuer, + "client_id": "resource-client", "sub": "user", "jti": "one", "iat": 0, + "exp": u64::MAX, "resource": resource}))?), + URL_SAFE_NO_PAD.encode(b"signature") + ); + for (endpoint, status, response) in [ + ( + "/mcp", + 200, + json!({"resource": resource, "authorization_servers": [issuer]}), + ), + ( + "/.well-known/oauth-authorization-server/as", + 200, + json!({"issuer": issuer, "token_endpoint": format!("{issuer}/token"), + "grant_types_supported": [JWT_BEARER_GRANT_TYPE], "token_endpoint_auth_methods_supported": ["none"]}), + ), + ( + "/idp/token", + if failure_source == EmaInvalidGrantSource::EnterpriseIdentity { + 400 + } else { + 200 + }, + if failure_source == EmaInvalidGrantSource::EnterpriseIdentity { + json!({"error": "invalid_grant"}) + } else { + json!({"access_token": assertion, "token_type": "N_A", "issued_token_type": "urn:ietf:params:oauth:token-type:id-jag"}) + }, + ), + ("/as/token", 400, json!({"error": "invalid_grant"})), + ] { + Mock::given(path(endpoint)) + .respond_with(ResponseTemplate::new(status).set_body_json(response)) + .mount(&server) + .await; + } + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + let tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": "ema-idp:revoked-test", "url": idp, "issuer": idp, "client_id": "enterprise-client", + "token_response": {"access_token": "unused", "token_type": "Bearer", "refresh_token": "old-grant"}, + }))?; + let credentials = StoredOAuthCredentialSnapshot::new(tokens.clone(), store); + let mut exchange = request(&resource, Some(&issuer)); + exchange.idp_issuer = &idp; + exchange.idp_identity = Box::pin(async { + Ok(EmaIdpIdentity { + token_endpoint: format!("{idp}/token"), + refresh_token: "old-grant".to_string(), + credentials: crate::EmaCredentialLease { credentials }, + }) + }); + let error = exchange_ema_auth_token(exchange) + .await + .expect_err("invalid grant"); + assert_eq!( + error.downcast_ref::(), + Some(&EmaAuthFailure::InvalidGrant { + grant_source: failure_source + }) + ); + let _exclusive = tokio::time::timeout( + std::time::Duration::from_secs(1), + RefreshCredentialLock::acquire_for_server(&tokens.server_name, &idp), + ) + .await + .expect("failed exchange must not retain the credential lock")?; + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/ema_exchange.rs b/codex-rs/rmcp-client/src/ema_exchange.rs index afe5252b11..f27cf17691 100644 --- a/codex-rs/rmcp-client/src/ema_exchange.rs +++ b/codex-rs/rmcp-client/src/ema_exchange.rs @@ -10,6 +10,7 @@ use rmcp::transport::auth::enterprise::EmaAuthorizationServer; use rmcp::transport::auth::enterprise::EmaError; use rmcp::transport::auth::enterprise::EmaExchangeRequest; use rmcp::transport::auth::enterprise::EmaExchangeStage; +use tokio::time::Instant; use crate::ema_auth_policy::EmaAuthFailure; use crate::ema_auth_policy::EmaInvalidGrantSource; @@ -26,6 +27,8 @@ pub(crate) const JWT_BEARER_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type pub struct EmaAccessToken { pub access_token: String, pub expires_in: Option, + /// Exchange completion time, before any asynchronous credential validation. + pub received_at: Instant, } impl std::fmt::Debug for EmaAccessToken { @@ -118,6 +121,7 @@ pub(crate) async fn exchange_id_jag( Ok(EmaAccessToken { access_token: token.access_token.secret().to_owned(), expires_in: token.expires_in, + received_at: Instant::now(), }) } diff --git a/codex-rs/rmcp-client/src/ema_exchange_tests.rs b/codex-rs/rmcp-client/src/ema_exchange_tests.rs index 1e7e1a3862..4c21b76d25 100644 --- a/codex-rs/rmcp-client/src/ema_exchange_tests.rs +++ b/codex-rs/rmcp-client/src/ema_exchange_tests.rs @@ -148,11 +148,13 @@ async fn public_client_round_trip_preserves_signed_narrowing() -> Result<()> { ); continue; } + let token = result?; assert_eq!( - result?, + token, EmaAccessToken { access_token: "resource-token".to_string(), expires_in: Some(Duration::from_secs(300)), + received_at: token.received_at, } ); let requests = server.received_requests().await.expect("requests"); diff --git a/codex-rs/rmcp-client/src/ema_http_client.rs b/codex-rs/rmcp-client/src/ema_http_client.rs new file mode 100644 index 0000000000..df778a1dce --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_http_client.rs @@ -0,0 +1,308 @@ +//! Renewable, endpoint-confined EMA bearers. An unauthorized MCP operation is never replayed. +//! Authentication shares the request deadline, and a rejected bearer cannot evict a newer one. + +use std::future::Future; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::PoisonError; +use std::time::Duration; + +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use futures::future::BoxFuture; +use tokio::sync::OwnedRwLockReadGuard; +use tokio::sync::RwLock; +use tokio::sync::Semaphore; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use url::Url; + +use crate::EmaAccessToken; +use crate::is_authentication_required_error; +use crate::validate_ema_auth_resource; + +/// Supplies resource tokens for one immutable endpoint and pinned enterprise login. +/// Credential ownership is checked when renewing a token, not for cached bearer requests. +pub trait EmaTokenSource: Send + Sync { + fn endpoint(&self) -> &Url; + /// Acquires and validates credentials for renewal. Implementations recheck ownership + /// and a saved grant's generation after exchange before publishing the bearer. + fn access_token(&self) -> impl Future> + Send; +} + +/// Revocable authority held while one EMA connection awaits HTTP responses. +#[derive(Clone, Default)] +pub struct EmaRequestAuthority { + revoked: CancellationToken, + responses: Arc>, +} + +impl EmaRequestAuthority { + pub fn revoke(&self) { + self.revoked.cancel(); + } + + pub fn is_revoked(&self) -> bool { + self.revoked.is_cancelled() + } + + /// Wait until response futures admitted before revocation complete or are cancelled. + pub async fn quiesce(&self) { + let _responses = self.responses.clone().write_owned().await; + } + + async fn acquire(&self) -> Result, ExecServerError> { + let response = self.responses.clone().read_owned().await; + if self.is_revoked() { + return Err(revoked_authority_error()); + } + Ok(response) + } + + async fn cancelled(&self) { + self.revoked.cancelled().await; + } +} + +/// Adds a renewable resource bearer to the existing runtime-selected HTTP capability. +pub struct EmaAuthenticatedHttpClient { + inner: Arc, + source: S, + authority: EmaRequestAuthority, + token: Mutex>, + renewal: Semaphore, +} + +struct CachedToken { + value: String, + refresh_at: Instant, +} + +impl CachedToken { + fn new(token: EmaAccessToken) -> Self { + // Bound reuse even when the authorization server omits an expiry. + let lifetime = token + .expires_in + .unwrap_or(Duration::from_secs(300)) + .min(Duration::from_secs(300)); + let skew = Duration::from_secs(30).min(lifetime / 2); + Self { + value: token.access_token, + refresh_at: token.received_at + lifetime.saturating_sub(skew), + } + } +} + +impl EmaAuthenticatedHttpClient { + pub fn new( + inner: Arc, + source: S, + access_token: EmaAccessToken, + authority: EmaRequestAuthority, + ) -> Self { + Self { + inner, + source, + authority, + token: Mutex::new(Some(CachedToken::new(access_token))), + renewal: Semaphore::new(/*permits*/ 1), + } + } + + fn cached_token(&self) -> MutexGuard<'_, Option> { + self.token.lock().unwrap_or_else(PoisonError::into_inner) + } + + async fn authenticate( + &self, + params: &mut HttpRequestParams, + ) -> Result { + let expected = self.source.endpoint(); + let target = Url::parse(¶ms.url).map_err(|_| { + ExecServerError::HttpRequest("invalid enterprise MCP destination".to_string()) + })?; + if expected.origin() != target.origin() + || expected.path() != target.path() + || expected.query() != target.query() + || validate_ema_auth_resource(¶ms.url, /*resource*/ None).is_err() + { + return Err(ExecServerError::HttpRequest( + "refusing to send enterprise bearer outside the configured MCP endpoint" + .to_string(), + )); + } + let current_token = || { + self.cached_token() + .as_ref() + .filter(|token| Instant::now() < token.refresh_at) + .map(|token| token.value.clone()) + }; + let value = if let Some(value) = current_token() { + value + } else { + let _renewal = self.renewal.acquire().await.map_err(|_| { + ExecServerError::HttpRequest("enterprise token renewal gate closed".to_string()) + })?; + if let Some(value) = current_token() { + value + } else { + let renewed = self.source.access_token().await.map_err(auth_error)?; + let token = CachedToken::new(renewed); + if Instant::now() >= token.refresh_at { + return Err(ExecServerError::HttpRequest( + "enterprise MCP bearer expired during credential validation".to_string(), + )); + } + let value = token.value.clone(); + *self.cached_token() = Some(token); + value + } + }; + params.redirect_policy = HttpRedirectPolicy::Stop; + params + .headers + .retain(|header| !header.name.eq_ignore_ascii_case("authorization")); + params.headers.push(HttpHeader { + name: "authorization".to_string(), + value: format!("Bearer {value}"), + value_env_var: None, + }); + Ok(value) + } + + fn clear_rejected_token(&self, rejected: &str) { + let mut token = self.cached_token(); + if token.as_ref().is_some_and(|token| token.value == rejected) { + *token = None; + } + } + + fn validate_response( + &self, + response: &HttpRequestResponse, + token: &str, + ) -> Result<(), ExecServerError> { + if response.status == 401 { + self.clear_rejected_token(token); + return Err(ExecServerError::AuthenticationRequired( + "enterprise MCP bearer was rejected; reconnect to renew authorization".to_string(), + )); + } + // An outer redirect wrapper may still follow a response even though the + // inner HTTP capability received Stop. Do not let it replay this operation. + if (300..400).contains(&response.status) { + return Err(ExecServerError::HttpRequest( + "enterprise MCP redirects are not allowed".to_string(), + )); + } + Ok(()) + } +} + +fn auth_error(error: anyhow::Error) -> ExecServerError { + if is_authentication_required_error(&error) { + ExecServerError::AuthenticationRequired( + "enterprise MCP authentication is required".to_string(), + ) + } else { + ExecServerError::HttpRequest( + "enterprise MCP token acquisition or credential validation failed".to_string(), + ) + } +} + +impl HttpClient for EmaAuthenticatedHttpClient { + fn http_request( + &self, + mut params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let deadline = params + .timeout_ms + .map(|timeout_ms| Instant::now() + Duration::from_millis(timeout_ms)); + let request = async { + let _authority = self.authority.acquire().await?; + let token = self.authenticate(&mut params).await?; + if let Some(deadline) = deadline { + params.timeout_ms = Some(remaining_timeout(deadline)?); + } + let response = self.inner.http_request(params).await?; + self.validate_response(&response, &token)?; + Ok(response) + }; + tokio::select! { + biased; + () = self.authority.cancelled() => Err(revoked_authority_error()), + () = deadline_elapsed(deadline) => Err(request_timeout_error()), + result = request => result, + } + }) + } + + fn http_request_stream( + &self, + mut params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + Box::pin(async move { + let deadline = params + .timeout_ms + .map(|timeout_ms| Instant::now() + Duration::from_millis(timeout_ms)); + let request = async { + let _authority = self.authority.acquire().await?; + let token = self.authenticate(&mut params).await?; + if let Some(deadline) = deadline { + params.timeout_ms = Some(remaining_timeout(deadline)?); + } + let response = self.inner.http_request_stream(params).await?; + self.validate_response(&response.0, &token)?; + Ok(( + response.0, + response.1.cancel_on(self.authority.revoked.clone()), + )) + }; + tokio::select! { + biased; + () = self.authority.cancelled() => Err(revoked_authority_error()), + () = deadline_elapsed(deadline) => Err(request_timeout_error()), + result = request => result, + } + }) + } +} + +async fn deadline_elapsed(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } +} + +fn remaining_timeout(deadline: Instant) -> Result { + let remaining = deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(request_timeout_error)?; + Ok(u64::try_from(remaining.as_millis()) + .unwrap_or(u64::MAX) + .max(1)) +} + +fn request_timeout_error() -> ExecServerError { + ExecServerError::HttpRequest("enterprise MCP request timed out".to_string()) +} + +fn revoked_authority_error() -> ExecServerError { + ExecServerError::AuthenticationRequired( + "enterprise MCP authorization changed; reconnect before using this server".to_string(), + ) +} + +#[cfg(test)] +#[path = "ema_http_client_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs b/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs new file mode 100644 index 0000000000..9a9d51430c --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_http_client_regression_tests.rs @@ -0,0 +1,387 @@ +//! Authentication, deadlines, and concurrent responses share the same request invariants. + +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; +use tokio::sync::oneshot; + +use super::*; +use crate::ema_auth_policy::EmaInvalidGrantSource; + +type Client = EmaAuthenticatedHttpClient>; + +#[derive(Clone, Copy)] +enum RequestMode { + Buffered, + Streaming, +} + +impl RequestMode { + fn start( + self, + client: Arc, + mut params: HttpRequestParams, + ) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + params.stream_response = matches!(self, Self::Streaming); + match self { + Self::Buffered => client.http_request(params).await, + Self::Streaming => client + .http_request_stream(params) + .await + .map(|(response, _body)| response), + } + }) + } +} + +#[derive(Debug)] +struct PendingRequest { + params: HttpRequestParams, + response: oneshot::Sender, +} + +struct ControlledHttpClient(mpsc::UnboundedSender); + +impl HttpClient for ControlledHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let (response, receiver) = oneshot::channel(); + self.0 + .send(PendingRequest { params, response }) + .expect("capture HTTP request"); + Ok(HttpRequestResponse { + status: receiver.await.expect("provide HTTP response"), + headers: Vec::new(), + body: Vec::new().into(), + }) + }) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + Box::pin(async move { + Ok(( + self.http_request(params).await?, + HttpResponseBodyStream::from_chunks(Vec::new()), + )) + }) + } +} + +fn controlled_fixture() -> ( + Arc, + Arc, + HttpRequestParams, + mpsc::UnboundedReceiver, +) { + controlled_fixture_for_endpoint("https://mcp.example/mcp") +} + +fn controlled_fixture_for_endpoint( + endpoint: &str, +) -> ( + Arc, + Arc, + HttpRequestParams, + mpsc::UnboundedReceiver, +) { + let (mut client, source, params) = fixture(endpoint); + let (sender, receiver) = mpsc::unbounded_channel(); + client.inner = Arc::new(ControlledHttpClient(sender)); + (Arc::new(client), source, params, receiver) +} + +#[tokio::test] +async fn exact_endpoint_and_renewal_credentials_are_checked_before_network_use() +-> anyhow::Result<()> { + let endpoint = "https://mcp.example/mcp?tenant=one"; + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + let (client, source, params, mut requests) = controlled_fixture_for_endpoint(endpoint); + for target in [ + endpoint.replace("one", "two"), + endpoint.replace("/mcp", "/other"), + "https://other.example/mcp?tenant=one".to_string(), + format!("{endpoint}#fragment"), + ] { + let mut request = params.clone(); + request.url = target; + assert!(matches!( + mode.start(Arc::clone(&client), request).await?, + Err(ExecServerError::HttpRequest(message)) + if message == "refusing to send enterprise bearer outside the configured MCP endpoint" + )); + } + source.valid.store(false, Ordering::SeqCst); + *client.cached_token() = None; + assert!(matches!( + mode.start(Arc::clone(&client), params.clone()).await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + source.valid.store(true, Ordering::SeqCst); + source.invalidate_on_refresh.store(true, Ordering::SeqCst); + *client.cached_token() = None; + assert!(matches!( + mode.start(client, params).await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + assert!(requests.try_recv().is_err()); + } + Ok(()) +} + +#[tokio::test] +async fn outgoing_requests_replace_all_authorization_headers_and_stop_redirects() +-> anyhow::Result<()> { + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + let (client, _source, mut params, mut requests) = controlled_fixture(); + params + .headers + .extend(["authorization", "aUtHoRiZaTiOn"].map(|name| HttpHeader { + name: name.to_string(), + value: "another-stale-credential".to_string(), + value_env_var: None, + })); + let response = mode.start(client, params); + let request = requests.recv().await.expect("authenticated request"); + assert_eq!(request.params.redirect_policy, HttpRedirectPolicy::Stop); + assert_eq!( + request.params.headers, + vec![HttpHeader { + name: "authorization".to_string(), + value: "Bearer initial".to_string(), + value_env_var: None, + }], + "replace every case variant with exactly one bearer" + ); + request.response.send(/*t*/ 200).expect("complete request"); + assert_eq!(response.await??.status, 200); + } + Ok(()) +} + +#[tokio::test] +async fn resource_bearers_never_follow_redirects() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(307).insert_header("location", "/mcp")) + .expect(2) + .mount(&server) + .await; + let (client, _source, params) = fixture(&format!("{}/mcp", server.uri())); + crate::oauth::test_support::warm_http_client(client.inner.as_ref()).await?; + let client = crate::http_client_redirect::SameOriginRedirectHttpClient::new(Arc::new(client)); + for result in [ + client.http_request(params.clone()).await.map(|_| ()), + client.http_request_stream(params).await.map(|_| ()), + ] { + assert!(matches!( + result, + Err(ExecServerError::HttpRequest(message)) + if message == "enterprise MCP redirects are not allowed" + )); + } + Ok(()) +} + +fn bearer(request: &PendingRequest) -> &str { + &request + .params + .headers + .iter() + .find(|header| header.name == "authorization") + .expect("bearer header") + .value +} + +#[tokio::test] +async fn credential_failures_are_classified_before_dispatch() -> anyhow::Result<()> { + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + for failure in [ + EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity, + }, + EmaAuthFailure::InsufficientUserAuthentication, + EmaAuthFailure::ReauthenticationRequired, + ] { + let (client, source, params, mut requests) = controlled_fixture(); + *client.cached_token() = None; + match failure { + EmaAuthFailure::ReauthenticationRequired => { + source.valid.store(/*val*/ false, Ordering::SeqCst); + } + failure => { + *source.renewal_error.lock().unwrap() = + Some(anyhow::Error::new(failure).context("token exchange failed")); + } + } + assert!(matches!( + mode.start(client, params).await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + assert!(requests.try_recv().is_err()); + } + for error in [ + anyhow::Error::new(EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::ResourceAuthorization, + }) + .context("resource authorization failed without invalidating the enterprise identity"), + anyhow::anyhow!("temporary network failure"), + ] { + let (client, source, params, mut requests) = controlled_fixture(); + *client.cached_token() = None; + *source.renewal_error.lock().unwrap() = Some(error); + assert!(matches!( + mode.start(client, params).await?, + Err(ExecServerError::HttpRequest(_)) + )); + assert!(requests.try_recv().is_err()); + } + } + Ok(()) +} + +#[tokio::test] +async fn delayed_unauthorized_response_preserves_the_renewed_bearer() -> anyhow::Result<()> { + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + let (client, source, params, mut requests) = controlled_fixture(); + let first = mode.start(Arc::clone(&client), params.clone()); + let first_request = requests.recv().await.expect("first request"); + let delayed = mode.start(Arc::clone(&client), params.clone()); + let delayed_request = requests.recv().await.expect("delayed request"); + assert_eq!( + (bearer(&first_request), bearer(&delayed_request)), + ("Bearer initial", "Bearer initial") + ); + + first_request.response.send(/*t*/ 401).expect("first 401"); + assert!(matches!( + first.await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + let renewed = mode.start(Arc::clone(&client), params.clone()); + let renewed_request = requests.recv().await.expect("renewed request"); + assert_eq!(bearer(&renewed_request), "Bearer renewed-1"); + renewed_request + .response + .send(/*t*/ 200) + .expect("renewed response"); + assert_eq!(renewed.await??.status, 200); + + delayed_request + .response + .send(/*t*/ 401) + .expect("delayed 401"); + assert!(matches!( + delayed.await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + let next = mode.start(Arc::clone(&client), params); + let next_request = requests.recv().await.expect("next request"); + assert_eq!(bearer(&next_request), "Bearer renewed-1"); + next_request + .response + .send(/*t*/ 200) + .expect("next response"); + assert_eq!(next.await??.status, 200); + assert_eq!(source.renewals.load(Ordering::SeqCst), 1); + assert!(requests.try_recv().is_err(), "no operation was replayed"); + } + Ok(()) +} + +#[tokio::test] +async fn authentication_waits_share_the_request_deadline() -> anyhow::Result<()> { + #[derive(Clone, Copy)] + enum Wait { + Authority, + Renewal, + Credentials, + Exchange, + } + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + for wait in [ + Wait::Authority, + Wait::Renewal, + Wait::Credentials, + Wait::Exchange, + ] { + let (client, source, mut params, mut requests) = controlled_fixture(); + let (gate, _entered, _release) = CredentialGate::new(); + let mut authority = None; + let mut renewal = None; + match wait { + Wait::Authority => authority = Some(client.authority.responses.write().await), + Wait::Renewal => { + *client.cached_token() = None; + renewal = Some(client.renewal.acquire().await?); + } + Wait::Credentials => { + *client.cached_token() = None; + *source.credential_gate.lock().unwrap() = Some(gate); + } + Wait::Exchange => { + *client.cached_token() = None; + *source.renewal_gate.lock().unwrap() = Some(gate); + } + } + params.timeout_ms = Some(20); + let response = tokio::time::timeout( + Duration::from_secs(/*secs*/ 1), + mode.start(Arc::clone(&client), params), + ) + .await??; + assert!( + matches!(response, Err(ExecServerError::HttpRequest(message)) + if message == "enterprise MCP request timed out") + ); + assert!(requests.try_recv().is_err()); + drop((authority, renewal)); + assert_eq!(source.active_leases.load(Ordering::SeqCst), 0); + } + let (client, source, mut params, mut requests) = controlled_fixture(); + params.timeout_ms = Some(0); + assert!(matches!( + mode.start(client, params).await?, + Err(ExecServerError::HttpRequest(_)) + )); + assert_eq!(source.renewals.load(Ordering::SeqCst), 0); + assert!(requests.try_recv().is_err()); + } + Ok(()) +} + +#[tokio::test] +async fn authentication_time_is_removed_from_the_inner_timeout() -> anyhow::Result<()> { + for mode in [RequestMode::Buffered, RequestMode::Streaming] { + for timeout in [Some(1000), None] { + let (client, source, mut params, mut requests) = controlled_fixture(); + let (gate, entered, release) = CredentialGate::new(); + *client.cached_token() = None; + *source.credential_gate.lock().unwrap() = Some(gate); + params.timeout_ms = timeout; + let response = mode.start(client, params); + entered.await?; + tokio::time::sleep(Duration::from_millis(/*millis*/ 30)).await; + release.send(()).expect("release credentials"); + let request = requests.recv().await.expect("authenticated request"); + match timeout { + Some(timeout) => assert!( + request + .params + .timeout_ms + .is_some_and(|remaining| remaining > 0 && remaining <= timeout - 30) + ), + None => assert_eq!(request.params.timeout_ms, None), + } + request.response.send(/*t*/ 200).expect("complete request"); + assert_eq!(response.await??.status, 200); + } + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/ema_http_client_tests.rs b/codex-rs/rmcp-client/src/ema_http_client_tests.rs new file mode 100644 index 0000000000..1ff154edf9 --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_http_client_tests.rs @@ -0,0 +1,437 @@ +//! Renewable sessions use cached bearers without credential reads and never replay operations. + +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use codex_exec_server::RouteAwareHttpClient; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; +use crate::ema_auth_policy::EmaAuthFailure; +use crate::oauth::RefreshCredentialLock; + +struct Source { + endpoint: Url, + renewals: AtomicUsize, + valid: AtomicBool, + invalidate_on_refresh: AtomicBool, + credential_gate: Mutex>, + renewal_gate: Mutex>, + renewal_error: Mutex>, + active_leases: AtomicUsize, + renewal_with_lease: AtomicBool, + lock_identity: Mutex>, +} + +struct CredentialGate { + entered: tokio::sync::oneshot::Sender<()>, + release: tokio::sync::oneshot::Receiver<()>, +} + +impl CredentialGate { + fn new() -> ( + Self, + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + ) { + let (entered, waiting) = tokio::sync::oneshot::channel(); + let (release, released) = tokio::sync::oneshot::channel(); + ( + Self { + entered, + release: released, + }, + waiting, + release, + ) + } + + async fn wait(self) { + let _ = self.entered.send(()); + let _ = self.release.await; + } +} + +struct CredentialLease { + source: Arc, + _lock: Option, +} + +impl Drop for CredentialLease { + fn drop(&mut self) { + self.source + .active_leases + .fetch_sub(/*val*/ 1, Ordering::SeqCst); + } +} + +impl EmaTokenSource for Arc { + fn endpoint(&self) -> &Url { + &self.endpoint + } + async fn access_token(&self) -> anyhow::Result { + if !self.valid.load(Ordering::SeqCst) { + return Err(anyhow::Error::new(EmaAuthFailure::ReauthenticationRequired) + .context("credentials were removed")); + } + let identity = self.lock_identity.lock().unwrap().clone(); + let lock = match identity { + Some((name, issuer)) => { + Some(RefreshCredentialLock::acquire_for_server(&name, &issuer).await?) + } + None => None, + }; + self.active_leases.fetch_add(/*val*/ 1, Ordering::SeqCst); + let _lease = CredentialLease { + source: Arc::clone(self), + _lock: lock, + }; + let gate = self.credential_gate.lock().unwrap().take(); + if let Some(gate) = gate { + gate.wait().await; + } + let renewal = self.renewals.fetch_add(/*val*/ 1, Ordering::SeqCst) + 1; + self.renewal_with_lease.fetch_or( + self.active_leases.load(Ordering::SeqCst) != 0, + Ordering::SeqCst, + ); + if let Some(error) = self.renewal_error.lock().unwrap().take() { + return Err(error); + } + let received_at = Instant::now(); + let gate = self.renewal_gate.lock().unwrap().take(); + if let Some(gate) = gate { + gate.wait().await; + } + tokio::task::yield_now().await; + if self.invalidate_on_refresh.load(Ordering::SeqCst) { + self.valid.store(false, Ordering::SeqCst); + return Err(anyhow::Error::new(EmaAuthFailure::ReauthenticationRequired)); + } + Ok(EmaAccessToken { + access_token: format!("renewed-{renewal}"), + expires_in: Some(Duration::from_secs(/*secs*/ 300)), + received_at, + }) + } +} + +fn fixture( + endpoint: &str, +) -> ( + EmaAuthenticatedHttpClient>, + Arc, + HttpRequestParams, +) { + let source = Arc::new(Source { + endpoint: Url::parse(endpoint).expect("endpoint"), + renewals: AtomicUsize::new(0), + valid: AtomicBool::new(true), + invalidate_on_refresh: AtomicBool::new(false), + credential_gate: Mutex::new(/*t*/ None), + renewal_gate: Mutex::new(/*t*/ None), + renewal_error: Mutex::new(/*t*/ None), + active_leases: AtomicUsize::new(/*v*/ 0), + renewal_with_lease: AtomicBool::new(/*v*/ false), + lock_identity: Mutex::new(None), + }); + let client = EmaAuthenticatedHttpClient::new( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + Arc::clone(&source), + EmaAccessToken { + access_token: "initial".to_string(), + expires_in: Some(Duration::from_secs(300)), + received_at: Instant::now(), + }, + EmaRequestAuthority::default(), + ); + let params = HttpRequestParams { + method: "POST".to_string(), + url: endpoint.to_string(), + headers: vec![HttpHeader { + name: "Authorization".to_string(), + value: "discard-me".to_string(), + value_env_var: None, + }], + body: Some(br#"{"method":"tools/call"}"#.to_vec().into()), + timeout_ms: Some(5000), + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "operation".to_string(), + stream_response: false, + }; + (client, source, params) +} + +#[tokio::test] +async fn cached_bearer_does_not_read_credentials_until_renewal() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let (client, source, params) = fixture(&format!("{}/mcp", server.uri())); + crate::oauth::test_support::warm_http_client(client.inner.as_ref()).await?; + assert_eq!(client.http_request(params.clone()).await?.status, 200); + source.valid.store(false, Ordering::SeqCst); + assert_eq!(client.http_request(params.clone()).await?.status, 200); + *client.cached_token() = None; + assert!(matches!( + client.http_request(params).await, + Err(ExecServerError::AuthenticationRequired(_)) + )); + assert_eq!(source.renewals.load(Ordering::SeqCst), 0); + let requests = server.received_requests().await.expect("recorded requests"); + let bearers = requests + .iter() + .map(|request| { + request.headers["authorization"] + .to_str() + .expect("bearer header") + .to_string() + }) + .collect::>(); + assert_eq!(bearers, ["Bearer initial", "Bearer initial"]); + Ok(()) +} + +#[tokio::test] +async fn revocation_during_renewal_prevents_dispatch() -> anyhow::Result<()> { + let server = MockServer::start().await; + let endpoint = format!("{}/mcp", server.uri()); + let (gate, validated_rx, release_tx) = CredentialGate::new(); + let (client, source, params) = fixture(&endpoint); + *client.cached_token() = None; + *source.credential_gate.lock().unwrap() = Some(gate); + let authority = client.authority.clone(); + let request = tokio::spawn(async move { client.http_request(params).await }); + validated_rx.await?; + authority.revoke(); + let _ = release_tx.send(()); + tokio::time::timeout(Duration::from_secs(1), authority.quiesce()).await?; + assert!(matches!( + request.await?, + Err(ExecServerError::AuthenticationRequired(_)) + )); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + Ok(()) +} + +#[tokio::test] +async fn revocation_ends_an_existing_response_stream() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_string("account-a-data")) + .mount(&server) + .await; + let (client, _source, params) = fixture(&format!("{}/mcp", server.uri())); + crate::oauth::test_support::warm_http_client(client.inner.as_ref()).await?; + let authority = client.authority.clone(); + let (_response, mut body) = client.http_request_stream(params).await?; + + authority.revoke(); + + assert_eq!(body.recv().await?, None); + Ok(()) +} + +#[tokio::test] +async fn revocation_interrupts_an_active_multichunk_response() -> anyhow::Result<()> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}/mcp", listener.local_addr()?); + let (chunks, receiver) = tokio::sync::mpsc::channel::>(2); + let receiver = Arc::new(Mutex::new(Some(receiver))); + let app = axum::Router::new().route( + "/mcp", + axum::routing::post(move || { + let receiver = receiver.lock().unwrap().take().expect("one request"); + async move { + let stream = futures::stream::unfold(receiver, |mut receiver| async { + receiver.recv().await.map(|chunk| (chunk, receiver)) + }); + ( + [("content-type", "text/event-stream")], + axum::body::Body::from_stream(stream), + ) + } + }), + ); + let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move { + axum::serve(listener, app).await + })); + let (client, _source, params) = fixture(&endpoint); + chunks + .send(Ok(bytes::Bytes::from_static(b"data: first\n\n"))) + .await?; + let (_, mut body) = client.http_request_stream(params).await?; + assert_eq!(body.recv().await?, Some(b"data: first\n\n".to_vec())); + let mut next = Box::pin(body.recv()); + assert!(futures::poll!(&mut next).is_pending()); + + client.authority.revoke(); + tokio::time::timeout(Duration::from_secs(1), client.authority.quiesce()).await?; + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), next).await??, + None + ); + // Even a later chunk cannot revive a stream whose authority was revoked. + let _ = chunks + .send(Ok(bytes::Bytes::from_static(b"data: second\n\n"))) + .await; + assert_eq!(body.recv().await?, None); + server.abort(); + let _ = server.await; + Ok(()) +} + +#[tokio::test] +async fn unauthorized_requests_are_not_replayed_and_later_requests_renew() -> anyhow::Result<()> { + for streaming in [false, true] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &wiremock::Request| { + ResponseTemplate::new(if request.headers["authorization"] == "Bearer initial" { + 401 + } else { + 200 + }) + }) + .mount(&server) + .await; + let (client, source, params) = fixture(&format!("{}/mcp", server.uri())); + crate::oauth::test_support::warm_http_client(client.inner.as_ref()).await?; + for (status, count) in [(401, 1), (200, 2), (200, 3)] { + let result = if streaming { + client + .http_request_stream(params.clone()) + .await + .map(|response| response.0.status) + } else { + client + .http_request(params.clone()) + .await + .map(|response| response.status) + }; + if status == 401 { + let error = result.expect_err("unauthorized operation fails without replay"); + assert!(is_authentication_required_error(&anyhow::Error::new(error))); + } else { + assert_eq!(result?, status); + } + assert_eq!( + server.received_requests().await.expect("requests").len(), + count + ); + } + assert_eq!(source.renewals.load(Ordering::SeqCst), 1); + // Concurrent callers of an expiring token share one renewal. + client + .cached_token() + .as_mut() + .expect("cached token") + .refresh_at = Instant::now(); + let results = + futures::future::join_all((0..8).map(|_| client.http_request(params.clone()))).await; + assert!(results.iter().all(Result::is_ok)); + assert_eq!(source.renewals.load(Ordering::SeqCst), 2); + } + Ok(()) +} + +#[path = "ema_http_client_regression_tests.rs"] +mod regressions; + +#[tokio::test] +async fn initial_bearer_aged_during_validation_is_renewed_before_dispatch() -> anyhow::Result<()> { + for streaming in [false, true] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let (client, source, params) = fixture(&format!("{}/mcp", server.uri())); + let client = EmaAuthenticatedHttpClient::new( + client.inner, + Arc::clone(&source), + EmaAccessToken { + access_token: "aged-in-validation".to_string(), + expires_in: Some(Duration::from_secs(300)), + received_at: Instant::now() - Duration::from_secs(301), + }, + client.authority, + ); + if streaming { + client.http_request_stream(params).await?; + } else { + client.http_request(params).await?; + } + let requests = server.received_requests().await.expect("requests"); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].headers["authorization"], "Bearer renewed-1"); + assert_eq!(source.renewals.load(Ordering::SeqCst), 1); + } + Ok(()) +} + +#[tokio::test] +async fn renewal_aged_during_validation_is_not_dispatched_or_cached() -> anyhow::Result<()> { + for streaming in [false, true] { + let server = MockServer::start().await; + let (client, source, mut params) = fixture(&format!("{}/mcp", server.uri())); + params.timeout_ms = None; + *client.cached_token() = None; + let client = Arc::new(client); + let (gate, waiting, release) = CredentialGate::new(); + *source.renewal_gate.lock().unwrap() = Some(gate); + tokio::time::pause(); + let request = { + let client = Arc::clone(&client); + tokio::spawn(async move { + if streaming { + client + .http_request_stream(params) + .await + .map(|(response, _)| response) + } else { + client.http_request(params).await + } + }) + }; + waiting.await?; + tokio::time::advance(Duration::from_secs(301)).await; + release.send(()).expect("release validation"); + let result = request.await?; + tokio::time::resume(); + assert!(matches!(result, Err(ExecServerError::HttpRequest(message)) + if message == "enterprise MCP bearer expired during credential validation")); + assert!(client.cached_token().is_none()); + assert!( + server + .received_requests() + .await + .expect("requests") + .is_empty() + ); + assert_eq!(source.renewals.load(Ordering::SeqCst), 1); + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/ema_identity.rs b/codex-rs/rmcp-client/src/ema_identity.rs index b320c8855e..64a1d6f2da 100644 --- a/codex-rs/rmcp-client/src/ema_identity.rs +++ b/codex-rs/rmcp-client/src/ema_identity.rs @@ -14,7 +14,9 @@ use oauth2::TokenResponse; use rmcp::transport::AuthorizationManager; use rmcp::transport::auth::AuthorizationMetadata; use rmcp::transport::auth::OAuthHttpClient; +use tracing::instrument::WithSubscriber; +use crate::EmaCredentialLease; use crate::ema_auth_policy::advertised_capability; use crate::ema_auth_policy::ema_reauthentication_required; use crate::ema_auth_policy::validate_ema_oauth_endpoint; @@ -23,7 +25,6 @@ use crate::ema_claims::ID_JAG_TOKEN_TYPE; use crate::ema_exchange::TOKEN_EXCHANGE_GRANT_TYPE; use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::oauth::RefreshCredentialLock; -use crate::oauth::StoredOAuthCredentialSnapshot; use crate::oauth::StoredOAuthTokens; use crate::oauth::stored_oidc_identity; use crate::oauth_http_client::OAuthHttpClientAdapter; @@ -32,16 +33,16 @@ use crate::utils::build_default_headers; pub struct EmaIdpIdentityRequest<'a> { pub issuer: &'a str, pub client_id: &'a str, - pub credentials: &'a StoredOAuthCredentialSnapshot, + pub credentials: EmaCredentialLease, pub http_client: Arc, pub redirect_mode: StreamableHttpRedirectMode, } -/// An opaque IdP refresh token whose credential lock is held through token exchange. +/// An opaque IdP refresh token paired with the credentials admitted for its exchange. pub struct EmaIdpIdentity { pub(crate) token_endpoint: String, pub(crate) refresh_token: String, - pub(crate) credential_lock: RefreshCredentialLock, + pub(crate) credentials: EmaCredentialLease, } /// Returns whether stored credentials contain a refresh token bound to the configured login. @@ -101,7 +102,7 @@ async fn resolve_ema_idp_identity_in( if request.issuer.trim().is_empty() || request.client_id.trim().is_empty() { bail!("ema_auth requires a non-empty enterprise IdP issuer and client ID"); } - let credentials = request.credentials.credentials(); + let credentials = request.credentials.credentials.credentials(); if credentials.url != request.issuer || credentials.bound_issuer() != Some(request.issuer) || credentials.client_id != request.client_id @@ -133,14 +134,21 @@ async fn resolve_ema_idp_identity_in( ); } } - let credential_lock = - RefreshCredentialLock::acquire_for_server(&credentials.server_name, &credentials.url) - .await?; - let snapshot = request.credentials.clone(); + let credentials = request.credentials; + let credential_lock = RefreshCredentialLock::acquire_for_server( + &credentials.credentials.credentials().server_name, + request.issuer, + ) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("failed to lock enterprise credentials"))?; let keyring_store = keyring_store.clone(); - // The worker retains the guard even if its caller stops waiting for the reread. + // The worker retains the short read lock if its caller stops waiting for the reread. tokio::task::spawn_blocking(move || { - let latest = snapshot.load_ema_credentials(&keyring_store)?; + let _credential_lock = credential_lock; + let latest = credentials + .credentials + .load_ema_credentials(&keyring_store)?; let refresh_token = latest .token_response .0 @@ -154,7 +162,7 @@ async fn resolve_ema_idp_identity_in( Ok(EmaIdpIdentity { token_endpoint: metadata.token_endpoint, refresh_token: refresh_token.secret().to_string(), - credential_lock, + credentials, }) }) .await diff --git a/codex-rs/rmcp-client/src/ema_identity_tests.rs b/codex-rs/rmcp-client/src/ema_identity_tests.rs index 56222f744d..21e57e2b5b 100644 --- a/codex-rs/rmcp-client/src/ema_identity_tests.rs +++ b/codex-rs/rmcp-client/src/ema_identity_tests.rs @@ -29,7 +29,9 @@ use wiremock::matchers::path; use super::*; use crate::EmaAuthFailure; use crate::WrappedOAuthTokenResponse; +use crate::oauth::RefreshCredentialLock; use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthCredentialSnapshot; use crate::oauth::test_support::TempCodexHome; fn credentials(issuer: &str, subject: &str, expires_at: u64) -> StoredOAuthTokens { @@ -66,7 +68,9 @@ fn request<'a>( EmaIdpIdentityRequest { issuer, client_id: "idp-client", - credentials, + credentials: EmaCredentialLease { + credentials: credentials.clone(), + }, http_client: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( OutboundProxyPolicy::ReqwestDefault, ))), @@ -288,14 +292,15 @@ async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() - "enterprise IdP credential reread task failed" ); } else { - assert!(error.to_string().contains("refusing file fallback")); - // The store's transparent wrapper exposes a platform error's source. - assert!(error.chain().any(|cause| { - matches!( - cause.downcast_ref::(), - Some(error) if error.kind() == io::ErrorKind::PermissionDenied - ) - })); + assert_eq!( + format!("{error:#}"), + "failed to read enterprise IdP credentials from keyring" + ); + assert!( + !error + .chain() + .any(::is::) + ); } } // Drain the detached read before TempCodexHome changes the process environment. @@ -323,15 +328,6 @@ async fn refresh_subject_rereads_pinned_credentials_after_id_token_expiry() -> R (&identity.token_endpoint, identity.refresh_token.as_str()), (&format!("{issuer}/token"), "stored-refresh") ); - assert!( - tokio::time::timeout( - Duration::from_millis(/*millis*/ 50), - RefreshCredentialLock::acquire_for_server(&stored.server_name, &issuer), - ) - .await - .is_err() - ); - drop(identity); let _released = RefreshCredentialLock::acquire_for_server(&stored.server_name, &issuer).await?; Ok(()) } diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs index a956420547..b4a5de157e 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_login.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs @@ -31,6 +31,7 @@ use crate::ema_identity::resolve_ema_idp_authorization_manager; use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::oauth::EnterpriseOAuthGeneration; use crate::oauth::EnterpriseOAuthGenerationFile; +use crate::oauth::EnterpriseOAuthGenerationKind; use crate::oauth::RefreshCredentialLock; use crate::oauth::delete_oauth_tokens_with_lock_held; use crate::oauth::save_oauth_tokens_with_lock_held; @@ -61,8 +62,13 @@ impl EnterpriseOAuthCredentialGuard { .with_subscriber(tracing::subscriber::NoSubscriber::default()) .await .map_err(|_| anyhow!("failed to lock enterprise credentials"))?; - let generation_file = EnterpriseOAuthGenerationFile::open(credential_name, issuer, &lock) - .map_err(|_| anyhow!("failed to open enterprise login generation"))?; + let generation_file = EnterpriseOAuthGenerationFile::open( + credential_name, + issuer, + EnterpriseOAuthGenerationKind::LoginAttempt, + &lock, + ) + .map_err(|_| anyhow!("failed to open enterprise login generation"))?; Ok(Self { credential_name: credential_name.to_owned(), issuer: issuer.to_owned(), diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs index 1487fd6446..7434dfe3c4 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs @@ -4,7 +4,9 @@ use std::any::Any; use std::collections::HashMap; use std::sync::Mutex; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::time::Duration; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; @@ -342,6 +344,67 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> { .map(|token| token.secret().as_str()), Some(SECRET) ); + // Exercise the production lease and keyring snapshot without installing another + // test backend. Use a non-sensitive name so routine keyring-read traces are safe. + let mut lease_tokens = stored.clone(); + lease_tokens.server_name = "ema-idp:lease-snapshot".to_string(); + crate::save_oauth_tokens( + &lease_tokens.server_name, + &lease_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await?; + let snapshot = crate::stored_oauth_credential_snapshot( + &lease_tokens.server_name, + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("keyring-backed snapshot"); + let pinned = snapshot.clone().pin_ema_credentials().await?; + let accesses = keyring.accesses.load(Ordering::SeqCst); + keyring.fail.store(true, Ordering::SeqCst); + let leases = futures::future::try_join_all( + (0..32).map(|_| pinned.acquire_current_ema_credentials()), + ) + .await?; + assert_eq!(keyring.accesses.load(Ordering::SeqCst), accesses); + drop(leases); + keyring.fail.store(false, Ordering::SeqCst); + + // Even an identical replacement, or a failed write, revokes the old version. + for fail in [false, true] { + let pinned = snapshot.clone().pin_ema_credentials().await?; + keyring.fail.store(fail, Ordering::SeqCst); + let saved = crate::save_oauth_tokens( + &lease_tokens.server_name, + &lease_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + .await; + keyring.fail.store(false, Ordering::SeqCst); + assert_eq!(saved.is_err(), fail); + assert!(pinned.acquire_current_ema_credentials().await.is_err()); + } + let pinned = snapshot.pin_ema_credentials().await?; + let first = pinned.acquire_current_ema_credentials().await?; + let second = pinned.acquire_current_ema_credentials().await?; + assert!( + tokio::time::timeout( + Duration::from_secs(5), + delete_enterprise_oauth_tokens( + &lease_tokens.server_name, + &issuer, + AuthKeyringBackendKind::Direct, + ), + ) + .await?? + ); + drop(first); + drop(second); + assert!(pinned.acquire_current_ema_credentials().await.is_err()); assert!( delete_enterprise_oauth_tokens( CREDENTIAL_NAME, @@ -436,6 +499,7 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> { struct TestKeyring { values: Arc>>>, fail: Arc, + accesses: Arc, } struct TestCredential(Arc); @@ -462,6 +526,7 @@ impl CredentialBuilderApi for TestKeyring { _service: &str, user: &str, ) -> keyring::Result> { + self.accesses.fetch_add(1, Ordering::SeqCst); if self.fail.load(Ordering::SeqCst) { return Err(keyring::Error::Invalid("account".into(), user.into())); } diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs index 3ca54d2120..3916c01b41 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs @@ -1,4 +1,4 @@ -//! Cross-process logout exercises staged login and the production keyring adapter. +//! Cross-process credential changes invalidate staged logins and cached EMA leases. use std::fs; use std::path::PathBuf; @@ -16,9 +16,80 @@ use tokio::sync::Notify; use super::*; +#[tokio::test] +async fn failed_persistence_fences_old_snapshots_but_allows_fresh_pins() -> Result<()> { + if isolated_process("enterprise_oauth_login::tests::logout::failed_persistence_fences_old_snapshots_but_allows_fresh_pins").await? { + return Ok(()); + } + let keyring = TestKeyring::default(); + keyring::set_default_credential_builder(Box::new(keyring.clone())); + let issuer = "https://idp.example"; + let assertion = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode( + json!({"iss":issuer,"sub":"user","aud":"enterprise-client","exp":4102444800_u64}) + .to_string() + ) + ); + let tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": CREDENTIAL_NAME, "url": issuer, "issuer": issuer, + "client_id": "enterprise-client", "token_response": { + "access_token": SECRET, "refresh_token": SECRET, "id_token": assertion, "token_type": "Bearer" + } + }))?; + let mode = OAuthCredentialsStoreMode::Keyring; + let backend = AuthKeyringBackendKind::Direct; + crate::save_oauth_tokens(CREDENTIAL_NAME, &tokens, mode, backend).await?; + for operation in ["save", "delete"] { + let snapshot = + crate::stored_oauth_credential_snapshot(CREDENTIAL_NAME, issuer, mode, backend)? + .expect("stored grant") + .pin_ema_credentials() + .await?; + drop(snapshot.acquire_current_ema_credentials().await?); + // Fail the actual keyring mutation, after the public API advances the generation. + keyring + .values + .lock() + .unwrap() + .values() + .next() + .expect("credential") + .set_error(keyring::Error::PlatformFailure(Box::new( + std::io::Error::other("injected persistence failure"), + ))); + let result = match operation { + "save" => crate::save_oauth_tokens(CREDENTIAL_NAME, &tokens, mode, backend).await, + "delete" => crate::delete_oauth_tokens(CREDENTIAL_NAME, issuer, mode, backend) + .await + .map(|_| ()), + _ => unreachable!(), + }; + assert!( + result.is_err(), + "{operation} must report the storage failure" + ); + assert!( + snapshot.acquire_current_ema_credentials().await.is_err(), + "{operation} must fence the old snapshot" + ); + assert_eq!(stored(issuer)?, Some(tokens.clone())); + // A generation is not a deletion tombstone: a fresh connection may use the surviving grant. + let fresh = + crate::stored_oauth_credential_snapshot(CREDENTIAL_NAME, issuer, mode, backend)? + .expect("surviving grant") + .pin_ema_credentials() + .await?; + drop(fresh.acquire_current_ema_credentials().await?); + } + Ok(()) +} + const TEST: &str = "enterprise_oauth_login::tests::logout::logout_invalidates_pending_login_across_processes"; const LOGOUT_ISSUER: &str = "CODEX_ENTERPRISE_LOGOUT_TEST_ISSUER"; +const REPLACE_CREDENTIAL: &str = "CODEX_ENTERPRISE_LOGOUT_TEST_REPLACE"; #[tokio::test] async fn logout_invalidates_pending_login_across_processes() -> Result<()> { @@ -30,8 +101,22 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> { fs::create_dir_all(&keyring.0)?; keyring::set_default_credential_builder(Box::new(keyring)); if let Ok(issuer) = std::env::var(LOGOUT_ISSUER) { - delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct) + if std::env::var_os(REPLACE_CREDENTIAL).is_some() { + crate::save_oauth_tokens( + CREDENTIAL_NAME, + &stored(&issuer)?.expect("stored grant"), + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) .await?; + } else { + delete_enterprise_oauth_tokens( + CREDENTIAL_NAME, + &issuer, + AuthKeyringBackendKind::Direct, + ) + .await?; + } return Ok(()); } let server = MockServer::start().await; @@ -139,6 +224,42 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> { assert_eq!(stored(&issuer)?, Some(fresh)); } + // Even a byte-identical replacement in another process retires the cached + // version. Its next lease must fail after that process exits, as must logout. + for operation in ["replace", "delete"] { + let before = stored(&issuer)?; + let cached = crate::stored_oauth_credential_snapshot( + CREDENTIAL_NAME, + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .expect("keyring-backed snapshot") + .pin_ema_credentials() + .await?; + drop(cached.acquire_current_ema_credentials().await?); + + let mut child = tokio::process::Command::new(std::env::current_exe()?); + child + .args(["--exact", TEST, "--nocapture"]) + .env(LOGOUT_ISSUER, &issuer); + if operation == "replace" { + child.env(REPLACE_CREDENTIAL, "1"); + } + let output = child.output().await?; + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!(cached.acquire_current_ema_credentials().await.is_err()); + assert_eq!( + stored(&issuer)?, + if operation == "replace" { before } else { None } + ); + } + let stale = complete_login(&issuer).await?; let generation_path = fs::read_dir(home.join("mcp-oauth-locks"))? .filter_map(std::result::Result::ok) @@ -164,6 +285,23 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> { // Metadata failures are logout errors and must leave the stored grant intact. let before = stored(&issuer)?; + let version_path = generation_path.with_extension("enterprise-credential-version"); + let version = fs::read(&version_path)?; + fs::remove_file(&version_path)?; + fs::create_dir(&version_path)?; + let mode = OAuthCredentialsStoreMode::Keyring; + let backend = AuthKeyringBackendKind::Direct; + let mut replacement = before.clone().expect("stored grant"); + replacement.client_id.push_str("-replacement"); + let saved = crate::save_oauth_tokens(CREDENTIAL_NAME, &replacement, mode, backend).await; + assert!(saved.is_err()); + assert_eq!(stored(&issuer)?, before); + let deleted = crate::delete_oauth_tokens(CREDENTIAL_NAME, &issuer, mode, backend).await; + assert!(deleted.is_err()); + assert_eq!(stored(&issuer)?, before); + fs::remove_dir(&version_path)?; + fs::write(&version_path, version)?; + fs::remove_file(&generation_path)?; fs::create_dir(&generation_path)?; assert!( diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index b4d65067ec..e1c0a52647 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod ema_auth; mod ema_auth_policy; mod ema_claims; mod ema_exchange; +mod ema_http_client; mod ema_identity; mod ema_resource; mod enterprise_oauth_login; @@ -49,6 +50,9 @@ pub use ema_auth_policy::EmaAuthFailure; pub use ema_auth_policy::EmaInvalidGrantSource; pub use ema_auth_policy::validate_ema_auth_resource; pub use ema_exchange::EmaAccessToken; +pub use ema_http_client::EmaAuthenticatedHttpClient; +pub use ema_http_client::EmaRequestAuthority; +pub use ema_http_client::EmaTokenSource; pub use ema_identity::EmaIdpIdentity; pub use ema_identity::EmaIdpIdentityRequest; pub use ema_identity::resolve_ema_idp_identity; @@ -63,6 +67,8 @@ pub use event_notification_transport::EventNotificationReceiver; pub use http_client_adapter::StreamableHttpRedirectMode; pub use http_headers::with_http_headers_helper; pub use in_process_transport::InProcessTransportFactory; +pub use oauth::EmaCredentialLease; +pub use oauth::EmaCredentialSnapshot; pub use oauth::StoredOAuthCredentialSnapshot; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index f3127d399a..a0773a768e 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -80,9 +80,13 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; pub(crate) use self::credential_store::OAuthCredentialStore; +pub use self::ema_identity::EmaCredentialLease; +pub use self::ema_identity::EmaCredentialSnapshot; pub(crate) use self::ema_identity::stored_oidc_identity; pub(crate) use self::enterprise_generation::EnterpriseOAuthGeneration; pub(crate) use self::enterprise_generation::EnterpriseOAuthGenerationFile; +pub(crate) use self::enterprise_generation::EnterpriseOAuthGenerationKind; +use self::enterprise_generation::invalidate_enterprise_credential_version; pub(crate) use self::issuer_binding::validate_authorization_server_endpoints; pub(crate) use self::issuer_binding::validate_refresh_token_issuer; pub(crate) use self::refresh_lock::RefreshCredentialLock; @@ -467,12 +471,13 @@ pub async fn save_oauth_tokens( /// Save while retaining the matching credential lock acquired by the caller. pub(crate) fn save_oauth_tokens_with_lock_held( - _lock: &RefreshCredentialLock, + lock: &RefreshCredentialLock, server_name: &str, tokens: &StoredOAuthTokens, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result<()> { + invalidate_enterprise_credential_version(server_name, &tokens.url, lock)?; let keyring_store = DefaultKeyringStore; let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Save); match store_mode { @@ -642,12 +647,13 @@ pub async fn delete_oauth_tokens( /// Delete while retaining the matching credential lock acquired by the caller. pub(crate) fn delete_oauth_tokens_with_lock_held( - _lock: &RefreshCredentialLock, + lock: &RefreshCredentialLock, server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { + invalidate_enterprise_credential_version(server_name, url, lock)?; let keyring_store = DefaultKeyringStore; delete_oauth_tokens_from_keyring_and_file( &keyring_store, diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity.rs b/codex-rs/rmcp-client/src/oauth/ema_identity.rs index bf4ea0d17a..9aaa122e9c 100644 --- a/codex-rs/rmcp-client/src/oauth/ema_identity.rs +++ b/codex-rs/rmcp-client/src/oauth/ema_identity.rs @@ -1,16 +1,53 @@ -//! Reread enterprise credentials without changing the connection's pinned identity. +//! Pin enterprise credentials to a non-secret version advanced by cooperating writers. +//! Exchanges reread the pinned keyring record and reject changes made outside this protocol. use anyhow::Result; +use anyhow::anyhow; use anyhow::bail; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use tracing::instrument::WithSubscriber; -use super::ResolvedOAuthCredentialStore; -use super::StoredOAuthCredentialSnapshot; -use super::StoredOAuthTokens; +use crate::ema_auth_policy::EmaAuthFailure; use crate::ema_auth_policy::ema_reauthentication_required; use crate::ema_claims::OidcClaims; use crate::ema_claims::oidc_identity; +use crate::oauth::EnterpriseOAuthGeneration; +use crate::oauth::EnterpriseOAuthGenerationFile; +use crate::oauth::EnterpriseOAuthGenerationKind; +use crate::oauth::RefreshCredentialLock; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthCredentialSnapshot; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::invalidate_enterprise_credential_version; + +/// One validated keyring record and its cross-process credential version. +#[derive(Clone)] +pub struct EmaCredentialSnapshot { + snapshot: StoredOAuthCredentialSnapshot, + version: EnterpriseOAuthGeneration, +} + +/// A credential snapshot admitted by a short generation check. +/// Logout may replace the grant while an admitted exchange finishes. +pub struct EmaCredentialLease { + pub(crate) credentials: StoredOAuthCredentialSnapshot, +} + +// Discard storage and lock error chains: keyring backends can include the stable +// account/workspace key. Keep only the sanitized authentication classification. +fn sanitize_ema_credential_error(error: anyhow::Error, message: &'static str) -> anyhow::Error { + match error.downcast::() { + Ok(failure) => anyhow::Error::new(failure).context(message), + Err(_) => anyhow!(message), + } +} + +impl std::fmt::Debug for EmaCredentialLease { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EmaCredentialLease").finish_non_exhaustive() + } +} pub(crate) fn stored_oidc_identity(tokens: &StoredOAuthTokens) -> Result { let assertion = tokens @@ -34,13 +71,154 @@ pub(crate) fn stored_oidc_identity(tokens: &StoredOAuthTokens) -> Result Result<()> { - self.load_ema_credentials(&DefaultKeyringStore).map(|_| ()) + /// Bind the record and its version under one lease, validating any earlier snapshot read. + pub async fn pin_ema_credentials(self) -> Result { + self.pin_ema_credentials_in(&DefaultKeyringStore).await } - /// Reread only the pinned keyring authority; exchange callers hold its credential lock. + async fn pin_ema_credentials_in( + self, + keyring_store: &K, + ) -> Result { + if !self.credentials.server_name.starts_with("ema-idp:") { + bail!("enterprise credentials require the enterprise namespace"); + } + let lock = RefreshCredentialLock::acquire_for_server( + &self.credentials.server_name, + &self.credentials.url, + ) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("failed to lock enterprise credentials"))?; + let keyring_store = keyring_store.clone(); + tokio::task::spawn_blocking(move || { + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + self.load_ema_credentials(&keyring_store)?; + let versions = EnterpriseOAuthGenerationFile::open( + &self.credentials.server_name, + &self.credentials.url, + EnterpriseOAuthGenerationKind::Credential, + &lock, + )?; + let version = match versions.current()? { + Some(version) => version, + None => versions.replace()?, + }; + Ok(EmaCredentialSnapshot { + snapshot: self, + version, + }) + }) + }) + .await + .map_err(|_| anyhow!("enterprise credential pinning task failed"))? + .map_err(|error: anyhow::Error| { + sanitize_ema_credential_error(error, "failed to pin enterprise credentials") + }) + } +} + +impl EmaCredentialSnapshot { + /// Validate the pinned version under the credential lock without reading the keyring. + pub async fn acquire_current_ema_credentials(&self) -> Result { + let credential_lock = RefreshCredentialLock::acquire_for_server( + &self.snapshot.credentials.server_name, + &self.snapshot.credentials.url, + ) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("failed to lock enterprise credentials"))?; + let pinned = self.clone(); + tokio::task::spawn_blocking(move || { + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + let credentials = pinned.snapshot.credentials(); + let versions = EnterpriseOAuthGenerationFile::open( + &credentials.server_name, + &credentials.url, + EnterpriseOAuthGenerationKind::Credential, + &credential_lock, + )?; + if versions.current()?.as_ref() != Some(&pinned.version) { + return Err(ema_reauthentication_required( + "enterprise credentials changed; reconnect", + )); + } + Ok(EmaCredentialLease { + credentials: pinned.snapshot, + }) + }) + }) + .await + .map_err(|_| anyhow!("enterprise credential validation task failed"))? + .map_err(|error: anyhow::Error| { + sanitize_ema_credential_error(error, "failed to validate enterprise credentials") + }) + } + + /// Exclusively lock and delete a rejected grant only if its pinned record is still current. + pub async fn invalidate_ema_credentials_if_current(&self) -> Result<()> { + self.invalidate_ema_credentials_if_current_in(&DefaultKeyringStore) + .await + } + + async fn invalidate_ema_credentials_if_current_in( + &self, + keyring_store: &K, + ) -> Result<()> { + let credential_lock = RefreshCredentialLock::acquire_for_server( + &self.snapshot.credentials.server_name, + &self.snapshot.credentials.url, + ) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("failed to lock enterprise credentials"))?; + let pinned = self.clone(); + let keyring_store = keyring_store.clone(); + tokio::task::spawn_blocking(move || { + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + let snapshot = &pinned.snapshot; + let previous = snapshot.credentials(); + let versions = EnterpriseOAuthGenerationFile::open( + &previous.server_name, + &previous.url, + EnterpriseOAuthGenerationKind::Credential, + &credential_lock, + )?; + if versions.current()?.as_ref() != Some(&pinned.version) { + return Ok(()); + } + if let Some(mut latest) = + snapshot + .store + .load(&keyring_store, &previous.server_name, &previous.url)? + { + latest.token_response.0.set_expires_in(None); + if latest == *previous { + invalidate_enterprise_credential_version( + &previous.server_name, + &previous.url, + &credential_lock, + )?; + snapshot.store.delete( + &keyring_store, + &previous.server_name, + &previous.url, + )?; + } + } + Ok(()) + }) + }) + .await + .map_err(|_| anyhow!("enterprise credential invalidation task failed"))? + .map_err(|error: anyhow::Error| { + sanitize_ema_credential_error(error, "failed to invalidate enterprise credentials") + }) + } +} + +impl StoredOAuthCredentialSnapshot { + /// Reread only the pinned keyring authority; a changed record rejects the exchange. pub(crate) fn load_ema_credentials( &self, keyring_store: &K, @@ -49,9 +227,12 @@ impl StoredOAuthCredentialSnapshot { bail!("enterprise IdP credentials require keyring storage"); } let previous = &self.credentials; - let mut latest = self - .store - .load(keyring_store, &previous.server_name, &previous.url)? + let mut latest = + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + self.store + .load(keyring_store, &previous.server_name, &previous.url) + }) + .map_err(|_| anyhow!("failed to read enterprise IdP credentials from keyring"))? .ok_or_else(|| { ema_reauthentication_required( "enterprise IdP credentials were removed; sign in again", diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs index bd804998d2..b146cc2733 100644 --- a/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs @@ -1,12 +1,247 @@ +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_config::types::AuthKeyringBackendKind; +use codex_keyring_store::CredentialStoreError; +use codex_keyring_store::KeyringStore; use codex_keyring_store::tests::MockKeyringStore; +use futures::FutureExt; +use pretty_assertions::assert_eq; use serde_json::json; +use tracing_test::traced_test; use super::*; +use crate::ema_auth_policy::EmaAuthFailure; use crate::oauth::RefreshCredentialLock; use crate::oauth::compute_store_key; use crate::oauth::test_support::TempCodexHome; +#[derive(Clone, Debug, Default)] +struct TracingKeyringStore(MockKeyringStore); + +impl KeyringStore for TracingKeyringStore { + fn load( + &self, + service: &str, + account: &str, + ) -> std::result::Result, CredentialStoreError> { + tracing::trace!(account, "reading enterprise keyring entry"); + self.0.load(service, account) + } + + fn save( + &self, + service: &str, + account: &str, + value: &str, + ) -> std::result::Result<(), CredentialStoreError> { + self.0.save(service, account, value) + } + + fn delete( + &self, + service: &str, + account: &str, + ) -> std::result::Result { + tracing::trace!(account, "deleting enterprise keyring entry"); + self.0.delete(service, account) + } +} + +#[tokio::test] +async fn invalidation_rereads_under_exclusive_lock_and_preserves_newer_grants() -> Result<()> { + let _home = TempCodexHome::new(); + for replacement in [ + None, + Some("refresh_token"), + Some("client_id"), + Some("version"), + ] { + let tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": "ema-idp:rejected-grant", "url": "https://idp.example", + "issuer": "https://idp.example", "client_id": "client", + "token_response": { + "access_token": "unused", "token_type": "Bearer", "refresh_token": "old-grant" + }, + }))?; + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + let keyring = MockKeyringStore::default(); + let writer = + RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?; + let versions = EnterpriseOAuthGenerationFile::open( + &tokens.server_name, + &tokens.url, + EnterpriseOAuthGenerationKind::Credential, + &writer, + )?; + let snapshot = EmaCredentialSnapshot { + snapshot: StoredOAuthCredentialSnapshot::new(tokens.clone(), store), + version: versions.replace()?, + }; + store.save(&keyring, &tokens.server_name, &tokens)?; + let mut invalidation = + Box::pin(snapshot.invalidate_ema_credentials_if_current_in(&keyring)); + assert!(invalidation.as_mut().now_or_never().is_none()); + assert_eq!( + store.load(&keyring, &tokens.server_name, &tokens.url)?, + Some(tokens.clone()) + ); + + // A login that wins the lock may change any part of the atomic credential record. + let mut latest = tokens.clone(); + match replacement { + Some("refresh_token") => latest + .token_response + .0 + .set_refresh_token(Some(oauth2::RefreshToken::new("new-grant".to_string()))), + Some("client_id") => latest.client_id = "new-client".to_string(), + // An identical token response can still belong to a newer login. + Some("version") => { + invalidate_enterprise_credential_version(&tokens.server_name, &tokens.url, &writer)? + } + None => {} + Some(_) => unreachable!(), + } + store.save(&keyring, &tokens.server_name, &latest)?; + drop(writer); + invalidation.await?; + assert_eq!( + store.load(&keyring, &tokens.server_name, &tokens.url)?, + replacement.map(|_| latest), + ); + let _reader = + RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?; + assert_eq!( + versions.current()?.as_ref() != Some(&snapshot.version), + replacement.is_none() || replacement == Some("version"), + "invalidation advances the version only when deleting its own grant", + ); + } + Ok(()) +} + +fn versioned_tokens() -> Result { + let assertion = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(br#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode( + json!({ + "iss":"https://idp.example", "aud":"client", "sub":"user", "exp":4102444800_u64, + }) + .to_string() + ), + ); + Ok(serde_json::from_value(json!({ + "server_name":"ema-idp:versioned", "url":"https://idp.example", + "issuer":"https://idp.example", "client_id":"client", + "token_response": { + "access_token":"unused", "token_type":"Bearer", + "refresh_token":"enterprise-refresh", "id_token":assertion, + }, + }))?) +} + +#[tokio::test] +async fn cached_credentials_require_the_current_version_after_waiting_for_a_writer() -> Result<()> { + let home = TempCodexHome::new(); + for change in ["replacement", "missing", "corrupt"] { + let tokens = versioned_tokens()?; + let keyring = MockKeyringStore::default(); + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + store.save(&keyring, &tokens.server_name, &tokens)?; + let snapshot = StoredOAuthCredentialSnapshot::new(tokens.clone(), store) + .pin_ema_credentials_in(&keyring) + .await?; + let version_path = std::fs::read_dir(home.path().join("mcp-oauth-locks"))? + .filter_map(std::result::Result::ok) + .map(|entry| entry.path()) + .find(|path| { + path.extension() + .is_some_and(|ext| ext == "enterprise-credential-version") + }) + .expect("credential version"); + let writer = + RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?; + let mut acquisition = Box::pin(snapshot.acquire_current_ema_credentials()); + assert!(acquisition.as_mut().now_or_never().is_none()); + match change { + "replacement" => { + invalidate_enterprise_credential_version(&tokens.server_name, &tokens.url, &writer)? + } + "missing" => std::fs::remove_file(&version_path)?, + "corrupt" => std::fs::write(&version_path, b"incomplete write")?, + _ => unreachable!(), + } + drop(writer); + let error = acquisition + .await + .expect_err("cached grant must not survive a change"); + if change != "corrupt" { + assert_eq!( + error.downcast_ref::(), + Some(&EmaAuthFailure::ReauthenticationRequired), + "{change}" + ); + } + assert_eq!( + store.load(&keyring, &tokens.server_name, &tokens.url)?, + Some(tokens) + ); + std::fs::remove_file(version_path)?; + } + Ok(()) +} + +#[tokio::test] +async fn pinning_does_not_bind_a_stale_record_to_a_new_credential_version() -> Result<()> { + let _home = TempCodexHome::new(); + let tokens = versioned_tokens()?; + let keyring = MockKeyringStore::default(); + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + store.save(&keyring, &tokens.server_name, &tokens)?; + let snapshot = StoredOAuthCredentialSnapshot::new(tokens.clone(), store); + let writer = + RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?; + let mut pinning = Box::pin(snapshot.pin_ema_credentials_in(&keyring)); + assert!(pinning.as_mut().now_or_never().is_none()); + let mut replacement = tokens; + replacement + .token_response + .0 + .set_refresh_token(Some(oauth2::RefreshToken::new( + "replacement-grant".to_string(), + ))); + invalidate_enterprise_credential_version(&replacement.server_name, &replacement.url, &writer)?; + store.save(&keyring, &replacement.server_name, &replacement)?; + drop(writer); + assert!(pinning.await.is_err()); + let current = StoredOAuthCredentialSnapshot::new(replacement, store) + .pin_ema_credentials_in(&keyring) + .await?; + current.acquire_current_ema_credentials().await?; + Ok(()) +} + +#[tokio::test] +async fn admitted_credentials_do_not_block_a_new_pin() -> Result<()> { + let _home = TempCodexHome::new(); + let tokens = versioned_tokens()?; + let keyring = MockKeyringStore::default(); + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + store.save(&keyring, &tokens.server_name, &tokens)?; + let snapshot = StoredOAuthCredentialSnapshot::new(tokens, store); + let pinned = snapshot.clone().pin_ema_credentials_in(&keyring).await?; + let _active_request = pinned.acquire_current_ema_credentials().await?; + + let concurrent = tokio::time::timeout( + std::time::Duration::from_secs(1), + snapshot.pin_ema_credentials_in(&keyring), + ) + .await??; + + assert!(concurrent.version == pinned.version); + Ok(()) +} + #[tokio::test] async fn keyring_failure_does_not_reuse_the_pinned_refresh_token() -> Result<()> { let _home = TempCodexHome::new(); @@ -35,7 +270,100 @@ async fn keyring_failure_does_not_reuse_the_pinned_refresh_token() -> Result<()> let error = snapshot .load_ema_credentials(&keyring) .expect_err("keyring failure must be terminal"); - assert!(error.to_string().contains("refusing file fallback")); + assert!( + error + .to_string() + .contains("failed to read enterprise IdP credentials from keyring") + ); + Ok(()) +} + +#[tokio::test] +#[traced_test] +async fn enterprise_credential_errors_and_keyring_traces_hide_account_key() -> Result<()> { + let home = TempCodexHome::new(); + let sentinel = "private-workspace-sentinel"; + let mut tokens = versioned_tokens()?; + tokens.server_name = format!("ema-idp:{sentinel}"); + let key = compute_store_key(&tokens.server_name, &tokens.url)?; + let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct); + let keyring = TracingKeyringStore::default(); + let snapshot = StoredOAuthCredentialSnapshot::new(tokens.clone(), store); + tracing::trace!("enterprise keyring privacy test capture enabled"); + assert!(logs_contain( + "enterprise keyring privacy test capture enabled" + )); + + keyring + .0 + .set_error(&key, keyring::Error::Invalid("account".into(), key.clone())); + let error = snapshot + .load_ema_credentials(&keyring) + .expect_err("backend read must fail"); + assert!(!format!("{error} {error:?} {error:#}").contains(sentinel)); + keyring + .0 + .set_error(&key, keyring::Error::Invalid("account".into(), key.clone())); + let error = snapshot + .clone() + .pin_ema_credentials_in(&keyring) + .await + .err() + .expect("backend pin must fail"); + assert!(!format!("{error} {error:?} {error:#}").contains(sentinel)); + assert_eq!(error.downcast_ref::(), None); + + let keyring = TracingKeyringStore::default(); + let removed = snapshot + .clone() + .pin_ema_credentials_in(&keyring) + .await + .err() + .expect("removed grant must fail"); + assert_eq!( + removed.downcast_ref::(), + Some(&EmaAuthFailure::ReauthenticationRequired), + ); + assert!(!format!("{removed} {removed:?} {removed:#}").contains(sentinel)); + store.save(&keyring, &tokens.server_name, &tokens)?; + let pinned = snapshot.clone().pin_ema_credentials_in(&keyring).await?; + keyring + .0 + .set_error(&key, keyring::Error::Invalid("account".into(), key.clone())); + let error = pinned + .invalidate_ema_credentials_if_current_in(&keyring) + .await + .unwrap_err(); + assert!(!format!("{error} {error:?} {error:#}").contains(sentinel)); + + let successful_keyring = TracingKeyringStore::default(); + store.save(&successful_keyring, &tokens.server_name, &tokens)?; + snapshot + .clone() + .pin_ema_credentials_in(&successful_keyring) + .await? + .invalidate_ema_credentials_if_current_in(&successful_keyring) + .await?; + assert_eq!(successful_keyring.0.saved_value(&key), None); + + let lock_dir = home.path().join("mcp-oauth-locks"); + std::fs::rename(&lock_dir, home.path().join("held-locks"))?; + std::fs::write(&lock_dir, b"not a directory")?; + for error in [ + snapshot + .pin_ema_credentials_in(&keyring) + .await + .err() + .expect("lock must fail"), + pinned.acquire_current_ema_credentials().await.unwrap_err(), + pinned + .invalidate_ema_credentials_if_current_in(&keyring) + .await + .unwrap_err(), + ] { + assert!(!format!("{error} {error:?} {error:#}").contains(sentinel)); + } + assert!(!logs_contain(sentinel)); Ok(()) } diff --git a/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs index 9845f79c98..f8c7f1ad1a 100644 --- a/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs +++ b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs @@ -1,5 +1,6 @@ -//! Persistent, non-secret generations for staged enterprise logins and logout. -//! All reads and writes require the credential lock; missing state never admits an old attempt. +//! Persistent, non-secret fences for login attempts and cached enterprise credentials. +//! Separate generations preserve pending-login semantics when a stored grant changes. +//! All reads and writes require the credential lock; missing state never admits old authority. use std::fs::File; use std::fs::OpenOptions; @@ -16,9 +17,14 @@ use sha2::Sha256; use super::RefreshCredentialLock; -#[derive(PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub(crate) struct EnterpriseOAuthGeneration([u8; 32]); +pub(crate) enum EnterpriseOAuthGenerationKind { + LoginAttempt, + Credential, +} + pub(crate) struct EnterpriseOAuthGenerationFile { file: File, } @@ -27,10 +33,15 @@ impl EnterpriseOAuthGenerationFile { pub(crate) fn open( credential_name: &str, issuer: &str, + kind: EnterpriseOAuthGenerationKind, _lock: &RefreshCredentialLock, ) -> Result { let key = super::compute_store_key(credential_name, issuer)?; - let name = format!("{:x}.enterprise-generation", Sha256::digest(key.as_bytes())); + let suffix = match kind { + EnterpriseOAuthGenerationKind::LoginAttempt => "enterprise-generation", + EnterpriseOAuthGenerationKind::Credential => "enterprise-credential-version", + }; + let name = format!("{:x}.{suffix}", Sha256::digest(key.as_bytes())); let path = find_codex_home()?.join("mcp-oauth-locks").join(name); let mut options = OpenOptions::new(); options.read(true).write(true).create(true).truncate(false); @@ -83,3 +94,21 @@ impl EnterpriseOAuthGenerationFile { Ok(generation) } } + +/// Invalidate cached grants before mutation, including when the following store write fails. +pub(crate) fn invalidate_enterprise_credential_version( + credential_name: &str, + issuer: &str, + lock: &RefreshCredentialLock, +) -> Result<()> { + if credential_name.starts_with("ema-idp:") { + EnterpriseOAuthGenerationFile::open( + credential_name, + issuer, + EnterpriseOAuthGenerationKind::Credential, + lock, + )? + .replace()?; + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs index e038dc2a1e..1e0e8b2401 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -3,6 +3,7 @@ //! The guard is intentionally acquired before the authoritative credential reread and retained //! through provider refresh and persistence. This prevents two processes from replaying the same //! rotating refresh token or observing a partially persisted transaction. +//! EMA uses the same lock only for short credential reads and generation checks. use anyhow::Context; use anyhow::Result; diff --git a/codex-rs/rmcp-client/src/startup_error.rs b/codex-rs/rmcp-client/src/startup_error.rs index 744e361e9b..486299322c 100644 --- a/codex-rs/rmcp-client/src/startup_error.rs +++ b/codex-rs/rmcp-client/src/startup_error.rs @@ -1,10 +1,13 @@ use anyhow::Error; +use codex_exec_server::ExecServerError; use rmcp::service::ClientInitializeError; use rmcp::service::ServiceError; use rmcp::transport::DynamicTransportError; use rmcp::transport::auth::AuthError; use rmcp::transport::streamable_http_client::StreamableHttpError; +use crate::ema_auth_policy::EmaAuthFailure; +use crate::ema_auth_policy::EmaInvalidGrantSource; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::rmcp_client::ClientOperationError; @@ -18,6 +21,22 @@ pub fn is_authentication_required_error(error: &Error) -> bool { source .downcast_ref::() .is_some_and(auth_error_requires_authentication) + || source + .downcast_ref::() + .is_some_and(|error| match error { + EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity, + } + | EmaAuthFailure::InsufficientUserAuthentication + | EmaAuthFailure::ReauthenticationRequired => true, + // Rejecting an ID-JAG does not invalidate the enterprise login itself. + EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::ResourceAuthorization, + } => false, + }) + || source + .downcast_ref::() + .is_some_and(exec_server_error_requires_authentication) || source .downcast_ref::() .is_some_and(|mut error| { @@ -49,10 +68,17 @@ fn transport_error_requires_authentication(error: &DynamicTransportError) -> boo .is_some_and(|error| match error { StreamableHttpError::AuthRequired(_) => true, StreamableHttpError::Auth(auth_error) => auth_error_requires_authentication(auth_error), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest(error)) => { + exec_server_error_requires_authentication(error) + } _ => false, }) } +fn exec_server_error_requires_authentication(error: &ExecServerError) -> bool { + matches!(error, ExecServerError::AuthenticationRequired(_)) +} + fn auth_error_requires_authentication(error: &AuthError) -> bool { matches!( error,