2862 lines · 102.8 KB · Rust
RawDownload
//! OAuth 2.0 authentication for MCP servers using the Authorization Code +//! PKCE flow, per the MCP spec's OAuth profile.//!//! The flow is split into two phases://!//! 1. **Discovery** ([`discover`]) fetches Protected Resource Metadata and//!    Authorization Server Metadata. This can happen early (e.g. on a 401//!    during server startup) because it doesn't need the redirect URI yet.//!//! 2. **Client registration** ([`resolve_client_registration`]) is separate//!    because DCR requires the actual loopback redirect URI, which includes an//!    ephemeral port that only exists once the callback server has started.//!//! After authentication, the full state is captured in [`OAuthSession`] which//! is persisted to the keychain. On next startup, the stored session feeds//! directly into [`McpOAuthTokenProvider`], giving a refresh-capable provider//! without requiring another browser flow. use anyhow::{Context as _, Result, anyhow, bail};use async_trait::async_trait;use base64::Engine as _;use futures::AsyncReadExt as _;use futures::FutureExt as _;use futures::channel::mpsc;use futures::future::BoxFuture;use http_client::{AsyncBody, HttpClient, Request};use parking_lot::Mutex as SyncMutex;use rand::Rng as _;use serde::{Deserialize, Serialize};use sha2::{Digest, Sha256}; use std::sync::Arc;use std::time::{Duration, SystemTime};use url::Url; /// The CIMD URL where Zed's OAuth client metadata document is hosted.pub const CIMD_URL: &str = "https://zed.dev/oauth/client-metadata.json"; /// Validate that a URL is safe to use as an OAuth endpoint.////// OAuth endpoints carry sensitive material (authorization codes, PKCE/// verifiers, tokens) and must use TLS. Plain HTTP is only permitted for/// loopback addresses, per RFC 8252 Section 8.3.fn require_https_or_loopback(url: &Url) -> Result<()> {    if url.scheme() == "https" {        return Ok(());    }    if url.scheme() == "http" {        if let Some(host) = url.host() {            match host {                url::Host::Ipv4(ip) if ip.is_loopback() => return Ok(()),                url::Host::Ipv6(ip) if ip.is_loopback() => return Ok(()),                url::Host::Domain(d) if d.eq_ignore_ascii_case("localhost") => return Ok(()),                _ => {}            }        }    }    bail!(        "OAuth endpoint must use HTTPS (got {}://{})",        url.scheme(),        url.host_str().unwrap_or("?")    )} /// Validate that a URL is safe to use as an OAuth endpoint, including SSRF/// protections against private/reserved IP ranges.////// This wraps [`require_https_or_loopback`] and adds IP-range checks to prevent/// an attacker-controlled MCP server from directing Zed to fetch internal/// network resources via metadata URLs.////// **Known limitation:** Domain-name URLs that resolve to private IPs are *not*/// blocked here — full mitigation requires resolver-level validation (e.g. a/// custom `Resolve` implementation). This function only blocks IP-literal URLs.fn validate_oauth_url(url: &Url) -> Result<()> {    require_https_or_loopback(url)?;     if let Some(host) = url.host() {        match host {            url::Host::Ipv4(ip) => {                // Loopback is already allowed by require_https_or_loopback.                if ip.is_private() || ip.is_link_local() || ip.is_broadcast() || ip.is_unspecified()                {                    bail!(                        "OAuth endpoint must not point to private/reserved IP: {}",                        ip                    );                }            }            url::Host::Ipv6(ip) => {                // Check for IPv4-mapped IPv6 addresses (::ffff:a.b.c.d) which                // could bypass the IPv4 checks above.                if let Some(mapped_v4) = ip.to_ipv4_mapped() {                    if mapped_v4.is_private()                        || mapped_v4.is_link_local()                        || mapped_v4.is_broadcast()                        || mapped_v4.is_unspecified()                    {                        bail!(                            "OAuth endpoint must not point to private/reserved IP: ::ffff:{}",                            mapped_v4                        );                    }                }                 if ip.is_unspecified() || ip.is_multicast() {                    bail!(                        "OAuth endpoint must not point to reserved IPv6 address: {}",                        ip                    );                }                // IPv6 Unique Local Addresses (fc00::/7). is_unique_local() is                // nightly-only, so check the prefix manually.                if (ip.segments()[0] & 0xfe00) == 0xfc00 {                    bail!(                        "OAuth endpoint must not point to IPv6 unique-local address: {}",                        ip                    );                }            }            url::Host::Domain(_) => {                // Domain-based SSRF prevention requires resolver-level checks.                // See known limitation in the doc comment above.            }        }    }     Ok(())} /// Parsed from the MCP server's WWW-Authenticate header or well-known endpoint/// per RFC 9728 (OAuth 2.0 Protected Resource Metadata).#[derive(Debug, Clone, Serialize, Deserialize)]pub struct ProtectedResourceMetadata {    pub resource: Url,    pub authorization_servers: Vec<Url>,    pub scopes_supported: Option<Vec<String>>,} /// Parsed from the authorization server's .well-known endpoint/// per RFC 8414 (OAuth 2.0 Authorization Server Metadata).#[derive(Debug, Clone, Serialize, Deserialize)]pub struct AuthServerMetadata {    pub issuer: Url,    pub authorization_endpoint: Url,    pub token_endpoint: Url,    pub registration_endpoint: Option<Url>,    pub scopes_supported: Option<Vec<String>>,    pub grant_types_supported: Option<Vec<String>>,    pub code_challenge_methods_supported: Option<Vec<String>>,    pub client_id_metadata_document_supported: bool,} /// The result of client registration — either CIMD or DCR.#[derive(Clone, Serialize, Deserialize)]pub struct OAuthClientRegistration {    pub client_id: String,    /// Only present for DCR-minted registrations.    pub client_secret: Option<String>,} impl std::fmt::Debug for OAuthClientRegistration {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("OAuthClientRegistration")            .field("client_id", &self.client_id)            .field(                "client_secret",                &self.client_secret.as_ref().map(|_| "[redacted]"),            )            .finish()    }} /// Access and refresh tokens obtained from the token endpoint.#[derive(Clone, Serialize, Deserialize)]pub struct OAuthTokens {    pub access_token: String,    pub refresh_token: Option<String>,    pub expires_at: Option<SystemTime>,} impl std::fmt::Debug for OAuthTokens {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("OAuthTokens")            .field("access_token", &"[redacted]")            .field(                "refresh_token",                &self.refresh_token.as_ref().map(|_| "[redacted]"),            )            .field("expires_at", &self.expires_at)            .finish()    }} /// Everything discovered before the browser flow starts. Client registration is/// resolved separately, once the real redirect URI is known.#[derive(Debug, Clone, Serialize, Deserialize)]pub struct OAuthDiscovery {    pub resource_metadata: ProtectedResourceMetadata,    pub auth_server_metadata: AuthServerMetadata,    pub scopes: Vec<String>,} /// The persisted OAuth session for a context server.////// Stored in the keychain so startup can restore a refresh-capable provider/// without another browser flow. Deliberately excludes the full discovery/// metadata to keep the serialized size well within keychain item limits.#[derive(Clone, Serialize, Deserialize)]pub struct OAuthSession {    pub token_endpoint: Url,    pub resource: Url,    pub client_registration: OAuthClientRegistration,    pub tokens: OAuthTokens,} impl std::fmt::Debug for OAuthSession {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("OAuthSession")            .field("token_endpoint", &self.token_endpoint)            .field("resource", &self.resource)            .field("client_registration", &self.client_registration)            .field("tokens", &self.tokens)            .finish()    }} /// Error codes defined by RFC 6750 Section 3.1 for Bearer token authentication.#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum BearerError {    /// The request is missing a required parameter, includes an unsupported    /// parameter or parameter value, or is otherwise malformed.    InvalidRequest,    /// The access token provided is expired, revoked, malformed, or invalid.    InvalidToken,    /// The request requires higher privileges than provided by the access token.    InsufficientScope,    /// An unrecognized error code (extension or future spec addition).    Other,} impl BearerError {    fn parse(value: &str) -> Self {        match value {            "invalid_request" => BearerError::InvalidRequest,            "invalid_token" => BearerError::InvalidToken,            "insufficient_scope" => BearerError::InsufficientScope,            _ => BearerError::Other,        }    }} /// Fields extracted from a `WWW-Authenticate: Bearer` header.////// Per RFC 9728 Section 5.1, MCP servers include `resource_metadata` to point/// at the Protected Resource Metadata document. The optional `scope` parameter/// (RFC 6750 Section 3) indicates scopes required for the request.#[derive(Debug, Clone, PartialEq, Eq)]pub struct WwwAuthenticate {    pub resource_metadata: Option<Url>,    pub scope: Option<Vec<String>>,    /// The parsed `error` parameter per RFC 6750 Section 3.1.    pub error: Option<BearerError>,    pub error_description: Option<String>,} /// Parse a `WWW-Authenticate` header value.////// Expects the `Bearer` scheme followed by comma-separated `key="value"` pairs./// Per RFC 6750 and RFC 9728, the relevant parameters are:/// - `resource_metadata` — URL of the Protected Resource Metadata document/// - `scope` — space-separated list of required scopes/// - `error` — error code (e.g. "insufficient_scope")/// - `error_description` — human-readable error descriptionpub fn parse_www_authenticate(header: &str) -> Result<WwwAuthenticate> {    let header = header.trim();     let params_str = if header.len() >= 6 && header[..6].eq_ignore_ascii_case("bearer") {        header[6..].trim()    } else {        bail!("WWW-Authenticate header does not use Bearer scheme");    };     if params_str.is_empty() {        return Ok(WwwAuthenticate {            resource_metadata: None,            scope: None,            error: None,            error_description: None,        });    }     let params = parse_auth_params(params_str);     let resource_metadata = params        .get("resource_metadata")        .map(|v| Url::parse(v))        .transpose()        .map_err(|e| anyhow!("invalid resource_metadata URL: {}", e))?;     let scope = params        .get("scope")        .map(|v| v.split_whitespace().map(String::from).collect());     let error = params.get("error").map(|v| BearerError::parse(v));    let error_description = params.get("error_description").cloned();     Ok(WwwAuthenticate {        resource_metadata,        scope,        error,        error_description,    })} /// Parse comma-separated `key="value"` or `key=token` parameters from an/// auth-param list (RFC 7235 Section 2.1).fn parse_auth_params(input: &str) -> collections::HashMap<String, String> {    let mut params = collections::HashMap::default();    let mut remaining = input.trim();     while !remaining.is_empty() {        // Skip leading whitespace and commas.        remaining = remaining.trim_start_matches(|c: char| c == ',' || c.is_whitespace());        if remaining.is_empty() {            break;        }         // Find the key (everything before '=').        let eq_pos = match remaining.find('=') {            Some(pos) => pos,            None => break,        };         let key = remaining[..eq_pos].trim().to_lowercase();        remaining = &remaining[eq_pos + 1..];        remaining = remaining.trim_start();         // Parse the value: either quoted or unquoted (token).        let value;        if remaining.starts_with('"') {            // Quoted string: find the closing quote, handling escaped chars.            remaining = &remaining[1..]; // skip opening quote            let mut val = String::new();            let mut chars = remaining.char_indices();            loop {                match chars.next() {                    Some((_, '\\')) => {                        // Escaped character — take the next char literally.                        if let Some((_, c)) = chars.next() {                            val.push(c);                        }                    }                    Some((i, '"')) => {                        remaining = &remaining[i + 1..];                        break;                    }                    Some((_, c)) => val.push(c),                    None => {                        remaining = "";                        break;                    }                }            }            value = val;        } else {            // Unquoted token: read until comma or whitespace.            let end = remaining                .find(|c: char| c == ',' || c.is_whitespace())                .unwrap_or(remaining.len());            value = remaining[..end].to_string();            remaining = &remaining[end..];        }         if !key.is_empty() {            params.insert(key, value);        }    }     params} /// Construct the well-known Protected Resource Metadata URIs for a given MCP/// server URL, per RFC 9728 Section 3.////// Returns URIs in priority order:/// 1. Path-specific: `https://<host>/.well-known/oauth-protected-resource/<path>`/// 2. Root: `https://<host>/.well-known/oauth-protected-resource`pub fn protected_resource_metadata_urls(server_url: &Url) -> Vec<Url> {    let mut urls = Vec::new();    let base = format!("{}://{}", server_url.scheme(), server_url.authority());     let path = server_url.path().trim_start_matches('/');    if !path.is_empty() {        if let Ok(url) = Url::parse(&format!(            "{}/.well-known/oauth-protected-resource/{}",            base, path        )) {            urls.push(url);        }    }     if let Ok(url) = Url::parse(&format!("{}/.well-known/oauth-protected-resource", base)) {        urls.push(url);    }     urls} /// Construct the well-known Authorization Server Metadata URIs for a given/// issuer URL, per RFC 8414 Section 3.1 and Section 5 (OIDC compat).////// Returns URIs in priority order, which differs depending on whether the/// issuer URL has a path component.pub fn auth_server_metadata_urls(issuer: &Url) -> Vec<Url> {    let mut urls = Vec::new();    let base = format!("{}://{}", issuer.scheme(), issuer.authority());    let path = issuer.path().trim_matches('/');     if !path.is_empty() {        // Issuer with path: try path-inserted variants first.        if let Ok(url) = Url::parse(&format!(            "{}/.well-known/oauth-authorization-server/{}",            base, path        )) {            urls.push(url);        }        if let Ok(url) = Url::parse(&format!(            "{}/.well-known/openid-configuration/{}",            base, path        )) {            urls.push(url);        }        if let Ok(url) = Url::parse(&format!(            "{}/{}/.well-known/openid-configuration",            base, path        )) {            urls.push(url);        }    } else {        // No path: standard well-known locations.        if let Ok(url) = Url::parse(&format!("{}/.well-known/oauth-authorization-server", base)) {            urls.push(url);        }        if let Ok(url) = Url::parse(&format!("{}/.well-known/openid-configuration", base)) {            urls.push(url);        }    }     urls} // -- Canonical server URI (RFC 8707) ----------------------------------------- /// Derive the canonical resource URI for an MCP server URL, suitable for the/// `resource` parameter in authorization and token requests per RFC 8707.////// Lowercases the scheme and host, preserves the path (without trailing slash),/// strips fragments and query strings.pub fn canonical_server_uri(server_url: &Url) -> String {    let mut uri = format!(        "{}://{}",        server_url.scheme().to_ascii_lowercase(),        server_url.host_str().unwrap_or("").to_ascii_lowercase(),    );    if let Some(port) = server_url.port() {        uri.push_str(&format!(":{}", port));    }    let path = server_url.path();    if path != "/" {        uri.push_str(path.trim_end_matches('/'));    }    uri} // -- Scope selection --------------------------------------------------------- /// Select scopes following the MCP spec's Scope Selection Strategy:/// 1. Use `scope` from the `WWW-Authenticate` challenge if present./// 2. Fall back to `scopes_supported` from Protected Resource Metadata./// 3. Return empty if neither is available.pub fn select_scopes(    www_authenticate: &WwwAuthenticate,    resource_metadata: &ProtectedResourceMetadata,) -> Vec<String> {    if let Some(ref scopes) = www_authenticate.scope {        if !scopes.is_empty() {            return scopes.clone();        }    }    resource_metadata        .scopes_supported        .clone()        .unwrap_or_default()} // -- Client registration strategy -------------------------------------------- /// The registration approach to use, determined from auth server metadata.#[derive(Debug, Clone, PartialEq, Eq)]

Lines 1–500 of 2862.