use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use anyhow::Result; use codex_exec_server::HttpClient; use codex_protocol::protocol::McpAuthStatus; use futures::FutureExt; use http::HeaderMap; use http::header::AUTHORIZATION; use rmcp::transport::AuthorizationManager; use rmcp::transport::auth::AuthError; use tracing::debug; use crate::http_client_adapter::StreamableHttpRedirectMode; use crate::oauth::StoredOAuthTokenStatus; use crate::oauth::oauth_token_status; use crate::oauth_callback::McpOAuthCallbackMode; use crate::oauth_callback::callback_mode; use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); /// Timeout policy for OAuth metadata discovery through a supplied HTTP client. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum OAuthDiscoveryTimeout { /// Preserve the timeout requested by the OAuth implementation. Requested, /// Cap OAuth discovery requests at the supplied duration. Capped(Duration), } impl OAuthDiscoveryTimeout { /// Preserves the existing timeout for local OAuth discovery. pub const LOCAL: Self = Self::Capped(DISCOVERY_TIMEOUT); } #[derive(Debug, Clone, PartialEq, Eq)] pub struct StreamableHttpOAuthDiscovery { pub scopes_supported: Option>, pub callback_mode: McpOAuthCallbackMode, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum McpLoginRequirement { Login, Reauthentication, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum McpAuthState { Unsupported, Unknown, LoggedOut(McpLoginRequirement), BearerToken, OAuth, } impl From for McpAuthStatus { fn from(value: McpAuthState) -> Self { match value { McpAuthState::Unsupported => Self::Unsupported, McpAuthState::Unknown => Self::Unknown, McpAuthState::LoggedOut(_) => Self::NotLoggedIn, McpAuthState::BearerToken => Self::BearerToken, McpAuthState::OAuth => Self::OAuth, } } } enum AuthStatusCheck { Complete(McpAuthState), Discover(HeaderMap), } /// Determine authentication status while routing OAuth discovery through the /// provided HTTP client. #[allow(clippy::too_many_arguments)] pub async fn determine_streamable_http_auth_status( server_name: &str, url: &str, bearer_token_env_var: Option<&str>, http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, discovery_timeout: OAuthDiscoveryTimeout, redirect_mode: StreamableHttpRedirectMode, ) -> Result { let has_configured_headers = has_configured_headers(&http_headers, &env_http_headers); let default_headers = match auth_status_before_discovery( server_name, url, bearer_token_env_var, http_headers, env_http_headers, store_mode, keyring_backend_kind, )? { AuthStatusCheck::Complete(status) => return Ok(status), AuthStatusCheck::Discover(default_headers) => default_headers, }; determine_auth_status_from_discovery( server_name, url, discover_streamable_http_oauth_with_headers_and_http_client( url, default_headers, http_client, discovery_timeout, has_configured_headers, redirect_mode, ) .await, ) } /// Determine authentication status using only configured and stored credentials. /// /// Returns `None` when determining the status would require OAuth metadata discovery. pub fn determine_streamable_http_auth_status_from_credentials( server_name: &str, url: &str, bearer_token_env_var: Option<&str>, http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result> { match auth_status_before_discovery( server_name, url, bearer_token_env_var, http_headers, env_http_headers, store_mode, keyring_backend_kind, )? { AuthStatusCheck::Complete(status) => Ok(Some(status)), AuthStatusCheck::Discover(_) => Ok(None), } } fn auth_status_before_discovery( server_name: &str, url: &str, bearer_token_env_var: Option<&str>, http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { if bearer_token_env_var.is_some() { return Ok(AuthStatusCheck::Complete(McpAuthState::BearerToken)); } let default_headers = build_default_headers(http_headers, env_http_headers)?; if default_headers.contains_key(AUTHORIZATION) { return Ok(AuthStatusCheck::Complete(McpAuthState::BearerToken)); } match oauth_token_status(server_name, url, store_mode, keyring_backend_kind)? { StoredOAuthTokenStatus::Usable => { return Ok(AuthStatusCheck::Complete(McpAuthState::OAuth)); } StoredOAuthTokenStatus::AuthorizationRequired => { return Ok(AuthStatusCheck::Complete(McpAuthState::LoggedOut( McpLoginRequirement::Reauthentication, ))); } StoredOAuthTokenStatus::Missing => {} } Ok(AuthStatusCheck::Discover(default_headers)) } fn determine_auth_status_from_discovery( server_name: &str, url: &str, discovery: Result>, ) -> Result { match discovery { Ok(Some(_)) => Ok(McpAuthState::LoggedOut(McpLoginRequirement::Login)), Ok(None) => Ok(McpAuthState::Unsupported), Err(error) => { debug!( "failed to detect OAuth support for MCP server `{server_name}` at {url}: {error:?}" ); Err(error) } } } pub async fn discover_streamable_http_oauth( url: &str, http_headers: Option>, env_http_headers: Option>, http_client: Arc, discovery_timeout: OAuthDiscoveryTimeout, redirect_mode: StreamableHttpRedirectMode, ) -> Result> { let has_configured_headers = has_configured_headers(&http_headers, &env_http_headers); let default_headers = build_default_headers(http_headers, env_http_headers)?; discover_streamable_http_oauth_with_headers_and_http_client( url, default_headers, http_client, discovery_timeout, has_configured_headers, redirect_mode, ) .await } async fn discover_streamable_http_oauth_with_headers_and_http_client( url: &str, default_headers: HeaderMap, http_client: Arc, discovery_timeout: OAuthDiscoveryTimeout, has_configured_headers: bool, redirect_mode: StreamableHttpRedirectMode, ) -> Result> { let oauth_http_client = match discovery_timeout { OAuthDiscoveryTimeout::Requested => OAuthHttpClientAdapter::new_with_redirect_mode( http_client, default_headers, url, has_configured_headers, redirect_mode, )?, OAuthDiscoveryTimeout::Capped(max_timeout) => { OAuthHttpClientAdapter::new_with_max_timeout_and_redirect_mode( http_client, default_headers, url, max_timeout, has_configured_headers, redirect_mode, )? } }; let mut authorization_manager = AuthorizationManager::new_with_oauth_http_client(url, Arc::new(oauth_http_client)).await?; authorization_manager.set_allow_missing_issuer(true); discover_streamable_http_oauth_with_manager(&authorization_manager).await } fn has_configured_headers( http_headers: &Option>, env_http_headers: &Option>, ) -> bool { http_headers .as_ref() .is_some_and(|headers| !headers.is_empty()) || env_http_headers .as_ref() .is_some_and(|headers| !headers.is_empty()) } async fn discover_streamable_http_oauth_with_manager( authorization_manager: &AuthorizationManager, ) -> Result> { match authorization_manager.resolve_metadata().boxed().await { Ok(resolution) if !resolution.source.is_discovered() => Ok(None), Ok(resolution) => { let metadata = resolution.metadata; Ok(Some(StreamableHttpOAuthDiscovery { callback_mode: callback_mode(&metadata) .unwrap_or(McpOAuthCallbackMode::CallbackSpecific), scopes_supported: normalize_scopes(metadata.scopes_supported), })) } Err(AuthError::NoAuthorizationSupport) => Ok(None), Err(err) => Err(err.into()), } } fn normalize_scopes(scopes_supported: Option>) -> Option> { let scopes_supported = scopes_supported?; let mut normalized = Vec::new(); for scope in scopes_supported { let scope = scope.trim(); if scope.is_empty() { continue; } let scope = scope.to_string(); if !normalized.contains(&scope) { normalized.push(scope); } } if normalized.is_empty() { None } else { Some(normalized) } } #[cfg(test)] mod tests { use super::*; use axum::Json; use axum::Router; use axum::http::StatusCode; use axum::http::header::WWW_AUTHENTICATE; use axum::routing::get; use codex_exec_server::ExecServerError; use codex_exec_server::HttpRedirectPolicy; 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 serial_test::serial; use std::collections::HashMap; use std::ffi::OsString; use std::sync::Mutex; use tokio::task::JoinHandle; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; struct TestServer { url: String, handle: JoinHandle<()>, } fn test_http_client() -> Arc { Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( OutboundProxyPolicy::ReqwestDefault, ))) } impl Drop for TestServer { fn drop(&mut self) { self.handle.abort(); } } #[derive(Default)] struct RecordingHttpClient { headers: Mutex>>, redirect_policy: Mutex>, timeout_ms: Mutex>>, } impl HttpClient for RecordingHttpClient { fn http_request( &self, _params: HttpRequestParams, ) -> BoxFuture<'_, Result> { Box::pin(async { Err(ExecServerError::HttpRequest( "unexpected buffered request".to_string(), )) }) } fn http_request_stream( &self, params: HttpRequestParams, ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { *self .headers .lock() .expect("header recorder lock should not be poisoned") = Some( params .headers .iter() .map(|header| (header.name.clone(), header.value.clone())) .collect(), ); *self .timeout_ms .lock() .expect("timeout recorder lock should not be poisoned") = Some(params.timeout_ms); *self .redirect_policy .lock() .expect("redirect policy recorder lock should not be poisoned") = Some(params.redirect_policy); Box::pin(async { Err(ExecServerError::HttpRequest( "expected discovery request failure".to_string(), )) }) } } fn assert_recorded_discovery_failure(discovery: Result>) { let error = discovery.expect_err("the recording HTTP client rejects OAuth discovery"); assert!( matches!( error.downcast_ref::(), Some(AuthError::MetadataError(reason)) if reason.contains("expected discovery request failure") ), "OAuth discovery must preserve the executor transport failure: {error:#}" ); } async fn spawn_oauth_discovery_server(metadata: serde_json::Value) -> TestServer { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let address = listener.local_addr().expect("listener should have address"); let mut metadata = metadata; if let Some(metadata) = metadata.as_object_mut() { metadata .entry("issuer") .or_insert_with(|| format!("http://{address}/mcp").into()); } let app = Router::new().route( "/.well-known/oauth-authorization-server/mcp", get({ let metadata = metadata.clone(); move || { let metadata = metadata.clone(); async move { Json(metadata) } } }), ); let handle = tokio::spawn(async move { axum::serve(listener, app).await.expect("server should run"); }); TestServer { url: format!("http://{address}/mcp"), handle, } } struct EnvVarGuard { key: String, original: Option, } impl EnvVarGuard { fn set(key: &str, value: &str) -> Self { let original = std::env::var_os(key); unsafe { std::env::set_var(key, value); } Self { key: key.to_string(), original, } } } impl Drop for EnvVarGuard { fn drop(&mut self) { if let Some(value) = &self.original { unsafe { std::env::set_var(&self.key, value); } } else { unsafe { std::env::remove_var(&self.key); } } } } #[tokio::test] async fn determine_auth_status_uses_bearer_token_when_authorization_header_present() { let status = determine_streamable_http_auth_status( "server", "not-a-url", /*bearer_token_env_var*/ None, Some(HashMap::from([( "Authorization".to_string(), "Bearer token".to_string(), )])), /*env_http_headers*/ None, OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::default(), test_http_client(), OAuthDiscoveryTimeout::Requested, StreamableHttpRedirectMode::Legacy, ) .await .expect("status should compute"); assert_eq!(status, McpAuthState::BearerToken); } #[tokio::test] #[serial(auth_status_env)] async fn determine_auth_status_uses_bearer_token_when_env_authorization_header_present() { let _guard = EnvVarGuard::set("CODEX_RMCP_CLIENT_AUTH_STATUS_TEST_TOKEN", "Bearer token"); let status = determine_streamable_http_auth_status( "server", "not-a-url", /*bearer_token_env_var*/ None, /*http_headers*/ None, Some(HashMap::from([( "Authorization".to_string(), "CODEX_RMCP_CLIENT_AUTH_STATUS_TEST_TOKEN".to_string(), )])), OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::default(), test_http_client(), OAuthDiscoveryTimeout::Requested, StreamableHttpRedirectMode::Legacy, ) .await .expect("status should compute"); assert_eq!(status, McpAuthState::BearerToken); } #[tokio::test] async fn oauth_metadata_preserves_login_without_probing_anonymous_tools() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let address = listener.local_addr().expect("listener should have address"); let metadata = serde_json::json!({ "issuer": format!("http://{address}/mcp"), "authorization_endpoint": format!("http://{address}/authorize"), "token_endpoint": format!("http://{address}/token"), }); let app = Router::new() .route( "/mcp", get(|| async { StatusCode::METHOD_NOT_ALLOWED }).post( |Json(request): Json| async move { let result = match request["method"].as_str() { Some("initialize") => serde_json::json!({ "protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "oauth", "version": "1"}, }), Some("tools/list") => serde_json::json!({"tools": []}), _ => serde_json::json!({}), }; Json(serde_json::json!({ "jsonrpc": "2.0", "id": request["id"], "result": result, })) }, ), ) .route( "/.well-known/oauth-authorization-server/mcp", get(move || async move { Json(metadata) }), ); let server = tokio::spawn(async move { axum::serve(listener, app).await.expect("server should run"); }); let url = format!("http://{address}/mcp"); let discovery = discover_streamable_http_oauth( &url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await; assert_eq!( determine_auth_status_from_discovery("server", &url, discovery) .expect("auth status should compute"), McpAuthState::LoggedOut(McpLoginRequirement::Login) ); server.abort(); } #[tokio::test] async fn oauth_discovery_does_not_follow_cross_origin_redirects() { let redirect_target = MockServer::start().await; let redirect_url = format!("{}/redirect-target", redirect_target.uri()); Mock::given(method("GET")) .and(path("/redirect-target")) .and(header("x-api-key", "sensitive-key")) .respond_with(ResponseTemplate::new(200)) .expect(0) .mount(&redirect_target) .await; let resource_server = MockServer::start().await; Mock::given(method("GET")) .and(path("/mcp")) .and(header("x-api-key", "sensitive-key")) .respond_with( ResponseTemplate::new(302).insert_header("location", redirect_url.clone()), ) .expect(1) .mount(&resource_server) .await; let error = discover_streamable_http_oauth( &format!("{}/mcp", resource_server.uri()), Some(HashMap::from([( "x-api-key".to_string(), "sensitive-key".to_string(), )])), /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect_err("cross-origin OAuth discovery redirects must be rejected"); assert!( matches!( error.downcast_ref::(), Some(AuthError::MetadataError(reason)) if reason.contains("OAuth discovery redirect to non-same-origin URL rejected") && reason.contains(&redirect_url) ), "OAuth discovery must preserve the cross-origin redirect rejection: {error:#}" ); redirect_target.verify().await; resource_server.verify().await; } #[tokio::test] async fn determine_auth_status_preserves_transient_http_errors() { for status in [ StatusCode::REQUEST_TIMEOUT, StatusCode::TOO_EARLY, StatusCode::TOO_MANY_REQUESTS, ] { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/mcp")) .respond_with(ResponseTemplate::new(status.as_u16())) .expect(1) .mount(&server) .await; let error = determine_streamable_http_auth_status( "transient-http-error", &format!("{}/mcp", server.uri()), /*bearer_token_env_var*/ None, /*http_headers*/ None, /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect_err("transient OAuth discovery failures must not become unsupported access"); assert!( matches!( error.downcast_ref::(), Some(AuthError::MetadataError(reason)) if reason.contains(status.as_str()) ), "auth-status discovery must preserve HTTP {status}: {error:#}" ); server.verify().await; } } #[tokio::test] async fn discover_streamable_http_oauth_returns_normalized_scopes() { let server = spawn_oauth_discovery_server(serde_json::json!({ "authorization_endpoint": "https://example.com/authorize", "token_endpoint": "https://example.com/token", "authorization_response_iss_parameter_supported": true, "scopes_supported": ["profile", " email ", "profile", "", " "], })) .await; let discovery = discover_streamable_http_oauth( &server.url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect("discovery should succeed") .expect("oauth support should be detected"); assert_eq!( discovery, StreamableHttpOAuthDiscovery { scopes_supported: Some(vec!["profile".to_string(), "email".to_string()]), callback_mode: McpOAuthCallbackMode::IssuerBound, } ); } #[tokio::test] async fn issuer_support_without_a_metadata_issuer_falls_back_to_distinct_callbacks() { let server = spawn_oauth_discovery_server(serde_json::json!({ "issuer": null, "authorization_endpoint": "https://example.com/authorize", "token_endpoint": "https://example.com/token", "authorization_response_iss_parameter_supported": true, })) .await; let discovery = discover_streamable_http_oauth( &server.url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect("discovery should succeed") .expect("oauth support should be detected"); assert_eq!( discovery, StreamableHttpOAuthDiscovery { scopes_supported: None, callback_mode: McpOAuthCallbackMode::CallbackSpecific, } ); } #[tokio::test] async fn routed_oauth_discovery_caps_local_discovery_timeout() { let http_client = Arc::new(RecordingHttpClient::default()); let discovery = discover_streamable_http_oauth( "http://example.com/mcp", /*http_headers*/ None, /*env_http_headers*/ None, http_client.clone(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await; assert_recorded_discovery_failure(discovery); assert_eq!( *http_client .timeout_ms .lock() .expect("timeout recorder lock should not be poisoned"), Some(Some( u64::try_from(DISCOVERY_TIMEOUT.as_millis()) .expect("discovery timeout should fit in u64") )) ); } #[tokio::test] async fn routed_oauth_discovery_preserves_requested_timeout() { let http_client = Arc::new(RecordingHttpClient::default()); let discovery = discover_streamable_http_oauth( "http://example.com/mcp", /*http_headers*/ None, /*env_http_headers*/ None, http_client.clone(), OAuthDiscoveryTimeout::Requested, StreamableHttpRedirectMode::Legacy, ) .await; assert_recorded_discovery_failure(discovery); assert_eq!( *http_client .timeout_ms .lock() .expect("timeout recorder lock should not be poisoned"), Some(Some(30_000)) ); } #[tokio::test] async fn routed_agent_plugin_oauth_discovery_stops_with_configured_headers() { let http_client = Arc::new(RecordingHttpClient::default()); let discovery = discover_streamable_http_oauth( "http://example.com/mcp", Some(HashMap::from([( "X-Mcp-Discovery".to_string(), "configured-value".to_string(), )])), /*env_http_headers*/ None, http_client.clone(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::AgentPluginV1, ) .await; assert_recorded_discovery_failure(discovery); let headers = http_client .headers .lock() .expect("header recorder lock should not be poisoned") .clone() .expect("discovery should issue an HTTP request"); assert_eq!( headers .iter() .find(|(name, _)| name.eq_ignore_ascii_case("x-mcp-discovery")) .map(|(_, value)| value.as_str()), Some("configured-value") ); assert_eq!( *http_client .redirect_policy .lock() .expect("redirect policy recorder lock should not be poisoned"), Some(HttpRedirectPolicy::Stop) ); } #[tokio::test] async fn discover_streamable_http_oauth_follows_protected_resource_metadata() { let authorization_server = spawn_oauth_discovery_server(serde_json::json!({ "authorization_endpoint": "https://example.com/authorize", "token_endpoint": "https://example.com/token", "scopes_supported": ["read", " write ", "read"], })) .await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let address = listener.local_addr().expect("listener should have address"); let resource_metadata_url = format!("http://{address}/oauth-resource"); let challenge = format!("Bearer resource_metadata=\"{resource_metadata_url}\""); let authorization_server_url = authorization_server.url.clone(); let app = Router::new() .route( "/mcp", get(move || { let challenge = challenge.clone(); async move { (StatusCode::UNAUTHORIZED, [(WWW_AUTHENTICATE, challenge)]) } }), ) .route( "/oauth-resource", get(move || { let authorization_server_url = authorization_server_url.clone(); async move { Json(serde_json::json!({ "resource": format!("http://{address}/mcp"), "authorization_servers": [authorization_server_url], })) } }), ); let handle = tokio::spawn(async move { axum::serve(listener, app).await.expect("server should run"); }); let resource_server = TestServer { url: format!("http://{address}/mcp"), handle, }; let discovery = discover_streamable_http_oauth( &resource_server.url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect("discovery should succeed") .expect("oauth support should be detected"); assert_eq!( discovery.scopes_supported, Some(vec!["read".to_string(), "write".to_string()]) ); } #[tokio::test] async fn discover_streamable_http_oauth_ignores_empty_scopes() { let server = spawn_oauth_discovery_server(serde_json::json!({ "authorization_endpoint": "https://example.com/authorize", "token_endpoint": "https://example.com/token", "scopes_supported": ["", " "], })) .await; let discovery = discover_streamable_http_oauth( &server.url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect("discovery should succeed") .expect("oauth support should be detected"); assert_eq!(discovery.scopes_supported, None); } #[tokio::test] async fn supports_oauth_login_does_not_require_scopes_supported() { let server = spawn_oauth_discovery_server(serde_json::json!({ "authorization_endpoint": "https://example.com/authorize", "token_endpoint": "https://example.com/token", })) .await; let supported = discover_streamable_http_oauth( &server.url, /*http_headers*/ None, /*env_http_headers*/ None, test_http_client(), OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) .await .expect("support check should succeed") .is_some(); assert!(supported); } }