fix: keep post probe eager to preserve scope hint

Author: DaleSeo

crates/rmcp/src/transport/auth.rs

@@ -1003,14 +1003,13 @@ pub struct AuthorizationManager {
 }
 
 /// Outcome of a resource metadata discovery GET issued by
-/// [`AuthorizationManager::fetch_resource_metadata_url`].
-enum ResourceMetadataUrlProbe {
+/// [`AuthorizationManager::probe_resource_metadata_url`].
+enum ResourceMetadataProbeOutcome {
     /// The probe located a resource metadata URL: either a 200 on the probed
     /// URL itself, or a 401 whose WWW-Authenticate header pointed at one.
     Found(Url),
     /// The probe got a 404/405, typical of streamable HTTP servers that
-    /// reject session-less GETs, so the POST probe fallback is worth trying
-    /// if well-known discovery also comes up empty.
+    /// reject session-less GETs, so the POST probe fallback is worth trying.
     PostProbeEligible,
     /// The probe yielded nothing usable.
     Unavailable,
@@ -2441,13 +2440,24 @@ impl AuthorizationManager {
     }
 
     async fn discover_resource_metadata_url(&self) -> Result<Option<Url>, AuthError> {
-        let should_post_probe = match self.fetch_resource_metadata_url(&self.base_url).await {
-            ResourceMetadataUrlProbe::Found(resource_metadata_url) => {
+        match self.probe_resource_metadata_url(&self.base_url).await {
+            ResourceMetadataProbeOutcome::Found(resource_metadata_url) => {
                 return Ok(Some(resource_metadata_url));
             }
-            ResourceMetadataUrlProbe::PostProbeEligible => true,
-            ResourceMetadataUrlProbe::Unavailable => false,
-        };
+            // The POST probe's 401 carries both the resource metadata URL and
+            // the scope hint from WWW-Authenticate, so it must run before
+            // well-known discovery: a well-known hit would otherwise hide the
+            // scope the server advertises only on that challenge.
+            ResourceMetadataProbeOutcome::PostProbeEligible => {
+                if let Some(resource_metadata_url) = self
+                    .post_probe_resource_metadata_url(&self.base_url)
+                    .await?
+                {
+                    return Ok(Some(resource_metadata_url));
+                }
+            }
+            ResourceMetadataProbeOutcome::Unavailable => {}
+        }
 
         // If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery.
         // https://www.rfc-editor.org/rfc/rfc9728.html#name-obtaining-protected-resourc
@@ -2458,56 +2468,47 @@ impl AuthorizationManager {
             discovery_url.set_query(None);
             discovery_url.set_fragment(None);
             discovery_url.set_path(&candidate_path);
-            if let ResourceMetadataUrlProbe::Found(resource_metadata_url) =
-                self.fetch_resource_metadata_url(&discovery_url).await
+            if let ResourceMetadataProbeOutcome::Found(resource_metadata_url) =
+                self.probe_resource_metadata_url(&discovery_url).await
             {
                 return Ok(Some(resource_metadata_url));
             }
         }
 
-        if should_post_probe {
-            return self
-                .fetch_resource_metadata_url_with_post_probe(&self.base_url)
-                .await;
-        }
-
         Ok(None)
     }
 
     /// Extract the resource metadata url from the 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) -> ResourceMetadataUrlProbe {
+    async fn probe_resource_metadata_url(&self, url: &Url) -> ResourceMetadataProbeOutcome {
         let response = match self.discovery_get(url).await {
             Ok(r) => r,
             Err(e) => {
                 debug!("resource metadata probe failed: {}", e);
-                return ResourceMetadataUrlProbe::Unavailable;
+                return ResourceMetadataProbeOutcome::Unavailable;
             }
         };
 
         match response.status() {
-            StatusCode::OK => ResourceMetadataUrlProbe::Found(url.clone()),
+            StatusCode::OK => ResourceMetadataProbeOutcome::Found(url.clone()),
             StatusCode::UNAUTHORIZED => self
                 .extract_resource_metadata_url_from_www_authenticate(&response)
                 .await
                 .map_or(
-                    ResourceMetadataUrlProbe::Unavailable,
-                    ResourceMetadataUrlProbe::Found,
+                    ResourceMetadataProbeOutcome::Unavailable,
+                    ResourceMetadataProbeOutcome::Found,
                 ),
             StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED => {
-                ResourceMetadataUrlProbe::PostProbeEligible
+                ResourceMetadataProbeOutcome::PostProbeEligible
             }
             status => {
                 debug!("resource metadata probe returned unexpected status: {status}");
-                ResourceMetadataUrlProbe::Unavailable
+                ResourceMetadataProbeOutcome::Unavailable
             }
         }
     }
 
-    async fn fetch_resource_metadata_url_with_post_probe(
-        &self,
-        url: &Url,
-    ) -> Result<Option<Url>, AuthError> {
+    async fn post_probe_resource_metadata_url(&self, url: &Url) -> Result<Option<Url>, AuthError> {
         let request = oauth2::http::Request::builder()
             .method("POST")
             .uri(url.as_str())
@@ -2546,7 +2547,7 @@ impl AuthorizationManager {
     }
 
     /// Best-effort cleanup of a session the server may have created for the
-    /// POST probe's synthetic `initialize` request (see issue #1048).
+    /// POST probe's synthetic `initialize` request.
     /// Failures are logged and swallowed: cleanup must never fail discovery.
     async fn delete_post_probe_session(&self, url: &Url, response: &HttpResponse) {
         let Some(session_id) = response.headers().get(HEADER_SESSION_ID) else {
@@ -3479,7 +3480,7 @@ impl OAuthState {
         )
     }
 
-    async fn placeholder(&self) -> Result<Self, AuthError> {
+    async fn placeholder_state(&self) -> Result<Self, AuthError> {
         let (http_client, refresh_redirect_policy) = self.oauth_http_client_config();
         Ok(OAuthState::Unauthorized(
             AuthorizationManager::new_inner(
@@ -3595,7 +3596,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;
@@ -3627,7 +3628,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(())
@@ -3637,7 +3638,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),
@@ -3657,7 +3658,7 @@ impl OAuthState {
         required_scope: &str,
         redirect_uri: &str,
     ) -> Result<String, AuthError> {
-        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;
@@ -3789,7 +3790,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(),
@@ -3971,12 +3972,20 @@ mod tests {
         );
     }
 
+    // The POST probe must run before well-known discovery: its 401 challenge
+    // is the only channel for the WWW-Authenticate scope hint, which a
+    // well-known hit would otherwise hide.
     #[tokio::test]
-    async fn well_known_resource_metadata_precedes_post_probe() {
-        let client = RecordingOAuthHttpClient::with_responses(vec![
-            empty_response(405),
-            empty_response(200),
-        ]);
+    async fn post_probe_precedes_well_known_discovery() {
+        let challenge = oauth2::http::Response::builder()
+            .status(401)
+            .header(
+                "www-authenticate",
+                r#"Bearer resource_metadata="https://mcp.example.com/custom/metadata/location.json""#,
+            )
+            .body(Vec::new())
+            .unwrap();
+        let client = RecordingOAuthHttpClient::with_responses(vec![empty_response(405), challenge]);
         let manager = AuthorizationManager::new_with_oauth_http_client(
             "https://mcp.example.com/mcp",
             Arc::new(client.clone()),
@@ -3996,15 +4005,10 @@ mod tests {
                     .collect::<Vec<_>>(),
             ),
             (
-                Some(
-                    "https://mcp.example.com/.well-known/oauth-protected-resource/mcp".to_string()
-                ),
+                Some("https://mcp.example.com/custom/metadata/location.json".to_string()),
                 vec![
                     ("GET", "https://mcp.example.com/mcp"),
-                    (
-                        "GET",
-                        "https://mcp.example.com/.well-known/oauth-protected-resource/mcp"
-                    ),
+                    ("POST", "https://mcp.example.com/mcp"),
                 ],
             )
         );
@@ -4019,11 +4023,11 @@ mod tests {
             .unwrap();
         let client = RecordingOAuthHttpClient::with_responses(vec![
             empty_response(405),
+            post_response,
+            empty_response(202),
             empty_response(404),
             empty_response(404),
             empty_response(404),
-            post_response,
-            empty_response(202),
         ]);
         let manager = AuthorizationManager::new_with_oauth_http_client(
             "https://mcp.example.com/mcp",
@@ -4053,6 +4057,12 @@ mod tests {
                 None,
                 vec![
                     ("GET", "https://mcp.example.com/mcp", None),
+                    ("POST", "https://mcp.example.com/mcp", None),
+                    (
+                        "DELETE",
+                        "https://mcp.example.com/mcp",
+                        Some("probe-session"),
+                    ),
                     (
                         "GET",
                         "https://mcp.example.com/.well-known/oauth-protected-resource/mcp",
@@ -4068,12 +4078,6 @@ mod tests {
                         "https://mcp.example.com/.well-known/oauth-protected-resource",
                         None,
                     ),
-                    ("POST", "https://mcp.example.com/mcp", None),
-                    (
-                        "DELETE",
-                        "https://mcp.example.com/mcp",
-                        Some("probe-session"),
-                    ),
                 ],
             )
         );
@@ -4092,9 +4096,6 @@ mod tests {
             .unwrap();
         let client = RecordingOAuthHttpClient::with_responses(vec![
             empty_response(405),
-            empty_response(404),
-            empty_response(404),
-            empty_response(404),
             post_response,
             empty_response(202),
         ]);
@@ -4126,21 +4127,6 @@ mod tests {
                 Some("https://mcp.example.com/custom/metadata/location.json".to_string()),
                 vec![
                     ("GET", "https://mcp.example.com/mcp", None),
-                    (
-                        "GET",
-                        "https://mcp.example.com/.well-known/oauth-protected-resource/mcp",
-                        None,
-                    ),
-                    (
-                        "GET",
-                        "https://mcp.example.com/mcp/.well-known/oauth-protected-resource",
-                        None,
-                    ),
-                    (
-                        "GET",
-                        "https://mcp.example.com/.well-known/oauth-protected-resource",
-                        None,
-                    ),
                     ("POST", "https://mcp.example.com/mcp", None),
                     (
                         "DELETE",
@@ -4341,9 +4327,6 @@ mod tests {
             .body(Vec::new())
             .unwrap();
         let client = RecordingOAuthHttpClient::with_responses(vec![
-            empty_response(404),
-            empty_response(404),
-            empty_response(404),
             empty_response(404),
             challenge,
             http_response(
@@ -4386,9 +4369,6 @@ mod tests {
                 "https://auth.example.com/tenant1/token",
                 vec![
                     "https://mcp.example.com/mcp",
-                    "https://mcp.example.com/.well-known/oauth-protected-resource/mcp",
-                    "https://mcp.example.com/mcp/.well-known/oauth-protected-resource",
-                    "https://mcp.example.com/.well-known/oauth-protected-resource",
                     "https://mcp.example.com/mcp",
                     "https://mcp.example.com/custom/metadata/location.json",
                     "https://auth.example.com/.well-known/oauth-authorization-server/tenant1",
@@ -4401,10 +4381,10 @@ mod tests {
             client
                 .requests()
                 .iter()
-                .take(5)
+                .take(2)
                 .map(|request| request.method.as_str())
                 .collect::<Vec<_>>(),
-            vec!["GET", "GET", "GET", "GET", "POST"]
+            vec!["GET", "POST"]
         );
     }
 
@@ -4445,8 +4425,8 @@ mod tests {
                 Some("https://legacy.example.com/register"),
                 vec![
                     "https://legacy.example.com/",
-                    "https://legacy.example.com/.well-known/oauth-protected-resource",
                     "https://legacy.example.com/",
+                    "https://legacy.example.com/.well-known/oauth-protected-resource",
                     "https://legacy.example.com/.well-known/oauth-authorization-server",
                     "https://legacy.example.com/.well-known/openid-configuration",
                 ],