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"))