From b090e901f8702f63a5fedbcdcf3f4d021d10a8a0 Mon Sep 17 00:00:00 2001 From: Nick Steele Date: Tue, 8 Sep 2026 15:33:06 +0000 Subject: [PATCH] Add staged enterprise OIDC login and coordinated logout (#43844) ## Why Enterprise sign-in needs to keep browser completion separate from credential storage so callers can recheck the active account and configuration before saving a grant. Logout must also prevent an earlier sign-in from restoring credentials, including from another process sharing `CODEX_HOME`. ## What changed - Add enterprise login APIs in `rmcp-client` that return an authorization URL and stage validated credentials for an explicit `commit_if` call. Store grants only in the keyring after rechecking caller authority under the credential lock. - Require a registered client ID, published metadata matching the configured issuer, HTTP loopback callbacks, a refresh token, and a valid OIDC identity assertion. Request `openid` and `offline_access` with `prompt=consent`, and omit MCP resource indicators from authorization and code exchange. - Persist a login generation under the credential lock so logout invalidates pending and staged sign-ins across processes, even when no grant is stored. - Keep credentials and account identifiers out of enterprise error chains and logs, and avoid logging callback payloads when the receiver has closed. ## Testing Add coverage for discovery validation, loopback callbacks, PKCE, staged keyring storage, cancellation, stale attempts, cross-process logout, and error/log privacy. Preserve ordinary MCP OAuth login without a refresh token, and adjust the terminal polling test deadline to include the minimum empty-poll wait. GitOrigin-RevId: 2a27b9a26505a2f6fdecce8877f6e2c21e148f72 --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/unified_exec/mod_tests.rs | 3 +- codex-rs/rmcp-client/Cargo.toml | 1 + codex-rs/rmcp-client/src/ema_identity.rs | 60 ++- .../rmcp-client/src/ema_identity_tests.rs | 84 +++- .../rmcp-client/src/enterprise_oauth_login.rs | 384 +++++++++++++++ .../src/enterprise_oauth_login_tests.rs | 466 ++++++++++++++++++ .../src/enterprise_oauth_logout_tests.rs | 195 ++++++++ codex-rs/rmcp-client/src/lib.rs | 7 + codex-rs/rmcp-client/src/oauth.rs | 29 +- .../src/oauth/enterprise_generation.rs | 85 ++++ .../rmcp-client/src/oauth/test_support.rs | 23 + .../rmcp-client/src/perform_oauth_login.rs | 178 +++++-- 13 files changed, 1435 insertions(+), 81 deletions(-) create mode 100644 codex-rs/rmcp-client/src/enterprise_oauth_login.rs create mode 100644 codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs create mode 100644 codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs create mode 100644 codex-rs/rmcp-client/src/oauth/enterprise_generation.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f642b4591b..fa30609770 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4327,6 +4327,7 @@ dependencies = [ "tiny_http", "tokio", "tracing", + "tracing-test", "url", "urlencoding", "webbrowser", diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index d090847a15..a6f488689b 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -951,8 +951,9 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()> assert!(original.interaction_lock().try_lock_owned().is_err()); } // Empty polling must complete without an approval response. + // The test deadline must allow the minimum empty-poll wait. tokio::time::timeout( - Duration::from_secs(/*secs*/ 5), + Duration::from_millis(MIN_EMPTY_YIELD_TIME_MS) + Duration::from_secs(/*secs*/ 5), write_stdin(&session, &turn, process_id, "", /*yield_time_ms*/ 250), ) .await??; diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 29890ff09a..fb3fd4a21b 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -72,6 +72,7 @@ codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serial_test = { workspace = true } tempfile = { workspace = true } +tracing-test = { workspace = true, features = ["no-env-filter"] } wiremock = { workspace = true } [target.'cfg(unix)'.dependencies] diff --git a/codex-rs/rmcp-client/src/ema_identity.rs b/codex-rs/rmcp-client/src/ema_identity.rs index 2fac801ab1..f906e553f7 100644 --- a/codex-rs/rmcp-client/src/ema_identity.rs +++ b/codex-rs/rmcp-client/src/ema_identity.rs @@ -1,4 +1,5 @@ -//! An enterprise IdP session is independent of Codex account authentication. +//! Shared enterprise IdP discovery and stored refresh-token resolution. +//! Login and token exchange require published metadata bound to the configured issuer. use std::sync::Arc; @@ -11,6 +12,8 @@ use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; use oauth2::TokenResponse; use rmcp::transport::AuthorizationManager; +use rmcp::transport::auth::AuthorizationMetadata; +use rmcp::transport::auth::OAuthHttpClient; use crate::ema_auth_policy::advertised_capability; use crate::ema_auth_policy::ema_reauthentication_required; @@ -55,6 +58,36 @@ pub fn stored_ema_identity_is_usable( && stored_oidc_identity(tokens).is_ok() } +pub(crate) async fn resolve_ema_idp_authorization_manager( + issuer: &str, + http_client: Arc, +) -> Result<(AuthorizationManager, AuthorizationMetadata)> { + validate_ema_oauth_endpoint(issuer, "enterprise IdP issuer")?; + let mut manager = AuthorizationManager::new_with_oauth_http_client(issuer, http_client) + .await + .context("failed to create enterprise IdP metadata discovery client")?; + manager.set_allow_missing_issuer(false); + let resolution = manager + .resolve_metadata() + .await + .context("failed to discover enterprise IdP authorization metadata")?; + if !resolution.source.is_discovered() { + bail!("enterprise IdP must publish authorization metadata"); + } + let metadata = resolution.metadata; + if metadata.issuer.as_deref() != Some(issuer) { + bail!("enterprise IdP authorization metadata issuer does not match configuration"); + } + validate_ema_oauth_endpoint(&metadata.token_endpoint, "enterprise IdP token endpoint")?; + validate_ema_public_client_auth( + metadata + .additional_fields + .get("token_endpoint_auth_methods_supported"), + "enterprise IdP", + )?; + Ok((manager, metadata)) +} + /// Resolve a stored refresh-token subject against the configured enterprise IdP metadata. pub async fn resolve_ema_idp_identity( request: EmaIdpIdentityRequest<'_>, @@ -69,7 +102,6 @@ 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"); } - 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) @@ -85,22 +117,8 @@ async fn resolve_ema_idp_identity_in( /*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")?; + let (_, metadata) = + resolve_ema_idp_authorization_manager(request.issuer, Arc::new(client)).await?; for (name, expected) in [ ( "identity_chaining_requested_token_types_supported", @@ -116,12 +134,6 @@ async fn resolve_ema_idp_identity_in( ); } } - 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?; diff --git a/codex-rs/rmcp-client/src/ema_identity_tests.rs b/codex-rs/rmcp-client/src/ema_identity_tests.rs index 8e0abb7bab..53c7ef497a 100644 --- a/codex-rs/rmcp-client/src/ema_identity_tests.rs +++ b/codex-rs/rmcp-client/src/ema_identity_tests.rs @@ -91,6 +91,83 @@ async fn discovery() -> (MockServer, String) { (server, issuer) } +#[tokio::test] +async fn enterprise_discovery_shares_idp_checks_but_keeps_login_endpoint_policy() -> Result<()> { + for case in [ + "valid", + "missing-discovery", + "issuer", + "missing-issuer", + "token-endpoint", + "public-client", + "authorization-endpoint", + ] { + let server = MockServer::start().await; + let issuer = format!("{}/idp", server.uri()); + // Login need not advertise ID-JAG exchange capabilities. + let mut metadata = json!({ + "issuer": issuer, "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint": format!("{issuer}/token"), + "token_endpoint_auth_methods_supported": ["none"], + "grant_types_supported": ["authorization_code"], + }); + match case { + "issuer" => metadata["issuer"] = json!("https://other.example"), + "missing-issuer" => { + metadata.as_object_mut().expect("metadata").remove("issuer"); + } + "token-endpoint" => metadata["token_endpoint"] = json!("http://unsafe.example/token"), + "public-client" => { + metadata["token_endpoint_auth_methods_supported"] = json!(["client_secret_basic"]) + } + "authorization-endpoint" => { + metadata["authorization_endpoint"] = json!("http://unsafe.example/authorize") + } + "valid" | "missing-discovery" => {} + _ => unreachable!(), + } + if case != "missing-discovery" { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/idp")) + .respond_with(ResponseTemplate::new(200).set_body_json(metadata)) + .mount(&server) + .await; + } + let client = Arc::new(OAuthHttpClientAdapter::new_with_redirect_mode( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + build_default_headers(/*http_headers*/ None, /*env_http_headers*/ None)?, + &issuer, + /*has_configured_headers*/ false, + StreamableHttpRedirectMode::Legacy, + )?); + let shared = resolve_ema_idp_authorization_manager(&issuer, client.clone()).await; + let login = crate::enterprise_oauth_login::resolve_enterprise_authorization_manager( + &issuer, client, + ) + .await; + assert_eq!( + (shared.is_ok(), login.is_ok()), + ( + matches!(case, "valid" | "authorization-endpoint"), + case == "valid" + ), + "{case}", + ); + if case == "missing-discovery" { + assert_eq!( + shared + .err() + .expect("reject synthesized metadata") + .to_string(), + "enterprise IdP must publish authorization metadata", + ); + } + } + Ok(()) +} + const REREAD_TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); const REREAD_PANIC_SENTINEL: &str = "credential-panic-payload-sentinel"; @@ -175,10 +252,9 @@ async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() - 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, - )); + let request = request(&issuer, &snapshot); + crate::oauth::test_support::warm_http_client(request.http_client.as_ref()).await?; + let mut identity = Box::pin(resolve_ema_idp_identity_in(request, &keyring)); tokio::select! { result = &mut identity => { result?; diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs new file mode 100644 index 0000000000..3ee6c6cc32 --- /dev/null +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs @@ -0,0 +1,384 @@ +//! Enterprise OIDC policy and staged, keyring-only credential commits. +//! Browser completion never persists a grant; its owner revalidates authority under the commit lock. + +use std::future::Future; +use std::net::IpAddr; +use std::net::Ipv4Addr; +use std::sync::Arc; + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::bail; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::HttpClient; +use http::Method; +use http::header::CONTENT_LENGTH; +use http::header::CONTENT_TYPE; +use oauth2::TokenResponse; +use rmcp::transport::AuthorizationManager; +use rmcp::transport::auth::AuthorizationMetadata; +use rmcp::transport::auth::OAuthHttpClient; +use rmcp::transport::auth::OAuthHttpClientFuture; +use rmcp::transport::auth::OAuthHttpRequest; +use tracing::instrument::WithSubscriber; +use url::Host; +use url::Url; + +use crate::StoredOAuthTokens; +use crate::ema_auth_policy::validate_ema_oauth_endpoint; +use crate::ema_claims::validate_oidc_identity_assertion; +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::RefreshCredentialLock; +use crate::oauth::delete_oauth_tokens_with_lock_held; +use crate::oauth::save_oauth_tokens_with_lock_held; +use crate::oauth::validate_authorization_server_endpoints; +use crate::oauth_client_registration::McpOAuthClientRegistration; +use crate::perform_oauth_login::OAuthHttpContext; +use crate::perform_oauth_login::OAuthLoginPurpose; +use crate::perform_oauth_login::OauthLoginFlow; + +/// An exclusive credential mutation guard. Hold it through primary account logout +/// so a competing process cannot commit between deleting the grant and signing out. +/// Coordination, like ordinary OAuth refresh, is scoped to the same CODEX_HOME. +pub struct EnterpriseOAuthCredentialGuard { + credential_name: String, + issuer: String, + keyring_backend: AuthKeyringBackendKind, + generation_file: EnterpriseOAuthGenerationFile, + _lock: RefreshCredentialLock, +} + +impl EnterpriseOAuthCredentialGuard { + pub async fn acquire( + credential_name: &str, + issuer: &str, + keyring_backend: AuthKeyringBackendKind, + ) -> Result { + let lock = RefreshCredentialLock::acquire_for_server(credential_name, issuer) + .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"))?; + Ok(Self { + credential_name: credential_name.to_owned(), + issuer: issuer.to_owned(), + keyring_backend, + generation_file, + _lock: lock, + }) + } + + /// Invalidate pending logins and delete the grant, if present. A false return + /// means no grant was stored; earlier login attempts are still invalidated. + pub fn delete_tokens(&self) -> Result { + // Invalidate attempts that have not stored a grant yet, including other processes. + // Persist before deletion so no successful logout can admit an earlier login. + self.generation_file + .replace() + .map_err(|_| anyhow!("failed to invalidate pending enterprise sign-ins"))?; + // Underlying keyring diagnostics include the account-scoped key. Suppress + // them and discard their error chain, not merely the outer error message. + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + delete_oauth_tokens_with_lock_held( + &self._lock, + &self.credential_name, + &self.issuer, + OAuthCredentialsStoreMode::Keyring, + self.keyring_backend, + ) + }) + .map_err(|_| anyhow!("failed to delete enterprise credentials")) + } +} + +/// Invalidate earlier enterprise logins across processes, then delete any stored grant. +pub async fn delete_enterprise_oauth_tokens( + credential_name: &str, + issuer: &str, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result { + EnterpriseOAuthCredentialGuard::acquire(credential_name, issuer, keyring_backend_kind) + .await? + .delete_tokens() +} + +/// A browser login with no detached persistence worker. Dropping this handle or +/// its wait future unblocks the callback listener and cannot write credentials. +pub struct EnterpriseOAuthLoginHandle { + flow: OauthLoginFlow, + keyring_backend: AuthKeyringBackendKind, + generation: EnterpriseOAuthGeneration, +} + +impl EnterpriseOAuthLoginHandle { + pub fn authorization_url(&self) -> String { + self.flow.authorization_url() + } + + pub async fn wait(self) -> Result { + let stored = self + .flow + .complete(/*emit_browser_url*/ false) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("enterprise IdP authorization failed"))?; + validate_enterprise_credentials(&stored)?; + Ok(EnterpriseOAuthCredentials { + stored, + keyring_backend: self.keyring_backend, + generation: self.generation, + }) + } +} + +/// A validated grant that is not yet stored. It intentionally does not expose tokens. +pub struct EnterpriseOAuthCredentials { + stored: StoredOAuthTokens, + keyring_backend: AuthKeyringBackendKind, + generation: EnterpriseOAuthGeneration, +} + +impl EnterpriseOAuthCredentials { + /// Revalidate the current attempt, account and configuration under the same + /// exclusive lock used by logout, then persist without an intervening await. + /// Return the authority proof so the caller can retire its attempt while still + /// holding the commit gate, before notifying clients or refreshing runtimes. + pub async fn commit_if(self, is_current: F) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let guard = EnterpriseOAuthCredentialGuard::acquire( + &self.stored.server_name, + &self.stored.url, + self.keyring_backend, + ) + .await?; + let generation = guard + .generation_file + .current() + .map_err(|_| anyhow!("failed to read enterprise login generation"))?; + if generation.as_ref() != Some(&self.generation) { + bail!("enterprise sign-in was invalidated by logout"); + } + let Some(authority) = is_current().await else { + bail!("enterprise sign-in no longer matches the active account or configuration"); + }; + tracing::subscriber::with_default(tracing::subscriber::NoSubscriber::default(), || { + save_oauth_tokens_with_lock_held( + &guard._lock, + &self.stored.server_name, + &self.stored, + OAuthCredentialsStoreMode::Keyring, + self.keyring_backend, + ) + }) + .map_err(|_| anyhow!("failed to store enterprise credentials"))?; + Ok(authority) + } +} + +/// Start a host-registered enterprise login without launching the browser or +/// persisting a credential. The caller owns both the attempt and its later commit. +pub async fn perform_enterprise_oauth_login_return_url( + request: EnterpriseOAuthLoginRequest<'_>, +) -> Result { + // Capture before discovery or browser setup, then release the lock while the user signs in. + let generation = { + let guard = EnterpriseOAuthCredentialGuard::acquire( + request.credential_name, + request.issuer, + request.keyring_backend_kind, + ) + .await?; + match guard.generation_file.current() { + Ok(Some(generation)) => generation, + Ok(None) => guard + .generation_file + .replace() + .map_err(|_| anyhow!("failed to initialize enterprise login generation"))?, + Err(_) => bail!("failed to read enterprise login generation"), + } + }; + let flow = OauthLoginFlow::new( + request.credential_name, + request.issuer, + OAuthCredentialsStoreMode::Keyring, + request.keyring_backend_kind, + OAuthHttpContext { + http_headers: None, + env_http_headers: None, + http_client: request.http_client, + redirect_mode: request.redirect_mode, + }, + &["openid".to_string(), "offline_access".to_string()], + Some(request.client_id), + OAuthLoginPurpose::EnterpriseIdp, + McpOAuthClientRegistration::Auto, + /*oauth_resource*/ None, + /*launch_browser*/ false, + request.callback_port, + request.callback_url, + /*global_callback_url*/ None, + request.timeout_secs, + ) + .with_subscriber(tracing::subscriber::NoSubscriber::default()) + .await + .map_err(|_| anyhow!("failed to start enterprise IdP authorization"))?; + Ok(EnterpriseOAuthLoginHandle { + flow, + keyring_backend: request.keyring_backend_kind, + generation, + }) +} + +pub(crate) fn enterprise_callback_settings( + issuer: &str, + client_id: Option<&str>, + callback_url: Option<&str>, + callback_port: Option, +) -> Result<(IpAddr, Option)> { + validate_ema_oauth_endpoint(issuer, "enterprise IdP issuer")?; + if client_id.is_none_or(|client_id| client_id.trim().is_empty()) { + bail!("enterprise IdP login requires its registered client ID"); + } + let ip = enterprise_callback_bind_ip(callback_url)?; + let registered_port = callback_url + .map(Url::parse) + .transpose()? + .and_then(|url| url.port()); + if callback_port + .zip(registered_port) + .is_some_and(|(configured, registered)| configured != registered) + { + bail!("enterprise IdP callback URL and listener specify different ports"); + } + Ok((ip, callback_port.or(registered_port))) +} + +fn enterprise_callback_bind_ip(callback_url: Option<&str>) -> Result { + let Some(callback_url) = callback_url else { + return Ok(Ipv4Addr::LOCALHOST.into()); + }; + validate_ema_oauth_endpoint(callback_url, "enterprise IdP callback URL")?; + let callback = Url::parse(callback_url)?; + if callback.scheme() == "http" { + match callback.host() { + Some(Host::Domain("localhost")) => return Ok(Ipv4Addr::LOCALHOST.into()), + Some(Host::Ipv4(ip)) if ip.is_loopback() => return Ok(ip.into()), + Some(Host::Ipv6(ip)) if ip.is_loopback() => return Ok(ip.into()), + _ => {} + } + } + bail!("enterprise IdP callback URL must use an HTTP loopback address") +} + +fn validate_enterprise_credentials(stored: &StoredOAuthTokens) -> Result<()> { + let credentials = &stored.token_response.0; + if credentials + .refresh_token() + .is_none_or(|refresh_token| refresh_token.secret().trim().is_empty()) + { + bail!("enterprise IdP login did not return a refresh token"); + } + let assertion = credentials + .extra_fields() + .0 + .get("id_token") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| anyhow!("enterprise IdP login did not return an OIDC identity assertion"))?; + validate_oidc_identity_assertion( + assertion, + stored.issuer.as_deref().unwrap_or(&stored.url), + &stored.client_id, + ) + .map_err(|_| anyhow!("enterprise IdP returned an invalid OIDC identity assertion")) +} + +pub(crate) async fn resolve_enterprise_authorization_manager( + issuer: &str, + http_client: Arc, +) -> Result<(AuthorizationManager, AuthorizationMetadata)> { + let (manager, metadata) = resolve_ema_idp_authorization_manager( + issuer, + Arc::new(EnterpriseOAuthHttpClient(http_client)), + ) + .await?; + validate_authorization_server_endpoints(&metadata)?; + validate_ema_oauth_endpoint( + &metadata.authorization_endpoint, + "enterprise IdP authorization endpoint", + )?; + Ok((manager, metadata)) +} + +pub(crate) fn enterprise_authorization_url(auth_url: &str) -> Result { + let mut url = Url::parse(auth_url)?; + let query = url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs( + url.query_pairs() + .filter(|(key, _)| key != "resource" && key != "prompt"), + ) + .append_pair("prompt", "consent") + .finish(); + url.set_query(Some(&query)); + Ok(url.to_string()) +} + +pub(crate) fn without_oauth_resource(encoded: &[u8]) -> String { + url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(url::form_urlencoded::parse(encoded).filter(|(key, _)| key != "resource")) + .finish() +} + +/// rmcp supplies a resource indicator for MCP OAuth, but the independent OIDC +/// login must not request the IdP issuer as a protected-resource audience. +pub(crate) struct EnterpriseOAuthHttpClient(pub(crate) Arc); + +impl OAuthHttpClient for EnterpriseOAuthHttpClient { + fn execute(&self, mut request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + if request.request.method() == Method::POST + && request + .request + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value.split(';').next().is_some_and(|mime| { + mime.trim() + .eq_ignore_ascii_case("application/x-www-form-urlencoded") + }) + }) + && url::form_urlencoded::parse(request.request.body()) + .any(|(key, value)| key == "grant_type" && value == "authorization_code") + { + let body = without_oauth_resource(request.request.body()).into_bytes(); + *request.request.body_mut() = body; + request.request.headers_mut().remove(CONTENT_LENGTH); + } + self.0.execute(request) + } +} + +/// Enterprise registration and host-owned login settings. The flow always uses +/// OpenID Connect, strict issuer validation, and keyring-only credential storage. +pub struct EnterpriseOAuthLoginRequest<'a> { + pub credential_name: &'a str, + pub issuer: &'a str, + pub client_id: &'a str, + pub keyring_backend_kind: AuthKeyringBackendKind, + pub callback_port: Option, + pub callback_url: Option<&'a str>, + pub timeout_secs: Option, + pub http_client: Arc, + pub redirect_mode: StreamableHttpRedirectMode, +} + +#[cfg(test)] +#[path = "enterprise_oauth_login_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs new file mode 100644 index 0000000000..20b4bde4c6 --- /dev/null +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs @@ -0,0 +1,466 @@ +//! Public enterprise login exercises the real callback and keyring adapter in an isolated process. + +use std::any::Any; +use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_exec_server::RouteAwareHttpClient; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use keyring::credential::Credential; +use keyring::credential::CredentialApi; +use keyring::credential::CredentialBuilderApi; +use keyring::credential::CredentialPersistence; +use keyring::mock::MockCredential; +use pretty_assertions::assert_eq; +use serde_json::json; +use sha2::Digest; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tracing_test::traced_test; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; + +const SECRET: &str = "enterprise-secret-sentinel"; +const CREDENTIAL_NAME: &str = "ema-idp:enterprise-secret-sentinel"; + +#[path = "enterprise_oauth_logout_tests.rs"] +mod logout; + +async fn isolated_process(test_name: &str) -> Result { + const CHILD: &str = "CODEX_ENTERPRISE_LOGIN_TEST_CHILD"; + if std::env::var_os(CHILD).is_some() { + return Ok(false); + } + let home = tempfile::tempdir()?; + let output = tokio::process::Command::new(std::env::current_exe()?) + .args(["--exact", test_name, "--nocapture"]) + .env(CHILD, "1") + .env("CODEX_HOME", home.path()) + .current_dir(home.path()) + .output() + .await?; + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + Ok(true) +} + +async fn login(issuer: &str, callback_url: Option<&str>) -> Result { + perform_enterprise_oauth_login_return_url(EnterpriseOAuthLoginRequest { + credential_name: CREDENTIAL_NAME, + issuer, + client_id: "enterprise-client", + keyring_backend_kind: AuthKeyringBackendKind::Direct, + callback_port: None, + callback_url, + timeout_secs: Some(5), + http_client: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + redirect_mode: StreamableHttpRedirectMode::Legacy, + }) + .await +} + +async fn metadata(server: &MockServer, issuer: &str) { + 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?prompt=none"), + "token_endpoint": format!("{issuer}/token"), "token_endpoint_auth_methods_supported": ["none"], + }))).mount(server).await; +} + +async fn callback(authorization_url: &str, issuer: &str, provider_error: bool) -> Result<()> { + let query = Url::parse(authorization_url)? + .query_pairs() + .into_owned() + .collect::>(); + assert_eq!(query.get("prompt").map(String::as_str), Some("consent")); + assert!(!query.contains_key("resource")); + let mut callback = Url::parse(&query["redirect_uri"])?; + let mut pairs = callback.query_pairs_mut(); + if provider_error { + pairs + .append_pair("error", SECRET) + .append_pair("error_description", SECRET); + } else { + pairs + .append_pair("code", SECRET) + .append_pair("state", &query["state"]) + .append_pair("iss", issuer); + } + drop(pairs); + let port = callback.port().expect("actual listener port"); + let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)).await?; + stream + .write_all( + format!( + "GET {}?{} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n", + callback.path(), + callback.query().expect("query") + ) + .as_bytes(), + ) + .await?; + stream.read_to_end(&mut Vec::new()).await?; + Ok(()) +} + +#[tokio::test] +#[traced_test] +async fn enterprise_callback_errors_and_sdk_logs_exclude_credentials() -> Result<()> { + if isolated_process( + "enterprise_oauth_login::tests::enterprise_callback_errors_and_sdk_logs_exclude_credentials", + ) + .await? + { + return Ok(()); + } + for (callback_error, token_success) in [(true, false), (false, false), (false, true)] { + let server = MockServer::start().await; + let issuer = format!("{}/idp", server.uri()); + metadata(&server, &issuer).await; + let response = if token_success { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": SECRET, "refresh_token": SECRET, "id_token": SECRET, "token_type": "Bearer", + })) + } else { + ResponseTemplate::new(400) + .set_body_json(json!({"error": "invalid_grant", "error_description": SECRET})) + }; + Mock::given(method("POST")) + .and(path("/idp/token")) + .respond_with(response) + .expect(u64::from(!callback_error)) + .mount(&server) + .await; + let login = login(&issuer, /*callback_url*/ None).await?; + callback(&login.authorization_url(), &issuer, callback_error).await?; + tracing::trace!(target: "rmcp::transport::auth", "SDK trace capture enabled"); + let error = login.wait().await.err().expect("reject provider response"); + assert!(!format!("{error} {error:?} {error:#}").contains(SECRET)); + assert!(logs_contain("SDK trace capture enabled")); + assert!(!logs_contain(SECRET)); + } + Ok(()) +} + +#[test] +fn enterprise_callback_requires_loopback() -> Result<()> { + for (callback, expected) in [ + (None, "127.0.0.1"), + (Some("http://localhost/callback"), "127.0.0.1"), + (Some("http://127.0.0.2/callback"), "127.0.0.2"), + (Some("http://[::1]/callback"), "::1"), + ] { + assert_eq!( + enterprise_callback_bind_ip(callback)?, + expected.parse::()? + ); + } + for callback in [ + "http://0.0.0.0/callback", + "http://[::]/callback", + "https://127.0.0.1/callback", + "http://remote.example/callback", + ] { + assert!(enterprise_callback_bind_ip(Some(callback)).is_err()); + } + Ok(()) +} + +#[tokio::test] +#[traced_test] +async fn enterprise_public_api_storage_and_privacy() -> Result<()> { + if isolated_process("enterprise_oauth_login::tests::enterprise_public_api_storage_and_privacy") + .await? + { + return Ok(()); + } + let keyring = TestKeyring::default(); + keyring::set_default_credential_builder(Box::new(keyring.clone())); + let server = MockServer::start().await; + let issuer = format!("{}/idp", server.uri()); + metadata(&server, &issuer).await; + assert!( + login(&issuer, Some(&format!("{}/callback", server.uri()))) + .await + .is_err(), + "an occupied registered callback port must not silently move" + ); + 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() + ) + ); + Mock::given(method("POST")) + .and(path("/idp/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token":SECRET,"refresh_token":SECRET,"id_token":assertion,"token_type":"Bearer", + }))) + .mount(&server) + .await; + + // Callback, exchange and DefaultKeyringStore all run through public entrypoints. + for (login_index, callback_url) in ["http://localhost/callback", "http://127.0.0.1/callback"] + .into_iter() + .enumerate() + { + let login = login(&issuer, Some(callback_url)).await?; + let authorization_url = login.authorization_url(); + let authorization_query = Url::parse(&authorization_url)? + .query_pairs() + .into_owned() + .collect::>(); + callback(&authorization_url, &issuer, /*provider_error*/ false).await?; + let credentials = login.wait().await?; + let token_requests = server + .received_requests() + .await + .expect("recorded OAuth requests") + .into_iter() + .filter(|request| { + request.method == http::Method::POST && request.url.path() == "/idp/token" + }) + .collect::>(); + assert_eq!(token_requests.len(), login_index + 1); + let token_form = url::form_urlencoded::parse(&token_requests[login_index].body) + .into_owned() + .collect::>(); + assert_eq!( + token_form.get("grant_type").map(String::as_str), + Some("authorization_code") + ); + assert!(!token_form.contains_key("resource")); + assert_eq!( + token_form.get("redirect_uri"), + Some( + authorization_query + .get("redirect_uri") + .expect("authorization redirect URI") + ) + ); + assert_eq!( + authorization_query + .get("code_challenge_method") + .map(String::as_str), + Some("S256") + ); + let verifier = token_form.get("code_verifier").expect("PKCE verifier"); + let challenge = URL_SAFE_NO_PAD.encode(sha2::Sha256::digest(verifier.as_bytes())); + assert_eq!(authorization_query.get("code_challenge"), Some(&challenge)); + assert!( + keyring + .values + .lock() + .unwrap() + .values() + .all(|value| value.get_secret().is_err()) + ); + let attempt = tokio::sync::Mutex::new(Some("active")); + { + let mut authority = credentials + .commit_if(|| async { Some(attempt.lock().await) }) + .await?; + assert!( + attempt.try_lock().is_err(), + "commit retains the attempt gate until the caller retires it" + ); + *authority = None; + } + assert_eq!(*attempt.lock().await, None); + let stored = tracing::subscriber::with_default( + tracing::subscriber::NoSubscriber::default(), + || { + crate::stored_oauth_credentials( + CREDENTIAL_NAME, + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) + }, + )? + .expect("stored grant"); + assert_eq!( + stored + .token_response + .0 + .refresh_token() + .map(|token| token.secret().as_str()), + Some(SECRET) + ); + assert!( + delete_enterprise_oauth_tokens( + CREDENTIAL_NAME, + &issuer, + AuthKeyringBackendKind::Direct + ) + .await? + ); + } + + // Cancellation while blocked on the actual credential lock cannot leave a + // detached persistence worker that writes after the other process releases it. + let canceled = login(&issuer, /*callback_url*/ None).await?; + callback( + &canceled.authorization_url(), + &issuer, + /*provider_error*/ false, + ) + .await?; + let canceled = canceled.wait().await?; + let guard = EnterpriseOAuthCredentialGuard::acquire( + CREDENTIAL_NAME, + &issuer, + AuthKeyringBackendKind::Direct, + ) + .await?; + let mut commit = Box::pin(canceled.commit_if(|| async { Some(()) })); + assert!(futures::poll!(&mut commit).is_pending()); + drop(commit); + drop(guard); + assert!( + keyring + .values + .lock() + .unwrap() + .values() + .all(|value| value.get_secret().is_err()) + ); + + // Rejected old attempts neither write nor delete a newer grant. + let old = login(&issuer, /*callback_url*/ None).await?; + callback( + &old.authorization_url(), + &issuer, + /*provider_error*/ false, + ) + .await?; + let old = old.wait().await?; + let winner = login(&issuer, /*callback_url*/ None).await?; + callback( + &winner.authorization_url(), + &issuer, + /*provider_error*/ false, + ) + .await?; + winner + .wait() + .await? + .commit_if(|| async { Some(()) }) + .await?; + assert!(old.commit_if(|| async { None::<()> }).await.is_err()); + assert!( + keyring + .values + .lock() + .unwrap() + .values() + .any(|value| value.get_secret().is_ok()) + ); + + // Inject raw account identifiers into the actual keyring adapter's error chain. + keyring.fail.store(true, Ordering::SeqCst); + let failed = login(&issuer, /*callback_url*/ None).await?; + callback( + &failed.authorization_url(), + &issuer, + /*provider_error*/ false, + ) + .await?; + let save_error = failed + .wait() + .await? + .commit_if(|| async { Some(()) }) + .await + .unwrap_err(); + let delete_error = + delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct) + .await + .unwrap_err(); + for error in [save_error, delete_error] { + assert!(!format!("{error} {error:?} {error:#}").contains(SECRET)); + } + keyring.fail.store(false, Ordering::SeqCst); + let home = std::path::PathBuf::from(std::env::var("CODEX_HOME")?); + assert!( + !home.join(".credentials.json").exists(), + "enterprise storage never falls back to plaintext" + ); + let locks = home.join("mcp-oauth-locks"); + std::fs::rename(&locks, home.join("held-locks"))?; + std::fs::write(&locks, b"not a directory")?; + let lock_error = + delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct) + .await + .unwrap_err(); + assert!(!format!("{lock_error} {lock_error:?} {lock_error:#}").contains(SECRET)); + assert!(!logs_contain(SECRET)); + Ok(()) +} + +#[derive(Clone, Default)] +struct TestKeyring { + values: Arc>>>, + fail: Arc, +} + +struct TestCredential(Arc); + +impl CredentialApi for TestCredential { + fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> { + self.0.set_secret(secret) + } + fn get_secret(&self) -> keyring::Result> { + self.0.get_secret() + } + fn delete_credential(&self) -> keyring::Result<()> { + self.0.delete_credential() + } + fn as_any(&self) -> &dyn Any { + self + } +} + +impl CredentialBuilderApi for TestKeyring { + fn build( + &self, + _target: Option<&str>, + _service: &str, + user: &str, + ) -> keyring::Result> { + if self.fail.load(Ordering::SeqCst) { + return Err(keyring::Error::Invalid("account".into(), user.into())); + } + Ok(Box::new(TestCredential( + self.values + .lock() + .unwrap() + .entry(user.into()) + .or_default() + .clone(), + ))) + } + fn as_any(&self) -> &dyn Any { + self + } + fn persistence(&self) -> CredentialPersistence { + CredentialPersistence::ProcessOnly + } +} diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs new file mode 100644 index 0000000000..8a74957e1d --- /dev/null +++ b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs @@ -0,0 +1,195 @@ +//! Cross-process logout exercises staged login and the production keyring adapter. + +use std::fs; +use std::path::PathBuf; + +use pretty_assertions::assert_eq; +use sha2::Digest; +use sha2::Sha256; + +use super::*; + +const TEST: &str = + "enterprise_oauth_login::tests::logout::logout_invalidates_pending_login_across_processes"; +const LOGOUT_ISSUER: &str = "CODEX_ENTERPRISE_LOGOUT_TEST_ISSUER"; + +#[tokio::test] +async fn logout_invalidates_pending_login_across_processes() -> Result<()> { + if isolated_process(TEST).await? { + return Ok(()); + } + let home = PathBuf::from(std::env::var("CODEX_HOME")?); + let keyring = FileKeyring(home.join("test-keyring")); + 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) + .await?; + return Ok(()); + } + let server = MockServer::start().await; + let issuer = format!("{}/idp", server.uri()); + metadata(&server, &issuer).await; + 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() + ) + ); + Mock::given(method("POST")) + .and(path("/idp/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token":SECRET,"refresh_token":SECRET,"id_token":assertion,"token_type":"Bearer", + }))) + .mount(&server) + .await; + + // The first logout has no grant to delete. The second deletes the fresh grant + // from the previous iteration. Both must invalidate pending callbacks and staged grants. + for _ in 0..2 { + let pending = login(&issuer, /*callback_url*/ None).await?; + let staged = complete_login(&issuer).await?; + let output = tokio::process::Command::new(std::env::current_exe()?) + .args(["--exact", TEST, "--nocapture"]) + .env(LOGOUT_ISSUER, &issuer) + .output() + .await?; + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!(stored(&issuer)?.is_none()); + // The logout process has exited; its invalidation must survive that exit. + assert!(staged.commit_if(|| async { Some(()) }).await.is_err()); + assert!(stored(&issuer)?.is_none()); + + // A fresh login after logout remains usable, including for the same account. + complete_login(&issuer) + .await? + .commit_if(|| async { Some(()) }) + .await?; + let fresh = stored(&issuer)?.expect("fresh grant"); + callback( + &pending.authorization_url(), + &issuer, + /*provider_error*/ false, + ) + .await?; + assert!( + pending + .wait() + .await? + .commit_if(|| async { Some(()) }) + .await + .is_err() + ); + assert_eq!(stored(&issuer)?, Some(fresh)); + } + + let stale = complete_login(&issuer).await?; + let generation_path = fs::read_dir(home.join("mcp-oauth-locks"))? + .filter_map(std::result::Result::ok) + .map(|entry| entry.path()) + .find(|path| { + path.extension() + .is_some_and(|ext| ext == "enterprise-generation") + }) + .expect("persistent generation"); + fs::write(&generation_path, b"incomplete write")?; + assert!(stale.commit_if(|| async { Some(()) }).await.is_err()); + assert!(login(&issuer, /*callback_url*/ None).await.is_err()); + + // A lost marker never revives a pre-logout attempt or a pre-reset generation. + fs::remove_file(&generation_path)?; + let stale = complete_login(&issuer).await?; + fs::remove_file(&generation_path)?; + complete_login(&issuer) + .await? + .commit_if(|| async { Some(()) }) + .await?; + assert!(stale.commit_if(|| async { Some(()) }).await.is_err()); + + // Metadata failures are logout errors and must leave the stored grant intact. + let before = stored(&issuer)?; + fs::remove_file(&generation_path)?; + fs::create_dir(&generation_path)?; + assert!( + delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct) + .await + .is_err() + ); + assert_eq!(stored(&issuer)?, before); + assert!(!home.join(".credentials.json").exists()); + Ok(()) +} + +async fn complete_login(issuer: &str) -> Result { + let handle = login(issuer, /*callback_url*/ None).await?; + callback( + &handle.authorization_url(), + issuer, + /*provider_error*/ false, + ) + .await?; + handle.wait().await +} + +fn stored(issuer: &str) -> Result> { + crate::stored_oauth_credentials( + CREDENTIAL_NAME, + issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + ) +} + +// Only synthetic fixture credentials are persisted here. Every process still uses +// DefaultKeyringStore and the production credential lock, serializer and logout API. +struct FileKeyring(PathBuf); +struct FileCredential(PathBuf); + +impl CredentialBuilderApi for FileKeyring { + fn build( + &self, + _target: Option<&str>, + _service: &str, + user: &str, + ) -> keyring::Result> { + let name = format!("{:x}", Sha256::digest(user.as_bytes())); + Ok(Box::new(FileCredential(self.0.join(name)))) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl CredentialApi for FileCredential { + fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> { + fs::write(&self.0, secret).map_err(keyring_error) + } + + fn get_secret(&self) -> keyring::Result> { + fs::read(&self.0).map_err(keyring_error) + } + + fn delete_credential(&self) -> keyring::Result<()> { + fs::remove_file(&self.0).map_err(keyring_error) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn keyring_error(error: std::io::Error) -> keyring::Error { + if error.kind() == std::io::ErrorKind::NotFound { + keyring::Error::NoEntry + } else { + keyring::Error::PlatformFailure(Box::new(error)) + } +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 7792e53806..56b278a580 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod ema_auth_policy; mod ema_claims; mod ema_exchange; mod ema_identity; +mod enterprise_oauth_login; mod event_notification_transport; mod executor_process_transport; mod http_client_adapter; @@ -49,6 +50,12 @@ 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 enterprise_oauth_login::EnterpriseOAuthCredentialGuard; +pub use enterprise_oauth_login::EnterpriseOAuthCredentials; +pub use enterprise_oauth_login::EnterpriseOAuthLoginHandle; +pub use enterprise_oauth_login::EnterpriseOAuthLoginRequest; +pub use enterprise_oauth_login::delete_enterprise_oauth_tokens; +pub use enterprise_oauth_login::perform_enterprise_oauth_login_return_url; 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 8d2882acc0..86822efa19 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -18,6 +18,7 @@ mod credential_store; mod ema_identity; +mod enterprise_generation; mod issuer_binding; mod refresh_lock; mod refresh_transaction; @@ -76,6 +77,8 @@ use codex_utils_home_dir::find_codex_home; pub(crate) use self::credential_store::OAuthCredentialStore; 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::issuer_binding::validate_authorization_server_endpoints; pub(crate) use self::issuer_binding::validate_refresh_token_issuer; pub(crate) use self::refresh_lock::RefreshCredentialLock; @@ -446,7 +449,18 @@ pub async fn save_oauth_tokens( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result<()> { - let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?; + let lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?; + save_oauth_tokens_with_lock_held(&lock, server_name, tokens, store_mode, keyring_backend_kind) +} + +/// Save while retaining the matching credential lock acquired by the caller. +pub(crate) fn save_oauth_tokens_with_lock_held( + _lock: &RefreshCredentialLock, + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { let keyring_store = DefaultKeyringStore; match store_mode { OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file( @@ -592,7 +606,18 @@ pub async fn delete_oauth_tokens( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { - let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?; + let lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?; + delete_oauth_tokens_with_lock_held(&lock, server_name, url, store_mode, keyring_backend_kind) +} + +/// Delete while retaining the matching credential lock acquired by the caller. +pub(crate) fn delete_oauth_tokens_with_lock_held( + _lock: &RefreshCredentialLock, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result { let keyring_store = DefaultKeyringStore; delete_oauth_tokens_from_keyring_and_file( &keyring_store, diff --git a/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs new file mode 100644 index 0000000000..cd227cfbeb --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs @@ -0,0 +1,85 @@ +//! Persistent, non-secret logout generations for staged enterprise logins. +//! All reads and writes require the credential lock; missing state never admits an old attempt. + +use std::fs::File; +use std::fs::OpenOptions; +use std::io::Read; +use std::io::Seek; +use std::io::Write; + +use anyhow::Result; +use anyhow::ensure; +use codex_utils_home_dir::find_codex_home; +use oauth2::CsrfToken; +use sha2::Digest; +use sha2::Sha256; + +use super::RefreshCredentialLock; + +#[derive(PartialEq, Eq)] +pub(crate) struct EnterpriseOAuthGeneration([u8; 32]); + +pub(crate) struct EnterpriseOAuthGenerationFile { + file: File, +} + +impl EnterpriseOAuthGenerationFile { + pub(crate) fn open( + credential_name: &str, + issuer: &str, + _lock: &RefreshCredentialLock, + ) -> Result { + let key = super::compute_store_key(credential_name, issuer)?; + let name = format!("{:x}.enterprise-generation", 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); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options + .mode(0o600) + .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK); + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); + } + let file = options.open(path)?; + ensure!( + file.metadata()?.is_file(), + "invalid enterprise generation file" + ); + Ok(Self { file }) + } + + pub(crate) fn current(&self) -> Result> { + let mut file = &self.file; + file.rewind()?; + match file.metadata()?.len() { + 0 => Ok(None), + 32 => { + let mut generation = [0; 32]; + file.read_exact(&mut generation)?; + Ok(Some(EnterpriseOAuthGeneration(generation))) + } + _ => anyhow::bail!("invalid enterprise generation file"), + } + } + + pub(crate) fn replace(&self) -> Result { + // Reuse OAuth's randomness without storing any credential. Random initialization + // also prevents an old attempt becoming valid if metadata is removed or truncated. + let random = CsrfToken::new_random_len(/*num_bytes*/ 32); + let generation = + EnterpriseOAuthGeneration(Sha256::digest(random.secret().as_bytes()).into()); + let mut file = &self.file; + file.set_len(/*size*/ 0)?; + file.rewind()?; + file.write_all(&generation.0)?; + file.sync_all()?; + Ok(generation) + } +} diff --git a/codex-rs/rmcp-client/src/oauth/test_support.rs b/codex-rs/rmcp-client/src/oauth/test_support.rs index 90176e0361..208d11b458 100644 --- a/codex-rs/rmcp-client/src/oauth/test_support.rs +++ b/codex-rs/rmcp-client/src/oauth/test_support.rs @@ -6,6 +6,29 @@ use std::sync::PoisonError; use tempfile::tempdir; +/// Build the actual client's Stop-policy route pool before measuring protocol deadlines. +/// Client construction loads platform/custom CA state and can be slow on developer hosts. +/// Use a separate local endpoint so warmup cannot affect the test server's request assertions; +/// the same production HTTP capability still applies all proxy and certificate policy. +pub(crate) async fn warm_http_client( + client: &dyn codex_exec_server::HttpClient, +) -> anyhow::Result<()> { + let server = wiremock::MockServer::start().await; + client + .http_request(codex_exec_server::HttpRequestParams { + method: "GET".to_string(), + url: server.uri(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: codex_exec_server::HttpRedirectPolicy::Stop, + request_id: "test-client-warmup".to_string(), + stream_response: false, + }) + .await?; + Ok(()) +} + /// Serializes tests that mutate process-wide CODEX_HOME. /// /// Keep OAuth tests on this one guard instead of defining per-module helpers; otherwise diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 67280978e2..22577db9e6 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::net::SocketAddr; use std::string::String; use std::sync::Arc; use std::time::Duration; @@ -23,6 +24,9 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; +use crate::enterprise_oauth_login::enterprise_authorization_url; +use crate::enterprise_oauth_login::enterprise_callback_settings; +use crate::enterprise_oauth_login::resolve_enterprise_authorization_manager; use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::oauth::compute_expires_at_millis; use crate::oauth::validate_authorization_server_endpoints; @@ -30,6 +34,7 @@ use crate::oauth_callback::McpOAuthCallbackMode; use crate::oauth_callback::append_callback_id_to_redirect_uri; use crate::oauth_callback::callback_id_from_server_url; use crate::oauth_callback::callback_mode; +use crate::oauth_callback::resolve_mcp_oauth_callback_url; use crate::oauth_callback::validate_callback_redirect; use crate::oauth_client_registration::McpOAuthClientRegistration; use crate::oauth_client_registration::PreparedOAuthLogin; @@ -40,11 +45,17 @@ use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; -struct OAuthHttpContext { - http_headers: Option>, - env_http_headers: Option>, - http_client: Arc, - redirect_mode: StreamableHttpRedirectMode, +#[derive(Clone, Copy)] +pub(crate) enum OAuthLoginPurpose { + Mcp, + EnterpriseIdp, +} + +pub(crate) struct OAuthHttpContext { + pub(crate) http_headers: Option>, + pub(crate) env_http_headers: Option>, + pub(crate) http_client: Arc, + pub(crate) redirect_mode: StreamableHttpRedirectMode, } struct CallbackServerGuard { @@ -197,6 +208,7 @@ async fn perform_oauth_login_with_browser_output( http_context, scopes, oauth_client_id, + OAuthLoginPurpose::Mcp, client_registration, oauth_resource, /*launch_browser*/ true, @@ -243,6 +255,7 @@ pub async fn perform_oauth_login_return_url( http_context, scopes, oauth_client_id, + OAuthLoginPurpose::Mcp, client_registration, oauth_resource, /*launch_browser*/ false, @@ -279,12 +292,15 @@ fn spawn_callback_server( if let Err(err) = request.respond(response) { eprintln!("Failed to respond to OAuth callback: {err}"); } - if let Err(err) = tx.send(CallbackResult::Success(OauthCallbackResult { - code, - state, - issuer, - })) { - eprintln!("Failed to send OAuth callback: {err:?}"); + if let Err(message) = send_oauth_callback( + tx, + CallbackResult::Success(OauthCallbackResult { + code, + state, + issuer, + }), + ) { + eprintln!("{message}"); } break; } @@ -293,8 +309,8 @@ fn spawn_callback_server( if let Err(err) = request.respond(response) { eprintln!("Failed to respond to OAuth callback: {err}"); } - if let Err(err) = tx.send(CallbackResult::Error(error)) { - eprintln!("Failed to send OAuth callback error: {err:?}"); + if let Err(message) = send_oauth_callback(tx, CallbackResult::Error(error)) { + eprintln!("{message}"); } break; } @@ -323,6 +339,14 @@ enum CallbackResult { Error(OAuthProviderError), } +fn send_oauth_callback( + tx: oneshot::Sender, + result: CallbackResult, +) -> std::result::Result<(), &'static str> { + tx.send(result) + .map_err(|_| "OAuth callback receiver closed") +} + #[derive(Debug, PartialEq, Eq)] enum CallbackOutcome { Success(OauthCallbackResult), @@ -405,7 +429,7 @@ impl OauthLoginHandle { } } -struct OauthLoginFlow { +pub(crate) struct OauthLoginFlow { auth_url: String, oauth_state: OAuthState, authorization_server_issuer: Option, @@ -502,7 +526,7 @@ fn callback_bind_host(callback_url: Option<&str>) -> &'static str { impl OauthLoginFlow { #[allow(clippy::too_many_arguments)] - async fn new( + pub(crate) async fn new( server_name: &str, server_url: &str, store_mode: OAuthCredentialsStoreMode, @@ -510,6 +534,7 @@ impl OauthLoginFlow { http_context: OAuthHttpContext, scopes: &[String], oauth_client_id: Option<&str>, + purpose: OAuthLoginPurpose, client_registration: McpOAuthClientRegistration, oauth_resource: Option<&str>, launch_browser: bool, @@ -521,6 +546,18 @@ impl OauthLoginFlow { const DEFAULT_OAUTH_TIMEOUT_SECS: i64 = 300; let callback_port = resolve_callback_port(callback_port)?; + let is_enterprise_idp = matches!(purpose, OAuthLoginPurpose::EnterpriseIdp); + let (enterprise_bind_ip, callback_port) = if is_enterprise_idp { + let (ip, port) = enterprise_callback_settings( + server_url, + oauth_client_id, + callback_url, + callback_port, + )?; + (Some(ip), port) + } else { + (None, callback_port) + }; let callback_id = callback_id_from_server_url(server_url)?; let oauth_client_id = oauth_client_id.filter(|client_id| !client_id.trim().is_empty()); let configured_callback = if oauth_client_id.is_some() { @@ -556,7 +593,10 @@ impl OauthLoginFlow { redirect_mode, )?); let registered_authorization = if oauth_client_id.is_some() { - Some(resolve_authorization_manager(server_url, Arc::clone(&oauth_http_client)).await?) + Some( + resolve_authorization_manager(server_url, Arc::clone(&oauth_http_client), purpose) + .await?, + ) } else { None }; @@ -564,8 +604,8 @@ impl OauthLoginFlow { .as_ref() .map(|(_, metadata)| callback_mode(metadata)) .transpose()?; - let use_legacy_fallback = registered_callback_mode - == Some(McpOAuthCallbackMode::CallbackSpecific) + let use_legacy_fallback = !is_enterprise_idp + && registered_callback_mode == Some(McpOAuthCallbackMode::CallbackSpecific) && configured_callback.as_ref().is_some_and(|callback_url| { callback_url .path_segments() @@ -584,18 +624,32 @@ impl OauthLoginFlow { callback_url }; - let bind_host = callback_bind_host(callback_url); + let bind_ip = match enterprise_bind_ip { + Some(ip) => ip, + None => callback_bind_host(callback_url).parse()?, + }; // Port zero asks the OS for a free ephemeral port; the resolved // redirect receives that port after the listener has been bound. - let bind_addr = match callback_port { - Some(port) => format!("{bind_host}:{port}"), - None => format!("{bind_host}:0"), - }; - let server = Arc::new(Server::http(&bind_addr).map_err(|err| anyhow!(err))?); + let bind_addr = SocketAddr::new(bind_ip, callback_port.unwrap_or(0)); + let server = Arc::new(Server::http(bind_addr).map_err(|err| anyhow!(err))?); let guard = CallbackServerGuard { server: Arc::clone(&server), }; let redirect_uri = resolve_redirect_uri(&server, callback_url)?; + let redirect_uri = if is_enterprise_idp { + let listener_port = server + .server_addr() + .to_ip() + .ok_or_else(|| anyhow!("unable to determine enterprise callback listener port"))? + .port(); + let mut redirect = Url::parse(&redirect_uri)?; + redirect + .set_port(Some(listener_port)) + .map_err(|()| anyhow!("invalid enterprise callback port"))?; + redirect.to_string() + } else { + redirect_uri + }; let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); let PreparedOAuthLogin { @@ -605,7 +659,13 @@ impl OauthLoginFlow { } = if let Some((oauth_client_id, (auth_manager, metadata))) = oauth_client_id.zip(registered_authorization) { - let redirect_uri = if callback_url.is_some() && !use_legacy_fallback { + let redirect_uri = if is_enterprise_idp { + resolve_mcp_oauth_callback_url( + server_url, + Some(&redirect_uri), + callback_mode(&metadata)?, + )? + } else if callback_url.is_some() && !use_legacy_fallback { redirect_uri } else { append_callback_id_to_redirect_uri(&redirect_uri, &callback_id)? @@ -617,6 +677,7 @@ impl OauthLoginFlow { &redirect_uri, &callback_id, oauth_client_id, + purpose, ) .await? } else { @@ -656,11 +717,24 @@ impl OauthLoginFlow { }) } - fn authorization_url(&self) -> String { + pub(crate) fn authorization_url(&self) -> String { self.auth_url.clone() } - async fn finish(mut self, emit_browser_url: bool) -> Result<()> { + async fn finish(self, emit_browser_url: bool) -> Result<()> { + let store_mode = self.store_mode; + let keyring_backend_kind = self.keyring_backend_kind; + let stored = self.complete(emit_browser_url).await?; + save_oauth_tokens( + &stored.server_name, + &stored, + store_mode, + keyring_backend_kind, + ) + .await + } + + pub(crate) async fn complete(mut self, emit_browser_url: bool) -> Result { if self.launch_browser { let server_name = &self.server_name; let auth_url = &self.auth_url; @@ -706,25 +780,15 @@ impl OauthLoginFlow { .context("failed to retrieve OAuth credentials")?; let credentials = credentials_opt .ok_or_else(|| anyhow!("OAuth provider did not return credentials"))?; - let expires_at = compute_expires_at_millis(&credentials); - let stored = StoredOAuthTokens { + Ok(StoredOAuthTokens { server_name: self.server_name.clone(), url: self.server_url.clone(), issuer: self.authorization_server_issuer.clone(), client_id, token_response: WrappedOAuthTokenResponse(credentials), expires_at, - }; - save_oauth_tokens( - &self.server_name, - &stored, - self.store_mode, - self.keyring_backend_kind, - ) - .await?; - - Ok(()) + }) } .await; @@ -733,16 +797,13 @@ impl OauthLoginFlow { } fn spawn(self) -> oneshot::Receiver> { - let server_name_for_logging = self.server_name.clone(); + let server_name = self.server_name.clone(); let (tx, rx) = oneshot::channel(); tokio::spawn(async move { let result = self.finish(/*emit_browser_url*/ false).await; - if let Err(err) = &result { - eprintln!( - "Failed to complete OAuth login for '{server_name_for_logging}': {err:#}" - ); + eprintln!("Failed to complete OAuth login for '{server_name}': {err:#}"); } let _ = tx.send(result); @@ -755,7 +816,11 @@ impl OauthLoginFlow { async fn resolve_authorization_manager( server_url: &str, http_client: Arc, + purpose: OAuthLoginPurpose, ) -> Result<(AuthorizationManager, AuthorizationMetadata)> { + if matches!(purpose, OAuthLoginPurpose::EnterpriseIdp) { + return resolve_enterprise_authorization_manager(server_url, http_client).await; + } let mut auth_manager = AuthorizationManager::new_with_oauth_http_client(server_url, http_client).await?; auth_manager.set_allow_missing_issuer(true); @@ -771,15 +836,21 @@ async fn start_authorization( redirect_uri: &str, callback_id: &str, oauth_client_id: &str, + purpose: OAuthLoginPurpose, ) -> Result { + let strict_enterprise_idp = matches!(purpose, OAuthLoginPurpose::EnterpriseIdp); let authorization_server_issuer = metadata.issuer.clone(); validate_callback_redirect(redirect_uri, callback_id, callback_mode(&metadata)?)?; auth_manager.set_metadata(metadata); - auth_manager.configure_client( - OAuthClientConfig::new(oauth_client_id, redirect_uri) - .with_scopes(scopes.iter().map(|scope| (*scope).to_string()).collect()), - )?; + let client_config = OAuthClientConfig::new(oauth_client_id, redirect_uri) + .with_scopes(scopes.iter().map(|scope| (*scope).to_string()).collect()); + auth_manager.configure_client(client_config)?; let auth_url = auth_manager.get_authorization_url(scopes).await?; + let auth_url = if strict_enterprise_idp { + enterprise_authorization_url(&auth_url)? + } else { + auth_url + }; Ok(PreparedOAuthLogin { oauth_state: OAuthState::Session(AuthorizationSession::for_scope_upgrade( @@ -835,6 +906,7 @@ mod tests { use codex_http_client::OutboundProxyPolicy; use futures::future::BoxFuture; use http::HeaderMap; + use oauth2::TokenResponse; use pretty_assertions::assert_eq; use serde_json::json; use tokio::net::TcpListener; @@ -844,6 +916,7 @@ mod tests { use super::McpOAuthClientRegistration; use super::OAuthHttpClientAdapter; use super::OAuthHttpContext; + use super::OAuthLoginPurpose; use super::OAuthProviderError; use super::OauthLoginFlow; use super::StreamableHttpRedirectMode; @@ -937,7 +1010,6 @@ mod tests { Json(json!({ "access_token": "test-access-token", "token_type": "Bearer", - "refresh_token": "test-refresh-token", })) }), ); @@ -981,7 +1053,7 @@ mod tests { } #[tokio::test] - async fn oauth_login_persists_discovered_issuer() -> anyhow::Result<()> { + async fn ordinary_oauth_login_persists_issuer_without_a_refresh_token() -> anyhow::Result<()> { let _env = TempCodexHome::new(); let (base_url, _registration_requests) = spawn_oauth_metadata_server().await; let server_url = format!("{base_url}/mcp"); @@ -1000,6 +1072,7 @@ mod tests { }, &[], Some("test-client"), + OAuthLoginPurpose::Mcp, McpOAuthClientRegistration::Auto, /*oauth_resource*/ None, /*launch_browser*/ false, @@ -1033,6 +1106,7 @@ mod tests { )? .expect("OAuth login should persist credentials"); assert_eq!(stored.issuer.as_deref(), Some(server_url.as_str())); + assert!(stored.token_response.0.refresh_token().is_none()); Ok(()) } @@ -1050,6 +1124,7 @@ mod tests { HeaderMap::new(), &format!("{base_url}/mcp"), )), + OAuthLoginPurpose::Mcp, ) .await .expect("resolve pre-registered OAuth metadata"); @@ -1060,6 +1135,7 @@ mod tests { redirect_uri, "configured-client", "eci-prd-pub-codex-123", + OAuthLoginPurpose::Mcp, ) .await .expect("start pre-registered OAuth authorization"); @@ -1157,6 +1233,7 @@ mod tests { HeaderMap::new(), &format!("{issuer}/mcp"), )), + OAuthLoginPurpose::Mcp, ) .await .expect("resolve issuer-aware authorization metadata"); @@ -1167,6 +1244,7 @@ mod tests { redirect_uri, "test-callback", "test-client", + OAuthLoginPurpose::Mcp, ) .await .expect("start issuer-aware authorization");