feat: add OAuth authentication support for remote MCP servers

2026-07-02 14:43:24 -06:00
parent cc50d39ab4
commit 16f324cefc
6 changed files with 582 additions and 19 deletions
@@ -0,0 +1,325 @@
+use crate::client::oauth::{OAuthProvider, load_oauth_tokens, run_oauth_flow};
+use crate::config::paths;
+use anyhow::{Context, Result, anyhow};
+use chrono::Utc;
+use inquire::Text;
+use log::warn;
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use std::fs;
+use std::net::TcpListener;
+use url::Url;
+
+#[derive(Debug, Deserialize)]
+struct ProtectedResourceMetadata {
+    #[serde(default)]
+    authorization_servers: Vec<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct OAuthServerMetadata {
+    authorization_endpoint: String,
+    token_endpoint: String,
+    #[serde(default)]
+    scopes_supported: Vec<String>,
+    registration_endpoint: Option<String>,
+}
+
+#[derive(Serialize, Deserialize)]
+struct McpRegistration {
+    client_id: String,
+}
+
+struct McpOAuthProvider {
+    client_id: String,
+    authorize_url: String,
+    token_url: String,
+    scopes: String,
+    fixed_redirect: String,
+}
+
+impl OAuthProvider for McpOAuthProvider {
+    fn provider_name(&self) -> &str {
+        "MCP"
+    }
+
+    fn client_id(&self) -> &str {
+        &self.client_id
+    }
+
+    fn authorize_url(&self) -> &str {
+        &self.authorize_url
+    }
+
+    fn token_url(&self) -> &str {
+        &self.token_url
+    }
+
+    fn redirect_uri(&self) -> &str {
+        ""
+    }
+
+    fn scopes(&self) -> &str {
+        &self.scopes
+    }
+
+    fn uses_localhost_redirect(&self) -> bool {
+        false
+    }
+
+    fn fixed_redirect_uri(&self) -> Option<String> {
+        Some(self.fixed_redirect.clone())
+    }
+}
+
+pub async fn run_mcp_oauth_flow(
+    server_name: &str,
+    server_url: &str,
+    configured_client_id: Option<&str>,
+) -> Result<()> {
+    let metadata = discover_oauth_metadata(server_url).await?;
+
+    let listener = TcpListener::bind("127.0.0.1:0")?;
+    let port = listener.local_addr()?.port();
+    drop(listener);
+    let redirect_uri = format!("http://127.0.0.1:{port}/callback");
+
+    let client_id = if let Some(id) = configured_client_id {
+        id.to_string()
+    } else if let Some(cached) = load_registered_client_id(server_name) {
+        cached
+    } else if let Some(reg_endpoint) = &metadata.registration_endpoint {
+        match register_client(reg_endpoint, &redirect_uri).await {
+            Ok(id) => {
+                let _ = save_registered_client_id(server_name, &id);
+                id
+            }
+            Err(e) => {
+                warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
+                Text::new("Enter the OAuth client ID for this MCP server:")
+                    .prompt()
+                    .context("Failed to read client ID")?
+            }
+        }
+    } else {
+        Text::new("Enter the OAuth client ID for this MCP server:")
+            .prompt()
+            .context("Failed to read client ID")?
+    };
+
+    let provider = McpOAuthProvider {
+        client_id,
+        authorize_url: metadata.authorization_endpoint,
+        token_url: metadata.token_endpoint,
+        scopes: metadata.scopes_supported.join(" "),
+        fixed_redirect: redirect_uri,
+    };
+
+    run_oauth_flow(&provider, &mcp_token_key(server_name)).await
+}
+
+pub fn load_valid_mcp_token(server_name: &str) -> Option<String> {
+    let tokens = load_oauth_tokens(&mcp_token_key(server_name))?;
+    if Utc::now().timestamp() < tokens.expires_at {
+        Some(tokens.access_token)
+    } else {
+        None
+    }
+}
+
+fn mcp_token_key(server_name: &str) -> String {
+    format!("mcp_{server_name}")
+}
+
+fn load_registered_client_id(server_name: &str) -> Option<String> {
+    let path = paths::oauth_tokens_path().join(format!("mcp_{server_name}_registration.json"));
+    let content = fs::read_to_string(path).ok()?;
+    let reg: McpRegistration = serde_json::from_str(&content).ok()?;
+
+    Some(reg.client_id)
+}
+
+fn save_registered_client_id(server_name: &str, client_id: &str) -> Result<()> {
+    let dir = paths::oauth_tokens_path();
+    fs::create_dir_all(&dir)?;
+
+    let path = dir.join(format!("mcp_{server_name}_registration.json"));
+    let reg = McpRegistration {
+        client_id: client_id.to_string(),
+    };
+
+    fs::write(path, serde_json::to_string_pretty(&reg)?)?;
+
+    Ok(())
+}
+
+async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
+    let body = serde_json::json!({
+        "client_name": "Coyote",
+        "redirect_uris": [redirect_uri],
+        "grant_types": ["authorization_code", "refresh_token"],
+        "response_types": ["code"],
+        "token_endpoint_auth_method": "none"
+    });
+
+    let response: serde_json::Value = Client::new()
+        .post(endpoint)
+        .json(&body)
+        .send()
+        .await
+        .context("Failed to reach registration endpoint")?
+        .json()
+        .await
+        .context("Failed to parse registration response")?;
+
+    response["client_id"]
+        .as_str()
+        .ok_or_else(|| anyhow!("Missing client_id in registration response: {response}"))
+        .map(|s| s.to_string())
+}
+
+async fn discover_oauth_metadata(server_url: &str) -> Result<OAuthServerMetadata> {
+    let base = extract_base_url(server_url)?;
+    let client = Client::new();
+
+    // RFC 9728: try protected resource metadata first; it points to the auth server
+    let pr_url = format!("{base}/.well-known/oauth-protected-resource");
+    if let Ok(resp) = client.get(&pr_url).send().await
+        && resp.status().is_success()
+        && let Ok(pr) = resp.json::<ProtectedResourceMetadata>().await
+        && let Some(auth_server) = pr.authorization_servers.first()
+    {
+        let as_url = format!("{auth_server}/.well-known/oauth-authorization-server");
+        if let Ok(resp) = client.get(&as_url).send().await
+            && resp.status().is_success()
+            && let Ok(meta) = resp.json::<OAuthServerMetadata>().await
+        {
+            return Ok(meta);
+        }
+    }
+
+    let as_url = format!("{base}/.well-known/oauth-authorization-server");
+    let resp = client
+        .get(&as_url)
+        .send()
+        .await
+        .with_context(|| format!("Failed to reach {as_url}"))?;
+
+    if resp.status().is_success() {
+        return resp
+            .json::<OAuthServerMetadata>()
+            .await
+            .with_context(|| format!("Failed to parse OAuth metadata from {as_url}"));
+    }
+
+    Err(anyhow!(
+        "Could not discover OAuth metadata for '{server_url}'.\n\
+         Tried:\n  {pr_url}\n  {as_url}\n\
+         Ensure the server supports MCP OAuth discovery, or consult its documentation."
+    ))
+}
+
+fn extract_base_url(url: &str) -> Result<String> {
+    let parsed = Url::parse(url).with_context(|| format!("Invalid URL: {url}"))?;
+    let scheme = parsed.scheme();
+    let host = parsed
+        .host_str()
+        .ok_or_else(|| anyhow!("No host in URL: {url}"))?;
+    let port = parsed.port().map(|p| format!(":{p}")).unwrap_or_default();
+
+    Ok(format!("{scheme}://{host}{port}"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::utils::get_env_name;
+    use serial_test::serial;
+    use std::{
+        env, fs,
+        time::{self, SystemTime},
+    };
+
+    fn with_temp_cache<F: FnOnce()>(f: F) {
+        let unique = SystemTime::now()
+            .duration_since(time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let root = env::temp_dir().join(format!("coyote-mcp-oauth-test-{unique}"));
+        fs::create_dir_all(&root).unwrap();
+        let env_key = get_env_name("cache_dir");
+        let prev = env::var_os(&env_key);
+        unsafe {
+            env::set_var(&env_key, &root);
+        }
+        f();
+        unsafe {
+            match prev {
+                Some(v) => env::set_var(&env_key, v),
+                None => env::remove_var(&env_key),
+            }
+        }
+        let _ = fs::remove_dir_all(&root);
+    }
+
+    #[test]
+    fn extract_base_url_strips_path_and_query() {
+        let result = extract_base_url("https://mcp.notion.com/mcp?foo=bar").unwrap();
+
+        assert_eq!(result, "https://mcp.notion.com");
+    }
+
+    #[test]
+    fn extract_base_url_preserves_explicit_port() {
+        let result = extract_base_url("http://localhost:8080/mcp").unwrap();
+
+        assert_eq!(result, "http://localhost:8080");
+    }
+
+    #[test]
+    fn extract_base_url_standard_port_omitted() {
+        let result = extract_base_url("https://example.com/mcp/v1").unwrap();
+
+        assert_eq!(result, "https://example.com");
+    }
+
+    #[test]
+    fn extract_base_url_rejects_invalid_url() {
+        assert!(extract_base_url("not-a-url").is_err());
+    }
+
+    #[test]
+    #[serial]
+    fn registered_client_id_roundtrip() {
+        with_temp_cache(|| {
+            save_registered_client_id("notion", "client-xyz-123").unwrap();
+
+            let loaded = load_registered_client_id("notion");
+
+            assert_eq!(loaded, Some("client-xyz-123".to_string()));
+        });
+    }
+
+    #[test]
+    #[serial]
+    fn load_registered_client_id_returns_none_for_missing() {
+        with_temp_cache(|| {
+            let loaded = load_registered_client_id("no-such-server");
+
+            assert!(loaded.is_none());
+        });
+    }
+
+    #[test]
+    #[serial]
+    fn registered_client_id_second_save_overwrites_first() {
+        with_temp_cache(|| {
+            save_registered_client_id("github", "first-id").unwrap();
+            save_registered_client_id("github", "second-id").unwrap();
+
+            let loaded = load_registered_client_id("github");
+
+            assert_eq!(loaded, Some("second-id".to_string()));
+        });
+    }
+}