mirror of
https://github.com/openai/codex.git
synced 2026-10-21 05:22:21 +00:00
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
This commit is contained in:
@@ -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<Duration>,
|
||||
) -> Result<HttpResponse, OAuthHttpClientError> {
|
||||
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::<AuthorizationMetadata>(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<dyn HttpClient>,
|
||||
@@ -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
|
||||
|
||||
@@ -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<AuthorizationManager> {
|
||||
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::<Vec<_>>(),
|
||||
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<HttpRequestResponse, ExecServerError>> {
|
||||
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";
|
||||
|
||||
@@ -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"))
|
||||
|
||||
Reference in New Issue
Block a user