From 3a18ab9fa7dac717650b311250b55428ce4b4807 Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Fri, 24 Jul 2026 18:51:31 -0400 Subject: [PATCH] refactor: make oauth discovery reactive --- conformance/src/bin/client.rs | 110 +++++-- crates/rmcp/src/service/client.rs | 25 ++ crates/rmcp/src/transport/auth.rs | 282 ++++++++++++------ .../common/auth/streamable_http_client.rs | 188 ++++++++---- .../src/transport/streamable_http_client.rs | 26 +- examples/clients/src/auth/oauth_client.rs | 191 ++++++++---- 6 files changed, 581 insertions(+), 241 deletions(-) diff --git a/conformance/src/bin/client.rs b/conformance/src/bin/client.rs index 852ca0917..387b45b12 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -195,25 +195,73 @@ const REDIRECT_URI: &str = "http://localhost:3000/callback"; const SCOPE_STEP_UP_INITIAL_SCOPES: &[&str] = &["mcp:basic"]; const SCOPE_STEP_UP_ESCALATED_SCOPES: &[&str] = &["mcp:basic", "mcp:write"]; -/// Perform the headless OAuth authorization-code flow. +/// Attempt the real connection unauthenticated and return the server's +/// `WWW-Authenticate` challenge from the 401 — the reactive discovery +/// trigger. /// -/// 1. Discover metadata, register (or use CIMD), get auth URL -/// 2. Fetch the auth URL with redirect:manual → extract code from Location header -/// 3. Exchange code for token -/// 4. Return an `AuthClient` wrapping `reqwest::Client` +/// `None` (server accepted the unauthenticated connection, which is then +/// closed cleanly) is a legitimate outcome, not an error: the scope-step-up +/// and scope-retry-limit mocks allow unauthenticated `initialize` and only +/// enforce authorization on tool calls. +async fn initialize_challenge( + server_url: &str, + lifecycle: ClientLifecycleMode, +) -> anyhow::Result> { + let transport = StreamableHttpClientTransport::from_uri(server_url); + match BasicClientHandler + .serve_with_lifecycle(transport, lifecycle) + .await + { + Ok(client) => { + client.cancel().await.ok(); + Ok(None) + } + Err(error) => match error.auth_challenge() { + Some(challenge) => Ok(Some(challenge.to_string())), + None => Err(error.into()), + }, + } +} + +fn with_optional_challenge( + request: AuthorizationRequest, + challenge: Option, +) -> AuthorizationRequest { + match challenge { + Some(challenge) => request.with_challenge(challenge), + None => request, + } +} + +/// Perform the headless OAuth authorization-code flow, reactively: +/// +/// 1. Attempt the real connection; take the 401's WWW-Authenticate challenge +/// 2. Discover from the challenge, register (or use CIMD), get auth URL +/// 3. Fetch the auth URL with redirect:manual → extract code from Location header +/// 4. Exchange code for token +/// 5. Return an `AuthClient` wrapping `reqwest::Client` async fn perform_oauth_flow( server_url: &str, _ctx: &ConformanceContext, ) -> anyhow::Result> { + // Always the discover lifecycle here (not `conformance_lifecycle()`): + // this flow serves `run_auth_client`, whose 2026-07-28 auth mocks require + // the per-request MCP-Protocol-Version negotiation. + let challenge = initialize_challenge( + server_url, + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + }, + ) + .await?; let mut oauth = OAuthState::new(server_url, None).await?; // Discover + register + get auth URL + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization( - AuthorizationRequest::new(REDIRECT_URI) - .with_client_name("conformance-client") - .with_client_metadata_url(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -255,14 +303,14 @@ async fn perform_oauth_flow_preregistered( client_id: &str, client_secret: &str, ) -> anyhow::Result> { + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; let mut oauth = OAuthState::new(server_url, None).await?; + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_preregistered_client(client_id) + .with_client_secret(client_secret); oauth - .start_authorization( - AuthorizationRequest::new(REDIRECT_URI) - .with_preregistered_client(client_id) - .with_client_secret(client_secret), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -325,14 +373,14 @@ async fn run_auth_scope_step_up_client( server_url: &str, _ctx: &ConformanceContext, ) -> anyhow::Result<()> { + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; let mut oauth = OAuthState::new(server_url, None).await?; + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied()) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization( - AuthorizationRequest::new(REDIRECT_URI) - .with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied()) - .with_client_name("conformance-client") - .with_client_metadata_url(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge)) .await?; let auth_url = oauth.get_authorization_url().await?; @@ -427,15 +475,15 @@ async fn run_auth_scope_retry_limit_client( ) -> anyhow::Result<()> { let max_retries = 3u32; let mut attempt = 0u32; + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; loop { let mut oauth = OAuthState::new(server_url, None).await?; + let request = AuthorizationRequest::new(REDIRECT_URI) + .with_client_name("conformance-client") + .with_client_metadata_url(CIMD_CLIENT_METADATA_URL); oauth - .start_authorization( - AuthorizationRequest::new(REDIRECT_URI) - .with_client_name("conformance-client") - .with_client_metadata_url(CIMD_CLIENT_METADATA_URL), - ) + .start_authorization(with_optional_challenge(request, challenge.clone())) .await?; let auth_url = oauth.get_authorization_url().await?; let callback = headless_authorize(&auth_url).await?; @@ -509,7 +557,10 @@ async fn migration_token( return Ok(manager.get_access_token().await?); } - let resolution = manager.resolve_metadata().await?; + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; + let resolution = manager + .resolve_metadata_from_challenge(challenge.as_deref()) + .await?; manager.set_metadata(resolution.metadata); manager .register_client("conformance-client", REDIRECT_URI, &[]) @@ -609,7 +660,10 @@ async fn run_client_credentials_basic( .unwrap_or("conformance-test-secret"); let mut manager = AuthorizationManager::new(server_url).await?; - let resolution = manager.resolve_metadata().await?; + let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?; + let resolution = manager + .resolve_metadata_from_challenge(challenge.as_deref()) + .await?; let token_endpoint = resolution.metadata.token_endpoint.clone(); manager.set_metadata(resolution.metadata); diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 40d410d5b..b23a1496d 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -86,6 +86,31 @@ impl ClientInitializeError { context: context.into(), } } + + /// The `WWW-Authenticate` challenge from the 401/403 the transport hit + /// during initialization, if that is why initialization failed. + /// + /// This is the trigger of the reactive OAuth flow: feed the challenge to + /// `AuthorizationRequest::with_challenge` to authorize, then reconnect. + #[cfg(feature = "transport-streamable-http-client")] + pub fn auth_challenge(&self) -> Option<&str> { + use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError}; + + let Self::TransportError { error, .. } = self else { + return None; + }; + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref()); + while let Some(current) = source { + if let Some(auth_required) = current.downcast_ref::() { + return Some(&auth_required.www_authenticate_header); + } + if let Some(insufficient_scope) = current.downcast_ref::() { + return Some(&insufficient_scope.www_authenticate_header); + } + source = current.source(); + } + None + } } /// Helper function to get the next message from the stream diff --git a/crates/rmcp/src/transport/auth.rs b/crates/rmcp/src/transport/auth.rs index 26e53785d..816344e2a 100644 --- a/crates/rmcp/src/transport/auth.rs +++ b/crates/rmcp/src/transport/auth.rs @@ -31,12 +31,6 @@ use crate::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION; const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30); const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024; const MAX_OAUTH_DISCOVERY_REDIRECTS: usize = 10; -const RESOURCE_METADATA_POST_PROBE_BODY: &str = concat!( - r#"{"jsonrpc":"2.0","id":"auth-discovery","method":"initialize","params":{"#, - r#""protocolVersion":"2024-11-05","capabilities":{},"#, - r#""clientInfo":{"name":"rmcp-auth-discovery","version":"0.0.0"}}"#, - r#"}"# -); const CLOUD_METADATA_HOSTS: &[&str] = &[ "metadata", "metadata.google.internal", @@ -634,6 +628,12 @@ pub struct WWWAuthenticateParams { } impl WWWAuthenticateParams { + /// Parse a `WWW-Authenticate` header value, resolving a relative + /// `resource_metadata` URL against `base_url`. + pub fn parse(header: &str, base_url: &Url) -> Self { + AuthorizationManager::extract_www_authenticate_params(header, base_url) + } + /// check if this is an insufficient_scope error pub fn is_insufficient_scope(&self) -> bool { self.error.as_deref() == Some("insufficient_scope") @@ -736,6 +736,11 @@ pub struct AuthorizationRequest { /// OIDC Dynamic Client Registration `application_type` (SEP-837), /// e.g. `"native"` or `"web"`. pub application_type: Option, + /// `WWW-Authenticate` header value from a real request's 401 response. + /// When set, discovery is seeded from the challenge (its + /// `resource_metadata` URL and `scope` hint) instead of probing the + /// server — the reactive discovery path. + pub challenge: Option, } impl AuthorizationRequest { @@ -750,6 +755,7 @@ impl AuthorizationRequest { client_secret: None, client_metadata_url: None, application_type: None, + challenge: None, } } @@ -805,6 +811,13 @@ impl AuthorizationRequest { self.application_type = Some(application_type.into()); self } + + /// Seed discovery from the `WWW-Authenticate` header value of a real + /// request's 401 response instead of probing the server. + pub fn with_challenge(mut self, www_authenticate: impl Into) -> Self { + self.challenge = Some(www_authenticate.into()); + self + } } // add type aliases for oauth2 types @@ -1406,6 +1419,52 @@ impl AuthorizationManager { }) } + /// Resolve authorization server metadata starting from the + /// `WWW-Authenticate` challenge of a real request's 401 response — the + /// reactive discovery path, matching the TypeScript and Python SDKs. + /// + /// Seeds scope selection with the challenge's `scope` hint and prefers + /// the challenge's `resource_metadata` URL; falls back to + /// [`resolve_metadata`](Self::resolve_metadata) when the challenge is + /// `None` or carries no usable metadata pointer. + pub async fn resolve_metadata_from_challenge( + &self, + www_authenticate: Option<&str>, + ) -> Result { + let Some(www_authenticate) = www_authenticate else { + return self.resolve_metadata().await; + }; + let params = WWWAuthenticateParams::parse(www_authenticate, &self.base_url); + + self.record_challenge_scope(¶ms).await; + + if let Some(resource_metadata_url) = ¶ms.resource_metadata_url + && let Some(metadata) = self + .discover_oauth_server_from_resource_metadata_url(resource_metadata_url) + .await? + { + return Ok(AuthorizationMetadataResolution { + metadata, + source: AuthorizationMetadataSource::ProtectedResourceMetadata, + }); + } + + self.resolve_metadata().await + } + + /// Store a challenge's `scope` hint for later scope selection. + async fn record_challenge_scope(&self, params: &WWWAuthenticateParams) { + let Some(scope) = ¶ms.scope else { + return; + }; + let scopes: Vec = scope.split_whitespace().map(str::to_string).collect(); + if scopes.is_empty() { + return; + } + debug!("WWW-Authenticate challenge contains scope: {scope}"); + *self.www_auth_scopes.write().await = scopes; + } + fn legacy_authorization_metadata(base_url: &Url) -> AuthorizationMetadata { let endpoint = |path: &str| { let mut url = base_url.clone(); @@ -2043,7 +2102,7 @@ impl AuthorizationManager { /// refresh token or the server rejected it, return `AuthorizationRequired` /// so the caller can re-prompt the user. Infrastructure errors (e.g. store /// I/O failures, misconfigured client) are propagated as-is. - async fn try_refresh_or_reauth(&self) -> Result { + pub(crate) async fn try_refresh_or_reauth(&self) -> Result { match self.refresh_token().await { Ok(new_creds) => { tracing::info!("Refreshed access token."); @@ -2311,12 +2370,19 @@ impl AuthorizationManager { async fn discover_oauth_server_via_resource_metadata( &self, ) -> Result, AuthError> { - let Some(resource_metadata_url) = self.discover_resource_metadata_url().await? else { + let Some(resource_metadata_url) = self.discover_resource_metadata_url().await else { return Ok(None); }; + self.discover_oauth_server_from_resource_metadata_url(&resource_metadata_url) + .await + } + async fn discover_oauth_server_from_resource_metadata_url( + &self, + resource_metadata_url: &Url, + ) -> Result, AuthError> { let Some(resource_metadata) = self - .fetch_resource_metadata_from_url(&resource_metadata_url) + .fetch_resource_metadata_from_url(resource_metadata_url) .await? else { return Ok(None); @@ -2426,11 +2492,10 @@ impl AuthorizationManager { && Self::is_same_origin(&root_resource, &path_resource) } - async fn discover_resource_metadata_url(&self) -> Result, AuthError> { - if let Ok(Some(resource_metadata_url)) = - self.fetch_resource_metadata_url(&self.base_url, true).await + async fn discover_resource_metadata_url(&self) -> Option { + if let Some(resource_metadata_url) = self.probe_resource_metadata_url(&self.base_url).await { - return Ok(Some(resource_metadata_url)); + return Some(resource_metadata_url); } // If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery. @@ -2442,84 +2507,40 @@ impl AuthorizationManager { discovery_url.set_query(None); discovery_url.set_fragment(None); discovery_url.set_path(&candidate_path); - if let Ok(Some(resource_metadata_url)) = self - .fetch_resource_metadata_url(&discovery_url, false) - .await + if let Some(resource_metadata_url) = + self.probe_resource_metadata_url(&discovery_url).await { - return Ok(Some(resource_metadata_url)); + return Some(resource_metadata_url); } } - Ok(None) + None } - /// Extract the resource metadata url from the WWW-Authenticate header value. + /// Probe `url` with a GET, extracting the resource metadata url from a + /// 200 (the url itself is the metadata document) or from a 401's + /// WWW-Authenticate header value. /// https://www.rfc-editor.org/rfc/rfc9728.html#name-use-of-www-authenticate-for - async fn fetch_resource_metadata_url( - &self, - url: &Url, - allow_post_probe: bool, - ) -> Result, AuthError> { + async fn probe_resource_metadata_url(&self, url: &Url) -> Option { let response = match self.discovery_get(url).await { Ok(r) => r, Err(e) => { debug!("resource metadata probe failed: {}", e); - return Ok(None); + return None; } }; match response.status() { - StatusCode::OK => Ok(Some(url.clone())), - StatusCode::UNAUTHORIZED => Ok(self - .extract_resource_metadata_url_from_www_authenticate(&response) - .await), - StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED if allow_post_probe => { - self.fetch_resource_metadata_url_with_post_probe(url).await + StatusCode::OK => Some(url.clone()), + StatusCode::UNAUTHORIZED => { + self.extract_resource_metadata_url_from_www_authenticate(&response) + .await } status => { debug!("resource metadata probe returned unexpected status: {status}"); - Ok(None) - } - } - } - - async fn fetch_resource_metadata_url_with_post_probe( - &self, - url: &Url, - ) -> Result, AuthError> { - let request = oauth2::http::Request::builder() - .method("POST") - .uri(url.as_str()) - .header(HEADER_MCP_PROTOCOL_VERSION, "2024-11-05") - .header(CONTENT_TYPE, "application/json") - .body(RESOURCE_METADATA_POST_PROBE_BODY.as_bytes().to_vec()) - .map_err(|error| AuthError::InternalError(error.to_string()))?; - let response = match self - .http_client - .execute(OAuthHttpRequest::new( - request, - OAuthHttpRedirectPolicy::Stop, - )) - .await - { - Ok(response) => response, - Err(error) => { - debug!("resource metadata POST probe failed: {}", error); - return Ok(None); + None } - }; - - if response.status() != StatusCode::UNAUTHORIZED { - debug!( - "resource metadata POST probe returned unexpected status: {}", - response.status() - ); - return Ok(None); } - - Ok(self - .extract_resource_metadata_url_from_www_authenticate(&response) - .await) } async fn extract_resource_metadata_url_from_www_authenticate( @@ -2532,14 +2553,9 @@ impl AuthorizationManager { continue; }; let params = Self::extract_www_authenticate_params(value_str, &self.base_url); - if let Some(url) = params.resource_metadata_url { - if let Some(scope) = ¶ms.scope { - debug!("WWW-Authenticate header contains scope: {}", scope); - let scopes: Vec = - scope.split_whitespace().map(|s| s.to_string()).collect(); - *self.www_auth_scopes.write().await = scopes; - } - parsed_url = Some(url); + if params.resource_metadata_url.is_some() { + self.record_challenge_scope(¶ms).await; + parsed_url = params.resource_metadata_url; break; } } @@ -3415,7 +3431,7 @@ impl OAuthState { ) } - async fn placeholder(&self) -> Result { + async fn placeholder_state(&self) -> Result { let (http_client, refresh_redirect_policy) = self.oauth_http_client_config(); Ok(OAuthState::Unauthorized( AuthorizationManager::new_inner( @@ -3531,7 +3547,7 @@ impl OAuthState { &mut self, request: AuthorizationRequest, ) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let old = std::mem::replace(self, placeholder); let OAuthState::Unauthorized(mut manager) = old else { *self = old; @@ -3540,7 +3556,10 @@ impl OAuthState { )); }; debug!("start discovery"); - let metadata = match manager.resolve_metadata().await { + let resolution = manager + .resolve_metadata_from_challenge(request.challenge.as_deref()) + .await; + let metadata = match resolution { Ok(resolution) => resolution.metadata, Err(e) => { *self = OAuthState::Unauthorized(manager); @@ -3563,7 +3582,7 @@ impl OAuthState { /// complete authorization pub async fn complete_authorization(&mut self) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; if let OAuthState::Session(session) = std::mem::replace(self, placeholder) { *self = OAuthState::Authorized(session.auth_manager); Ok(()) @@ -3573,7 +3592,7 @@ impl OAuthState { } /// covert to authorized http client pub async fn to_authorized_http_client(&mut self) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; if let OAuthState::Authorized(manager) = std::mem::replace(self, placeholder) { *self = OAuthState::AuthorizedHttpClient(AuthorizedHttpClient::new( Arc::new(manager), @@ -3593,7 +3612,7 @@ impl OAuthState { required_scope: &str, redirect_uri: &str, ) -> Result { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let old = std::mem::replace(self, placeholder); let OAuthState::Authorized(manager) = old else { *self = old; @@ -3725,7 +3744,7 @@ impl OAuthState { &mut self, config: ClientCredentialsConfig, ) -> Result<(), AuthError> { - let placeholder = self.placeholder().await?; + let placeholder = self.placeholder_state().await?; let OAuthState::Unauthorized(mut manager) = std::mem::replace(self, placeholder) else { return Err(AuthError::InternalError( "Client credentials flow requires Unauthorized state".to_string(), @@ -4087,7 +4106,6 @@ mod tests { .body(Vec::new()) .unwrap(); let client = RecordingOAuthHttpClient::with_responses(vec![ - empty_response(404), challenge, http_response( 200, @@ -4128,7 +4146,6 @@ mod tests { ( "https://auth.example.com/tenant1/token", vec![ - "https://mcp.example.com/mcp", "https://mcp.example.com/mcp", "https://mcp.example.com/custom/metadata/location.json", "https://auth.example.com/.well-known/oauth-authorization-server/tenant1", @@ -4137,14 +4154,12 @@ mod tests { ], ) ); - assert_eq!( + assert!( client .requests() .iter() - .take(2) - .map(|request| request.method.as_str()) - .collect::>(), - vec!["GET", "POST"] + .all(|request| request.method == "GET"), + "discovery must not send non-GET requests" ); } @@ -4155,7 +4170,6 @@ mod tests { empty_response(404), empty_response(404), empty_response(404), - empty_response(404), ]); let manager = AuthorizationManager::new_with_oauth_http_client( "https://legacy.example.com/", @@ -4184,7 +4198,6 @@ mod tests { "https://legacy.example.com/token", Some("https://legacy.example.com/register"), vec![ - "https://legacy.example.com/", "https://legacy.example.com/", "https://legacy.example.com/.well-known/oauth-protected-resource", "https://legacy.example.com/.well-known/oauth-authorization-server", @@ -4243,6 +4256,85 @@ mod tests { ); } + #[tokio::test] + async fn resolve_metadata_from_challenge_uses_challenge_pointer_and_scope() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + http_response( + 200, + serde_json::json!({ + "resource": "https://mcp.example.com/mcp", + "authorization_servers": ["https://auth.example.com"] + }), + ), + http_response( + 200, + serde_json::json!({ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }), + ), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(client.clone()), + ) + .await + .unwrap(); + + let resolution = manager + .resolve_metadata_from_challenge(Some( + r#"Bearer resource_metadata="https://mcp.example.com/custom/prm.json", scope="mcp:read mcp:write""#, + )) + .await + .unwrap(); + + assert_eq!( + ( + resolution.source, + resolution.metadata.token_endpoint.as_str(), + manager.select_scopes(None, &[]), + client.requests().first().map(|request| request.uri.clone()), + ), + ( + AuthorizationMetadataSource::ProtectedResourceMetadata, + "https://auth.example.com/token", + vec!["mcp:read".to_string(), "mcp:write".to_string()], + // discovery starts at the challenge's pointer: no probing + Some("https://mcp.example.com/custom/prm.json".to_string()), + ) + ); + } + + #[tokio::test] + async fn resolve_metadata_from_challenge_falls_back_without_metadata_pointer() { + let client = RecordingOAuthHttpClient::with_responses(vec![ + empty_response(404), + empty_response(404), + empty_response(404), + empty_response(404), + ]); + let manager = AuthorizationManager::new_with_oauth_http_client( + "https://legacy.example.com/", + Arc::new(client), + ) + .await + .unwrap(); + + let resolution = manager + .resolve_metadata_from_challenge(Some(r#"Bearer scope="mcp:basic""#)) + .await + .unwrap(); + + assert_eq!( + (resolution.source, manager.select_scopes(None, &[]),), + ( + AuthorizationMetadataSource::LegacyEndpointFallback, + vec!["mcp:basic".to_string()], + ) + ); + } + #[tokio::test] async fn resolve_metadata_reports_authorization_server_metadata() { let client = RecordingOAuthHttpClient::with_responses(vec![ diff --git a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index 2069e6d31..2053920d2 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -1,11 +1,70 @@ use std::collections::HashMap; use http::{HeaderName, HeaderValue}; +use tracing::debug; use crate::transport::{ - auth::AuthClient, + auth::{AuthClient, AuthError}, streamable_http_client::{StreamableHttpClient, StreamableHttpError}, }; + +impl AuthClient +where + C: StreamableHttpClient + Send + Sync, +{ + /// Run `call` with a token when one is available, reacting to the + /// server's auth verdict instead of requiring credentials up front: + /// + /// - no usable credentials → the request goes out unauthenticated, and a + /// 401 propagates as [`StreamableHttpError::AuthRequired`] carrying the + /// `WWW-Authenticate` challenge for the caller to authorize with; + /// - a token the server rejects (e.g. revoked) → one silent refresh, one + /// retry, then the challenge propagates. + async fn call_reacting_to_challenges( + &self, + auth_token: Option, + call: F, + ) -> Result> + where + F: Fn(Option) -> Fut, + Fut: Future>>, + { + // Missing credentials are not an error in the reactive model: the + // request goes out unauthenticated and the server's 401 challenge + // drives authorization. + let auth_token = match auth_token { + None => match self.get_access_token().await { + Ok(token) => Some(token), + Err(AuthError::AuthorizationRequired) => None, + Err(error) => return Err(error.into()), + }, + token => token, + }; + match call(auth_token.clone()).await { + Err(StreamableHttpError::AuthRequired(challenge)) => { + // One silent recovery attempt: refresh the rejected token and + // retry only when the refresh actually produced a new one. + let Some(sent_token) = auth_token else { + return Err(StreamableHttpError::AuthRequired(challenge)); + }; + let refreshed = { + let manager = self.auth_manager.lock().await; + manager.try_refresh_or_reauth().await + }; + match refreshed { + Ok(fresh_token) if fresh_token != sent_token => call(Some(fresh_token)).await, + Ok(_) => Err(StreamableHttpError::AuthRequired(challenge)), + Err(error) => { + debug!("token refresh after server rejection failed: {error}"); + Err(StreamableHttpError::AuthRequired(challenge)) + } + } + } + result => result, + } + } +} + impl StreamableHttpClient for AuthClient where C: StreamableHttpClient + Send + Sync, @@ -16,16 +75,21 @@ where &self, uri: std::sync::Arc, session_id: std::sync::Arc, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError> { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .delete_session(uri, session_id, auth_token, custom_headers) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .delete_session(uri, session_id, token, custom_headers) + .await + } + }) + .await } async fn get_stream( @@ -33,18 +97,24 @@ where uri: std::sync::Arc, session_id: Option>, last_event_id: Option, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result< futures::stream::BoxStream<'static, Result>, crate::transport::streamable_http_client::StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let last_event_id = last_event_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .get_stream(uri, session_id, last_event_id, token, custom_headers) + .await + } + }) + .await } async fn get_stream_with_max_sse_event_size( @@ -52,26 +122,32 @@ where uri: std::sync::Arc, session_id: Option>, last_event_id: Option, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, ) -> Result< futures::stream::BoxStream<'static, Result>, crate::transport::streamable_http_client::StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .get_stream_with_max_sse_event_size( - uri, - session_id, - last_event_id, - auth_token, - custom_headers, - max_sse_event_size, - ) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let session_id = session_id.clone(); + let last_event_id = last_event_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + token, + custom_headers, + max_sse_event_size, + ) + .await + } + }) + .await } async fn post_message( @@ -79,18 +155,24 @@ where uri: std::sync::Arc, message: crate::model::ClientJsonRpcMessage, session_id: Option>, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, ) -> Result< crate::transport::streamable_http_client::StreamableHttpPostResponse, StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .post_message(uri, message, session_id, auth_token, custom_headers) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let message = message.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .post_message(uri, message, session_id, token, custom_headers) + .await + } + }) + .await } async fn post_message_with_max_sse_event_size( @@ -98,25 +180,31 @@ where uri: std::sync::Arc, message: crate::model::ClientJsonRpcMessage, session_id: Option>, - mut auth_token: Option, + auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, ) -> Result< crate::transport::streamable_http_client::StreamableHttpPostResponse, StreamableHttpError, > { - if auth_token.is_none() { - auth_token = Some(self.get_access_token().await?); - } - self.http_client - .post_message_with_max_sse_event_size( - uri, - message, - session_id, - auth_token, - custom_headers, - max_sse_event_size, - ) - .await + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let message = message.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .post_message_with_max_sse_event_size( + uri, + message, + session_id, + token, + custom_headers, + max_sse_event_size, + ) + .await + } + }) + .await } } diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 5194ba960..c43fb7dd5 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -124,7 +124,8 @@ fn negotiate_version_headers( (version, headers) } -#[derive(Debug)] +#[derive(Debug, Error)] +#[error("authorization required: {www_authenticate_header}")] #[non_exhaustive] pub struct AuthRequiredError { pub www_authenticate_header: String, @@ -139,7 +140,8 @@ impl AuthRequiredError { } } -#[derive(Debug)] +#[derive(Debug, Error)] +#[error("insufficient scope: {www_authenticate_header}")] #[non_exhaustive] pub struct InsufficientScopeError { pub www_authenticate_header: String, @@ -197,15 +199,31 @@ pub enum StreamableHttpError { #[error("Auth error: {0}")] Auth(#[from] crate::transport::auth::AuthError), #[error("Auth required")] - AuthRequired(AuthRequiredError), + AuthRequired(#[source] AuthRequiredError), #[error("Insufficient scope")] - InsufficientScope(InsufficientScopeError), + InsufficientScope(#[source] InsufficientScopeError), #[error("Header name '{0}' is reserved and conflicts with default headers")] ReservedHeaderConflict(String), #[error("Session expired (HTTP 404)")] SessionExpired, } +impl StreamableHttpError { + /// The `WWW-Authenticate` challenge carried by this error, when the + /// server answered 401 ([`AuthRequired`](Self::AuthRequired)) or 403 + /// ([`InsufficientScope`](Self::InsufficientScope)). Feed it to + /// [`AuthorizationRequest::with_challenge`](crate::transport::auth::AuthorizationRequest::with_challenge) + /// to authorize reactively. + #[cfg(feature = "auth")] + pub fn auth_challenge(&self) -> Option<&str> { + match self { + Self::AuthRequired(error) => Some(&error.www_authenticate_header), + Self::InsufficientScope(error) => Some(&error.www_authenticate_header), + _ => None, + } + } +} + #[derive(Debug, Clone, Error)] #[non_exhaustive] pub enum StreamableHttpProtocolError { diff --git a/examples/clients/src/auth/oauth_client.rs b/examples/clients/src/auth/oauth_client.rs index 58565fb72..1ab931331 100644 --- a/examples/clients/src/auth/oauth_client.rs +++ b/examples/clients/src/auth/oauth_client.rs @@ -8,8 +8,9 @@ use axum::{ routing::get, }; use rmcp::{ - ServiceExt, + RoleClient, ServiceExt, model::ClientInfo, + service::RunningService, transport::{ StreamableHttpClientTransport, auth::{AuthClient, AuthorizationRequest, OAuthState}, @@ -55,6 +56,109 @@ async fn callback_handler( Html(CALLBACK_HTML.to_string()) } +enum ConnectOutcome { + /// The server accepted the unauthenticated connection. + Connected(RunningService), + /// The server answered 401; authorize with this `WWW-Authenticate` + /// challenge and reconnect. + AuthRequired(String), +} + +/// Attempt the real connection unauthenticated — the reactive discovery +/// trigger (matching the TypeScript and Python SDKs). The server's 401 +/// challenge, not a probe, tells us whether and how to authorize. +async fn try_connect(http_client: reqwest::Client, server_url: &str) -> Result { + let transport = StreamableHttpClientTransport::with_client( + http_client, + StreamableHttpClientTransportConfig::with_uri(server_url), + ); + match ClientInfo::default().serve(transport).await { + Ok(client) => Ok(ConnectOutcome::Connected(client)), + Err(error) => match error.auth_challenge() { + Some(challenge) => Ok(ConnectOutcome::AuthRequired(challenge.to_string())), + None => Err(error.into()), + }, + } +} + +/// Run the browser OAuth flow seeded by the server's challenge, then +/// reconnect with the authorized transport. +async fn authorize_and_connect( + challenge: String, + oauth_http_client: reqwest::Client, + server_url: &str, + client_metadata_url: &str, + code_receiver: oneshot::Receiver, + output: &mut BufWriter, +) -> Result> { + tracing::info!("Server requires authorization: {challenge}"); + + // initialize oauth state machine + let mut oauth_state = OAuthState::new(server_url, Some(oauth_http_client)) + .await + .context("Failed to initialize oauth state machine")?; + // Seed discovery from the server's challenge, and use CIMD (SEP-991) + // with client metadata URL. Passing no scopes lets the SDK auto-select + // from the challenge's scope hint, Protected Resource Metadata, or AS + // metadata. + oauth_state + .start_authorization( + AuthorizationRequest::new(MCP_REDIRECT_URI) + .with_client_name("Test MCP Client") + .with_client_metadata_url(client_metadata_url) + .with_challenge(challenge), + ) + .await + .context("Failed to start authorization")?; + + // Output authorization URL to user + output + .write_all(b"Please open the following URL in your browser to authorize:\n\n") + .await?; + output + .write_all(oauth_state.get_authorization_url().await?.as_bytes()) + .await?; + output + .write_all(b"\n\nWaiting for browser callback, please do not close this window...\n") + .await?; + output.flush().await?; + + // Wait for authorization code + tracing::info!("Waiting for authorization code..."); + let CallbackParams { + code: auth_code, + state: csrf_token, + iss, + } = code_receiver + .await + .context("Failed to get authorization code")?; + tracing::info!("Received authorization code: {}", auth_code); + // Exchange code for access token + tracing::info!("Exchanging authorization code for access token..."); + oauth_state + .handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref()) + .await + .context("Failed to handle callback")?; + tracing::info!("Successfully obtained access token"); + + output + .write_all(b"\nAuthorization successful! Access token obtained.\n\n") + .await?; + output.flush().await?; + + // Reconnect with the authorized transport + tracing::info!("Establishing authorized connection to MCP server..."); + let am = oauth_state + .into_authorization_manager() + .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; + let auth_client = AuthClient::new(reqwest::Client::default(), am); + let transport = StreamableHttpClientTransport::with_client( + auth_client, + StreamableHttpClientTransportConfig::with_uri(server_url), + ); + Ok(ClientInfo::default().serve(transport).await?) +} + #[tokio::main] async fn main() -> Result<()> { // Initialize logging @@ -123,73 +227,32 @@ async fn main() -> Result<()> { .build() .context("Failed to build OAuth HTTP client")?; - // initialize oauth state machine - let mut oauth_state = OAuthState::new(&server_url, Some(oauth_http_client)) - .await - .context("Failed to initialize oauth state machine")?; - // use CIMD (SEP-991) with client metadata URL. - // passing no scopes lets the SDK auto-select from the server's - // WWW-Authenticate header, Protected Resource Metadata, or AS metadata. - oauth_state - .start_authorization( - AuthorizationRequest::new(MCP_REDIRECT_URI) - .with_client_name("Test MCP Client") - .with_client_metadata_url(&client_metadata_url), - ) - .await - .context("Failed to start authorization")?; - - // Output authorization URL to user let mut output = BufWriter::new(tokio::io::stdout()); output.write_all(b"\n=== MCP OAuth Client ===\n\n").await?; - output - .write_all(b"Please open the following URL in your browser to authorize:\n\n") - .await?; - output - .write_all(oauth_state.get_authorization_url().await?.as_bytes()) - .await?; - output - .write_all(b"\n\nWaiting for browser callback, please do not close this window...\n") - .await?; - output.flush().await?; - - // Wait for authorization code - tracing::info!("Waiting for authorization code..."); - let CallbackParams { - code: auth_code, - state: csrf_token, - iss, - } = code_receiver - .await - .context("Failed to get authorization code")?; - tracing::info!("Received authorization code: {}", auth_code); - // Exchange code for access token - tracing::info!("Exchanging authorization code for access token..."); - oauth_state - .handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref()) - .await - .context("Failed to handle callback")?; - tracing::info!("Successfully obtained access token"); - - output - .write_all(b"\nAuthorization successful! Access token obtained.\n\n") - .await?; output.flush().await?; - // Create authorized transport, this transport is authorized by the oauth state machine - tracing::info!("Establishing authorized connection to MCP server..."); - let am = oauth_state - .into_authorization_manager() - .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; - let client = AuthClient::new(reqwest::Client::default(), am); - let transport = StreamableHttpClientTransport::with_client( - client, - StreamableHttpClientTransportConfig::with_uri(server_url.as_str()), - ); - - // Create client and connect to MCP server - let client_service = ClientInfo::default(); - let client = client_service.serve(transport).await?; + // Reactive discovery: attempt the real connection first. The server's + // 401 challenge — not a probe — tells us whether and how to authorize. + // The transport gets a default client: `oauth_http_client`'s request + // timeout would cut long-lived SSE streams short. + tracing::info!("Attempting connection to MCP server..."); + let client = match try_connect(reqwest::Client::default(), &server_url).await? { + ConnectOutcome::Connected(client) => { + tracing::info!("Server accepted the connection without authorization"); + client + } + ConnectOutcome::AuthRequired(challenge) => { + authorize_and_connect( + challenge, + oauth_http_client, + &server_url, + &client_metadata_url, + code_receiver, + &mut output, + ) + .await? + } + }; tracing::info!("Successfully connected to MCP server"); // Test API requests