From 36512a34ffe71a48bd9aab66fc81c3b909ef2dfc Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 20 Oct 2026 21:54:46 +0000 Subject: [PATCH] Add account credential support to EMA token exchange helpers (#55685) ## What changed - Add an account IdP identity resolver that checks ID-token bindings against trusted issuer and client provenance before discovering the refresh-token endpoint. An expired ID token does not by itself invalidate the refresh grant. - Generalize `EmaIdpIdentity` and token exchange requests to carry caller-owned credential leases. Invoke a `finish_idp` callback after the IdP exchange and before resource authorization server I/O so callers can release their refresh guards at that boundary. - Reject empty expected OIDC issuers and client IDs, and expose an unverified issuer hint solely for explaining legacy login migration. GitOrigin-RevId: 34473fed82a9a6e58997664cd5e6e67049b5ee7e --- codex-rs/codex-mcp/src/ema_auth.rs | 49 ++++---- codex-rs/rmcp-client/src/ema_auth.rs | 48 ++++---- codex-rs/rmcp-client/src/ema_auth_tests.rs | 4 +- codex-rs/rmcp-client/src/ema_claims.rs | 9 ++ codex-rs/rmcp-client/src/ema_exchange.rs | 75 +++++++----- .../rmcp-client/src/ema_exchange_tests.rs | 64 +++++----- codex-rs/rmcp-client/src/ema_identity.rs | 114 +++++++++++++----- codex-rs/rmcp-client/src/lib.rs | 3 + 8 files changed, 235 insertions(+), 131 deletions(-) diff --git a/codex-rs/codex-mcp/src/ema_auth.rs b/codex-rs/codex-mcp/src/ema_auth.rs index 3039c7d25b..821eb53144 100644 --- a/codex-rs/codex-mcp/src/ema_auth.rs +++ b/codex-rs/codex-mcp/src/ema_auth.rs @@ -301,28 +301,33 @@ impl EmaTokenSource for EnterpriseTokenSource { async fn access_token(&self) -> Result { self.auth.check_account().await?; let idp = self.registration.idp(); - let exchange = exchange_ema_auth_token(EmaAuthTokenExchangeRequest { - server_url: self.registration.server_url(), - resource: self.registration.resource(), - scopes: self.registration.scopes(), - mcp_client_id: self.registration.client_id(), - expected_authorization_server_issuer: self.registration.authorization_server_issuer(), - idp_issuer: &idp.issuer, - idp_client_id: &idp.client_id, - idp_identity: Box::pin(async { - self.auth.check_account().await?; - resolve_ema_idp_identity(EmaIdpIdentityRequest { - issuer: &idp.issuer, - client_id: &idp.client_id, - credentials: self.credentials.acquire_current_ema_credentials().await?, - http_client: Arc::clone(&self.http_client), - redirect_mode: StreamableHttpRedirectMode::Legacy, - }) - .await - }), - idp_http_client: Arc::clone(&self.http_client), - resource_http_client: Arc::clone(&self.http_client), - }); + let exchange = exchange_ema_auth_token( + EmaAuthTokenExchangeRequest { + server_url: self.registration.server_url(), + resource: self.registration.resource(), + scopes: self.registration.scopes(), + mcp_client_id: self.registration.client_id(), + expected_authorization_server_issuer: self + .registration + .authorization_server_issuer(), + idp_issuer: &idp.issuer, + idp_client_id: &idp.client_id, + idp_identity: Box::pin(async { + self.auth.check_account().await?; + resolve_ema_idp_identity(EmaIdpIdentityRequest { + issuer: &idp.issuer, + client_id: &idp.client_id, + credentials: self.credentials.acquire_current_ema_credentials().await?, + http_client: Arc::clone(&self.http_client), + redirect_mode: StreamableHttpRedirectMode::Legacy, + }) + .await + }), + idp_http_client: Arc::clone(&self.http_client), + resource_http_client: Arc::clone(&self.http_client), + }, + std::convert::identity, + ); let result = tokio::select! { biased; () = self.auth.config.mcp_ema_cancellation.cancelled() => { diff --git a/codex-rs/rmcp-client/src/ema_auth.rs b/codex-rs/rmcp-client/src/ema_auth.rs index 02f1697e85..6d18f8df2d 100644 --- a/codex-rs/rmcp-client/src/ema_auth.rs +++ b/codex-rs/rmcp-client/src/ema_auth.rs @@ -41,7 +41,7 @@ struct EmaAuthorizationMetadata { } /// Inputs for non-interactive enterprise-managed MCP authorization. -pub struct EmaAuthTokenExchangeRequest<'a> { +pub struct EmaAuthTokenExchangeRequest<'a, L = EmaCredentialLease> { pub server_url: &'a str, pub resource: Option<&'a str>, pub scopes: &'a [String], @@ -49,13 +49,13 @@ pub struct EmaAuthTokenExchangeRequest<'a> { pub expected_authorization_server_issuer: Option<&'a str>, pub idp_issuer: &'a str, pub idp_client_id: &'a str, - pub idp_identity: BoxFuture<'a, Result>, + pub idp_identity: BoxFuture<'a, Result>>, pub idp_http_client: Arc, pub resource_http_client: Arc, } -impl<'a> From<&EmaAuthTokenExchangeRequest<'a>> for EmaDiscoveryRequest<'a> { - fn from(request: &EmaAuthTokenExchangeRequest<'a>) -> Self { +impl<'a, L> From<&EmaAuthTokenExchangeRequest<'a, L>> for EmaDiscoveryRequest<'a> { + fn from(request: &EmaAuthTokenExchangeRequest<'a, L>) -> Self { Self { server_url: request.server_url, resource: request.resource, @@ -66,9 +66,10 @@ 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<(EmaAccessToken, EmaCredentialLease)> { +pub async fn exchange_ema_auth_token( + request: EmaAuthTokenExchangeRequest<'_, L>, + finish_idp: impl FnOnce(L) -> R + Send, +) -> Result<(EmaAccessToken, R)> { validate_ema_auth_resource(request.server_url, request.resource)?; let client_id = request .mcp_client_id @@ -89,21 +90,24 @@ pub async fn exchange_ema_auth_token( refresh_token, credentials, } = request.idp_identity.await?; - let token = exchange_id_jag(EmaIdJagExchangeRequest { - resource: &resource, - scopes: request.scopes, - mcp_client_id: client_id, - authorization_server_issuer: issuer, - authorization_server_token_endpoint: &metadata.token_endpoint, - idp_token_endpoint: &token_endpoint, - idp_issuer: request.idp_issuer, - idp_client_id: request.idp_client_id, - refresh_token, - idp_http_client: request.idp_http_client, - resource_http_client: request.resource_http_client, - }) - .await?; - Ok((token, credentials)) + exchange_id_jag( + EmaIdJagExchangeRequest { + resource: &resource, + scopes: request.scopes, + mcp_client_id: client_id, + authorization_server_issuer: issuer, + authorization_server_token_endpoint: &metadata.token_endpoint, + idp_token_endpoint: &token_endpoint, + idp_issuer: request.idp_issuer, + idp_client_id: request.idp_client_id, + refresh_token, + idp_http_client: request.idp_http_client, + resource_http_client: request.resource_http_client, + }, + credentials, + finish_idp, + ) + .await } 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 cf263a3352..740f290e3c 100644 --- a/codex-rs/rmcp-client/src/ema_auth_tests.rs +++ b/codex-rs/rmcp-client/src/ema_auth_tests.rs @@ -133,7 +133,7 @@ async fn discovery_pins_resource_issuer_and_public_jwt_bearer_before_identity() polled.store(true, Ordering::SeqCst); bail!("identity resolver reached") }); - exchange_ema_auth_token(exchange) + exchange_ema_auth_token(exchange, std::convert::identity) .await .expect_err("stop before token exchange"); assert_eq!(polled.load(Ordering::SeqCst), valid, "{additions}"); @@ -452,7 +452,7 @@ async fn rejected_grants_preserve_their_source_and_release_the_credential_lease( credentials: crate::EmaCredentialLease { credentials }, }) }); - let error = exchange_ema_auth_token(exchange) + let error = exchange_ema_auth_token(exchange, std::convert::identity) .await .expect_err("invalid grant"); assert_eq!( diff --git a/codex-rs/rmcp-client/src/ema_claims.rs b/codex-rs/rmcp-client/src/ema_claims.rs index 61a65cc34e..56444b85fd 100644 --- a/codex-rs/rmcp-client/src/ema_claims.rs +++ b/codex-rs/rmcp-client/src/ema_claims.rs @@ -59,6 +59,9 @@ pub(crate) fn oidc_identity( expected_issuer: &str, expected_audience: &str, ) -> Result { + if expected_issuer.trim().is_empty() || expected_audience.trim().is_empty() { + bail!("OIDC identity assertion requires a non-empty expected issuer and client"); + } let claims: OidcClaims = signed_jwt(assertion)?; if claims.iss != expected_issuer || claims.sub.trim().is_empty() { bail!("OIDC identity assertion issuer or subject does not match the enterprise IdP"); @@ -82,6 +85,12 @@ pub(crate) fn oidc_identity( Ok(claims) } +/// Unverified issuer claim used only to explain a legacy login migration. +/// Never use this hint to choose where a refresh token is sent. +pub(crate) fn oidc_issuer_hint(assertion: &str, expected_issuer: &str) -> bool { + signed_jwt::(assertion).is_ok_and(|claims| claims.iss == expected_issuer) +} + pub(crate) fn validate_oidc_identity_assertion( assertion: &str, expected_issuer: &str, diff --git a/codex-rs/rmcp-client/src/ema_exchange.rs b/codex-rs/rmcp-client/src/ema_exchange.rs index f27cf17691..1c611810c2 100644 --- a/codex-rs/rmcp-client/src/ema_exchange.rs +++ b/codex-rs/rmcp-client/src/ema_exchange.rs @@ -58,9 +58,11 @@ pub(crate) struct EmaIdJagExchangeRequest<'a> { } /// Exchanges an enterprise IdP credential for a resource-bound MCP bearer token. -pub(crate) async fn exchange_id_jag( +pub(crate) async fn exchange_id_jag( request: EmaIdJagExchangeRequest<'_>, -) -> Result { + credentials: L, + finish_idp: impl FnOnce(L) -> R + Send, +) -> Result<(EmaAccessToken, R)> { let idp_http = OAuthHttpClientAdapter::new_with_redirect_mode( request.idp_http_client, build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?, @@ -78,7 +80,7 @@ pub(crate) async fn exchange_id_jag( let refresh_token = RefreshToken::new(request.refresh_token); // Discovery has already verified both public-client registrations. The SDK owns // token validation and the two exchanges; Codex retains routing and credential policy. - let token = EmaExchangeRequest::new( + let assertion = EmaExchangeRequest::new( EmaAuthorizationServer::new( request.idp_issuer, request.idp_token_endpoint, @@ -93,36 +95,47 @@ pub(crate) async fn exchange_id_jag( &refresh_token, ) .with_scopes(request.scopes.iter().cloned()) - .exchange(&idp_http, &resource_http) + .exchange_id_jag(&idp_http) .await - .map_err(|error| { - let failure = match &error { - EmaError::InvalidGrant(EmaExchangeStage::IdentityProvider) => { - Some(EmaAuthFailure::InvalidGrant { - grant_source: EmaInvalidGrantSource::EnterpriseIdentity, - }) - } - EmaError::InvalidGrant(EmaExchangeStage::ResourceAuthorizationServer) => { - Some(EmaAuthFailure::InvalidGrant { - grant_source: EmaInvalidGrantSource::ResourceAuthorization, - }) - } - EmaError::InsufficientUserAuthentication(_) => { - Some(EmaAuthFailure::InsufficientUserAuthentication) - } - // Future SDK stages must not invalidate the shared enterprise credential. - _ => None, - }; - match failure { - Some(failure) => anyhow::Error::new(failure).context(error), - None => anyhow::Error::new(error), + .map_err(exchange_error)?; + drop(refresh_token); + // Only the IdP uses the refresh token. Release its lease before resource-AS I/O. + let receipt = finish_idp(credentials); + let token = assertion + .exchange(&resource_http) + .await + .map_err(exchange_error)?; + Ok(( + EmaAccessToken { + access_token: token.access_token.secret().to_owned(), + expires_in: token.expires_in, + received_at: Instant::now(), + }, + receipt, + )) +} + +fn exchange_error(error: EmaError) -> anyhow::Error { + let failure = match &error { + EmaError::InvalidGrant(EmaExchangeStage::IdentityProvider) => { + Some(EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::EnterpriseIdentity, + }) } - })?; - Ok(EmaAccessToken { - access_token: token.access_token.secret().to_owned(), - expires_in: token.expires_in, - received_at: Instant::now(), - }) + EmaError::InvalidGrant(EmaExchangeStage::ResourceAuthorizationServer) => { + Some(EmaAuthFailure::InvalidGrant { + grant_source: EmaInvalidGrantSource::ResourceAuthorization, + }) + } + EmaError::InsufficientUserAuthentication(_) => { + Some(EmaAuthFailure::InsufficientUserAuthentication) + } + _ => None, + }; + match failure { + Some(failure) => anyhow::Error::new(failure).context(error), + None => anyhow::Error::new(error), + } } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/ema_exchange_tests.rs b/codex-rs/rmcp-client/src/ema_exchange_tests.rs index 4c21b76d25..40d22f96df 100644 --- a/codex-rs/rmcp-client/src/ema_exchange_tests.rs +++ b/codex-rs/rmcp-client/src/ema_exchange_tests.rs @@ -75,20 +75,25 @@ fn token_response() -> Value { async fn exchange(server: &MockServer, refresh_token: &str) -> Result { let issuer = format!("{}/idp", server.uri()); let audience = format!("{}/as", server.uri()); - exchange_id_jag(EmaIdJagExchangeRequest { - resource: &format!("{}/mcp", server.uri()), - scopes: &[], - mcp_client_id: "mcp-client", - authorization_server_issuer: &audience, - authorization_server_token_endpoint: &format!("{audience}/token"), - idp_token_endpoint: &format!("{issuer}/token"), - idp_issuer: &issuer, - idp_client_id: "idp-client", - refresh_token: refresh_token.to_string(), - idp_http_client: http_client(), - resource_http_client: http_client(), - }) + exchange_id_jag( + EmaIdJagExchangeRequest { + resource: &format!("{}/mcp", server.uri()), + scopes: &[], + mcp_client_id: "mcp-client", + authorization_server_issuer: &audience, + authorization_server_token_endpoint: &format!("{audience}/token"), + idp_token_endpoint: &format!("{issuer}/token"), + idp_issuer: &issuer, + idp_client_id: "idp-client", + refresh_token: refresh_token.to_string(), + idp_http_client: http_client(), + resource_http_client: http_client(), + }, + (), + std::convert::identity, + ) .await + .map(|(token, ())| token) } #[tokio::test] @@ -122,21 +127,26 @@ async fn public_client_round_trip_preserves_signed_narrowing() -> Result<()> { .await; } let expected_subject = refresh_token.to_string(); - let result = exchange_id_jag(EmaIdJagExchangeRequest { - resource: &resource, - scopes, - mcp_client_id: "mcp-client", - authorization_server_issuer: &audience, - authorization_server_token_endpoint: &format!("{audience}/token"), - idp_token_endpoint: &format!("{issuer}/token"), - idp_issuer: &issuer, - idp_client_id: "idp-client", - refresh_token: refresh_token.to_string(), - idp_http_client: Arc::clone(&client), - resource_http_client: Arc::clone(&client), - }) + let result = exchange_id_jag( + EmaIdJagExchangeRequest { + resource: &resource, + scopes, + mcp_client_id: "mcp-client", + authorization_server_issuer: &audience, + authorization_server_token_endpoint: &format!("{audience}/token"), + idp_token_endpoint: &format!("{issuer}/token"), + idp_issuer: &issuer, + idp_client_id: "idp-client", + refresh_token: refresh_token.to_string(), + idp_http_client: Arc::clone(&client), + resource_http_client: Arc::clone(&client), + }, + (), + std::convert::identity, + ) .boxed() - .await; + .await + .map(|(token, ())| token); if !valid { assert!(result.is_err(), "invalid subject must fail before HTTP"); assert!( diff --git a/codex-rs/rmcp-client/src/ema_identity.rs b/codex-rs/rmcp-client/src/ema_identity.rs index 64a1d6f2da..937dcb3093 100644 --- a/codex-rs/rmcp-client/src/ema_identity.rs +++ b/codex-rs/rmcp-client/src/ema_identity.rs @@ -39,10 +39,60 @@ pub struct EmaIdpIdentityRequest<'a> { } /// An opaque IdP refresh token paired with the credentials admitted for its exchange. -pub struct EmaIdpIdentity { +pub struct EmaIdpIdentity { pub(crate) token_endpoint: String, pub(crate) refresh_token: String, - pub(crate) credentials: EmaCredentialLease, + pub(crate) credentials: L, +} + +impl EmaIdpIdentity { + /// Preserve an opaque credential lease when selecting an EMA credential source. + /// The mapping must retain its owned refresh guard until the IdP exchange finishes. + pub fn map_credentials(self, map: impl FnOnce(L) -> U) -> EmaIdpIdentity { + EmaIdpIdentity { + token_endpoint: self.token_endpoint, + refresh_token: self.refresh_token, + credentials: map(self.credentials), + } + } +} + +/// Resolve a refresh token issued by the account's validated OAuth authority. +/// The caller supplies trusted persisted issuer/client provenance and an ID/refresh-token +/// pair from one admitted account record, retaining its refresh guard in `credentials`. +/// ID-token claims corroborate that record; this does not verify a JWT signature or +/// authenticate an expired ID token. Only the refresh token is exchanged with the IdP. +pub async fn resolve_ema_account_idp_identity( + issuer: &str, + client_id: &str, + id_token: &str, + refresh_token: String, + credentials: L, + http_client: Arc, + redirect_mode: StreamableHttpRedirectMode, +) -> Result> { + if refresh_token.trim().is_empty() { + return Err(ema_reauthentication_required( + "Codex account has no refresh token; sign in again", + )); + } + crate::ema_claims::oidc_identity(id_token, issuer, client_id)?; + let token_endpoint = resolve_ema_idp_token_endpoint(issuer, http_client, redirect_mode).await?; + Ok(EmaIdpIdentity { + token_endpoint, + refresh_token, + credentials, + }) +} + +/// Validate account ID-token bindings without treating an expired ID token as an expired grant. +pub fn account_oidc_identity_is_usable(id_token: &str, issuer: &str, client_id: &str) -> bool { + crate::ema_claims::oidc_identity(id_token, issuer, client_id).is_ok() +} + +/// Only an explanatory legacy-login hint, never authority for RT exchange. +pub fn account_oidc_issuer_hint(id_token: &str, issuer: &str) -> bool { + crate::ema_claims::oidc_issuer_hint(id_token, issuer) } /// Returns whether stored credentials contain a refresh token bound to the configured login. @@ -110,30 +160,9 @@ async fn resolve_ema_idp_identity_in( bail!("stored enterprise IdP credentials do not match the configured issuer and client"); } stored_oidc_identity(credentials)?; - let client = OAuthHttpClientAdapter::new_with_redirect_mode( - request.http_client, - build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?, - request.issuer, - /*has_configured_headers*/ false, - request.redirect_mode, - )?; - let (_, metadata) = - resolve_ema_idp_authorization_manager(request.issuer, Arc::new(client)).await?; - for (name, expected) in [ - ( - "identity_chaining_requested_token_types_supported", - ID_JAG_TOKEN_TYPE, - ), - ("grant_types_supported", TOKEN_EXCHANGE_GRANT_TYPE), - ] { - if advertised_capability(metadata.additional_fields.get(name), expected, name)? - == Some(false) - { - bail!( - "enterprise IdP does not advertise the required ID-JAG token exchange capability" - ); - } - } + let token_endpoint = + resolve_ema_idp_token_endpoint(request.issuer, request.http_client, request.redirect_mode) + .await?; let credentials = request.credentials; let credential_lock = RefreshCredentialLock::acquire_for_server( &credentials.credentials.credentials().server_name, @@ -160,7 +189,7 @@ async fn resolve_ema_idp_identity_in( ) })?; Ok(EmaIdpIdentity { - token_endpoint: metadata.token_endpoint, + token_endpoint, refresh_token: refresh_token.secret().to_string(), credentials, }) @@ -169,6 +198,37 @@ async fn resolve_ema_idp_identity_in( .map_err(|_| anyhow!("enterprise IdP credential reread task failed"))? } +async fn resolve_ema_idp_token_endpoint( + issuer: &str, + http_client: Arc, + redirect_mode: StreamableHttpRedirectMode, +) -> Result { + let client = OAuthHttpClientAdapter::new_with_redirect_mode( + http_client, + build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?, + issuer, + /*has_configured_headers*/ false, + redirect_mode, + )?; + let (_, metadata) = resolve_ema_idp_authorization_manager(issuer, Arc::new(client)).await?; + for (name, expected) in [ + ( + "identity_chaining_requested_token_types_supported", + ID_JAG_TOKEN_TYPE, + ), + ("grant_types_supported", TOKEN_EXCHANGE_GRANT_TYPE), + ] { + if advertised_capability(metadata.additional_fields.get(name), expected, name)? + == Some(false) + { + bail!( + "enterprise IdP does not advertise the required ID-JAG token exchange capability" + ); + } + } + Ok(metadata.token_endpoint) +} + #[cfg(test)] #[path = "ema_identity_tests.rs"] mod tests; diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 615f64280e..469da7472c 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -56,6 +56,9 @@ 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::account_oidc_identity_is_usable; +pub use ema_identity::account_oidc_issuer_hint; +pub use ema_identity::resolve_ema_account_idp_identity; pub use ema_identity::resolve_ema_idp_identity; pub use ema_identity::stored_ema_identity_is_usable; pub use enterprise_oauth_login::EnterpriseOAuthCredentialGuard;