diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 44d3f85302..4716f03494 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -260,10 +260,13 @@ impl McpServerConfig { self.environment_id == DEFAULT_MCP_SERVER_ENVIRONMENT_ID } - /// Keeps local OAuth credentials compatible while isolating executor-owned servers. + /// Keeps local OAuth credentials compatible while reserving managed credential namespaces. pub fn oauth_credential_name<'a>(&self, server_name: &'a str) -> Cow<'a, str> { if self.is_local_environment() { - if server_name.starts_with("executor:") || server_name.starts_with("local:") { + if server_name.starts_with("executor:") + || server_name.starts_with("local:") + || server_name.starts_with("ema-idp:") + { Cow::Owned(format!("local:{server_name}")) } else { Cow::Borrowed(server_name) diff --git a/codex-rs/rmcp-client/src/ema_auth_policy.rs b/codex-rs/rmcp-client/src/ema_auth_policy.rs index 99d0688a70..42057b1f78 100644 --- a/codex-rs/rmcp-client/src/ema_auth_policy.rs +++ b/codex-rs/rmcp-client/src/ema_auth_policy.rs @@ -2,7 +2,9 @@ use anyhow::Context; use anyhow::Result; +use anyhow::anyhow; use anyhow::bail; +use serde_json::Value; use url::Host; use url::Url; @@ -23,6 +25,10 @@ pub enum EmaInvalidGrantSource { ResourceAuthorization, } +pub(crate) fn ema_reauthentication_required(message: &'static str) -> anyhow::Error { + anyhow::Error::new(EmaAuthFailure::ReauthenticationRequired).context(message) +} + pub(crate) fn safe_oauth_error_code(code: Option<&str>) -> &str { code.filter(|code| { matches!( @@ -43,6 +49,24 @@ pub(crate) fn safe_oauth_error_code(code: Option<&str>) -> &str { .unwrap_or("OAuth token request rejected") } +pub(crate) fn validate_ema_public_client_auth( + advertised_methods: Option<&Value>, + issuer_description: &str, +) -> Result<()> { + let advertised_methods = advertised_methods.ok_or_else(|| { + anyhow!( + "{issuer_description} does not explicitly advertise public-client token endpoint authentication" + ) + })?; + let methods = advertised_methods.as_array().ok_or_else(|| { + anyhow!("{issuer_description} advertised malformed token endpoint authentication methods") + })?; + if !methods.iter().any(|method| method.as_str() == Some("none")) { + bail!("{issuer_description} does not support public-client token endpoint authentication"); + } + Ok(()) +} + pub(crate) fn validate_ema_oauth_endpoint(endpoint: &str, description: &str) -> Result<()> { let url = Url::parse(endpoint).with_context(|| format!("{description} is not a valid URL"))?; validate_credential_destination(&url, description) @@ -63,3 +87,48 @@ fn validate_credential_destination(url: &Url, description: &str) -> Result<()> { } Ok(()) } + +pub(crate) fn advertised_capability( + value: Option<&Value>, + expected: &str, + description: &str, +) -> Result> { + value + .map(|value| { + let values = value + .as_array() + .ok_or_else(|| anyhow!("{description} is malformed"))?; + Ok(values.iter().any(|value| value.as_str() == Some(expected))) + }) + .transpose() +} + +/// Resource indicators must describe the configured MCP origin, query and path. +pub fn validate_ema_auth_resource(server_url: &str, resource: Option<&str>) -> Result<()> { + let server = Url::parse(server_url).context("enterprise MCP server URL is invalid")?; + validate_credential_destination(&server, "enterprise MCP server URL")?; + let Some(resource) = resource.filter(|resource| !resource.trim().is_empty()) else { + return Ok(()); + }; + let resource = Url::parse(resource).context("enterprise MCP resource indicator is invalid")?; + validate_credential_destination(&resource, "enterprise MCP resource indicator")?; + if resource.origin() != server.origin() || resource.query() != server.query() { + bail!( + "enterprise MCP resource indicator must match the configured MCP server origin and query" + ); + } + let resource_path = resource.path().trim_end_matches('/'); + let server_path = server.path().trim_end_matches('/'); + if server_path != resource_path + && !server_path + .strip_prefix(resource_path) + .is_some_and(|suffix| suffix.starts_with('/')) + { + bail!("enterprise MCP resource indicator path must contain the configured MCP server path"); + } + Ok(()) +} + +#[cfg(test)] +#[path = "ema_auth_policy_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/ema_auth_policy_tests.rs b/codex-rs/rmcp-client/src/ema_auth_policy_tests.rs new file mode 100644 index 0000000000..3bdaed367d --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_auth_policy_tests.rs @@ -0,0 +1,51 @@ +use pretty_assertions::assert_eq; +use serde_json::json; + +use super::*; + +#[test] +fn resource_origin_query_and_path_are_bound() { + let server = "https://mcp.example/enterprise/tools?tenant=one"; + for (resource, valid) in [ + ("https://mcp.example/enterprise?tenant=one", true), + ("https://other.example/enterprise?tenant=one", false), + ("https://mcp.example/enterprise-admin?tenant=one", false), + ("https://mcp.example/enterprise?tenant=two", false), + ("https://mcp.example/enterprise", false), + ("http://localhost:4000/enterprise?tenant=one", false), + ] { + assert_eq!( + validate_ema_auth_resource(server, Some(resource)).is_ok(), + valid, + "{resource}" + ); + } + for (server, valid) in [ + ("https://mcp.example/mcp", true), + ("http://localhost:4000/mcp", true), + ("http://127.0.0.1:4000/mcp", true), + ("http://mcp.example/mcp", false), + ] { + assert_eq!( + validate_ema_auth_resource(server, /*resource*/ None).is_ok(), + valid, + "{server}" + ); + } +} + +#[test] +fn public_clients_require_an_explicitly_advertised_auth_method() { + for (advertised, valid) in [ + (None, false), + (Some(json!(["none"])), true), + (Some(json!(["client_secret_basic", "none"])), true), + (Some(json!(["private_key_jwt"])), false), + (Some(json!("none")), false), + ] { + assert_eq!( + validate_ema_public_client_auth(advertised.as_ref(), "IdP").is_ok(), + valid + ); + } +} diff --git a/codex-rs/rmcp-client/src/ema_claims.rs b/codex-rs/rmcp-client/src/ema_claims.rs index f12c6f1ff3..886f148a32 100644 --- a/codex-rs/rmcp-client/src/ema_claims.rs +++ b/codex-rs/rmcp-client/src/ema_claims.rs @@ -59,6 +59,56 @@ fn signed_jwt(token: &str) -> Result<(JwtHeader, T)> { Ok((header, claims)) } +#[derive(Deserialize)] +pub(crate) struct OidcClaims { + iss: String, + sub: String, + aud: OAuthResource, + azp: Option, + exp: u64, +} + +pub(crate) fn oidc_identity( + assertion: &str, + expected_issuer: &str, + expected_audience: &str, +) -> Result { + 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"); + } + let (audience_matches, multiple_audiences) = match &claims.aud { + OAuthResource::Single(value) => (value == expected_audience, false), + OAuthResource::Multiple(values) => ( + values.iter().any(|value| value == expected_audience), + values.len() > 1, + ), + }; + if !audience_matches + || claims + .azp + .as_deref() + .is_some_and(|party| party != expected_audience) + || multiple_audiences && claims.azp.as_deref() != Some(expected_audience) + { + bail!("OIDC identity assertion audience or authorized party does not match the IdP client"); + } + Ok(claims) +} + +pub fn validate_oidc_identity_assertion( + assertion: &str, + expected_issuer: &str, + expected_audience: &str, +) -> Result<()> { + let claims = oidc_identity(assertion, expected_issuer, expected_audience)?; + let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); + if claims.exp <= now { + bail!("OIDC identity assertion is expired"); + } + Ok(()) +} + #[derive(Deserialize)] struct IdJagClaims { iss: String, diff --git a/codex-rs/rmcp-client/src/ema_exchange_tests.rs b/codex-rs/rmcp-client/src/ema_exchange_tests.rs index f23bdb75c3..4671790710 100644 --- a/codex-rs/rmcp-client/src/ema_exchange_tests.rs +++ b/codex-rs/rmcp-client/src/ema_exchange_tests.rs @@ -19,6 +19,7 @@ use wiremock::matchers::method; use wiremock::matchers::path; use super::*; +use crate::ema_claims::validate_oidc_identity_assertion; fn http_client() -> Arc { Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( @@ -348,6 +349,40 @@ fn signed_claims_and_resource_tokens_cannot_widen_authority() -> Result<()> { Ok(()) } +#[test] +fn identity_and_credential_destinations_are_bound() { + for endpoint in [ + "http://idp.example/token", + "https://user:pass@idp.example/token", + "https://idp.example/token#fragment", + ] { + assert!( + validate_ema_oauth_endpoint(endpoint, "IdP").is_err(), + "accepted {endpoint}" + ); + } + let original = claims("https://idp.example", "idp-client", "https://mcp.example"); + assert!( + validate_oidc_identity_assertion(&jwt(&original), "https://idp.example", "idp-client") + .is_ok() + ); + for (field, value) in [ + ("iss", json!("https://other.example")), + ("aud", json!(["idp-client", "other"])), + ("azp", json!("other")), + ("exp", json!(0)), + ("sub", json!("")), + ] { + let mut changed = original.clone(); + changed[field] = value; + assert!( + validate_oidc_identity_assertion(&jwt(&changed), "https://idp.example", "idp-client") + .is_err(), + "accepted changed {field}" + ); + } +} + #[tokio::test] async fn provider_errors_cannot_reflect_credentials() { const SENTINEL: &str = "secret-assertion-sentinel"; diff --git a/codex-rs/rmcp-client/src/ema_identity.rs b/codex-rs/rmcp-client/src/ema_identity.rs new file mode 100644 index 0000000000..2fac801ab1 --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_identity.rs @@ -0,0 +1,155 @@ +//! An enterprise IdP session is independent of Codex account authentication. + +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use anyhow::bail; +use codex_exec_server::HttpClient; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use oauth2::TokenResponse; +use rmcp::transport::AuthorizationManager; + +use crate::ema_auth_policy::advertised_capability; +use crate::ema_auth_policy::ema_reauthentication_required; +use crate::ema_auth_policy::validate_ema_oauth_endpoint; +use crate::ema_auth_policy::validate_ema_public_client_auth; +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; +use crate::utils::build_default_headers; + +pub struct EmaIdpIdentityRequest<'a> { + pub issuer: &'a str, + pub client_id: &'a str, + pub credentials: &'a StoredOAuthCredentialSnapshot, + pub http_client: Arc, + pub redirect_mode: StreamableHttpRedirectMode, +} + +/// An opaque IdP refresh token whose credential lock is held through token exchange. +#[allow(dead_code)] +pub struct EmaIdpIdentity { + pub(crate) token_endpoint: String, + pub(crate) refresh_token: String, + pub(crate) credential_lock: RefreshCredentialLock, +} + +/// Returns whether stored credentials contain a refresh token bound to the configured login. +pub fn stored_ema_identity_is_usable( + tokens: &StoredOAuthTokens, + issuer: &str, + client_id: &str, +) -> bool { + tokens.url == issuer + && tokens.bound_issuer() == Some(issuer) + && tokens.client_id == client_id + && tokens.has_refresh_token() + && stored_oidc_identity(tokens).is_ok() +} + +/// Resolve a stored refresh-token subject against the configured enterprise IdP metadata. +pub async fn resolve_ema_idp_identity( + request: EmaIdpIdentityRequest<'_>, +) -> Result { + resolve_ema_idp_identity_in(request, &DefaultKeyringStore).await +} + +async fn resolve_ema_idp_identity_in( + request: EmaIdpIdentityRequest<'_>, + keyring_store: &K, +) -> Result { + 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"); + } + validate_ema_oauth_endpoint(request.issuer, "enterprise IdP issuer")?; + let credentials = request.credentials.credentials(); + if credentials.url != request.issuer + || credentials.bound_issuer() != Some(request.issuer) + || credentials.client_id != request.client_id + { + 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 mut manager = AuthorizationManager::new_with_oauth_http_client( + request.issuer.to_string(), + Arc::new(client), + ) + .await + .context("failed to create enterprise IdP metadata discovery client")?; + manager.set_allow_missing_issuer(false); + let metadata = manager + .resolve_metadata() + .await + .context("failed to discover enterprise IdP authorization metadata")? + .metadata; + if metadata.issuer.as_deref() != Some(request.issuer) { + bail!("enterprise IdP authorization metadata issuer does not match configuration"); + } + validate_ema_oauth_endpoint(&metadata.token_endpoint, "enterprise IdP token endpoint")?; + 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" + ); + } + } + validate_ema_public_client_auth( + metadata + .additional_fields + .get("token_endpoint_auth_methods_supported"), + "enterprise IdP", + )?; + let credential_lock = + RefreshCredentialLock::acquire_for_server(&credentials.server_name, &credentials.url) + .await?; + let snapshot = request.credentials.clone(); + let keyring_store = keyring_store.clone(); + // The worker retains the guard even if its caller stops waiting for the reread. + tokio::task::spawn_blocking(move || { + let latest = snapshot.load_ema_credentials(&keyring_store)?; + let refresh_token = latest + .token_response + .0 + .refresh_token() + .filter(|token| !token.secret().trim().is_empty()) + .ok_or_else(|| { + ema_reauthentication_required( + "enterprise IdP session has no refresh token; sign in again", + ) + })?; + Ok(EmaIdpIdentity { + token_endpoint: metadata.token_endpoint, + refresh_token: refresh_token.secret().to_string(), + credential_lock, + }) + }) + .await + .map_err(|_| anyhow!("enterprise IdP credential reread task failed"))? +} + +#[cfg(test)] +#[path = "ema_identity_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/ema_identity_tests.rs b/codex-rs/rmcp-client/src/ema_identity_tests.rs new file mode 100644 index 0000000000..8e0abb7bab --- /dev/null +++ b/codex-rs/rmcp-client/src/ema_identity_tests.rs @@ -0,0 +1,359 @@ +use std::io; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +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; +use codex_keyring_store::CredentialStoreError; +use codex_keyring_store::tests::MockKeyringStore; +use futures::FutureExt; +use pretty_assertions::assert_eq; +use pretty_assertions::assert_ne; +use serde_json::json; +use tokio::sync::oneshot; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; +use crate::EmaAuthFailure; +use crate::WrappedOAuthTokenResponse; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::test_support::TempCodexHome; + +fn credentials(issuer: &str, subject: &str, expires_at: u64) -> StoredOAuthTokens { + let assertion = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(br#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode( + serde_json::to_vec(&json!({ + "iss":issuer,"aud":"idp-client","sub":subject,"exp":expires_at, + })) + .expect("claims") + ) + ); + StoredOAuthTokens { + server_name: "ema-idp:enterprise-test".to_string(), + url: issuer.to_string(), + issuer: Some(issuer.to_string()), + client_id: "idp-client".to_string(), + token_response: WrappedOAuthTokenResponse( + serde_json::from_value(json!({ + "access_token":"unused","token_type":"Bearer","id_token":assertion, + "refresh_token":"stored-refresh","scope":"openid offline_access", + })) + .expect("credentials"), + ), + expires_at: None, + } +} + +fn request<'a>( + issuer: &'a str, + credentials: &'a StoredOAuthCredentialSnapshot, +) -> EmaIdpIdentityRequest<'a> { + EmaIdpIdentityRequest { + issuer, + client_id: "idp-client", + credentials, + http_client: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + redirect_mode: StreamableHttpRedirectMode::Legacy, + } +} + +async fn discovery() -> (MockServer, String) { + let server = MockServer::start().await; + let issuer = format!("{}/idp", server.uri()); + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/idp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "issuer":issuer,"authorization_endpoint":format!("{issuer}/authorize"), + "token_endpoint":format!("{issuer}/token"), + "identity_chaining_requested_token_types_supported":[ID_JAG_TOKEN_TYPE], + "grant_types_supported":[TOKEN_EXCHANGE_GRANT_TYPE], + "token_endpoint_auth_methods_supported":["none"], + }))) + .mount(&server) + .await; + (server, issuer) +} + +const REREAD_TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); +const REREAD_PANIC_SENTINEL: &str = "credential-panic-payload-sentinel"; + +#[derive(Clone, Copy, Debug)] +enum ReadOutcome { + Stored, + BackendError, + Panic, +} + +#[derive(Clone, Debug)] +struct GatedKeyringStore { + inner: MockKeyringStore, + executor_thread: thread::ThreadId, + entered: Arc>>>, + release: Arc>>, +} + +impl KeyringStore for GatedKeyringStore { + fn load(&self, service: &str, account: &str) -> Result, CredentialStoreError> { + // Fail before blocking if a regression puts this read on the current-thread executor. + assert_ne!(thread::current().id(), self.executor_thread); + if let Some(entered) = self + .entered + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + { + let _ = entered.send(()); + } + let outcome = match self + .release + .lock() + .unwrap_or_else(PoisonError::into_inner) + .recv_timeout(REREAD_TEST_TIMEOUT) + { + Ok(outcome) => outcome, + Err(mpsc::RecvTimeoutError::Disconnected) => ReadOutcome::Stored, + Err(mpsc::RecvTimeoutError::Timeout) => panic!("credential reread gate timed out"), + }; + match outcome { + ReadOutcome::Stored => self.inner.load(service, account), + ReadOutcome::BackendError => Err(CredentialStoreError::new( + keyring::Error::PlatformFailure(Box::new(io::Error::new( + io::ErrorKind::PermissionDenied, + "credential backend unavailable", + ))), + )), + ReadOutcome::Panic => panic!("{REREAD_PANIC_SENTINEL}"), + } + } + + fn save(&self, service: &str, account: &str, value: &str) -> Result<(), CredentialStoreError> { + self.inner.save(service, account, value) + } + + fn delete(&self, service: &str, account: &str) -> Result { + self.inner.delete(service, account) + } +} + +#[tokio::test(flavor = "current_thread")] +async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() -> Result<()> { + let _home = TempCodexHome::new(); + let (_server, issuer) = discovery().await; + let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct); + let stored = credentials(&issuer, "user", /*expires_at*/ 0); + let snapshot = StoredOAuthCredentialSnapshot::new(stored.clone(), store); + for outcome in [ + ReadOutcome::Stored, + ReadOutcome::BackendError, + ReadOutcome::Panic, + ] { + let inner = MockKeyringStore::default(); + store.save(&inner, &stored.server_name, &stored)?; + let (entered_tx, entered_rx) = oneshot::channel(); + // Dropping the sole sender releases the worker on any early return or panic. + let (release_tx, release_rx) = mpsc::channel(); + let keyring = GatedKeyringStore { + inner, + executor_thread: thread::current().id(), + entered: Arc::new(Mutex::new(Some(entered_tx))), + release: Arc::new(Mutex::new(release_rx)), + }; + let mut identity = Box::pin(resolve_ema_idp_identity_in( + request(&issuer, &snapshot), + &keyring, + )); + tokio::select! { + result = &mut identity => { + result?; + bail!("credential reread completed before its gate was released"); + } + entered = tokio::time::timeout(REREAD_TEST_TIMEOUT, entered_rx) => { entered??; } + } + + if matches!(outcome, ReadOutcome::Stored) { + // The caller times out while the already-started blocking read remains gated. + assert!( + tokio::time::timeout(Duration::ZERO, identity) + .await + .is_err() + ); + assert!( + RefreshCredentialLock::acquire_for_server(&stored.server_name, &issuer) + .now_or_never() + .is_none() + ); + release_tx.send(outcome)?; + } else { + release_tx.send(outcome)?; + let error = tokio::time::timeout(REREAD_TEST_TIMEOUT, identity) + .await? + .err() + .context("credential reread should fail")?; + if matches!(outcome, ReadOutcome::Panic) { + assert_eq!( + format!("{error:#}"), + "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 + ) + })); + } + } + // Drain the detached read before TempCodexHome changes the process environment. + let _released = tokio::time::timeout( + REREAD_TEST_TIMEOUT, + RefreshCredentialLock::acquire_for_server(&stored.server_name, &issuer), + ) + .await??; + } + Ok(()) +} + +#[tokio::test] +async fn refresh_subject_rereads_pinned_credentials_after_id_token_expiry() -> Result<()> { + let _home = TempCodexHome::new(); + let (_server, issuer) = discovery().await; + let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct); + let stored = credentials(&issuer, "user", /*expires_at*/ 0); + let snapshot = StoredOAuthCredentialSnapshot::new(stored.clone(), store); + let keyring = MockKeyringStore::default(); + store.save(&keyring, &stored.server_name, &stored)?; + + let identity = resolve_ema_idp_identity_in(request(&issuer, &snapshot), &keyring).await?; + assert_eq!( + (&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(()) +} + +#[tokio::test] +async fn refresh_subject_rejects_removed_replaced_or_missing_credentials() -> Result<()> { + let _home = TempCodexHome::new(); + let (server, issuer) = discovery().await; + let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); + let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct); + let original = credentials(&issuer, "user", /*expires_at*/ 0); + for change in [ + "deleted", + "subject", + "issuer", + "client", + "refresh", + "login", + "missing-refresh", + "file", + ] { + let keyring = MockKeyringStore::default(); + let snapshot = StoredOAuthCredentialSnapshot::new( + original.clone(), + if change == "file" { + ResolvedOAuthCredentialStore::File + } else { + store + }, + ); + let mut latest = original.clone(); + match change { + "subject" => latest = credentials(&issuer, "other-user", /*expires_at*/ 0), + "issuer" => latest.issuer = Some("https://other.example".to_string()), + "client" => latest.client_id = "other-client".to_string(), + "refresh" => { + latest + .token_response + .0 + .set_refresh_token(Some(oauth2::RefreshToken::new( + "other-users-refresh".to_string(), + ))) + } + "login" => latest = credentials(&issuer, "user", now + 3600), + "missing-refresh" => latest.token_response.0.set_refresh_token(None), + "deleted" | "file" => {} + _ => panic!("unexpected change"), + } + if change != "deleted" { + store.save(&keyring, &original.server_name, &latest)?; + } + let error = resolve_ema_idp_identity_in(request(&issuer, &snapshot), &keyring) + .await + .err() + .expect("must not reuse the stale snapshot or fall back to its ID token"); + if change == "file" { + assert!(error.to_string().contains("require keyring storage")); + } else { + assert_eq!( + error.downcast_ref::(), + Some(&EmaAuthFailure::ReauthenticationRequired), + "{change}" + ); + } + } + assert!( + server + .received_requests() + .await + .expect("requests") + .iter() + .all(|request| request.method.as_str() == "GET") + ); + Ok(()) +} + +#[test] +fn stored_identity_usability_requires_a_bound_refresh_token_not_a_current_id_token() { + let issuer = "https://idp.example"; + let expired = credentials(issuer, "user", /*expires_at*/ 0); + assert!(stored_ema_identity_is_usable( + &expired, + issuer, + "idp-client" + )); + for change in ["url", "issuer", "client", "refresh"] { + let mut changed = expired.clone(); + match change { + "url" => changed.url = "https://other.example".to_string(), + "issuer" => changed.issuer = Some("https://other.example".to_string()), + "client" => changed.client_id = "other-client".to_string(), + "refresh" => changed.token_response.0.set_refresh_token(None), + _ => panic!("unexpected change"), + } + assert!(!stored_ema_identity_is_usable( + &changed, + issuer, + "idp-client" + )); + } +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 623b519abd..06ba74a7e0 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -3,6 +3,7 @@ mod elicitation_client_service; mod ema_auth_policy; mod ema_claims; mod ema_exchange; +mod ema_identity; mod event_notification_transport; mod executor_process_transport; mod http_client_adapter; @@ -33,9 +34,15 @@ pub use auth_status::discover_streamable_http_oauth; pub use codex_protocol::protocol::McpAuthStatus; pub use ema_auth_policy::EmaAuthFailure; pub use ema_auth_policy::EmaInvalidGrantSource; +pub use ema_auth_policy::validate_ema_auth_resource; +pub use ema_claims::validate_oidc_identity_assertion; pub use ema_exchange::EmaAccessToken; pub use ema_exchange::EmaIdJagExchangeRequest; pub use ema_exchange::exchange_id_jag; +pub use ema_identity::EmaIdpIdentity; +pub use ema_identity::EmaIdpIdentityRequest; +pub use ema_identity::resolve_ema_idp_identity; +pub use ema_identity::stored_ema_identity_is_usable; pub use event_notification_transport::EventNotificationReceiver; pub use http_client_adapter::StreamableHttpRedirectMode; pub use http_headers::with_http_headers_helper; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 6411551261..6b7eaae501 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,7 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +mod ema_identity; mod issuer_binding; mod refresh_lock; mod refresh_transaction; @@ -71,8 +72,10 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; +pub(crate) use self::ema_identity::stored_oidc_identity; 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; pub(crate) use self::refresh_transaction::install_tokens_in_manager; pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; pub(crate) use self::resolved_store::ResolvedOAuthTokens; @@ -136,6 +139,18 @@ impl PartialEq for StoredOAuthCredentialSnapshot { } impl StoredOAuthCredentialSnapshot { + pub(crate) fn new( + mut credentials: StoredOAuthTokens, + store: ResolvedOAuthCredentialStore, + ) -> Self { + credentials.token_response.0.set_expires_in(None); + Self { + credentials, + store, + store_was_contended: false, + } + } + /// Returns the normalized credentials originally read from the selected store. pub fn credentials(&self) -> &StoredOAuthTokens { &self.credentials @@ -219,7 +234,7 @@ pub struct WrappedOAuthTokenResponse(pub OAuthTokenResponse); impl PartialEq for WrappedOAuthTokenResponse { fn eq(&self, other: &Self) -> bool { - match (serde_json::to_string(self), serde_json::to_string(other)) { + match (serde_json::to_value(self), serde_json::to_value(other)) { (Ok(s1), Ok(s2)) => s1 == s2, _ => false, } @@ -283,13 +298,10 @@ pub fn stored_oauth_credential_snapshot( else { return Ok(None); }; - let mut credentials = resolved.tokens; - credentials.token_response.0.set_expires_in(None); - Ok(Some(StoredOAuthCredentialSnapshot { - credentials, - store: resolved.store, - store_was_contended: false, - })) + Ok(Some(StoredOAuthCredentialSnapshot::new( + resolved.tokens, + resolved.store, + ))) } fn oauth_store_is_contended(error: &Error) -> bool { @@ -965,6 +977,7 @@ fn token_needs_refresh(expires_at: Option) -> bool { fn compute_store_key(server_name: &str, server_url: &str) -> Result { let executor_owned = server_name.starts_with("executor:"); + let enterprise_owned = server_name.starts_with("ema-idp:"); let server_name = server_name.strip_prefix("local:").unwrap_or(server_name); let mut payload = JsonMap::new(); payload.insert( @@ -973,8 +986,21 @@ fn compute_store_key(server_name: &str, server_url: &str) -> Result { ); payload.insert("url".to_string(), Value::String(server_url.to_string())); payload.insert("headers".to_string(), Value::Object(JsonMap::new())); - - let truncated = sha_256_prefix(&Value::Object(payload))?; + let payload = if enterprise_owned { + // The OS keyring is shared across homes. Keep enterprise sessions + // isolated by Codex profile as well as authenticated user and workspace. + let codex_home = find_codex_home()?; + fs::create_dir_all(&codex_home)?; + payload.insert( + "codex_home".to_string(), + serde_json::to_value(codex_home.as_path().canonicalize()?)?, + ); + // Different binaries can enable different serde_json ordering features. + serde_json::to_value(payload.into_iter().collect::>())? + } else { + Value::Object(payload) + }; + let truncated = sha_256_prefix(&payload)?; let separator = if executor_owned { ':' } else { '|' }; Ok(format!("{server_name}{separator}{truncated}")) } diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity.rs b/codex-rs/rmcp-client/src/oauth/ema_identity.rs new file mode 100644 index 0000000000..6f804a245f --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/ema_identity.rs @@ -0,0 +1,78 @@ +//! Reread enterprise credentials without changing the connection's pinned identity. + +use anyhow::Result; +use anyhow::bail; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; + +use super::ResolvedOAuthCredentialStore; +use super::StoredOAuthCredentialSnapshot; +use super::StoredOAuthTokens; +use crate::ema_auth_policy::ema_reauthentication_required; +use crate::ema_claims::OidcClaims; +use crate::ema_claims::oidc_identity; + +pub(crate) fn stored_oidc_identity(tokens: &StoredOAuthTokens) -> Result { + let assertion = tokens + .token_response + .0 + .extra_fields() + .0 + .get("id_token") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + ema_reauthentication_required( + "enterprise IdP session has no OIDC ID token; sign in again", + ) + })?; + // The ID token binds the login identity. Its expiry does not determine + // whether the IdP will accept the independently valid refresh token. + oidc_identity(assertion, &tokens.url, &tokens.client_id).map_err(|error| { + ema_reauthentication_required("stored enterprise IdP identity is invalid; sign in again") + .context(error.to_string()) + }) +} + +impl StoredOAuthCredentialSnapshot { + /// Reject deletion or replacement before using this session's cached resource bearer. + /// Call from a blocking task: the pinned keyring backend may perform blocking I/O. + pub fn validate_current_ema_credentials(&self) -> Result<()> { + self.load_ema_credentials(&DefaultKeyringStore).map(|_| ()) + } + + /// Reread only the pinned keyring authority; exchange callers hold its credential lock. + pub(crate) fn load_ema_credentials( + &self, + keyring_store: &K, + ) -> Result { + if !matches!(self.store, ResolvedOAuthCredentialStore::Keyring(_)) { + bail!("enterprise IdP credentials require keyring storage"); + } + let previous = &self.credentials; + let mut latest = self + .store + .load(keyring_store, &previous.server_name, &previous.url)? + .ok_or_else(|| { + ema_reauthentication_required( + "enterprise IdP credentials were removed; sign in again", + ) + })?; + // There is no refresh-token rotation writer in the supported EMA profile. + // Pin the whole atomic login record, not just the claims of its ID token. + latest.token_response.0.set_expires_in(None); + if latest != *previous + || previous.bound_issuer() != Some(previous.url.as_str()) + || !latest.has_refresh_token() + { + return Err(ema_reauthentication_required( + "enterprise IdP identity changed; sign in again and reconnect", + )); + } + stored_oidc_identity(&latest)?; + Ok(latest) + } +} + +#[cfg(test)] +#[path = "ema_identity_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs new file mode 100644 index 0000000000..c66721f395 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs @@ -0,0 +1,85 @@ +use codex_config::types::AuthKeyringBackendKind; +use codex_keyring_store::tests::MockKeyringStore; +use serde_json::json; + +use super::*; +use crate::oauth::RefreshCredentialLock; +use crate::oauth::compute_store_key; +use crate::oauth::test_support::TempCodexHome; + +#[tokio::test] +async fn keyring_failure_does_not_reuse_the_pinned_refresh_token() -> Result<()> { + let _home = TempCodexHome::new(); + let tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": "ema-idp:keyring-failure", + "url": "https://idp.example", + "issuer": "https://idp.example", + "client_id": "client", + "token_response": { + "access_token": "unused", + "token_type": "Bearer", + "refresh_token": "stale-refresh", + }, + }))?; + let key = compute_store_key(&tokens.server_name, &tokens.url)?; + let _lock = RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?; + let snapshot = StoredOAuthCredentialSnapshot::new( + tokens, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + ); + let keyring = MockKeyringStore::default(); + keyring.set_error( + &key, + keyring::Error::Invalid("backend".into(), "unavailable".into()), + ); + let error = snapshot + .load_ema_credentials(&keyring) + .expect_err("keyring failure must be terminal"); + assert!(error.to_string().contains("refusing file fallback")); + Ok(()) +} + +#[test] +fn ordinary_oauth_names_cannot_alias_enterprise_credential_keys() -> Result<()> { + let _home = TempCodexHome::new(); + let issuer = "https://idp.example"; + let enterprise_name = "ema-idp:synthetic-identity"; + let ordinary: codex_config::McpServerConfig = serde_json::from_value(json!({ + "url": issuer, + "oauth": {"client_id": "idp-client"}, + }))?; + let ordinary_name = ordinary.oauth_credential_name(enterprise_name); + + pretty_assertions::assert_ne!( + compute_store_key(&ordinary_name, issuer)?, + compute_store_key(enterprise_name, issuer)?, + "an ordinary server name must not select the enterprise credential namespace" + ); + let legacy_key = compute_store_key("ordinary-server", issuer)?; + let legacy_hash = legacy_key.split_once('|').expect("legacy key separator").1; + pretty_assertions::assert_eq!( + compute_store_key(&ordinary_name, issuer)?, + format!("{enterprise_name}|{legacy_hash}"), + "escaping the reserved prefix preserves the pre-EMA ordinary credential key" + ); + + let keyring = MockKeyringStore::default(); + let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct); + let enterprise_tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": enterprise_name, "url": issuer, "issuer": issuer, + "client_id": "idp-client", "token_response": { + "access_token": "unused", "token_type": "Bearer", "refresh_token": "enterprise-refresh" + } + }))?; + store.save(&keyring, enterprise_name, &enterprise_tokens)?; + let mut ordinary_tokens = enterprise_tokens.clone(); + ordinary_tokens.server_name = ordinary_name.to_string(); + store.save(&keyring, &ordinary_name, &ordinary_tokens)?; + assert!(store.delete(&keyring, &ordinary_name, issuer)?); + pretty_assertions::assert_eq!( + store.load(&keyring, enterprise_name, issuer)?, + Some(enterprise_tokens), + "ordinary OAuth save/logout must not overwrite or remove the enterprise entry" + ); + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs index 9d1ecf9e60..e038dc2a1e 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -25,12 +25,12 @@ const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(/*millis*/ 50); // WouldBlock contention from a contender that merely started late and observed persisted tokens. const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; -pub(super) struct RefreshCredentialLock { +pub(crate) struct RefreshCredentialLock { _file: File, } impl RefreshCredentialLock { - pub(super) async fn acquire_for_server(server_name: &str, url: &str) -> Result { + pub(crate) async fn acquire_for_server(server_name: &str, url: &str) -> Result { let store_key = super::compute_store_key(server_name, url)?; let codex_home = find_codex_home()?; Self::acquire_in(&codex_home, &store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT)