From 8e2afc09126c0cea4c282725fe68af43adad73d7 Mon Sep 17 00:00:00 2001 From: malsamiri-oai Date: Thu, 10 Sep 2026 20:38:21 +0000 Subject: [PATCH] Recover OAuth metadata discovery from 503 responses via OIDC (#44636) ## Why A `503` from an issuer's OAuth metadata endpoint can block discovery even when its OIDC metadata is available, preventing expired OAuth tokens from being refreshed during MCP startup. ## What changed Try the same issuer's OIDC discovery endpoints when an OAuth authorization-server metadata GET returns `503`. Support issuers with path components and keep all attempts within the original request's timeout. Reuse existing header, origin, and response-size checks, leave issuer validation to RMCP, and do not follow fallback redirects. Preserve discovery errors when fallback fails instead of enabling legacy endpoint fallback, and propagate terminal HTTP failures. ## Testing Add tests for OIDC candidate order, malformed metadata, issuer mismatches, header isolation, redirects, terminal failures, unaffected requests, and the shared timeout. Add startup coverage for refreshing expired persisted tokens after OIDC fallback in both legacy and coordinated refresh modes. GitOrigin-RevId: 6d94fd09e9ba26cb83e029171449e91410a0f6cb --- codex-rs/rmcp-client/src/oauth_http_client.rs | 92 +++++- .../src/oauth_http_client_security_tests.rs | 274 ++++++++++++++++++ .../tests/streamable_http_oauth_startup.rs | 27 ++ 3 files changed, 391 insertions(+), 2 deletions(-) diff --git a/codex-rs/rmcp-client/src/oauth_http_client.rs b/codex-rs/rmcp-client/src/oauth_http_client.rs index 9e4e21e93d..f827e1720f 100644 --- a/codex-rs/rmcp-client/src/oauth_http_client.rs +++ b/codex-rs/rmcp-client/src/oauth_http_client.rs @@ -74,6 +74,91 @@ pub(crate) struct OAuthHttpClientAdapter { } impl OAuthHttpClientAdapter { + /// Recover only a candidate-local 503 without turning failed discovery into + /// missing metadata (which would enable RMCP's legacy endpoint fallback). + async fn execute_with_metadata_fallback( + &self, + request: HttpRequest, + redirect_policy: OAuthHttpRedirectPolicy, + timeout: Option, + ) -> Result { + let issuer_path = request + .uri() + .path() + .strip_prefix("/.well-known/oauth-authorization-server") + .filter(|suffix| suffix.is_empty() || suffix.starts_with('/')); + let Some(issuer_path) = issuer_path.filter(|_| { + request.method() == Method::GET + && matches!(redirect_policy, OAuthHttpRedirectPolicy::Stop) + }) else { + return self + .execute_request(request, redirect_policy, timeout) + .await; + }; + let mut candidates = vec![format!("/.well-known/openid-configuration{issuer_path}")]; + if !issuer_path.is_empty() { + candidates.push(format!("{issuer_path}/.well-known/openid-configuration")); + } + let mut candidate_url = + Url::parse(&request.uri().to_string()).map_err(oauth_http_client_error)?; + candidate_url.set_query(None); + candidate_url.set_fragment(None); + let operation = async { + let original = self + .execute_request(request.clone(), redirect_policy, timeout) + .await?; + if original.status() != StatusCode::SERVICE_UNAVAILABLE { + return Ok(original); + } + // These are the same issuer's OIDC candidates, in RMCP's discovery + // order. Reuse the adapter's header, origin and body-size checks. + // Do not follow redirects here: RMCP owns discovery redirect policy. + for path in candidates { + candidate_url.set_path(&path); + let mut candidate = request.clone(); + *candidate.uri_mut() = candidate_url + .as_str() + .parse() + .map_err(oauth_http_client_error)?; + let response = self + .execute_request(candidate, OAuthHttpRedirectPolicy::Stop, timeout) + .await?; + match response.status() { + StatusCode::OK => { + if serde_json::from_slice::(response.body()).is_ok() + { + // RMCP still validates the exact expected issuer + // before using these endpoints or saved credentials. + return Ok(response); + } + } + StatusCode::NOT_FOUND + | StatusCode::METHOD_NOT_ALLOWED + | StatusCode::SERVICE_UNAVAILABLE => {} + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_EARLY + | StatusCode::TOO_MANY_REQUESTS => return Ok(response), + status if status.is_server_error() => return Ok(response), + _ => return Ok(original), + } + } + Ok(original) + }; + // Additional candidates share the original request's time budget. + let timeout = match self.timeout { + OAuthDiscoveryTimeout::Requested => timeout, + OAuthDiscoveryTimeout::Capped(cap) => { + Some(timeout.map_or(cap, |timeout| timeout.min(cap))) + } + }; + match timeout { + Some(timeout) => tokio::time::timeout(timeout, operation) + .await + .map_err(|_| oauth_http_client_error(OAuthHttpClientAdapterError::TimedOut))?, + None => operation.await, + } + } + #[cfg(test)] pub(crate) fn new( http_client: Arc, @@ -330,8 +415,11 @@ fn oauth_redirect_policy( impl OAuthHttpClient for OAuthHttpClientAdapter { fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { Box::pin(async move { - let operation = - self.execute_request(request.request, request.redirect_policy, request.timeout); + let operation = self.execute_with_metadata_fallback( + request.request, + request.redirect_policy, + request.timeout, + ); match PROACTIVE_REFRESH_TIMEOUT.try_with(|duration| *duration) { Ok(duration) => tokio::time::timeout(duration, operation) .await diff --git a/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs b/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs index 736ffc8d9b..d4b8a64444 100644 --- a/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs +++ b/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs @@ -3,11 +3,20 @@ use std::sync::Arc; use std::time::Duration; use anyhow::Result; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; use codex_exec_server::RouteAwareHttpClient; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::AuthorizationMetadata; use rmcp::transport::auth::OAuthHttpRedirectPolicy; +use serde_json::json; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -21,6 +30,271 @@ use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::utils::MCP_USER_AGENT; use crate::utils::build_default_headers; +async fn metadata_manager( + resource: &MockServer, + authorization: &MockServer, + issuer_path: &str, +) -> Result { + let resource_url = format!("{}/mcp", resource.uri()); + Mock::given(method("GET")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(401).insert_header( + "www-authenticate", + format!( + "Bearer resource_metadata=\"{}/resource-metadata\"", + resource.uri() + ), + )) + .mount(resource) + .await; + Mock::given(method("GET")) + .and(path("/resource-metadata")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "resource": resource_url, + "authorization_servers": [format!("{}{issuer_path}", authorization.uri())] + }))) + .mount(resource) + .await; + let adapter = OAuthHttpClientAdapter::new( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + build_default_headers( + Some(HashMap::from([( + "X-Api-Key".to_string(), + "resource-secret".to_string(), + )])), + /*env_http_headers*/ None, + )?, + &resource_url, + ); + Ok(AuthorizationManager::new_with_oauth_http_client(resource_url, Arc::new(adapter)).await?) +} + +#[tokio::test] +async fn oauth_metadata_503_falls_back_to_oidc() -> Result<()> { + for (issuer_path, first_oidc_response) in [ + ("", None), + ("/adfs", Some(ResponseTemplate::new(404))), + ("/adfs", Some(ResponseTemplate::new(405))), + ("/adfs", Some(ResponseTemplate::new(503))), + ( + "/adfs", + Some(ResponseTemplate::new(200).set_body_string("not metadata")), + ), + ] { + let resource = MockServer::start().await; + let authorization = MockServer::start().await; + let manager = metadata_manager(&resource, &authorization, issuer_path).await?; + let oauth_path = format!("/.well-known/oauth-authorization-server{issuer_path}"); + let oidc_path = format!("{issuer_path}/.well-known/openid-configuration"); + let metadata = json!({ + "issuer": format!("{}{issuer_path}", authorization.uri()), + "authorization_endpoint": format!("{}/authorize", authorization.uri()), + "token_endpoint": format!("{}/token", authorization.uri()) + }); + Mock::given(method("GET")) + .and(path(oauth_path.clone())) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&authorization) + .await; + if let Some(response) = first_oidc_response { + Mock::given(method("GET")) + .and(path(format!( + "/.well-known/openid-configuration{issuer_path}" + ))) + .respond_with(response) + .expect(1) + .mount(&authorization) + .await; + } + Mock::given(method("GET")) + .and(path(oidc_path.clone())) + .respond_with(ResponseTemplate::new(200).set_body_json(&metadata)) + .expect(1) + .mount(&authorization) + .await; + + let resolved = manager.resolve_metadata().await?; + let expected: AuthorizationMetadata = serde_json::from_value(metadata)?; + assert_eq!( + serde_json::to_value(resolved.metadata)?, + serde_json::to_value(expected)? + ); + let requests = authorization.received_requests().await.unwrap(); + let mut expected_paths = vec![oauth_path]; + if !issuer_path.is_empty() { + expected_paths.push(format!("/.well-known/openid-configuration{issuer_path}")); + } + expected_paths.push(oidc_path); + assert_eq!( + requests + .iter() + .map(|request| request.url.path().to_string()) + .collect::>(), + expected_paths + ); + assert!( + requests + .iter() + .all(|request| !request.headers.contains_key("x-api-key")) + ); + authorization.verify().await; + } + Ok(()) +} + +#[tokio::test] +async fn oauth_metadata_fallback_preserves_terminal_failures() -> Result<()> { + let resource = MockServer::start().await; + let authorization = MockServer::start().await; + let manager = metadata_manager(&resource, &authorization, "/adfs").await?; + let redirect_target = MockServer::start().await; + let wrong_issuer = json!({ + "issuer": format!("{}/other-tenant", authorization.uri()), + "authorization_endpoint": format!("{}/authorize", authorization.uri()), + "token_endpoint": format!("{}/token", authorization.uri()) + }); + for (fallback, expected_error, request_count) in [ + (ResponseTemplate::new(404), "503", 3), + ( + ResponseTemplate::new(200).set_body_json(wrong_issuer), + "issuer mismatch", + 2, + ), + ( + ResponseTemplate::new(307).insert_header("location", redirect_target.uri()), + "503", + 2, + ), + (ResponseTemplate::new(408), "408", 2), + (ResponseTemplate::new(425), "425", 2), + (ResponseTemplate::new(429), "429", 2), + (ResponseTemplate::new(500), "500", 2), + ] { + authorization.reset().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/adfs")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&authorization) + .await; + Mock::given(method("GET")) + .and(path("/.well-known/openid-configuration/adfs")) + .respond_with(fallback) + .expect(1) + .mount(&authorization) + .await; + let error = manager + .resolve_metadata() + .await + .expect_err("discovery must fail closed"); + assert!(error.to_string().contains(expected_error), "{error}"); + assert_eq!( + authorization.received_requests().await.unwrap().len(), + request_count + ); + authorization.verify().await; + } + assert_eq!(redirect_target.received_requests().await.unwrap().len(), 0); + Ok(()) +} + +#[tokio::test] +async fn oauth_metadata_fallback_does_not_retry_other_requests_or_statuses() -> Result<()> { + for (verb, request_path, status) in [ + ("POST", "/.well-known/oauth-authorization-server/adfs", 503), + ("GET", "/.well-known/oauth-protected-resource/mcp", 503), + ("GET", "/.well-known/oauth-authorization-server-other", 503), + ("GET", "/.well-known/oauth-authorization-server/adfs", 500), + ] { + let server = MockServer::start().await; + Mock::given(method(verb)) + .and(path(request_path)) + .respond_with(ResponseTemplate::new(status)) + .mount(&server) + .await; + let adapter = OAuthHttpClientAdapter::new( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + Default::default(), + &format!("{}/mcp", server.uri()), + ); + let response = adapter + .execute_with_metadata_fallback( + oauth2::http::Request::builder() + .method(verb) + .uri(format!("{}{request_path}", server.uri())) + .body(Vec::new())?, + OAuthHttpRedirectPolicy::Stop, + Some(Duration::from_secs(/*secs*/ 5)), + ) + .await + .map_err(|error| anyhow::anyhow!(error))?; + assert_eq!(response.status().as_u16(), status); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } + Ok(()) +} + +struct DelayedMetadataHttpClient; + +impl HttpClient for DelayedMetadataHttpClient { + fn http_request( + &self, + _params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + Box::pin(async { panic!("OAuth requests must stream responses") }) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + Box::pin(async move { + // Each response fits the per-request timeout. Only the shared + // deadline can prevent the fallback from finishing. + tokio::time::sleep(Duration::from_millis(/*millis*/ 100)).await; + Ok(( + HttpRequestResponse { + status: if params.url.contains("oauth-authorization-server") { + 503 + } else { + 404 + }, + headers: Vec::new(), + body: Vec::new().into(), + }, + HttpResponseBodyStream::from_chunks(Vec::new()), + )) + }) + } +} + +#[tokio::test] +async fn oauth_metadata_fallback_shares_the_original_timeout() -> Result<()> { + let adapter = OAuthHttpClientAdapter::new( + Arc::new(DelayedMetadataHttpClient), + Default::default(), + "https://resource.example/mcp", + ); + let error = adapter + .execute_with_metadata_fallback( + oauth2::http::Request::builder() + .method("GET") + .uri("https://issuer.example/.well-known/oauth-authorization-server/adfs") + .body(Vec::new())?, + OAuthHttpRedirectPolicy::Stop, + Some(Duration::from_millis(/*millis*/ 150)), + ) + .await + .expect_err("fallback requests must share the original deadline"); + assert!(error.to_string().contains("timed out"), "{error}"); + Ok(()) +} + #[tokio::test] async fn oauth_registration_redirects_never_forward_resource_only_headers() -> Result<()> { const RESOURCE_API_KEY: &str = "resource-api-key-secret"; diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index e41986541e..374e3232c2 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -68,6 +68,7 @@ const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; #[derive(Clone, Copy)] enum OAuthStartupScenario { DirectAuthorizationMetadata, + OidcMetadataAfter503, GatewayHeadersHelper, SameOriginGatewayHeadersHelper, ProtectedResourceMetadata, @@ -255,6 +256,18 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result .await } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refreshes_expired_persisted_token_after_oidc_fallback() -> anyhow::Result<()> { + for refresh_mode in [ + McpOAuthRefreshMode::Legacy, + McpOAuthRefreshMode::Coordinated, + ] { + assert_expired_token_refresh(OAuthStartupScenario::OidcMetadataAfter503, refresh_mode) + .await?; + } + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_oauth_with_gateway_headers_helper() -> anyhow::Result<()> { assert_expired_token_refresh( @@ -311,8 +324,22 @@ async fn assert_expired_token_refresh( "/mcp/", "/.well-known/oauth-authorization-server", ), + OAuthStartupScenario::OidcMetadataAfter503 => ( + resource_url.clone(), + "/mcp", + "/mcp/.well-known/openid-configuration", + ), }; + if matches!(scenario, OAuthStartupScenario::OidcMetadataAfter503) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(503)) + .expect(2) + .mount(&server) + .await; + } + if matches!(scenario, OAuthStartupScenario::ProtectedResourceMetadata) { let resource_metadata_url = format!("{}/resource-metadata", server.uri()); Mock::given(method("GET"))