Skip to repository content2863 lines · 105.3 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T01:50:23.258Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
oauth.rs
1//! OAuth 2.0 authentication for MCP servers using the Authorization Code +
2//! PKCE flow, per the MCP spec's OAuth profile.
3//!
4//! The flow is split into two phases:
5//!
6//! 1. **Discovery** ([`discover`]) fetches Protected Resource Metadata and
7//! Authorization Server Metadata. This can happen early (e.g. on a 401
8//! during server startup) because it doesn't need the redirect URI yet.
9//!
10//! 2. **Client registration** ([`resolve_client_registration`]) is separate
11//! because DCR requires the actual loopback redirect URI, which includes an
12//! ephemeral port that only exists once the callback server has started.
13//!
14//! After authentication, the full state is captured in [`OAuthSession`] which
15//! is persisted to the keychain. On next startup, the stored session feeds
16//! directly into [`McpOAuthTokenProvider`], giving a refresh-capable provider
17//! without requiring another browser flow.
18
19use anyhow::{Context as _, Result, anyhow, bail};
20use async_trait::async_trait;
21use base64::Engine as _;
22use futures::AsyncReadExt as _;
23use futures::FutureExt as _;
24use futures::channel::mpsc;
25use futures::future::BoxFuture;
26use http_client::{AsyncBody, HttpClient, Request};
27use parking_lot::Mutex as SyncMutex;
28use rand::Rng as _;
29use serde::{Deserialize, Serialize};
30use sha2::{Digest, Sha256};
31
32use std::sync::Arc;
33use std::time::{Duration, SystemTime};
34use url::Url;
35
36/// The CIMD URL where Zed's OAuth client metadata document is hosted.
37pub const CIMD_URL: &str = "https://zed.dev/oauth/client-metadata.json";
38
39/// Validate that a URL is safe to use as an OAuth endpoint.
40///
41/// OAuth endpoints carry sensitive material (authorization codes, PKCE
42/// verifiers, tokens) and must use TLS. Plain HTTP is only permitted for
43/// loopback addresses, per RFC 8252 Section 8.3.
44fn require_https_or_loopback(url: &Url) -> Result<()> {
45 if url.scheme() == "https" {
46 return Ok(());
47 }
48 if url.scheme() == "http" {
49 if let Some(host) = url.host() {
50 match host {
51 url::Host::Ipv4(ip) if ip.is_loopback() => return Ok(()),
52 url::Host::Ipv6(ip) if ip.is_loopback() => return Ok(()),
53 url::Host::Domain(d) if d.eq_ignore_ascii_case("localhost") => return Ok(()),
54 _ => {}
55 }
56 }
57 }
58 bail!(
59 "OAuth endpoint must use HTTPS (got {}://{})",
60 url.scheme(),
61 url.host_str().unwrap_or("?")
62 )
63}
64
65/// Validate that a URL is safe to use as an OAuth endpoint, including SSRF
66/// protections against private/reserved IP ranges.
67///
68/// This wraps [`require_https_or_loopback`] and adds IP-range checks to prevent
69/// an attacker-controlled MCP server from directing Zed to fetch internal
70/// network resources via metadata URLs.
71///
72/// **Known limitation:** Domain-name URLs that resolve to private IPs are *not*
73/// blocked here — full mitigation requires resolver-level validation (e.g. a
74/// custom `Resolve` implementation). This function only blocks IP-literal URLs.
75fn validate_oauth_url(url: &Url) -> Result<()> {
76 require_https_or_loopback(url)?;
77
78 if let Some(host) = url.host() {
79 match host {
80 url::Host::Ipv4(ip) => {
81 // Loopback is already allowed by require_https_or_loopback.
82 if ip.is_private() || ip.is_link_local() || ip.is_broadcast() || ip.is_unspecified()
83 {
84 bail!(
85 "OAuth endpoint must not point to private/reserved IP: {}",
86 ip
87 );
88 }
89 }
90 url::Host::Ipv6(ip) => {
91 // Check for IPv4-mapped IPv6 addresses (::ffff:a.b.c.d) which
92 // could bypass the IPv4 checks above.
93 if let Some(mapped_v4) = ip.to_ipv4_mapped() {
94 if mapped_v4.is_private()
95 || mapped_v4.is_link_local()
96 || mapped_v4.is_broadcast()
97 || mapped_v4.is_unspecified()
98 {
99 bail!(
100 "OAuth endpoint must not point to private/reserved IP: ::ffff:{}",
101 mapped_v4
102 );
103 }
104 }
105
106 if ip.is_unspecified() || ip.is_multicast() {
107 bail!(
108 "OAuth endpoint must not point to reserved IPv6 address: {}",
109 ip
110 );
111 }
112 // IPv6 Unique Local Addresses (fc00::/7). is_unique_local() is
113 // nightly-only, so check the prefix manually.
114 if (ip.segments()[0] & 0xfe00) == 0xfc00 {
115 bail!(
116 "OAuth endpoint must not point to IPv6 unique-local address: {}",
117 ip
118 );
119 }
120 }
121 url::Host::Domain(_) => {
122 // Domain-based SSRF prevention requires resolver-level checks.
123 // See known limitation in the doc comment above.
124 }
125 }
126 }
127
128 Ok(())
129}
130
131/// Parsed from the MCP server's WWW-Authenticate header or well-known endpoint
132/// per RFC 9728 (OAuth 2.0 Protected Resource Metadata).
133#[derive(Debug, Clone, Serialize, Deserialize)]
134pub struct ProtectedResourceMetadata {
135 pub resource: Url,
136 pub authorization_servers: Vec<Url>,
137 pub scopes_supported: Option<Vec<String>>,
138}
139
140/// Parsed from the authorization server's .well-known endpoint
141/// per RFC 8414 (OAuth 2.0 Authorization Server Metadata).
142#[derive(Debug, Clone, Serialize, Deserialize)]
143pub struct AuthServerMetadata {
144 pub issuer: Url,
145 pub authorization_endpoint: Url,
146 pub token_endpoint: Url,
147 pub registration_endpoint: Option<Url>,
148 pub scopes_supported: Option<Vec<String>>,
149 pub grant_types_supported: Option<Vec<String>>,
150 pub code_challenge_methods_supported: Option<Vec<String>>,
151 pub client_id_metadata_document_supported: bool,
152}
153
154/// The result of client registration — either CIMD or DCR.
155#[derive(Clone, Serialize, Deserialize)]
156pub struct OAuthClientRegistration {
157 pub client_id: String,
158 /// Only present for DCR-minted registrations.
159 pub client_secret: Option<String>,
160}
161
162impl std::fmt::Debug for OAuthClientRegistration {
163 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164 f.debug_struct("OAuthClientRegistration")
165 .field("client_id", &self.client_id)
166 .field(
167 "client_secret",
168 &self.client_secret.as_ref().map(|_| "[redacted]"),
169 )
170 .finish()
171 }
172}
173
174/// Access and refresh tokens obtained from the token endpoint.
175#[derive(Clone, Serialize, Deserialize)]
176pub struct OAuthTokens {
177 pub access_token: String,
178 pub refresh_token: Option<String>,
179 pub expires_at: Option<SystemTime>,
180}
181
182impl std::fmt::Debug for OAuthTokens {
183 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184 f.debug_struct("OAuthTokens")
185 .field("access_token", &"[redacted]")
186 .field(
187 "refresh_token",
188 &self.refresh_token.as_ref().map(|_| "[redacted]"),
189 )
190 .field("expires_at", &self.expires_at)
191 .finish()
192 }
193}
194
195/// Everything discovered before the browser flow starts. Client registration is
196/// resolved separately, once the real redirect URI is known.
197#[derive(Debug, Clone, Serialize, Deserialize)]
198pub struct OAuthDiscovery {
199 pub resource_metadata: ProtectedResourceMetadata,
200 pub auth_server_metadata: AuthServerMetadata,
201 pub scopes: Vec<String>,
202}
203
204/// The persisted OAuth session for a context server.
205///
206/// Stored in the keychain so startup can restore a refresh-capable provider
207/// without another browser flow. Deliberately excludes the full discovery
208/// metadata to keep the serialized size well within keychain item limits.
209#[derive(Clone, Serialize, Deserialize)]
210pub struct OAuthSession {
211 pub token_endpoint: Url,
212 pub resource: Url,
213 pub client_registration: OAuthClientRegistration,
214 pub tokens: OAuthTokens,
215}
216
217impl std::fmt::Debug for OAuthSession {
218 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219 f.debug_struct("OAuthSession")
220 .field("token_endpoint", &self.token_endpoint)
221 .field("resource", &self.resource)
222 .field("client_registration", &self.client_registration)
223 .field("tokens", &self.tokens)
224 .finish()
225 }
226}
227
228/// Error codes defined by RFC 6750 Section 3.1 for Bearer token authentication.
229#[derive(Debug, Clone, Copy, PartialEq, Eq)]
230pub enum BearerError {
231 /// The request is missing a required parameter, includes an unsupported
232 /// parameter or parameter value, or is otherwise malformed.
233 InvalidRequest,
234 /// The access token provided is expired, revoked, malformed, or invalid.
235 InvalidToken,
236 /// The request requires higher privileges than provided by the access token.
237 InsufficientScope,
238 /// An unrecognized error code (extension or future spec addition).
239 Other,
240}
241
242impl BearerError {
243 fn parse(value: &str) -> Self {
244 match value {
245 "invalid_request" => BearerError::InvalidRequest,
246 "invalid_token" => BearerError::InvalidToken,
247 "insufficient_scope" => BearerError::InsufficientScope,
248 _ => BearerError::Other,
249 }
250 }
251}
252
253/// Fields extracted from a `WWW-Authenticate: Bearer` header.
254///
255/// Per RFC 9728 Section 5.1, MCP servers include `resource_metadata` to point
256/// at the Protected Resource Metadata document. The optional `scope` parameter
257/// (RFC 6750 Section 3) indicates scopes required for the request.
258#[derive(Debug, Clone, PartialEq, Eq)]
259pub struct WwwAuthenticate {
260 pub resource_metadata: Option<Url>,
261 pub scope: Option<Vec<String>>,
262 /// The parsed `error` parameter per RFC 6750 Section 3.1.
263 pub error: Option<BearerError>,
264 pub error_description: Option<String>,
265}
266
267/// Parse a `WWW-Authenticate` header value.
268///
269/// Expects the `Bearer` scheme followed by comma-separated `key="value"` pairs.
270/// Per RFC 6750 and RFC 9728, the relevant parameters are:
271/// - `resource_metadata` — URL of the Protected Resource Metadata document
272/// - `scope` — space-separated list of required scopes
273/// - `error` — error code (e.g. "insufficient_scope")
274/// - `error_description` — human-readable error description
275pub fn parse_www_authenticate(header: &str) -> Result<WwwAuthenticate> {
276 let header = header.trim();
277
278 let params_str = if header.len() >= 6 && header[..6].eq_ignore_ascii_case("bearer") {
279 header[6..].trim()
280 } else {
281 bail!("WWW-Authenticate header does not use Bearer scheme");
282 };
283
284 if params_str.is_empty() {
285 return Ok(WwwAuthenticate {
286 resource_metadata: None,
287 scope: None,
288 error: None,
289 error_description: None,
290 });
291 }
292
293 let params = parse_auth_params(params_str);
294
295 let resource_metadata = params
296 .get("resource_metadata")
297 .map(|v| Url::parse(v))
298 .transpose()
299 .map_err(|e| anyhow!("invalid resource_metadata URL: {}", e))?;
300
301 let scope = params
302 .get("scope")
303 .map(|v| v.split_whitespace().map(String::from).collect());
304
305 let error = params.get("error").map(|v| BearerError::parse(v));
306 let error_description = params.get("error_description").cloned();
307
308 Ok(WwwAuthenticate {
309 resource_metadata,
310 scope,
311 error,
312 error_description,
313 })
314}
315
316/// Parse comma-separated `key="value"` or `key=token` parameters from an
317/// auth-param list (RFC 7235 Section 2.1).
318fn parse_auth_params(input: &str) -> collections::HashMap<String, String> {
319 let mut params = collections::HashMap::default();
320 let mut remaining = input.trim();
321
322 while !remaining.is_empty() {
323 // Skip leading whitespace and commas.
324 remaining = remaining.trim_start_matches(|c: char| c == ',' || c.is_whitespace());
325 if remaining.is_empty() {
326 break;
327 }
328
329 // Find the key (everything before '=').
330 let eq_pos = match remaining.find('=') {
331 Some(pos) => pos,
332 None => break,
333 };
334
335 let key = remaining[..eq_pos].trim().to_lowercase();
336 remaining = &remaining[eq_pos + 1..];
337 remaining = remaining.trim_start();
338
339 // Parse the value: either quoted or unquoted (token).
340 let value;
341 if remaining.starts_with('"') {
342 // Quoted string: find the closing quote, handling escaped chars.
343 remaining = &remaining[1..]; // skip opening quote
344 let mut val = String::new();
345 let mut chars = remaining.char_indices();
346 loop {
347 match chars.next() {
348 Some((_, '\\')) => {
349 // Escaped character — take the next char literally.
350 if let Some((_, c)) = chars.next() {
351 val.push(c);
352 }
353 }
354 Some((i, '"')) => {
355 remaining = &remaining[i + 1..];
356 break;
357 }
358 Some((_, c)) => val.push(c),
359 None => {
360 remaining = "";
361 break;
362 }
363 }
364 }
365 value = val;
366 } else {
367 // Unquoted token: read until comma or whitespace.
368 let end = remaining
369 .find(|c: char| c == ',' || c.is_whitespace())
370 .unwrap_or(remaining.len());
371 value = remaining[..end].to_string();
372 remaining = &remaining[end..];
373 }
374
375 if !key.is_empty() {
376 params.insert(key, value);
377 }
378 }
379
380 params
381}
382
383/// Construct the well-known Protected Resource Metadata URIs for a given MCP
384/// server URL, per RFC 9728 Section 3.
385///
386/// Returns URIs in priority order:
387/// 1. Path-specific: `https://<host>/.well-known/oauth-protected-resource/<path>`
388/// 2. Root: `https://<host>/.well-known/oauth-protected-resource`
389pub fn protected_resource_metadata_urls(server_url: &Url) -> Vec<Url> {
390 let mut urls = Vec::new();
391 let base = format!("{}://{}", server_url.scheme(), server_url.authority());
392
393 let path = server_url.path().trim_start_matches('/');
394 if !path.is_empty() {
395 if let Ok(url) = Url::parse(&format!(
396 "{}/.well-known/oauth-protected-resource/{}",
397 base, path
398 )) {
399 urls.push(url);
400 }
401 }
402
403 if let Ok(url) = Url::parse(&format!("{}/.well-known/oauth-protected-resource", base)) {
404 urls.push(url);
405 }
406
407 urls
408}
409
410/// Construct the well-known Authorization Server Metadata URIs for a given
411/// issuer URL, per RFC 8414 Section 3.1 and Section 5 (OIDC compat).
412///
413/// Returns URIs in priority order, which differs depending on whether the
414/// issuer URL has a path component.
415pub fn auth_server_metadata_urls(issuer: &Url) -> Vec<Url> {
416 let mut urls = Vec::new();
417 let base = format!("{}://{}", issuer.scheme(), issuer.authority());
418 let path = issuer.path().trim_matches('/');
419
420 if !path.is_empty() {
421 // Issuer with path: try path-inserted variants first.
422 if let Ok(url) = Url::parse(&format!(
423 "{}/.well-known/oauth-authorization-server/{}",
424 base, path
425 )) {
426 urls.push(url);
427 }
428 if let Ok(url) = Url::parse(&format!(
429 "{}/.well-known/openid-configuration/{}",
430 base, path
431 )) {
432 urls.push(url);
433 }
434 if let Ok(url) = Url::parse(&format!(
435 "{}/{}/.well-known/openid-configuration",
436 base, path
437 )) {
438 urls.push(url);
439 }
440 } else {
441 // No path: standard well-known locations.
442 if let Ok(url) = Url::parse(&format!("{}/.well-known/oauth-authorization-server", base)) {
443 urls.push(url);
444 }
445 if let Ok(url) = Url::parse(&format!("{}/.well-known/openid-configuration", base)) {
446 urls.push(url);
447 }
448 }
449
450 urls
451}
452
453// -- Canonical server URI (RFC 8707) -----------------------------------------
454
455/// Derive the canonical resource URI for an MCP server URL, suitable for the
456/// `resource` parameter in authorization and token requests per RFC 8707.
457///
458/// Lowercases the scheme and host, preserves the path (without trailing slash),
459/// strips fragments and query strings.
460pub fn canonical_server_uri(server_url: &Url) -> String {
461 let mut uri = format!(
462 "{}://{}",
463 server_url.scheme().to_ascii_lowercase(),
464 server_url.host_str().unwrap_or("").to_ascii_lowercase(),
465 );
466 if let Some(port) = server_url.port() {
467 uri.push_str(&format!(":{}", port));
468 }
469 let path = server_url.path();
470 if path != "/" {
471 uri.push_str(path.trim_end_matches('/'));
472 }
473 uri
474}
475
476// -- Scope selection ---------------------------------------------------------
477
478/// Select scopes following the MCP spec's Scope Selection Strategy:
479/// 1. Use `scope` from the `WWW-Authenticate` challenge if present.
480/// 2. Fall back to `scopes_supported` from Protected Resource Metadata.
481/// 3. Return empty if neither is available.
482pub fn select_scopes(
483 www_authenticate: &WwwAuthenticate,
484 resource_metadata: &ProtectedResourceMetadata,
485) -> Vec<String> {
486 if let Some(ref scopes) = www_authenticate.scope {
487 if !scopes.is_empty() {
488 return scopes.clone();
489 }
490 }
491 resource_metadata
492 .scopes_supported
493 .clone()
494 .unwrap_or_default()
495}
496
497// -- Client registration strategy --------------------------------------------
498
499/// The registration approach to use, determined from auth server metadata.
500#[derive(Debug, Clone, PartialEq, Eq)]
501pub enum ClientRegistrationStrategy {
502 /// The auth server supports CIMD. Use the CIMD URL as client_id directly.
503 Cimd { client_id: String },
504 /// The auth server has a registration endpoint. Caller must POST to it.
505 Dcr { registration_endpoint: Url },
506 /// No supported registration mechanism.
507 Unavailable,
508}
509
510/// Determine how to register with the authorization server, following the
511/// spec's recommended priority: CIMD first, DCR fallback.
512pub fn determine_registration_strategy(
513 auth_server_metadata: &AuthServerMetadata,
514) -> ClientRegistrationStrategy {
515 if auth_server_metadata.client_id_metadata_document_supported {
516 ClientRegistrationStrategy::Cimd {
517 client_id: CIMD_URL.to_string(),
518 }
519 } else if let Some(ref endpoint) = auth_server_metadata.registration_endpoint {
520 ClientRegistrationStrategy::Dcr {
521 registration_endpoint: endpoint.clone(),
522 }
523 } else {
524 ClientRegistrationStrategy::Unavailable
525 }
526}
527
528// -- PKCE (RFC 7636) ---------------------------------------------------------
529
530/// A PKCE code verifier and its S256 challenge.
531#[derive(Clone)]
532pub struct PkceChallenge {
533 pub verifier: String,
534 pub challenge: String,
535}
536
537impl std::fmt::Debug for PkceChallenge {
538 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539 f.debug_struct("PkceChallenge")
540 .field("verifier", &"[redacted]")
541 .field("challenge", &self.challenge)
542 .finish()
543 }
544}
545
546/// Generate a PKCE code verifier and S256 challenge per RFC 7636.
547///
548/// The verifier is 43 base64url characters derived from 32 random bytes.
549/// The challenge is `BASE64URL(SHA256(verifier))`.
550pub fn generate_pkce_challenge() -> PkceChallenge {
551 let mut random_bytes = [0u8; 32];
552 rand::rng().fill(&mut random_bytes);
553 let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
554 let verifier = engine.encode(&random_bytes);
555
556 let digest = Sha256::digest(verifier.as_bytes());
557 let challenge = engine.encode(digest);
558
559 PkceChallenge {
560 verifier,
561 challenge,
562 }
563}
564
565// -- Authorization URL construction ------------------------------------------
566
567/// Build the authorization URL for the OAuth Authorization Code + PKCE flow.
568pub fn build_authorization_url(
569 auth_server_metadata: &AuthServerMetadata,
570 client_id: &str,
571 redirect_uri: &str,
572 scopes: &[String],
573 resource: &str,
574 pkce: &PkceChallenge,
575 state: &str,
576) -> Url {
577 let mut url = auth_server_metadata.authorization_endpoint.clone();
578 {
579 let mut query = url.query_pairs_mut();
580 query.append_pair("response_type", "code");
581 query.append_pair("client_id", client_id);
582 query.append_pair("redirect_uri", redirect_uri);
583 if !scopes.is_empty() {
584 query.append_pair("scope", &scopes.join(" "));
585 }
586 query.append_pair("resource", resource);
587 query.append_pair("code_challenge", &pkce.challenge);
588 query.append_pair("code_challenge_method", "S256");
589 query.append_pair("state", state);
590 }
591 url
592}
593
594// -- Token endpoint request bodies -------------------------------------------
595
596/// The JSON body returned by the token endpoint on success.
597#[derive(Deserialize)]
598pub struct TokenResponse {
599 pub access_token: String,
600 #[serde(default)]
601 pub refresh_token: Option<String>,
602 #[serde(default)]
603 pub expires_in: Option<u64>,
604 #[serde(default)]
605 pub token_type: Option<String>,
606}
607
608impl std::fmt::Debug for TokenResponse {
609 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
610 f.debug_struct("TokenResponse")
611 .field("access_token", &"[redacted]")
612 .field(
613 "refresh_token",
614 &self.refresh_token.as_ref().map(|_| "[redacted]"),
615 )
616 .field("expires_in", &self.expires_in)
617 .field("token_type", &self.token_type)
618 .finish()
619 }
620}
621
622impl TokenResponse {
623 /// Convert into `OAuthTokens`, computing `expires_at` from `expires_in`.
624 pub fn into_tokens(self) -> OAuthTokens {
625 let expires_at = self
626 .expires_in
627 .map(|secs| SystemTime::now() + Duration::from_secs(secs));
628 OAuthTokens {
629 access_token: self.access_token,
630 refresh_token: self.refresh_token,
631 expires_at,
632 }
633 }
634}
635
636/// An OAuth token error response (RFC 6749 Section 5.2).
637#[derive(Debug, Deserialize, PartialEq)]
638pub struct OAuthTokenError {
639 pub error: String,
640 #[serde(default)]
641 pub error_description: Option<String>,
642}
643
644impl std::fmt::Display for OAuthTokenError {
645 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
646 write!(f, "OAuth token error: {}", self.error)?;
647 if let Some(description) = &self.error_description {
648 write!(f, " ({description})")?;
649 }
650 Ok(())
651 }
652}
653
654impl std::error::Error for OAuthTokenError {}
655
656/// Build the form-encoded body for an authorization code token exchange.
657pub fn token_exchange_params(
658 code: &str,
659 client_id: &str,
660 redirect_uri: &str,
661 code_verifier: &str,
662 resource: &str,
663 client_secret: Option<&str>,
664) -> Vec<(&'static str, String)> {
665 let mut params = vec![
666 ("grant_type", "authorization_code".to_string()),
667 ("code", code.to_string()),
668 ("redirect_uri", redirect_uri.to_string()),
669 ("client_id", client_id.to_string()),
670 ("code_verifier", code_verifier.to_string()),
671 ("resource", resource.to_string()),
672 ];
673 if let Some(secret) = client_secret {
674 params.push(("client_secret", secret.to_string()));
675 }
676 params
677}
678
679/// Build the form-encoded body for a token refresh request.
680pub fn token_refresh_params(
681 refresh_token: &str,
682 client_id: &str,
683 resource: &str,
684 client_secret: Option<&str>,
685) -> Vec<(&'static str, String)> {
686 let mut params = vec![
687 ("grant_type", "refresh_token".to_string()),
688 ("refresh_token", refresh_token.to_string()),
689 ("client_id", client_id.to_string()),
690 ("resource", resource.to_string()),
691 ];
692 if let Some(secret) = client_secret {
693 params.push(("client_secret", secret.to_string()));
694 }
695 params
696}
697
698// -- DCR request body (RFC 7591) ---------------------------------------------
699
700/// Build the JSON body for a Dynamic Client Registration request.
701///
702/// The `redirect_uri` should be the actual loopback URI with the ephemeral
703/// port (e.g. `http://127.0.0.1:12345/callback`). Some auth servers do strict
704/// redirect URI matching even for loopback addresses, so we register the
705/// exact URI we intend to use.
706/// The grant types Zed can use. Intersected with the server's
707/// `grant_types_supported` to build the DCR request.
708const SUPPORTED_GRANT_TYPES: &[&str] = &["authorization_code", "refresh_token"];
709
710pub fn dcr_registration_body(
711 redirect_uri: &str,
712 server_grant_types: Option<&[String]>,
713) -> serde_json::Value {
714 // Use the intersection of what we support and what the server advertises.
715 // When the server doesn't advertise grant_types_supported, send all of
716 // ours — the server will reject what it doesn't like.
717 let grant_types: Vec<&str> = match server_grant_types {
718 Some(server) => SUPPORTED_GRANT_TYPES
719 .iter()
720 .copied()
721 .filter(|gt| server.iter().any(|s| s == *gt))
722 .collect(),
723 None => SUPPORTED_GRANT_TYPES.to_vec(),
724 };
725
726 serde_json::json!({
727 "client_name": "Zed",
728 "redirect_uris": [redirect_uri],
729 "grant_types": grant_types,
730 "response_types": ["code"],
731 "token_endpoint_auth_method": "none"
732 })
733}
734
735// -- Discovery (async, hits real endpoints) ----------------------------------
736
737/// Fetch Protected Resource Metadata from the MCP server.
738///
739/// Tries the `resource_metadata` URL from the `WWW-Authenticate` header first,
740/// then falls back to well-known URIs constructed from `server_url`.
741pub async fn fetch_protected_resource_metadata(
742 http_client: &Arc<dyn HttpClient>,
743 server_url: &Url,
744 www_authenticate: &WwwAuthenticate,
745) -> Result<ProtectedResourceMetadata> {
746 let candidate_urls = match &www_authenticate.resource_metadata {
747 Some(url) if url.origin() == server_url.origin() => {
748 // Try the header-provided URL first (per MCP spec: "use the resource
749 // metadata URL from the parsed WWW-Authenticate headers when present"),
750 // then fall back to RFC 9728 well-known URIs in case the header URL is
751 // wrong (e.g. a buggy server that doubles the path component).
752 let mut urls = vec![url.clone()];
753 for fallback in protected_resource_metadata_urls(server_url) {
754 if !urls.contains(&fallback) {
755 urls.push(fallback);
756 }
757 }
758 urls
759 }
760 Some(url) => {
761 log::warn!(
762 "Ignoring cross-origin resource_metadata URL {} \
763 (server origin: {})",
764 url,
765 server_url.origin().unicode_serialization()
766 );
767 protected_resource_metadata_urls(server_url)
768 }
769 None => protected_resource_metadata_urls(server_url),
770 };
771
772 for url in &candidate_urls {
773 match fetch_json::<ProtectedResourceMetadataResponse>(http_client, url).await {
774 Ok(response) => {
775 if response.authorization_servers.is_empty() {
776 bail!(
777 "Protected Resource Metadata at {} has no authorization_servers",
778 url
779 );
780 }
781 return Ok(ProtectedResourceMetadata {
782 resource: response.resource.unwrap_or_else(|| server_url.clone()),
783 authorization_servers: response.authorization_servers,
784 scopes_supported: response.scopes_supported,
785 });
786 }
787 Err(err) => {
788 log::debug!(
789 "Failed to fetch Protected Resource Metadata from {}: {}",
790 url,
791 err
792 );
793 }
794 }
795 }
796
797 bail!(
798 "Could not fetch Protected Resource Metadata for {}",
799 server_url
800 )
801}
802
803/// Fetch Authorization Server Metadata, trying RFC 8414 and OIDC Discovery
804/// endpoints in the priority order specified by the MCP spec.
805pub async fn fetch_auth_server_metadata(
806 http_client: &Arc<dyn HttpClient>,
807 issuer: &Url,
808) -> Result<AuthServerMetadata> {
809 let candidate_urls = auth_server_metadata_urls(issuer);
810
811 for url in &candidate_urls {
812 match fetch_json::<AuthServerMetadataResponse>(http_client, url).await {
813 Ok(response) => {
814 let reported_issuer = response.issuer.unwrap_or_else(|| issuer.clone());
815
816 if reported_issuer != *issuer {
817 bail!(
818 "Auth server metadata issuer mismatch: expected {}, got {}",
819 issuer,
820 reported_issuer
821 );
822 }
823
824 return Ok(AuthServerMetadata {
825 issuer: reported_issuer,
826 grant_types_supported: response.grant_types_supported,
827 authorization_endpoint: response
828 .authorization_endpoint
829 .ok_or_else(|| anyhow!("missing authorization_endpoint"))?,
830 token_endpoint: response
831 .token_endpoint
832 .ok_or_else(|| anyhow!("missing token_endpoint"))?,
833 registration_endpoint: response.registration_endpoint,
834 scopes_supported: response.scopes_supported,
835 code_challenge_methods_supported: response.code_challenge_methods_supported,
836 client_id_metadata_document_supported: response
837 .client_id_metadata_document_supported
838 .unwrap_or(false),
839 });
840 }
841 Err(err) => {
842 log::debug!("Failed to fetch Auth Server Metadata from {}: {}", url, err);
843 }
844 }
845 }
846
847 bail!(
848 "Could not fetch Authorization Server Metadata for {}",
849 issuer
850 )
851}
852
853/// Run the full discovery flow: fetch resource metadata, then auth server
854/// metadata, then select scopes. Client registration is resolved separately,
855/// once the real redirect URI is known.
856pub async fn discover(
857 http_client: &Arc<dyn HttpClient>,
858 server_url: &Url,
859 www_authenticate: &WwwAuthenticate,
860) -> Result<OAuthDiscovery> {
861 let resource_metadata =
862 fetch_protected_resource_metadata(http_client, server_url, www_authenticate).await?;
863
864 let auth_server_url = resource_metadata
865 .authorization_servers
866 .first()
867 .ok_or_else(|| anyhow!("no authorization servers in resource metadata"))?;
868
869 let auth_server_metadata = fetch_auth_server_metadata(http_client, auth_server_url).await?;
870
871 // Verify PKCE S256 support (spec requirement).
872 match &auth_server_metadata.code_challenge_methods_supported {
873 Some(methods) if methods.iter().any(|m| m == "S256") => {}
874 Some(_) => bail!("authorization server does not support S256 PKCE"),
875 None => bail!("authorization server does not advertise code_challenge_methods_supported"),
876 }
877
878 let scopes = select_scopes(www_authenticate, &resource_metadata);
879
880 Ok(OAuthDiscovery {
881 resource_metadata,
882 auth_server_metadata,
883 scopes,
884 })
885}
886
887/// Resolve the OAuth client registration for an authorization flow.
888///
889/// CIMD uses the static client metadata document directly. For DCR, a fresh
890/// registration is performed each time because the loopback redirect URI
891/// includes an ephemeral port that changes every flow.
892pub async fn resolve_client_registration(
893 http_client: &Arc<dyn HttpClient>,
894 discovery: &OAuthDiscovery,
895 redirect_uri: &str,
896) -> Result<OAuthClientRegistration> {
897 match determine_registration_strategy(&discovery.auth_server_metadata) {
898 ClientRegistrationStrategy::Cimd { client_id } => Ok(OAuthClientRegistration {
899 client_id,
900 client_secret: None,
901 }),
902 ClientRegistrationStrategy::Dcr {
903 registration_endpoint,
904 } => {
905 perform_dcr(
906 http_client,
907 ®istration_endpoint,
908 redirect_uri,
909 discovery
910 .auth_server_metadata
911 .grant_types_supported
912 .as_deref(),
913 )
914 .await
915 }
916 ClientRegistrationStrategy::Unavailable => {
917 bail!("authorization server supports neither CIMD nor DCR")
918 }
919 }
920}
921
922// -- Dynamic Client Registration (RFC 7591) ----------------------------------
923
924/// Perform Dynamic Client Registration with the authorization server.
925pub async fn perform_dcr(
926 http_client: &Arc<dyn HttpClient>,
927 registration_endpoint: &Url,
928 redirect_uri: &str,
929 server_grant_types: Option<&[String]>,
930) -> Result<OAuthClientRegistration> {
931 validate_oauth_url(registration_endpoint)?;
932
933 let body = dcr_registration_body(redirect_uri, server_grant_types);
934 let body_bytes = serde_json::to_vec(&body)?;
935
936 let request = Request::builder()
937 .method(http_client::http::Method::POST)
938 .uri(registration_endpoint.as_str())
939 .header("Content-Type", "application/json")
940 .header("Accept", "application/json")
941 .body(AsyncBody::from(body_bytes))?;
942
943 let mut response = http_client.send(request).await?;
944
945 if !response.status().is_success() {
946 let mut error_body = String::new();
947 response.body_mut().read_to_string(&mut error_body).await?;
948 bail!(
949 "DCR failed with status {}: {}",
950 response.status(),
951 error_body
952 );
953 }
954
955 let mut response_body = String::new();
956 response
957 .body_mut()
958 .read_to_string(&mut response_body)
959 .await?;
960
961 let dcr_response: DcrResponse =
962 serde_json::from_str(&response_body).context("failed to parse DCR response")?;
963
964 Ok(OAuthClientRegistration {
965 client_id: dcr_response.client_id,
966 client_secret: dcr_response.client_secret,
967 })
968}
969
970// -- Token exchange and refresh (async) --------------------------------------
971
972/// Exchange an authorization code for tokens at the token endpoint.
973pub async fn exchange_code(
974 http_client: &Arc<dyn HttpClient>,
975 auth_server_metadata: &AuthServerMetadata,
976 code: &str,
977 client_id: &str,
978 redirect_uri: &str,
979 code_verifier: &str,
980 resource: &str,
981 client_secret: Option<&str>,
982) -> Result<OAuthTokens> {
983 let params = token_exchange_params(
984 code,
985 client_id,
986 redirect_uri,
987 code_verifier,
988 resource,
989 client_secret,
990 );
991 post_token_request(http_client, &auth_server_metadata.token_endpoint, ¶ms).await
992}
993
994/// Refresh tokens using a refresh token.
995pub async fn refresh_tokens(
996 http_client: &Arc<dyn HttpClient>,
997 token_endpoint: &Url,
998 refresh_token: &str,
999 client_id: &str,
1000 resource: &str,
1001 client_secret: Option<&str>,
1002) -> Result<OAuthTokens> {
1003 let params = token_refresh_params(refresh_token, client_id, resource, client_secret);
1004 post_token_request(http_client, token_endpoint, ¶ms).await
1005}
1006
1007/// POST form-encoded parameters to a token endpoint and parse the response.
1008async fn post_token_request(
1009 http_client: &Arc<dyn HttpClient>,
1010 token_endpoint: &Url,
1011 params: &[(&str, String)],
1012) -> Result<OAuthTokens> {
1013 validate_oauth_url(token_endpoint)?;
1014
1015 let body = url::form_urlencoded::Serializer::new(String::new())
1016 .extend_pairs(params.iter().map(|(k, v)| (*k, v.as_str())))
1017 .finish();
1018
1019 let request = Request::builder()
1020 .method(http_client::http::Method::POST)
1021 .uri(token_endpoint.as_str())
1022 .header("Content-Type", "application/x-www-form-urlencoded")
1023 .header("Accept", "application/json")
1024 .body(AsyncBody::from(body.into_bytes()))?;
1025
1026 let mut response = http_client.send(request).await?;
1027
1028 if !response.status().is_success() {
1029 let mut error_body = String::new();
1030 response.body_mut().read_to_string(&mut error_body).await?;
1031 let status = response.status();
1032 // Try to parse as an OAuth error response (RFC 6749 Section 5.2).
1033 if let Ok(token_error) = serde_json::from_str::<OAuthTokenError>(&error_body) {
1034 return Err(token_error.into());
1035 }
1036 bail!("token request failed with status {status}: {error_body}");
1037 }
1038
1039 let mut response_body = String::new();
1040 response
1041 .body_mut()
1042 .read_to_string(&mut response_body)
1043 .await?;
1044
1045 let token_response: TokenResponse =
1046 serde_json::from_str(&response_body).context("failed to parse token response")?;
1047
1048 Ok(token_response.into_tokens())
1049}
1050
1051// -- Loopback HTTP callback server -------------------------------------------
1052
1053/// An OAuth authorization callback received via the loopback HTTP server.
1054pub struct OAuthCallback {
1055 pub code: String,
1056 pub state: String,
1057}
1058
1059impl std::fmt::Debug for OAuthCallback {
1060 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1061 f.debug_struct("OAuthCallback")
1062 .field("code", &"[redacted]")
1063 .field("state", &"[redacted]")
1064 .finish()
1065 }
1066}
1067
1068impl OAuthCallback {
1069 /// Parse the query string from a callback URL like
1070 /// `http://127.0.0.1:<port>/callback?code=...&state=...`.
1071 pub fn parse_query(query: &str) -> Result<Self> {
1072 let params = oauth_callback_server::OAuthCallbackParams::parse_query(query)?;
1073 Ok(Self {
1074 code: params.code,
1075 state: params.state,
1076 })
1077 }
1078}
1079
1080/// Start a loopback HTTP server to receive the OAuth authorization callback.
1081///
1082/// Binds to an ephemeral loopback port for each flow.
1083///
1084/// Returns `(redirect_uri, callback_future)`. The caller should use the
1085/// redirect URI in the authorization request, open the browser, then await
1086/// the future to receive the callback.
1087///
1088/// The server accepts exactly one request on `/callback`, validates that it
1089/// contains `code` and `state` query parameters, responds with a minimal
1090/// HTML page telling the user they can close the tab, and shuts down.
1091///
1092/// The callback server shuts down when the returned future is dropped (e.g.
1093/// because the authentication task was cancelled), or after a timeout.
1094pub fn start_callback_server() -> Result<(String, BoxFuture<'static, Result<OAuthCallback>>)> {
1095 let (redirect_uri, rx) = oauth_callback_server::start_oauth_callback_server()?;
1096 let future = async move {
1097 match rx.await {
1098 Ok(Ok(params)) => Ok(OAuthCallback {
1099 code: params.code,
1100 state: params.state,
1101 }),
1102 Ok(Err(e)) => Err(e),
1103 Err(_) => Err(anyhow!(
1104 "OAuth callback server was shut down before receiving a response"
1105 )),
1106 }
1107 }
1108 .boxed();
1109 Ok((redirect_uri, future))
1110}
1111
1112// -- JSON fetch helper -------------------------------------------------------
1113
1114async fn fetch_json<T: serde::de::DeserializeOwned>(
1115 http_client: &Arc<dyn HttpClient>,
1116 url: &Url,
1117) -> Result<T> {
1118 validate_oauth_url(url)?;
1119
1120 let request = Request::builder()
1121 .method(http_client::http::Method::GET)
1122 .uri(url.as_str())
1123 .header("Accept", "application/json")
1124 .body(AsyncBody::default())?;
1125
1126 let mut response = http_client.send(request).await?;
1127
1128 if !response.status().is_success() {
1129 bail!("HTTP {} fetching {}", response.status(), url);
1130 }
1131
1132 let mut body = String::new();
1133 response.body_mut().read_to_string(&mut body).await?;
1134 serde_json::from_str(&body).with_context(|| format!("failed to parse JSON from {}", url))
1135}
1136
1137// -- Serde response types for discovery --------------------------------------
1138
1139#[derive(Debug, Deserialize)]
1140struct ProtectedResourceMetadataResponse {
1141 #[serde(default)]
1142 resource: Option<Url>,
1143 #[serde(default)]
1144 authorization_servers: Vec<Url>,
1145 #[serde(default)]
1146 scopes_supported: Option<Vec<String>>,
1147}
1148
1149#[derive(Debug, Deserialize)]
1150struct AuthServerMetadataResponse {
1151 #[serde(default)]
1152 issuer: Option<Url>,
1153 #[serde(default)]
1154 authorization_endpoint: Option<Url>,
1155 #[serde(default)]
1156 token_endpoint: Option<Url>,
1157 #[serde(default)]
1158 registration_endpoint: Option<Url>,
1159 #[serde(default)]
1160 scopes_supported: Option<Vec<String>>,
1161 #[serde(default)]
1162 grant_types_supported: Option<Vec<String>>,
1163 #[serde(default)]
1164 code_challenge_methods_supported: Option<Vec<String>>,
1165 #[serde(default)]
1166 client_id_metadata_document_supported: Option<bool>,
1167}
1168
1169#[derive(Debug, Deserialize)]
1170struct DcrResponse {
1171 client_id: String,
1172 #[serde(default)]
1173 client_secret: Option<String>,
1174}
1175
1176/// Provides OAuth tokens to the HTTP transport layer.
1177///
1178/// The transport calls `access_token()` before each request. On a 401 response
1179/// it calls `try_refresh()` and retries once if the refresh succeeds.
1180#[async_trait]
1181pub trait OAuthTokenProvider: Send + Sync {
1182 /// Returns the current access token, if one is available.
1183 fn access_token(&self) -> Option<String>;
1184
1185 /// Attempts to refresh the access token. Returns `true` if a new token was
1186 /// obtained and the request should be retried.
1187 async fn try_refresh(&self) -> Result<bool>;
1188}
1189
1190/// Concrete `OAuthTokenProvider` backed by a full persisted OAuth session and
1191/// an HTTP client for token refresh. The same provider type is used both after
1192/// an interactive authentication flow and when restoring a saved session from
1193/// the keychain on startup.
1194pub struct McpOAuthTokenProvider {
1195 session: SyncMutex<OAuthSession>,
1196 http_client: Arc<dyn HttpClient>,
1197 token_refresh_tx: Option<mpsc::UnboundedSender<OAuthSession>>,
1198}
1199
1200impl McpOAuthTokenProvider {
1201 pub fn new(
1202 session: OAuthSession,
1203 http_client: Arc<dyn HttpClient>,
1204 token_refresh_tx: Option<mpsc::UnboundedSender<OAuthSession>>,
1205 ) -> Self {
1206 Self {
1207 session: SyncMutex::new(session),
1208 http_client,
1209 token_refresh_tx,
1210 }
1211 }
1212
1213 fn access_token_is_expired(tokens: &OAuthTokens) -> bool {
1214 tokens.expires_at.is_some_and(|expires_at| {
1215 SystemTime::now()
1216 .checked_add(Duration::from_secs(30))
1217 .is_some_and(|now_with_buffer| expires_at <= now_with_buffer)
1218 })
1219 }
1220}
1221
1222#[async_trait]
1223impl OAuthTokenProvider for McpOAuthTokenProvider {
1224 fn access_token(&self) -> Option<String> {
1225 let session = self.session.lock();
1226 if Self::access_token_is_expired(&session.tokens) {
1227 return None;
1228 }
1229 Some(session.tokens.access_token.clone())
1230 }
1231
1232 async fn try_refresh(&self) -> Result<bool> {
1233 let (refresh_token, token_endpoint, resource, client_id, client_secret) = {
1234 let session = self.session.lock();
1235 match session.tokens.refresh_token.clone() {
1236 Some(refresh_token) => (
1237 refresh_token,
1238 session.token_endpoint.clone(),
1239 session.resource.clone(),
1240 session.client_registration.client_id.clone(),
1241 session.client_registration.client_secret.clone(),
1242 ),
1243 None => return Ok(false),
1244 }
1245 };
1246
1247 let resource_str = canonical_server_uri(&resource);
1248
1249 match refresh_tokens(
1250 &self.http_client,
1251 &token_endpoint,
1252 &refresh_token,
1253 &client_id,
1254 &resource_str,
1255 client_secret.as_deref(),
1256 )
1257 .await
1258 {
1259 Ok(mut new_tokens) => {
1260 if new_tokens.refresh_token.is_none() {
1261 new_tokens.refresh_token = Some(refresh_token);
1262 }
1263
1264 {
1265 let mut session = self.session.lock();
1266 session.tokens = new_tokens;
1267
1268 if let Some(ref tx) = self.token_refresh_tx {
1269 tx.unbounded_send(session.clone()).ok();
1270 }
1271 }
1272
1273 Ok(true)
1274 }
1275 Err(err) => {
1276 log::warn!("OAuth token refresh failed: {}", err);
1277 Ok(false)
1278 }
1279 }
1280 }
1281}
1282
1283#[cfg(test)]
1284mod tests {
1285 use super::*;
1286 use http_client::Response;
1287
1288 // -- require_https_or_loopback tests ------------------------------------
1289
1290 #[test]
1291 fn test_require_https_or_loopback_accepts_https() {
1292 let url = Url::parse("https://auth.example.com/token").unwrap();
1293 assert!(require_https_or_loopback(&url).is_ok());
1294 }
1295
1296 #[test]
1297 fn test_require_https_or_loopback_rejects_http_remote() {
1298 let url = Url::parse("http://auth.example.com/token").unwrap();
1299 assert!(require_https_or_loopback(&url).is_err());
1300 }
1301
1302 #[test]
1303 fn test_require_https_or_loopback_accepts_http_127_0_0_1() {
1304 let url = Url::parse("http://127.0.0.1:8080/callback").unwrap();
1305 assert!(require_https_or_loopback(&url).is_ok());
1306 }
1307
1308 #[test]
1309 fn test_require_https_or_loopback_accepts_http_ipv6_loopback() {
1310 let url = Url::parse("http://[::1]:8080/callback").unwrap();
1311 assert!(require_https_or_loopback(&url).is_ok());
1312 }
1313
1314 #[test]
1315 fn test_require_https_or_loopback_accepts_http_localhost() {
1316 let url = Url::parse("http://localhost:8080/callback").unwrap();
1317 assert!(require_https_or_loopback(&url).is_ok());
1318 }
1319
1320 #[test]
1321 fn test_require_https_or_loopback_accepts_http_localhost_case_insensitive() {
1322 let url = Url::parse("http://LOCALHOST:8080/callback").unwrap();
1323 assert!(require_https_or_loopback(&url).is_ok());
1324 }
1325
1326 #[test]
1327 fn test_require_https_or_loopback_rejects_http_non_loopback_ip() {
1328 let url = Url::parse("http://192.168.1.1:8080/token").unwrap();
1329 assert!(require_https_or_loopback(&url).is_err());
1330 }
1331
1332 #[test]
1333 fn test_require_https_or_loopback_rejects_ftp() {
1334 let url = Url::parse("ftp://auth.example.com/token").unwrap();
1335 assert!(require_https_or_loopback(&url).is_err());
1336 }
1337
1338 // -- validate_oauth_url (SSRF) tests ------------------------------------
1339
1340 #[test]
1341 fn test_validate_oauth_url_accepts_https_public() {
1342 let url = Url::parse("https://auth.example.com/token").unwrap();
1343 assert!(validate_oauth_url(&url).is_ok());
1344 }
1345
1346 #[test]
1347 fn test_validate_oauth_url_rejects_private_ipv4_10() {
1348 let url = Url::parse("https://10.0.0.1/token").unwrap();
1349 assert!(validate_oauth_url(&url).is_err());
1350 }
1351
1352 #[test]
1353 fn test_validate_oauth_url_rejects_private_ipv4_172() {
1354 let url = Url::parse("https://172.16.0.1/token").unwrap();
1355 assert!(validate_oauth_url(&url).is_err());
1356 }
1357
1358 #[test]
1359 fn test_validate_oauth_url_rejects_private_ipv4_192() {
1360 let url = Url::parse("https://192.168.1.1/token").unwrap();
1361 assert!(validate_oauth_url(&url).is_err());
1362 }
1363
1364 #[test]
1365 fn test_validate_oauth_url_rejects_link_local() {
1366 let url = Url::parse("https://169.254.169.254/latest/meta-data/").unwrap();
1367 assert!(validate_oauth_url(&url).is_err());
1368 }
1369
1370 #[test]
1371 fn test_validate_oauth_url_rejects_ipv6_ula() {
1372 let url = Url::parse("https://[fd12:3456:789a::1]/token").unwrap();
1373 assert!(validate_oauth_url(&url).is_err());
1374 }
1375
1376 #[test]
1377 fn test_validate_oauth_url_rejects_ipv6_unspecified() {
1378 let url = Url::parse("https://[::]/token").unwrap();
1379 assert!(validate_oauth_url(&url).is_err());
1380 }
1381
1382 #[test]
1383 fn test_validate_oauth_url_rejects_ipv4_mapped_ipv6_private() {
1384 let url = Url::parse("https://[::ffff:10.0.0.1]/token").unwrap();
1385 assert!(validate_oauth_url(&url).is_err());
1386 }
1387
1388 #[test]
1389 fn test_validate_oauth_url_rejects_ipv4_mapped_ipv6_link_local() {
1390 let url = Url::parse("https://[::ffff:169.254.169.254]/token").unwrap();
1391 assert!(validate_oauth_url(&url).is_err());
1392 }
1393
1394 #[test]
1395 fn test_validate_oauth_url_allows_http_loopback() {
1396 // Loopback is permitted (it's our callback server).
1397 let url = Url::parse("http://127.0.0.1:8080/callback").unwrap();
1398 assert!(validate_oauth_url(&url).is_ok());
1399 }
1400
1401 #[test]
1402 fn test_validate_oauth_url_allows_https_public_ip() {
1403 let url = Url::parse("https://93.184.216.34/token").unwrap();
1404 assert!(validate_oauth_url(&url).is_ok());
1405 }
1406
1407 // -- parse_www_authenticate tests ----------------------------------------
1408
1409 #[test]
1410 fn test_parse_www_authenticate_with_resource_metadata_and_scope() {
1411 let header = r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", scope="files:read user:profile""#;
1412 let result = parse_www_authenticate(header).unwrap();
1413
1414 assert_eq!(
1415 result.resource_metadata.as_ref().map(|u| u.as_str()),
1416 Some("https://mcp.example.com/.well-known/oauth-protected-resource")
1417 );
1418 assert_eq!(
1419 result.scope,
1420 Some(vec!["files:read".to_string(), "user:profile".to_string()])
1421 );
1422 assert_eq!(result.error, None);
1423 }
1424
1425 #[test]
1426 fn test_parse_www_authenticate_resource_metadata_only() {
1427 let header = r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#;
1428 let result = parse_www_authenticate(header).unwrap();
1429
1430 assert_eq!(
1431 result.resource_metadata.as_ref().map(|u| u.as_str()),
1432 Some("https://mcp.example.com/.well-known/oauth-protected-resource")
1433 );
1434 assert_eq!(result.scope, None);
1435 }
1436
1437 #[test]
1438 fn test_parse_www_authenticate_bare_bearer() {
1439 let result = parse_www_authenticate("Bearer").unwrap();
1440 assert_eq!(result.resource_metadata, None);
1441 assert_eq!(result.scope, None);
1442 }
1443
1444 #[test]
1445 fn test_parse_www_authenticate_with_error() {
1446 let header = r#"Bearer error="insufficient_scope", scope="files:read files:write", resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", error_description="Additional file write permission required""#;
1447 let result = parse_www_authenticate(header).unwrap();
1448
1449 assert_eq!(result.error, Some(BearerError::InsufficientScope));
1450 assert_eq!(
1451 result.error_description.as_deref(),
1452 Some("Additional file write permission required")
1453 );
1454 assert_eq!(
1455 result.scope,
1456 Some(vec!["files:read".to_string(), "files:write".to_string()])
1457 );
1458 assert!(result.resource_metadata.is_some());
1459 }
1460
1461 #[test]
1462 fn test_parse_www_authenticate_invalid_token_error() {
1463 let header =
1464 r#"Bearer error="invalid_token", error_description="The access token expired""#;
1465 let result = parse_www_authenticate(header).unwrap();
1466 assert_eq!(result.error, Some(BearerError::InvalidToken));
1467 }
1468
1469 #[test]
1470 fn test_parse_www_authenticate_invalid_request_error() {
1471 let header = r#"Bearer error="invalid_request""#;
1472 let result = parse_www_authenticate(header).unwrap();
1473 assert_eq!(result.error, Some(BearerError::InvalidRequest));
1474 }
1475
1476 #[test]
1477 fn test_parse_www_authenticate_unknown_error() {
1478 let header = r#"Bearer error="some_future_error""#;
1479 let result = parse_www_authenticate(header).unwrap();
1480 assert_eq!(result.error, Some(BearerError::Other));
1481 }
1482
1483 #[test]
1484 fn test_parse_www_authenticate_rejects_non_bearer() {
1485 let result = parse_www_authenticate("Basic realm=\"example\"");
1486 assert!(result.is_err());
1487 }
1488
1489 #[test]
1490 fn test_parse_www_authenticate_case_insensitive_scheme() {
1491 let header = r#"bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource""#;
1492 let result = parse_www_authenticate(header).unwrap();
1493 assert!(result.resource_metadata.is_some());
1494 }
1495
1496 #[test]
1497 fn test_parse_www_authenticate_multiline_style() {
1498 // Some servers emit the header spread across multiple lines joined by
1499 // whitespace, as shown in the spec examples.
1500 let header = "Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\",\n scope=\"files:read\"";
1501 let result = parse_www_authenticate(header).unwrap();
1502 assert!(result.resource_metadata.is_some());
1503 assert_eq!(result.scope, Some(vec!["files:read".to_string()]));
1504 }
1505
1506 #[test]
1507 fn test_protected_resource_metadata_urls_with_path() {
1508 let server_url = Url::parse("https://api.example.com/v1/mcp").unwrap();
1509 let urls = protected_resource_metadata_urls(&server_url);
1510
1511 assert_eq!(urls.len(), 2);
1512 assert_eq!(
1513 urls[0].as_str(),
1514 "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
1515 );
1516 assert_eq!(
1517 urls[1].as_str(),
1518 "https://api.example.com/.well-known/oauth-protected-resource"
1519 );
1520 }
1521
1522 #[test]
1523 fn test_protected_resource_metadata_urls_without_path() {
1524 let server_url = Url::parse("https://mcp.example.com").unwrap();
1525 let urls = protected_resource_metadata_urls(&server_url);
1526
1527 assert_eq!(urls.len(), 1);
1528 assert_eq!(
1529 urls[0].as_str(),
1530 "https://mcp.example.com/.well-known/oauth-protected-resource"
1531 );
1532 }
1533
1534 #[test]
1535 fn test_auth_server_metadata_urls_with_path() {
1536 let issuer = Url::parse("https://auth.example.com/tenant1").unwrap();
1537 let urls = auth_server_metadata_urls(&issuer);
1538
1539 assert_eq!(urls.len(), 3);
1540 assert_eq!(
1541 urls[0].as_str(),
1542 "https://auth.example.com/.well-known/oauth-authorization-server/tenant1"
1543 );
1544 assert_eq!(
1545 urls[1].as_str(),
1546 "https://auth.example.com/.well-known/openid-configuration/tenant1"
1547 );
1548 assert_eq!(
1549 urls[2].as_str(),
1550 "https://auth.example.com/tenant1/.well-known/openid-configuration"
1551 );
1552 }
1553
1554 #[test]
1555 fn test_auth_server_metadata_urls_without_path() {
1556 let issuer = Url::parse("https://auth.example.com").unwrap();
1557 let urls = auth_server_metadata_urls(&issuer);
1558
1559 assert_eq!(urls.len(), 2);
1560 assert_eq!(
1561 urls[0].as_str(),
1562 "https://auth.example.com/.well-known/oauth-authorization-server"
1563 );
1564 assert_eq!(
1565 urls[1].as_str(),
1566 "https://auth.example.com/.well-known/openid-configuration"
1567 );
1568 }
1569
1570 // -- Canonical server URI tests ------------------------------------------
1571
1572 #[test]
1573 fn test_canonical_server_uri_simple() {
1574 let url = Url::parse("https://mcp.example.com").unwrap();
1575 assert_eq!(canonical_server_uri(&url), "https://mcp.example.com");
1576 }
1577
1578 #[test]
1579 fn test_canonical_server_uri_with_path() {
1580 let url = Url::parse("https://mcp.example.com/v1/mcp").unwrap();
1581 assert_eq!(canonical_server_uri(&url), "https://mcp.example.com/v1/mcp");
1582 }
1583
1584 #[test]
1585 fn test_canonical_server_uri_strips_trailing_slash() {
1586 let url = Url::parse("https://mcp.example.com/").unwrap();
1587 assert_eq!(canonical_server_uri(&url), "https://mcp.example.com");
1588 }
1589
1590 #[test]
1591 fn test_canonical_server_uri_preserves_port() {
1592 let url = Url::parse("https://mcp.example.com:8443").unwrap();
1593 assert_eq!(canonical_server_uri(&url), "https://mcp.example.com:8443");
1594 }
1595
1596 #[test]
1597 fn test_canonical_server_uri_lowercases() {
1598 let url = Url::parse("HTTPS://MCP.Example.COM/Server/MCP").unwrap();
1599 assert_eq!(
1600 canonical_server_uri(&url),
1601 "https://mcp.example.com/Server/MCP"
1602 );
1603 }
1604
1605 // -- Scope selection tests -----------------------------------------------
1606
1607 #[test]
1608 fn test_select_scopes_prefers_www_authenticate() {
1609 let www_auth = WwwAuthenticate {
1610 resource_metadata: None,
1611 scope: Some(vec!["files:read".into()]),
1612 error: None,
1613 error_description: None,
1614 };
1615 let resource_meta = ProtectedResourceMetadata {
1616 resource: Url::parse("https://example.com").unwrap(),
1617 authorization_servers: vec![],
1618 scopes_supported: Some(vec!["files:read".into(), "files:write".into()]),
1619 };
1620 assert_eq!(select_scopes(&www_auth, &resource_meta), vec!["files:read"]);
1621 }
1622
1623 #[test]
1624 fn test_select_scopes_falls_back_to_resource_metadata() {
1625 let www_auth = WwwAuthenticate {
1626 resource_metadata: None,
1627 scope: None,
1628 error: None,
1629 error_description: None,
1630 };
1631 let resource_meta = ProtectedResourceMetadata {
1632 resource: Url::parse("https://example.com").unwrap(),
1633 authorization_servers: vec![],
1634 scopes_supported: Some(vec!["admin".into()]),
1635 };
1636 assert_eq!(select_scopes(&www_auth, &resource_meta), vec!["admin"]);
1637 }
1638
1639 #[test]
1640 fn test_select_scopes_empty_when_nothing_available() {
1641 let www_auth = WwwAuthenticate {
1642 resource_metadata: None,
1643 scope: None,
1644 error: None,
1645 error_description: None,
1646 };
1647 let resource_meta = ProtectedResourceMetadata {
1648 resource: Url::parse("https://example.com").unwrap(),
1649 authorization_servers: vec![],
1650 scopes_supported: None,
1651 };
1652 assert!(select_scopes(&www_auth, &resource_meta).is_empty());
1653 }
1654
1655 // -- Client registration strategy tests ----------------------------------
1656
1657 #[test]
1658 fn test_registration_strategy_prefers_cimd() {
1659 let metadata = AuthServerMetadata {
1660 issuer: Url::parse("https://auth.example.com").unwrap(),
1661 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
1662 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
1663 registration_endpoint: Some(Url::parse("https://auth.example.com/register").unwrap()),
1664 scopes_supported: None,
1665 code_challenge_methods_supported: Some(vec!["S256".into()]),
1666 client_id_metadata_document_supported: true,
1667 grant_types_supported: None,
1668 };
1669 assert_eq!(
1670 determine_registration_strategy(&metadata),
1671 ClientRegistrationStrategy::Cimd {
1672 client_id: CIMD_URL.to_string(),
1673 }
1674 );
1675 }
1676
1677 #[test]
1678 fn test_registration_strategy_falls_back_to_dcr() {
1679 let reg_endpoint = Url::parse("https://auth.example.com/register").unwrap();
1680 let metadata = AuthServerMetadata {
1681 issuer: Url::parse("https://auth.example.com").unwrap(),
1682 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
1683 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
1684 registration_endpoint: Some(reg_endpoint.clone()),
1685 scopes_supported: None,
1686 code_challenge_methods_supported: Some(vec!["S256".into()]),
1687 client_id_metadata_document_supported: false,
1688 grant_types_supported: None,
1689 };
1690 assert_eq!(
1691 determine_registration_strategy(&metadata),
1692 ClientRegistrationStrategy::Dcr {
1693 registration_endpoint: reg_endpoint,
1694 }
1695 );
1696 }
1697
1698 #[test]
1699 fn test_registration_strategy_unavailable() {
1700 let metadata = AuthServerMetadata {
1701 issuer: Url::parse("https://auth.example.com").unwrap(),
1702 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
1703 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
1704 registration_endpoint: None,
1705 scopes_supported: None,
1706 code_challenge_methods_supported: Some(vec!["S256".into()]),
1707 client_id_metadata_document_supported: false,
1708 grant_types_supported: None,
1709 };
1710 assert_eq!(
1711 determine_registration_strategy(&metadata),
1712 ClientRegistrationStrategy::Unavailable,
1713 );
1714 }
1715
1716 // -- PKCE tests ----------------------------------------------------------
1717
1718 #[test]
1719 fn test_pkce_challenge_verifier_length() {
1720 let pkce = generate_pkce_challenge();
1721 // 32 random bytes → 43 base64url chars (no padding).
1722 assert_eq!(pkce.verifier.len(), 43);
1723 }
1724
1725 #[test]
1726 fn test_pkce_challenge_is_valid_base64url() {
1727 let pkce = generate_pkce_challenge();
1728 for c in pkce.verifier.chars().chain(pkce.challenge.chars()) {
1729 assert!(
1730 c.is_ascii_alphanumeric() || c == '-' || c == '_',
1731 "invalid base64url character: {}",
1732 c
1733 );
1734 }
1735 }
1736
1737 #[test]
1738 fn test_pkce_challenge_is_s256_of_verifier() {
1739 let pkce = generate_pkce_challenge();
1740 let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
1741 let expected_digest = Sha256::digest(pkce.verifier.as_bytes());
1742 let expected_challenge = engine.encode(expected_digest);
1743 assert_eq!(pkce.challenge, expected_challenge);
1744 }
1745
1746 #[test]
1747 fn test_pkce_challenges_are_unique() {
1748 let a = generate_pkce_challenge();
1749 let b = generate_pkce_challenge();
1750 assert_ne!(a.verifier, b.verifier);
1751 }
1752
1753 // -- Authorization URL tests ---------------------------------------------
1754
1755 #[test]
1756 fn test_build_authorization_url() {
1757 let metadata = AuthServerMetadata {
1758 issuer: Url::parse("https://auth.example.com").unwrap(),
1759 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
1760 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
1761 registration_endpoint: None,
1762 scopes_supported: None,
1763 code_challenge_methods_supported: Some(vec!["S256".into()]),
1764 client_id_metadata_document_supported: true,
1765 grant_types_supported: None,
1766 };
1767 let pkce = PkceChallenge {
1768 verifier: "test_verifier".into(),
1769 challenge: "test_challenge".into(),
1770 };
1771 let url = build_authorization_url(
1772 &metadata,
1773 "https://zed.dev/oauth/client-metadata.json",
1774 "http://127.0.0.1:12345/callback",
1775 &["files:read".into(), "files:write".into()],
1776 "https://mcp.example.com",
1777 &pkce,
1778 "random_state_123",
1779 );
1780
1781 let pairs: std::collections::HashMap<_, _> = url.query_pairs().collect();
1782 assert_eq!(pairs.get("response_type").unwrap(), "code");
1783 assert_eq!(
1784 pairs.get("client_id").unwrap(),
1785 "https://zed.dev/oauth/client-metadata.json"
1786 );
1787 assert_eq!(
1788 pairs.get("redirect_uri").unwrap(),
1789 "http://127.0.0.1:12345/callback"
1790 );
1791 assert_eq!(pairs.get("scope").unwrap(), "files:read files:write");
1792 assert_eq!(pairs.get("resource").unwrap(), "https://mcp.example.com");
1793 assert_eq!(pairs.get("code_challenge").unwrap(), "test_challenge");
1794 assert_eq!(pairs.get("code_challenge_method").unwrap(), "S256");
1795 assert_eq!(pairs.get("state").unwrap(), "random_state_123");
1796 }
1797
1798 #[test]
1799 fn test_build_authorization_url_omits_empty_scope() {
1800 let metadata = AuthServerMetadata {
1801 issuer: Url::parse("https://auth.example.com").unwrap(),
1802 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
1803 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
1804 registration_endpoint: None,
1805 scopes_supported: None,
1806 code_challenge_methods_supported: Some(vec!["S256".into()]),
1807 client_id_metadata_document_supported: false,
1808 grant_types_supported: None,
1809 };
1810 let pkce = PkceChallenge {
1811 verifier: "v".into(),
1812 challenge: "c".into(),
1813 };
1814 let url = build_authorization_url(
1815 &metadata,
1816 "client_123",
1817 "http://127.0.0.1:9999/callback",
1818 &[],
1819 "https://mcp.example.com",
1820 &pkce,
1821 "state",
1822 );
1823
1824 let pairs: std::collections::HashMap<_, _> = url.query_pairs().collect();
1825 assert!(!pairs.contains_key("scope"));
1826 }
1827
1828 // -- Token exchange / refresh param tests --------------------------------
1829
1830 #[test]
1831 fn test_token_exchange_params() {
1832 let params = token_exchange_params(
1833 "auth_code_abc",
1834 "client_xyz",
1835 "http://127.0.0.1:5555/callback",
1836 "verifier_123",
1837 "https://mcp.example.com",
1838 None,
1839 );
1840 let map: std::collections::HashMap<&str, &str> =
1841 params.iter().map(|(k, v)| (*k, v.as_str())).collect();
1842
1843 assert_eq!(map["grant_type"], "authorization_code");
1844 assert_eq!(map["code"], "auth_code_abc");
1845 assert_eq!(map["redirect_uri"], "http://127.0.0.1:5555/callback");
1846 assert_eq!(map["client_id"], "client_xyz");
1847 assert_eq!(map["code_verifier"], "verifier_123");
1848 assert_eq!(map["resource"], "https://mcp.example.com");
1849 }
1850
1851 #[test]
1852 fn test_token_refresh_params() {
1853 let params = token_refresh_params(
1854 "refresh_token_abc",
1855 "client_xyz",
1856 "https://mcp.example.com",
1857 None,
1858 );
1859 let map: std::collections::HashMap<&str, &str> =
1860 params.iter().map(|(k, v)| (*k, v.as_str())).collect();
1861
1862 assert_eq!(map["grant_type"], "refresh_token");
1863 assert_eq!(map["refresh_token"], "refresh_token_abc");
1864 assert_eq!(map["client_id"], "client_xyz");
1865 assert_eq!(map["resource"], "https://mcp.example.com");
1866 }
1867
1868 // -- Token response tests ------------------------------------------------
1869
1870 #[test]
1871 fn test_token_response_into_tokens_with_expiry() {
1872 let response: TokenResponse = serde_json::from_str(
1873 r#"{"access_token": "at_123", "refresh_token": "rt_456", "expires_in": 3600, "token_type": "Bearer"}"#,
1874 )
1875 .unwrap();
1876
1877 let tokens = response.into_tokens();
1878 assert_eq!(tokens.access_token, "at_123");
1879 assert_eq!(tokens.refresh_token.as_deref(), Some("rt_456"));
1880 assert!(tokens.expires_at.is_some());
1881 }
1882
1883 #[test]
1884 fn test_token_response_into_tokens_minimal() {
1885 let response: TokenResponse =
1886 serde_json::from_str(r#"{"access_token": "at_789"}"#).unwrap();
1887
1888 let tokens = response.into_tokens();
1889 assert_eq!(tokens.access_token, "at_789");
1890 assert_eq!(tokens.refresh_token, None);
1891 assert_eq!(tokens.expires_at, None);
1892 }
1893
1894 // -- DCR body test -------------------------------------------------------
1895
1896 #[test]
1897 fn test_dcr_registration_body_without_server_metadata() {
1898 // When server metadata is unavailable, include all supported grant types.
1899 let body = dcr_registration_body("http://127.0.0.1:12345/callback", None);
1900 assert_eq!(body["client_name"], "Zed");
1901 assert_eq!(body["redirect_uris"][0], "http://127.0.0.1:12345/callback");
1902 assert_eq!(body["grant_types"][0], "authorization_code");
1903 assert_eq!(body["grant_types"][1], "refresh_token");
1904 assert_eq!(body["response_types"][0], "code");
1905 assert_eq!(body["token_endpoint_auth_method"], "none");
1906 }
1907
1908 #[test]
1909 fn test_dcr_registration_body_mirrors_server_grant_types() {
1910 // When the server only supports authorization_code, omit refresh_token.
1911 let server_types = vec!["authorization_code".to_string()];
1912 let body = dcr_registration_body("http://127.0.0.1:12345/callback", Some(&server_types));
1913 assert_eq!(body["grant_types"][0], "authorization_code");
1914 assert!(body["grant_types"].as_array().unwrap().len() == 1);
1915
1916 // When the server supports both, include both.
1917 let server_types = vec![
1918 "authorization_code".to_string(),
1919 "refresh_token".to_string(),
1920 ];
1921 let body = dcr_registration_body("http://127.0.0.1:12345/callback", Some(&server_types));
1922 assert_eq!(body["grant_types"][0], "authorization_code");
1923 assert_eq!(body["grant_types"][1], "refresh_token");
1924 }
1925
1926 // -- Test helpers for async/HTTP tests -----------------------------------
1927
1928 fn make_fake_http_client(
1929 handler: impl Fn(
1930 http_client::Request<AsyncBody>,
1931 ) -> std::pin::Pin<
1932 Box<dyn std::future::Future<Output = anyhow::Result<Response<AsyncBody>>> + Send>,
1933 > + Send
1934 + Sync
1935 + 'static,
1936 ) -> Arc<dyn HttpClient> {
1937 http_client::FakeHttpClient::create(handler) as Arc<dyn HttpClient>
1938 }
1939
1940 fn json_response(status: u16, body: &str) -> anyhow::Result<Response<AsyncBody>> {
1941 Ok(Response::builder()
1942 .status(status)
1943 .header("Content-Type", "application/json")
1944 .body(AsyncBody::from(body.as_bytes().to_vec()))
1945 .unwrap())
1946 }
1947
1948 // -- Discovery integration tests -----------------------------------------
1949
1950 #[test]
1951 fn test_fetch_protected_resource_metadata() {
1952 gpui::block_on(async {
1953 let client = make_fake_http_client(|req| {
1954 Box::pin(async move {
1955 let uri = req.uri().to_string();
1956 if uri.contains(".well-known/oauth-protected-resource") {
1957 json_response(
1958 200,
1959 r#"{
1960 "resource": "https://mcp.example.com",
1961 "authorization_servers": ["https://auth.example.com"],
1962 "scopes_supported": ["read", "write"]
1963 }"#,
1964 )
1965 } else {
1966 json_response(404, "{}")
1967 }
1968 })
1969 });
1970
1971 let server_url = Url::parse("https://mcp.example.com").unwrap();
1972 let www_auth = WwwAuthenticate {
1973 resource_metadata: None,
1974 scope: None,
1975 error: None,
1976 error_description: None,
1977 };
1978
1979 let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
1980 .await
1981 .unwrap();
1982
1983 assert_eq!(metadata.resource.as_str(), "https://mcp.example.com/");
1984 assert_eq!(metadata.authorization_servers.len(), 1);
1985 assert_eq!(
1986 metadata.authorization_servers[0].as_str(),
1987 "https://auth.example.com/"
1988 );
1989 assert_eq!(
1990 metadata.scopes_supported,
1991 Some(vec!["read".to_string(), "write".to_string()])
1992 );
1993 });
1994 }
1995
1996 #[test]
1997 fn test_fetch_protected_resource_metadata_prefers_www_authenticate_url() {
1998 gpui::block_on(async {
1999 let client = make_fake_http_client(|req| {
2000 Box::pin(async move {
2001 let uri = req.uri().to_string();
2002 if uri == "https://mcp.example.com/custom-resource-metadata" {
2003 json_response(
2004 200,
2005 r#"{
2006 "resource": "https://mcp.example.com",
2007 "authorization_servers": ["https://auth.example.com"]
2008 }"#,
2009 )
2010 } else {
2011 json_response(500, r#"{"error": "should not be called"}"#)
2012 }
2013 })
2014 });
2015
2016 let server_url = Url::parse("https://mcp.example.com").unwrap();
2017 let www_auth = WwwAuthenticate {
2018 resource_metadata: Some(
2019 Url::parse("https://mcp.example.com/custom-resource-metadata").unwrap(),
2020 ),
2021 scope: None,
2022 error: None,
2023 error_description: None,
2024 };
2025
2026 let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
2027 .await
2028 .unwrap();
2029
2030 assert_eq!(metadata.authorization_servers.len(), 1);
2031 });
2032 }
2033
2034 #[test]
2035 fn test_fetch_protected_resource_metadata_falls_back_when_header_url_fails() {
2036 // Reproduces the Pydantic Logfire case: the server's WWW-Authenticate
2037 // header contains a resource_metadata URL with a doubled path (e.g.
2038 // /mcp/mcp), which returns HTML instead of JSON. The client should
2039 // fall back to the RFC 9728 well-known URL, which works correctly.
2040 gpui::block_on(async {
2041 let client = make_fake_http_client(|req| {
2042 Box::pin(async move {
2043 let uri = req.uri().to_string();
2044 if uri
2045 == "https://mcp.example.com/.well-known/oauth-protected-resource/api/mcp/mcp"
2046 {
2047 // Buggy header URL returns HTML (like a SPA catch-all).
2048 Ok(Response::builder()
2049 .status(200)
2050 .header("Content-Type", "text/html")
2051 .body(AsyncBody::from(b"<!doctype html><html></html>".to_vec()))
2052 .unwrap())
2053 } else if uri
2054 == "https://mcp.example.com/.well-known/oauth-protected-resource/api/mcp"
2055 {
2056 // Correct well-known URL returns valid metadata.
2057 json_response(
2058 200,
2059 r#"{
2060 "resource": "https://mcp.example.com/api/mcp",
2061 "authorization_servers": ["https://auth.example.com"]
2062 }"#,
2063 )
2064 } else {
2065 json_response(404, "{}")
2066 }
2067 })
2068 });
2069
2070 let server_url = Url::parse("https://mcp.example.com/api/mcp").unwrap();
2071 let www_auth = WwwAuthenticate {
2072 resource_metadata: Some(
2073 // Buggy URL with doubled path component.
2074 Url::parse(
2075 "https://mcp.example.com/.well-known/oauth-protected-resource/api/mcp/mcp",
2076 )
2077 .unwrap(),
2078 ),
2079 scope: None,
2080 error: None,
2081 error_description: None,
2082 };
2083
2084 let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
2085 .await
2086 .unwrap();
2087
2088 assert_eq!(
2089 metadata.resource.as_str(),
2090 "https://mcp.example.com/api/mcp"
2091 );
2092 assert_eq!(
2093 metadata.authorization_servers[0].as_str(),
2094 "https://auth.example.com/"
2095 );
2096 });
2097 }
2098
2099 #[test]
2100 fn test_fetch_protected_resource_metadata_rejects_cross_origin_url() {
2101 gpui::block_on(async {
2102 let client = make_fake_http_client(|req| {
2103 Box::pin(async move {
2104 let uri = req.uri().to_string();
2105 // The cross-origin URL should NOT be fetched; only the
2106 // well-known fallback at the server's own origin should be.
2107 if uri.contains("attacker.example.com") {
2108 panic!("should not fetch cross-origin resource_metadata URL");
2109 } else if uri.contains(".well-known/oauth-protected-resource") {
2110 json_response(
2111 200,
2112 r#"{
2113 "resource": "https://mcp.example.com",
2114 "authorization_servers": ["https://auth.example.com"]
2115 }"#,
2116 )
2117 } else {
2118 json_response(404, "{}")
2119 }
2120 })
2121 });
2122
2123 let server_url = Url::parse("https://mcp.example.com").unwrap();
2124 let www_auth = WwwAuthenticate {
2125 resource_metadata: Some(
2126 Url::parse("https://attacker.example.com/fake-metadata").unwrap(),
2127 ),
2128 scope: None,
2129 error: None,
2130 error_description: None,
2131 };
2132
2133 let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
2134 .await
2135 .unwrap();
2136
2137 // Should have used the fallback well-known URL, not the attacker's.
2138 assert_eq!(metadata.resource.as_str(), "https://mcp.example.com/");
2139 });
2140 }
2141
2142 #[test]
2143 fn test_fetch_auth_server_metadata() {
2144 gpui::block_on(async {
2145 let client = make_fake_http_client(|req| {
2146 Box::pin(async move {
2147 let uri = req.uri().to_string();
2148 if uri.contains(".well-known/oauth-authorization-server") {
2149 json_response(
2150 200,
2151 r#"{
2152 "issuer": "https://auth.example.com",
2153 "authorization_endpoint": "https://auth.example.com/authorize",
2154 "token_endpoint": "https://auth.example.com/token",
2155 "registration_endpoint": "https://auth.example.com/register",
2156 "code_challenge_methods_supported": ["S256"],
2157 "client_id_metadata_document_supported": true
2158 }"#,
2159 )
2160 } else {
2161 json_response(404, "{}")
2162 }
2163 })
2164 });
2165
2166 let issuer = Url::parse("https://auth.example.com").unwrap();
2167 let metadata = fetch_auth_server_metadata(&client, &issuer).await.unwrap();
2168
2169 assert_eq!(metadata.issuer.as_str(), "https://auth.example.com/");
2170 assert_eq!(
2171 metadata.authorization_endpoint.as_str(),
2172 "https://auth.example.com/authorize"
2173 );
2174 assert_eq!(
2175 metadata.token_endpoint.as_str(),
2176 "https://auth.example.com/token"
2177 );
2178 assert!(metadata.registration_endpoint.is_some());
2179 assert!(metadata.client_id_metadata_document_supported);
2180 assert_eq!(
2181 metadata.code_challenge_methods_supported,
2182 Some(vec!["S256".to_string()])
2183 );
2184 });
2185 }
2186
2187 #[test]
2188 fn test_fetch_auth_server_metadata_falls_back_to_oidc() {
2189 gpui::block_on(async {
2190 let client = make_fake_http_client(|req| {
2191 Box::pin(async move {
2192 let uri = req.uri().to_string();
2193 if uri.contains("openid-configuration") {
2194 json_response(
2195 200,
2196 r#"{
2197 "issuer": "https://auth.example.com",
2198 "authorization_endpoint": "https://auth.example.com/authorize",
2199 "token_endpoint": "https://auth.example.com/token",
2200 "code_challenge_methods_supported": ["S256"]
2201 }"#,
2202 )
2203 } else {
2204 json_response(404, "{}")
2205 }
2206 })
2207 });
2208
2209 let issuer = Url::parse("https://auth.example.com").unwrap();
2210 let metadata = fetch_auth_server_metadata(&client, &issuer).await.unwrap();
2211
2212 assert_eq!(
2213 metadata.authorization_endpoint.as_str(),
2214 "https://auth.example.com/authorize"
2215 );
2216 assert!(!metadata.client_id_metadata_document_supported);
2217 });
2218 }
2219
2220 #[test]
2221 fn test_fetch_auth_server_metadata_rejects_issuer_mismatch() {
2222 gpui::block_on(async {
2223 let client = make_fake_http_client(|req| {
2224 Box::pin(async move {
2225 let uri = req.uri().to_string();
2226 if uri.contains(".well-known/oauth-authorization-server") {
2227 // Response claims to be a different issuer.
2228 json_response(
2229 200,
2230 r#"{
2231 "issuer": "https://evil.example.com",
2232 "authorization_endpoint": "https://evil.example.com/authorize",
2233 "token_endpoint": "https://evil.example.com/token",
2234 "code_challenge_methods_supported": ["S256"]
2235 }"#,
2236 )
2237 } else {
2238 json_response(404, "{}")
2239 }
2240 })
2241 });
2242
2243 let issuer = Url::parse("https://auth.example.com").unwrap();
2244 let result = fetch_auth_server_metadata(&client, &issuer).await;
2245
2246 assert!(result.is_err());
2247 let err_msg = result.unwrap_err().to_string();
2248 assert!(
2249 err_msg.contains("issuer mismatch"),
2250 "unexpected error: {}",
2251 err_msg
2252 );
2253 });
2254 }
2255
2256 // -- Full discover integration tests -------------------------------------
2257
2258 #[test]
2259 fn test_full_discover_with_cimd() {
2260 gpui::block_on(async {
2261 let client = make_fake_http_client(|req| {
2262 Box::pin(async move {
2263 let uri = req.uri().to_string();
2264 if uri.contains("oauth-protected-resource") {
2265 json_response(
2266 200,
2267 r#"{
2268 "resource": "https://mcp.example.com",
2269 "authorization_servers": ["https://auth.example.com"],
2270 "scopes_supported": ["mcp:read"]
2271 }"#,
2272 )
2273 } else if uri.contains("oauth-authorization-server") {
2274 json_response(
2275 200,
2276 r#"{
2277 "issuer": "https://auth.example.com",
2278 "authorization_endpoint": "https://auth.example.com/authorize",
2279 "token_endpoint": "https://auth.example.com/token",
2280 "code_challenge_methods_supported": ["S256"],
2281 "client_id_metadata_document_supported": true
2282 }"#,
2283 )
2284 } else {
2285 json_response(404, "{}")
2286 }
2287 })
2288 });
2289
2290 let server_url = Url::parse("https://mcp.example.com").unwrap();
2291 let www_auth = WwwAuthenticate {
2292 resource_metadata: None,
2293 scope: None,
2294 error: None,
2295 error_description: None,
2296 };
2297
2298 let discovery = discover(&client, &server_url, &www_auth).await.unwrap();
2299 let registration =
2300 resolve_client_registration(&client, &discovery, "http://127.0.0.1:12345/callback")
2301 .await
2302 .unwrap();
2303
2304 assert_eq!(registration.client_id, CIMD_URL);
2305 assert_eq!(registration.client_secret, None);
2306 assert_eq!(discovery.scopes, vec!["mcp:read"]);
2307 });
2308 }
2309
2310 #[test]
2311 fn test_full_discover_with_dcr_fallback() {
2312 gpui::block_on(async {
2313 let client = make_fake_http_client(|req| {
2314 Box::pin(async move {
2315 let uri = req.uri().to_string();
2316 if uri.contains("oauth-protected-resource") {
2317 json_response(
2318 200,
2319 r#"{
2320 "resource": "https://mcp.example.com",
2321 "authorization_servers": ["https://auth.example.com"]
2322 }"#,
2323 )
2324 } else if uri.contains("oauth-authorization-server") {
2325 json_response(
2326 200,
2327 r#"{
2328 "issuer": "https://auth.example.com",
2329 "authorization_endpoint": "https://auth.example.com/authorize",
2330 "token_endpoint": "https://auth.example.com/token",
2331 "registration_endpoint": "https://auth.example.com/register",
2332 "code_challenge_methods_supported": ["S256"],
2333 "client_id_metadata_document_supported": false
2334 }"#,
2335 )
2336 } else if uri.contains("/register") {
2337 json_response(
2338 201,
2339 r#"{
2340 "client_id": "dcr-minted-id-123",
2341 "client_secret": "dcr-secret-456"
2342 }"#,
2343 )
2344 } else {
2345 json_response(404, "{}")
2346 }
2347 })
2348 });
2349
2350 let server_url = Url::parse("https://mcp.example.com").unwrap();
2351 let www_auth = WwwAuthenticate {
2352 resource_metadata: None,
2353 scope: Some(vec!["files:read".into()]),
2354 error: None,
2355 error_description: None,
2356 };
2357
2358 let discovery = discover(&client, &server_url, &www_auth).await.unwrap();
2359 let registration =
2360 resolve_client_registration(&client, &discovery, "http://127.0.0.1:9999/callback")
2361 .await
2362 .unwrap();
2363
2364 assert_eq!(registration.client_id, "dcr-minted-id-123");
2365 assert_eq!(
2366 registration.client_secret.as_deref(),
2367 Some("dcr-secret-456")
2368 );
2369 assert_eq!(discovery.scopes, vec!["files:read"]);
2370 });
2371 }
2372
2373 #[test]
2374 fn test_discover_fails_without_pkce_support() {
2375 gpui::block_on(async {
2376 let client = make_fake_http_client(|req| {
2377 Box::pin(async move {
2378 let uri = req.uri().to_string();
2379 if uri.contains("oauth-protected-resource") {
2380 json_response(
2381 200,
2382 r#"{
2383 "resource": "https://mcp.example.com",
2384 "authorization_servers": ["https://auth.example.com"]
2385 }"#,
2386 )
2387 } else if uri.contains("oauth-authorization-server") {
2388 json_response(
2389 200,
2390 r#"{
2391 "issuer": "https://auth.example.com",
2392 "authorization_endpoint": "https://auth.example.com/authorize",
2393 "token_endpoint": "https://auth.example.com/token"
2394 }"#,
2395 )
2396 } else {
2397 json_response(404, "{}")
2398 }
2399 })
2400 });
2401
2402 let server_url = Url::parse("https://mcp.example.com").unwrap();
2403 let www_auth = WwwAuthenticate {
2404 resource_metadata: None,
2405 scope: None,
2406 error: None,
2407 error_description: None,
2408 };
2409
2410 let result = discover(&client, &server_url, &www_auth).await;
2411 assert!(result.is_err());
2412 let err_msg = result.unwrap_err().to_string();
2413 assert!(
2414 err_msg.contains("code_challenge_methods_supported"),
2415 "unexpected error: {}",
2416 err_msg
2417 );
2418 });
2419 }
2420
2421 // -- Token exchange integration tests ------------------------------------
2422
2423 #[test]
2424 fn test_exchange_code_success() {
2425 gpui::block_on(async {
2426 let client = make_fake_http_client(|req| {
2427 Box::pin(async move {
2428 let uri = req.uri().to_string();
2429 if uri.contains("/token") {
2430 json_response(
2431 200,
2432 r#"{
2433 "access_token": "new_access_token",
2434 "refresh_token": "new_refresh_token",
2435 "expires_in": 3600,
2436 "token_type": "Bearer"
2437 }"#,
2438 )
2439 } else {
2440 json_response(404, "{}")
2441 }
2442 })
2443 });
2444
2445 let metadata = AuthServerMetadata {
2446 issuer: Url::parse("https://auth.example.com").unwrap(),
2447 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
2448 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
2449 registration_endpoint: None,
2450 scopes_supported: None,
2451 code_challenge_methods_supported: Some(vec!["S256".into()]),
2452 client_id_metadata_document_supported: true,
2453 grant_types_supported: None,
2454 };
2455
2456 let tokens = exchange_code(
2457 &client,
2458 &metadata,
2459 "auth_code_123",
2460 CIMD_URL,
2461 "http://127.0.0.1:9999/callback",
2462 "verifier_abc",
2463 "https://mcp.example.com",
2464 None,
2465 )
2466 .await
2467 .unwrap();
2468
2469 assert_eq!(tokens.access_token, "new_access_token");
2470 assert_eq!(tokens.refresh_token.as_deref(), Some("new_refresh_token"));
2471 assert!(tokens.expires_at.is_some());
2472 });
2473 }
2474
2475 #[test]
2476 fn test_refresh_tokens_success() {
2477 gpui::block_on(async {
2478 let client = make_fake_http_client(|req| {
2479 Box::pin(async move {
2480 let uri = req.uri().to_string();
2481 if uri.contains("/token") {
2482 json_response(
2483 200,
2484 r#"{
2485 "access_token": "refreshed_token",
2486 "expires_in": 1800,
2487 "token_type": "Bearer"
2488 }"#,
2489 )
2490 } else {
2491 json_response(404, "{}")
2492 }
2493 })
2494 });
2495
2496 let token_endpoint = Url::parse("https://auth.example.com/token").unwrap();
2497
2498 let tokens = refresh_tokens(
2499 &client,
2500 &token_endpoint,
2501 "old_refresh_token",
2502 CIMD_URL,
2503 "https://mcp.example.com",
2504 None,
2505 )
2506 .await
2507 .unwrap();
2508
2509 assert_eq!(tokens.access_token, "refreshed_token");
2510 assert_eq!(tokens.refresh_token, None);
2511 assert!(tokens.expires_at.is_some());
2512 });
2513 }
2514
2515 #[test]
2516 fn test_exchange_code_failure() {
2517 gpui::block_on(async {
2518 let client = make_fake_http_client(|_req| {
2519 Box::pin(async move { json_response(400, r#"{"error": "invalid_grant"}"#) })
2520 });
2521
2522 let metadata = AuthServerMetadata {
2523 issuer: Url::parse("https://auth.example.com").unwrap(),
2524 authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
2525 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
2526 registration_endpoint: None,
2527 scopes_supported: None,
2528 code_challenge_methods_supported: Some(vec!["S256".into()]),
2529 client_id_metadata_document_supported: true,
2530 grant_types_supported: None,
2531 };
2532
2533 let result = exchange_code(
2534 &client,
2535 &metadata,
2536 "bad_code",
2537 "client",
2538 "http://127.0.0.1:1/callback",
2539 "verifier",
2540 "https://mcp.example.com",
2541 None,
2542 )
2543 .await;
2544
2545 let err = result.unwrap_err();
2546 let token_error = err
2547 .downcast_ref::<OAuthTokenError>()
2548 .expect("expected OAuthTokenError");
2549 assert_eq!(
2550 *token_error,
2551 OAuthTokenError {
2552 error: "invalid_grant".into(),
2553 error_description: None,
2554 }
2555 );
2556 });
2557 }
2558
2559 // -- DCR integration tests -----------------------------------------------
2560
2561 #[test]
2562 fn test_perform_dcr() {
2563 gpui::block_on(async {
2564 let client = make_fake_http_client(|_req| {
2565 Box::pin(async move {
2566 json_response(
2567 201,
2568 r#"{
2569 "client_id": "dynamic-client-001",
2570 "client_secret": "dynamic-secret-001"
2571 }"#,
2572 )
2573 })
2574 });
2575
2576 let endpoint = Url::parse("https://auth.example.com/register").unwrap();
2577 let registration =
2578 perform_dcr(&client, &endpoint, "http://127.0.0.1:9999/callback", None)
2579 .await
2580 .unwrap();
2581
2582 assert_eq!(registration.client_id, "dynamic-client-001");
2583 assert_eq!(
2584 registration.client_secret.as_deref(),
2585 Some("dynamic-secret-001")
2586 );
2587 });
2588 }
2589
2590 #[test]
2591 fn test_perform_dcr_failure() {
2592 gpui::block_on(async {
2593 let client = make_fake_http_client(|_req| {
2594 Box::pin(
2595 async move { json_response(403, r#"{"error": "registration_not_allowed"}"#) },
2596 )
2597 });
2598
2599 let endpoint = Url::parse("https://auth.example.com/register").unwrap();
2600 let result =
2601 perform_dcr(&client, &endpoint, "http://127.0.0.1:9999/callback", None).await;
2602
2603 assert!(result.is_err());
2604 assert!(result.unwrap_err().to_string().contains("403"));
2605 });
2606 }
2607
2608 // -- OAuthCallback parse tests -------------------------------------------
2609
2610 #[test]
2611 fn test_oauth_callback_parse_query() {
2612 let callback = OAuthCallback::parse_query("code=test_auth_code&state=test_state").unwrap();
2613 assert_eq!(callback.code, "test_auth_code");
2614 assert_eq!(callback.state, "test_state");
2615 }
2616
2617 #[test]
2618 fn test_oauth_callback_parse_query_reversed_order() {
2619 let callback = OAuthCallback::parse_query("state=test_state&code=test_auth_code").unwrap();
2620 assert_eq!(callback.code, "test_auth_code");
2621 assert_eq!(callback.state, "test_state");
2622 }
2623
2624 #[test]
2625 fn test_oauth_callback_parse_query_with_extra_params() {
2626 let callback =
2627 OAuthCallback::parse_query("code=test_auth_code&state=test_state&extra=ignored")
2628 .unwrap();
2629 assert_eq!(callback.code, "test_auth_code");
2630 assert_eq!(callback.state, "test_state");
2631 }
2632
2633 #[test]
2634 fn test_oauth_callback_parse_query_missing_code() {
2635 let result = OAuthCallback::parse_query("state=test_state");
2636 assert!(result.is_err());
2637 assert!(result.unwrap_err().to_string().contains("code"));
2638 }
2639
2640 #[test]
2641 fn test_oauth_callback_parse_query_missing_state() {
2642 let result = OAuthCallback::parse_query("code=test_auth_code");
2643 assert!(result.is_err());
2644 assert!(result.unwrap_err().to_string().contains("state"));
2645 }
2646
2647 #[test]
2648 fn test_oauth_callback_parse_query_empty_code() {
2649 let result = OAuthCallback::parse_query("code=&state=test_state");
2650 assert!(result.is_err());
2651 }
2652
2653 #[test]
2654 fn test_oauth_callback_parse_query_empty_state() {
2655 let result = OAuthCallback::parse_query("code=test_auth_code&state=");
2656 assert!(result.is_err());
2657 }
2658
2659 #[test]
2660 fn test_oauth_callback_parse_query_url_encoded_values() {
2661 let callback = OAuthCallback::parse_query("code=abc%20def&state=test%3Dstate").unwrap();
2662 assert_eq!(callback.code, "abc def");
2663 assert_eq!(callback.state, "test=state");
2664 }
2665
2666 #[test]
2667 fn test_oauth_callback_parse_query_error_response() {
2668 let result = OAuthCallback::parse_query(
2669 "error=access_denied&error_description=User%20denied%20access&state=abc",
2670 );
2671 assert!(result.is_err());
2672 let err_msg = result.unwrap_err().to_string();
2673 assert!(
2674 err_msg.contains("access_denied"),
2675 "unexpected error: {}",
2676 err_msg
2677 );
2678 assert!(
2679 err_msg.contains("User denied access"),
2680 "unexpected error: {}",
2681 err_msg
2682 );
2683 }
2684
2685 #[test]
2686 fn test_oauth_callback_parse_query_error_without_description() {
2687 let result = OAuthCallback::parse_query("error=server_error&state=abc");
2688 assert!(result.is_err());
2689 let err_msg = result.unwrap_err().to_string();
2690 assert!(
2691 err_msg.contains("server_error"),
2692 "unexpected error: {}",
2693 err_msg
2694 );
2695 assert!(
2696 err_msg.contains("no description"),
2697 "unexpected error: {}",
2698 err_msg
2699 );
2700 }
2701
2702 // -- McpOAuthTokenProvider tests -----------------------------------------
2703
2704 fn make_test_session(
2705 access_token: &str,
2706 refresh_token: Option<&str>,
2707 expires_at: Option<SystemTime>,
2708 ) -> OAuthSession {
2709 OAuthSession {
2710 token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
2711 resource: Url::parse("https://mcp.example.com").unwrap(),
2712 client_registration: OAuthClientRegistration {
2713 client_id: "test-client".into(),
2714 client_secret: None,
2715 },
2716 tokens: OAuthTokens {
2717 access_token: access_token.into(),
2718 refresh_token: refresh_token.map(String::from),
2719 expires_at,
2720 },
2721 }
2722 }
2723
2724 #[test]
2725 fn test_mcp_oauth_provider_returns_none_when_token_expired() {
2726 let expired = SystemTime::now() - Duration::from_secs(60);
2727 let session = make_test_session("stale-token", Some("rt"), Some(expired));
2728 let provider = McpOAuthTokenProvider::new(
2729 session,
2730 make_fake_http_client(|_| Box::pin(async { unreachable!() })),
2731 None,
2732 );
2733
2734 assert_eq!(provider.access_token(), None);
2735 }
2736
2737 #[test]
2738 fn test_mcp_oauth_provider_returns_token_when_not_expired() {
2739 let far_future = SystemTime::now() + Duration::from_secs(3600);
2740 let session = make_test_session("valid-token", Some("rt"), Some(far_future));
2741 let provider = McpOAuthTokenProvider::new(
2742 session,
2743 make_fake_http_client(|_| Box::pin(async { unreachable!() })),
2744 None,
2745 );
2746
2747 assert_eq!(provider.access_token().as_deref(), Some("valid-token"));
2748 }
2749
2750 #[test]
2751 fn test_mcp_oauth_provider_returns_token_when_no_expiry() {
2752 let session = make_test_session("no-expiry-token", Some("rt"), None);
2753 let provider = McpOAuthTokenProvider::new(
2754 session,
2755 make_fake_http_client(|_| Box::pin(async { unreachable!() })),
2756 None,
2757 );
2758
2759 assert_eq!(provider.access_token().as_deref(), Some("no-expiry-token"));
2760 }
2761
2762 #[test]
2763 fn test_mcp_oauth_provider_refresh_without_refresh_token_returns_false() {
2764 gpui::block_on(async {
2765 let session = make_test_session("token", None, None);
2766 let provider = McpOAuthTokenProvider::new(
2767 session,
2768 make_fake_http_client(|_| {
2769 Box::pin(async { unreachable!("no HTTP call expected") })
2770 }),
2771 None,
2772 );
2773
2774 let refreshed = provider.try_refresh().await.unwrap();
2775 assert!(!refreshed);
2776 });
2777 }
2778
2779 #[test]
2780 fn test_mcp_oauth_provider_refresh_updates_session_and_notifies_channel() {
2781 gpui::block_on(async {
2782 let session = make_test_session("old-access", Some("my-refresh-token"), None);
2783 let (tx, mut rx) = futures::channel::mpsc::unbounded();
2784
2785 let http_client = make_fake_http_client(|_req| {
2786 Box::pin(async {
2787 json_response(
2788 200,
2789 r#"{
2790 "access_token": "new-access",
2791 "refresh_token": "new-refresh",
2792 "expires_in": 1800
2793 }"#,
2794 )
2795 })
2796 });
2797
2798 let provider = McpOAuthTokenProvider::new(session, http_client, Some(tx));
2799
2800 let refreshed = provider.try_refresh().await.unwrap();
2801 assert!(refreshed);
2802 assert_eq!(provider.access_token().as_deref(), Some("new-access"));
2803
2804 let notified_session = rx.try_recv().expect("channel should have a session");
2805 assert_eq!(notified_session.tokens.access_token, "new-access");
2806 assert_eq!(
2807 notified_session.tokens.refresh_token.as_deref(),
2808 Some("new-refresh")
2809 );
2810 });
2811 }
2812
2813 #[test]
2814 fn test_mcp_oauth_provider_refresh_preserves_old_refresh_token_when_server_omits_it() {
2815 gpui::block_on(async {
2816 let session = make_test_session("old-access", Some("original-refresh"), None);
2817 let (tx, mut rx) = futures::channel::mpsc::unbounded();
2818
2819 let http_client = make_fake_http_client(|_req| {
2820 Box::pin(async {
2821 json_response(
2822 200,
2823 r#"{
2824 "access_token": "new-access",
2825 "expires_in": 900
2826 }"#,
2827 )
2828 })
2829 });
2830
2831 let provider = McpOAuthTokenProvider::new(session, http_client, Some(tx));
2832
2833 let refreshed = provider.try_refresh().await.unwrap();
2834 assert!(refreshed);
2835
2836 let notified_session = rx.try_recv().expect("channel should have a session");
2837 assert_eq!(notified_session.tokens.access_token, "new-access");
2838 assert_eq!(
2839 notified_session.tokens.refresh_token.as_deref(),
2840 Some("original-refresh"),
2841 );
2842 });
2843 }
2844
2845 #[test]
2846 fn test_mcp_oauth_provider_refresh_returns_false_on_http_error() {
2847 gpui::block_on(async {
2848 let session = make_test_session("old-access", Some("my-refresh"), None);
2849
2850 let http_client = make_fake_http_client(|_req| {
2851 Box::pin(async { json_response(401, r#"{"error": "invalid_grant"}"#) })
2852 });
2853
2854 let provider = McpOAuthTokenProvider::new(session, http_client, None);
2855
2856 let refreshed = provider.try_refresh().await.unwrap();
2857 assert!(!refreshed);
2858 // The old token should still be in place.
2859 assert_eq!(provider.access_token().as_deref(), Some("old-access"));
2860 });
2861 }
2862}
2863