Skip to main content

prism_mcp_rs/auth/
discovery.rs

1//! Authorization Server Discovery
2//!
3//! Module implements discovery mechanisms for OAuth 2.0 authorization servers
4//! including Protected Resource Metadata (RFC 9728), Authorization Server Metadata
5//! (RFC 8414), and OpenID Connect Discovery.
6
7use reqwest::Client;
8use url::Url;
9
10use crate::auth::types::*;
11use crate::core::error::{McpError, McpResult};
12
13/// Discovery client for authorization server metadata
14pub struct DiscoveryClient {
15    http_client: Client,
16}
17
18impl DiscoveryClient {
19    /// Create a new discovery client
20    pub fn new() -> Self {
21        Self {
22            http_client: Client::new(),
23        }
24    }
25
26    /// Create with custom HTTP client
27    pub fn with_client(client: Client) -> Self {
28        Self {
29            http_client: client,
30        }
31    }
32
33    /// Discover authorization server from Protected Resource Metadata
34    ///
35    /// This follows RFC 9728 to discover the authorization server(s)
36    /// for a protected resource (MCP server).
37    pub async fn discover_from_resource(
38        &self,
39        _resource_url: &str,
40    ) -> McpResult<ProtectedResourceMetadata> {
41        // Try to fetch protected resource metadata
42        let metadata_url = self.build_resource_metadata_url(_resource_url)?;
43
44        let response = self
45            .http_client
46            .get(&metadata_url)
47            .send()
48            .await
49            .map_err(|e| McpError::Auth(format!("Failed to fetch resource metadata: {e}")))?;
50
51        if !response.status().is_success() {
52            return Err(McpError::Auth(format!(
53                "Failed to fetch resource metadata: HTTP {}",
54                response.status()
55            )));
56        }
57
58        let metadata: ProtectedResourceMetadata = response
59            .json()
60            .await
61            .map_err(|e| McpError::Auth(format!("Invalid resource metadata: {e}")))?;
62
63        // Validate metadata
64        if metadata.authorization_servers.is_empty() {
65            return Err(McpError::Auth(
66                "Resource metadata does not specify any authorization servers".to_string(),
67            ));
68        }
69
70        Ok(metadata)
71    }
72
73    /// Parse WWW-Authenticate header and extract optional resource metadata URL.
74    ///
75    /// `resource_metadata` is optional in newer MCP auth guidance; callers should
76    /// fall back to protected-resource discovery from the MCP resource URL when absent.
77    pub fn parse_www_authenticate(&self, header_value: &str) -> McpResult<Option<String>> {
78        let challenge = AuthChallenge::parse(header_value)
79            .ok_or_else(|| McpError::Auth("Invalid WWW-Authenticate header".to_string()))?;
80        Ok(challenge.resource_metadata)
81    }
82
83    /// Discover authorization server metadata
84    ///
85    /// This tries multiple discovery endpoints in priority order:
86    /// 1. OAuth 2.0 Authorization Server Metadata (RFC 8414)
87    /// 2. OpenID Connect Discovery
88    pub async fn discover_auth_server(
89        &self,
90        issuer_url: &str,
91    ) -> McpResult<AuthorizationServerMetadata> {
92        let issuer = Url::parse(issuer_url)
93            .map_err(|e| McpError::Auth(format!("Invalid issuer URL: {e}")))?;
94
95        // Build discovery URLs based on issuer format
96        let discovery_urls = self.build_discovery_urls(&issuer)?;
97
98        let mut last_error = None;
99
100        // Try each discovery URL in order
101        for url in discovery_urls {
102            match self.fetch_auth_server_metadata(&url).await {
103                Ok(metadata) => {
104                    // Validate issuer matches
105                    if metadata.issuer != issuer_url {
106                        continue; // Try next URL
107                    }
108                    return Ok(metadata);
109                }
110                Err(e) => {
111                    last_error = Some(e);
112                    continue; // Try next URL
113                }
114            }
115        }
116
117        Err(last_error.unwrap_or_else(|| {
118            McpError::Auth("Failed to discover authorization server metadata".to_string())
119        }))
120    }
121
122    /// Build Protected Resource Metadata URL
123    fn build_resource_metadata_url(&self, resource_url: &str) -> McpResult<String> {
124        let base = Url::parse(resource_url)
125            .map_err(|e| McpError::Auth(format!("Invalid resource URL: {e}")))?;
126
127        // RFC 9728: /.well-known/oauth-protected-resource
128        let metadata_url = base
129            .join("/.well-known/oauth-protected-resource")
130            .map_err(|e| McpError::Auth(format!("Failed to build metadata URL: {e}")))?;
131
132        Ok(metadata_url.to_string())
133    }
134
135    /// Build discovery URLs for authorization server
136    fn build_discovery_urls(&self, issuer: &Url) -> McpResult<Vec<String>> {
137        let mut urls = Vec::new();
138
139        // Get the path component (excluding leading slash)
140        let path = issuer.path();
141        let has_path = path != "/" && !path.is_empty();
142
143        if has_path {
144            // For issuer URLs with path components
145            let path_component = path.trim_start_matches('/');
146
147            // 1. OAuth 2.0 with path insertion
148            let oauth_url = format!(
149                "{}://{}/{}/.well-known/oauth-authorization-server/{}",
150                issuer.scheme(),
151                issuer.host_str().unwrap_or(""),
152                issuer.port().map(|p| format!(":{p}")).unwrap_or_default(),
153                path_component
154            );
155            urls.push(oauth_url);
156
157            // 2. OpenID Connect with path insertion
158            let oidc_insert_url = format!(
159                "{}://{}/{}/.well-known/openid-configuration/{}",
160                issuer.scheme(),
161                issuer.host_str().unwrap_or(""),
162                issuer.port().map(|p| format!(":{p}")).unwrap_or_default(),
163                path_component
164            );
165            urls.push(oidc_insert_url);
166
167            // 3. OpenID Connect with path appending
168            let oidc_append_url = format!(
169                "{}/.well-known/openid-configuration",
170                issuer.as_str().trim_end_matches('/')
171            );
172            urls.push(oidc_append_url);
173        } else {
174            // For issuer URLs without path components
175
176            // 1. OAuth 2.0 Authorization Server Metadata
177            let oauth_url = format!(
178                "{}/.well-known/oauth-authorization-server",
179                issuer.as_str().trim_end_matches('/')
180            );
181            urls.push(oauth_url);
182
183            // 2. OpenID Connect Discovery
184            let oidc_url = format!(
185                "{}/.well-known/openid-configuration",
186                issuer.as_str().trim_end_matches('/')
187            );
188            urls.push(oidc_url);
189        }
190
191        Ok(urls)
192    }
193
194    /// Fetch authorization server metadata from a discovery URL
195    async fn fetch_auth_server_metadata(
196        &self,
197        url: &str,
198    ) -> McpResult<AuthorizationServerMetadata> {
199        let response = self
200            .http_client
201            .get(url)
202            .send()
203            .await
204            .map_err(|e| McpError::Auth(format!("Failed to fetch metadata: {e}")))?;
205
206        if !response.status().is_success() {
207            return Err(McpError::Auth(format!(
208                "Failed to fetch metadata from {}: HTTP {}",
209                url,
210                response.status()
211            )));
212        }
213
214        // Try to parse as OAuth 2.0 metadata first, then as OpenID Connect
215        let json: serde_json::Value = response
216            .json()
217            .await
218            .map_err(|e| McpError::Auth(format!("Invalid metadata JSON: {e}")))?;
219
220        // Convert OpenID metadata to OAuth metadata if needed
221        if let Ok(oidc_metadata) = serde_json::from_value::<OpenIDProviderMetadata>(json.clone()) {
222            Ok(self.convert_oidc_to_oauth(oidc_metadata))
223        } else {
224            serde_json::from_value::<AuthorizationServerMetadata>(json)
225                .map_err(|e| McpError::Auth(format!("Invalid authorization server metadata: {e}")))
226        }
227    }
228
229    /// Convert OpenID Connect metadata to OAuth 2.0 metadata
230    fn convert_oidc_to_oauth(&self, oidc: OpenIDProviderMetadata) -> AuthorizationServerMetadata {
231        AuthorizationServerMetadata {
232            issuer: oidc.issuer,
233            authorization_endpoint: oidc.authorization_endpoint,
234            token_endpoint: oidc.token_endpoint,
235            registration_endpoint: oidc.registration_endpoint,
236            scopes_supported: oidc.scopes_supported,
237            response_types_supported: oidc.response_types_supported,
238            response_modes_supported: None,
239            grant_types_supported: None,
240            token_endpoint_auth_methods_supported: None,
241            code_challenge_methods_supported: oidc.code_challenge_methods_supported,
242            revocation_endpoint: None,
243            introspection_endpoint: None,
244            additional: oidc.additional,
245        }
246    }
247}
248
249impl Default for DiscoveryClient {
250    fn default() -> Self {
251        Self::new()
252    }
253}
254
255/// Perform full discovery flow from an authorization challenge response.
256pub async fn discover_from_auth_challenge(
257    client: &Client,
258    www_authenticate: &str,
259    resource_url: &str,
260) -> McpResult<(ProtectedResourceMetadata, AuthorizationServerMetadata)> {
261    let discovery = DiscoveryClient::with_client(client.clone());
262
263    // Parse WWW-Authenticate header
264    let metadata_url = discovery.parse_www_authenticate(www_authenticate)?;
265
266    // Fetch resource metadata (fallback to RFC 9728 well-known discovery when
267    // resource_metadata is not present in challenge parameters).
268    let resource_metadata = match metadata_url {
269        Some(url) => discovery.discover_from_resource(&url).await?,
270        None => discovery.discover_from_resource(resource_url).await?,
271    };
272
273    // Select first authorization server (client should implement selection logic)
274    let auth_server_url = resource_metadata
275        .authorization_servers
276        .first()
277        .ok_or_else(|| {
278            McpError::Auth("No authorization servers specified in resource metadata".to_string())
279        })?
280        .clone();
281
282    // Discover authorization server metadata
283    let auth_metadata = discovery.discover_auth_server(&auth_server_url).await?;
284
285    Ok((resource_metadata, auth_metadata))
286}
287
288/// Backward-compatible alias for callers that still model this as a 401-only flow.
289pub async fn discover_from_401(
290    client: &Client,
291    www_authenticate: &str,
292    resource_url: &str,
293) -> McpResult<(ProtectedResourceMetadata, AuthorizationServerMetadata)> {
294    discover_from_auth_challenge(client, www_authenticate, resource_url).await
295}
296
297/// Check if an authorization server supports required MCP features
298pub fn validate_auth_server_for_mcp(metadata: &AuthorizationServerMetadata) -> McpResult<()> {
299    // Check for PKCE support (required)
300    if metadata.code_challenge_methods_supported.is_none()
301        || metadata
302            .code_challenge_methods_supported
303            .as_ref()
304            .unwrap()
305            .is_empty()
306    {
307        return Err(McpError::Auth(
308            "Authorization server does not support PKCE (required for MCP)".to_string(),
309        ));
310    }
311
312    // Check for authorization code flow support
313    if !metadata
314        .response_types_supported
315        .contains(&"code".to_string())
316        && !metadata
317            .response_types_supported
318            .contains(&"code id_token".to_string())
319    {
320        return Err(McpError::Auth(
321            "Authorization server does not support authorization code flow".to_string(),
322        ));
323    }
324
325    Ok(())
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    #[test]
333    fn test_parse_www_authenticate() {
334        let header = r#"Bearer realm="example", resource_metadata="https://example.com/.well-known/oauth-protected-resource", error="invalid_token""#;
335
336        let discovery = DiscoveryClient::new();
337        let parsed = discovery.parse_www_authenticate(header).unwrap();
338        assert_eq!(
339            parsed,
340            Some("https://example.com/.well-known/oauth-protected-resource".to_string())
341        );
342    }
343
344    #[test]
345    fn test_parse_www_authenticate_without_resource_metadata() {
346        let header = r#"Bearer realm="example", error="insufficient_scope""#;
347        let discovery = DiscoveryClient::new();
348        let parsed = discovery.parse_www_authenticate(header).unwrap();
349        assert_eq!(parsed, None);
350    }
351
352    #[test]
353    fn test_build_discovery_urls_no_path() {
354        let client = DiscoveryClient::new();
355        let issuer = Url::parse("https://auth.example.com").unwrap();
356
357        let urls = client.build_discovery_urls(&issuer).unwrap();
358        assert_eq!(urls.len(), 2);
359        assert_eq!(
360            urls[0],
361            "https://auth.example.com/.well-known/oauth-authorization-server"
362        );
363        assert_eq!(
364            urls[1],
365            "https://auth.example.com/.well-known/openid-configuration"
366        );
367    }
368
369    #[test]
370    fn test_build_discovery_urls_with_path() {
371        let client = DiscoveryClient::new();
372        let issuer = Url::parse("https://auth.example.com/tenant1").unwrap();
373
374        let urls = client.build_discovery_urls(&issuer).unwrap();
375        assert_eq!(urls.len(), 3);
376        assert!(urls[0].contains("/.well-known/oauth-authorization-server/tenant1"));
377        assert!(urls[1].contains("/.well-known/openid-configuration/tenant1"));
378        assert!(urls[2].contains("/tenant1/.well-known/openid-configuration"));
379    }
380
381    #[test]
382    fn test_validate_auth_server() {
383        let mut metadata = AuthorizationServerMetadata {
384            issuer: "https://auth.example.com".to_string(),
385            authorization_endpoint: "https://auth.example.com/authorize".to_string(),
386            token_endpoint: "https://auth.example.com/token".to_string(),
387            registration_endpoint: None,
388            scopes_supported: None,
389            response_types_supported: vec!["code".to_string()],
390            response_modes_supported: None,
391            grant_types_supported: None,
392            token_endpoint_auth_methods_supported: None,
393            code_challenge_methods_supported: Some(vec!["S256".to_string()]),
394            revocation_endpoint: None,
395            introspection_endpoint: None,
396            additional: Default::default(),
397        };
398
399        // Should pass validation
400        assert!(validate_auth_server_for_mcp(&metadata).is_ok());
401
402        // Should fail without PKCE
403        metadata.code_challenge_methods_supported = None;
404        assert!(validate_auth_server_for_mcp(&metadata).is_err());
405
406        // Should fail without code flow
407        metadata.code_challenge_methods_supported = Some(vec!["S256".to_string()]);
408        metadata.response_types_supported = vec!["token".to_string()];
409        assert!(validate_auth_server_for_mcp(&metadata).is_err());
410    }
411}