From 18b59ae9a79b11ca782090255b643e4bd792d9e3 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 3 Mar 2026 09:08:40 -0800 Subject: [PATCH] feat: add OAuth support for WASM tools in web gateway (#489) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add OAuth support for WASM tools in web gateway Extract reusable OAuth functions (build_oauth_url, exchange_oauth_code, store_oauth_tokens, validate_oauth_token) from CLI into shared oauth_defaults module, then wire them into the web gateway's ExtensionManager. Key changes: - Install auto-activates WASM tools (no separate Activate button) - Configure button triggers OAuth flow via save_setup_secrets - Scope merging: installing a second Google tool triggers re-auth with merged scopes from all tools sharing the same secret_name - Cancel-and-retry: aborting stale OAuth listeners prevents port conflicts - Post-auth validation: wrong account detected via validation_endpoint - Reconfigure always re-auths (deletes old token before starting fresh) - UI shows error toast on OAuth failure, refreshes extension list Flow: Install → Active → Configure (enter client_id/secret) → Save → OAuth popup → authorize → done. Second Google tool install auto-triggers scope expansion OAuth. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: address PR review comments - Add custom headers support to ValidationEndpointSchema (fixes missing Notion-Version header regression) - Guard activate handler auth check with status == "awaiting_authorization" to prevent unexpected OAuth popups - Add window dimensions to OAuth popup in activateExtension() - Simplify UTF-8 truncation boundary check Co-Authored-By: Claude Opus 4.6 (1M context) * fix: address Copilot PR review comments (security, UX, bugs) - Add CSRF state parameter to OAuth flow (random state in auth URL, validated in callback) - Restore MCP server Activate button in web UI (was hidden for all non-channel extensions) - Abort JoinHandle in cleanup_expired_auths to prevent port 9876 conflicts - Fix Google-specific error message for non-Google OAuth providers - Add has_auth field to ExtensionInfo API response (fixes Configure button visibility) - Use oauth_defaults::callback_url() instead of hardcoded redirect_uri (both CLI and manager) - Update auth check comment to match actual behavior (scope expansion + first-time auth) - Add unit tests for build_oauth_url (basic, PKCE, extra params, state uniqueness) - Check all required setup secrets (client_id + client_secret) before starting OAuth Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/channels/web/handlers/extensions.rs | 7 +- src/channels/web/server.rs | 49 ++- src/channels/web/static/app.js | 43 +- src/channels/web/types.rs | 3 + src/cli/oauth_defaults.rs | 453 +++++++++++++++++++- src/cli/tool.rs | 242 +++-------- src/extensions/manager.rs | 529 +++++++++++++++++++++++- src/extensions/mod.rs | 3 + src/llm/session.rs | 2 +- src/tools/mcp/auth.rs | 5 +- src/tools/wasm/capabilities_schema.rs | 5 + 11 files changed, 1109 insertions(+), 232 deletions(-) diff --git a/src/channels/web/handlers/extensions.rs b/src/channels/web/handlers/extensions.rs index 888ffa99..8199b63c 100644 --- a/src/channels/web/handlers/extensions.rs +++ b/src/channels/web/handlers/extensions.rs @@ -59,6 +59,7 @@ pub async fn extensions_list_handler( active: ext.active, tools: ext.tools, needs_setup: ext.needs_setup, + has_auth: ext.has_auth, activation_status, activation_error: ext.activation_error, } @@ -123,7 +124,11 @@ pub async fn extensions_activate_handler( ))?; match ext_mgr.activate(&name).await { - Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Ok(result) => { + // Activation just loads the WASM module. Auth (OAuth/manual) is + // triggered separately via save_setup_secrets or the auth endpoint. + Ok(Json(ActionResponse::ok(result.message))) + } Err(activate_err) => { let err_str = activate_err.to_string(); let needs_auth = err_str.contains("authentication") diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 18b5f473..4f891678 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -1244,6 +1244,7 @@ async fn extensions_list_handler( active: ext.active, tools: ext.tools, needs_setup: ext.needs_setup, + has_auth: ext.has_auth, activation_status, activation_error: ext.activation_error, } @@ -1313,7 +1314,37 @@ async fn extensions_install_handler( .install(&req.name, req.url.as_deref(), kind_hint) .await { - Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Ok(result) => { + let mut resp = ActionResponse::ok(result.message); + + // Auto-activate WASM tools after install (install = active). + if result.kind == crate::extensions::ExtensionKind::WasmTool { + if let Err(e) = ext_mgr.activate(&req.name).await { + tracing::debug!( + extension = %req.name, + error = %e, + "Auto-activation after install failed" + ); + } + + // Check auth after activation. This may initiate OAuth both for scope + // expansion and for first-time auth when credentials are already + // configured (e.g., built-in providers). We only surface an auth_url + // when the extension reports it is awaiting authorization. + match ext_mgr.auth(&req.name, None).await { + Ok(auth_result) + if auth_result.auth_url.is_some() + && auth_result.status == "awaiting_authorization" => + { + // Scope expansion or initial OAuth: user needs to authorize + resp.auth_url = auth_result.auth_url; + } + _ => {} + } + } + + Ok(Json(resp)) + } Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), } } @@ -1328,7 +1359,20 @@ async fn extensions_activate_handler( ))?; match ext_mgr.activate(&name).await { - Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Ok(result) => { + // Activation loaded the WASM module. Check if the tool needs + // OAuth scope expansion (e.g., adding google-docs when gmail + // already has a token but missing the documents scope). + // Initial OAuth setup is triggered via save_setup_secrets. + let mut resp = ActionResponse::ok(result.message); + if let Ok(auth_result) = ext_mgr.auth(&name, None).await + && auth_result.auth_url.is_some() + && auth_result.status == "awaiting_authorization" + { + resp.auth_url = auth_result.auth_url; + } + Ok(Json(resp)) + } Err(activate_err) => { let err_str = activate_err.to_string(); let needs_auth = err_str.contains("authentication") @@ -1550,6 +1594,7 @@ async fn extensions_setup_submit_handler( Ok(result) => { let mut resp = ActionResponse::ok(result.message); resp.activated = Some(result.activated); + resp.auth_url = result.auth_url; Ok(Json(resp)) } Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 7e653408..738f6dde 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -228,7 +228,13 @@ function connectSSE() { eventSource.addEventListener('auth_completed', (e) => { const data = JSON.parse(e.data); removeAuthCard(data.extension_name); - showToast(data.message, 'success'); + if (data.success) { + showToast(data.message, 'success'); + } else { + showToast(data.message, 'error'); + } + // Refresh extensions list so status indicators update + if (currentTab === 'extensions') loadExtensions(); enableChatInput(); }); @@ -1760,6 +1766,11 @@ function renderAvailableExtensionCard(entry) { }).then(function(res) { if (res.success) { showToast('Installed ' + entry.display_name, 'success'); + // OAuth popup if auth started during install (builtin creds) + if (res.auth_url) { + showToast('Opening authentication for ' + entry.display_name, 'info'); + window.open(res.auth_url, '_blank', 'width=600,height=700'); + } loadExtensions(); // Auto-open configure for WASM channels if (entry.kind === 'wasm_channel') { @@ -1961,24 +1972,25 @@ function renderExtensionCard(ext) { actions.appendChild(setupBtn); } } else { - // Non-WASM-channel extensions: original behavior - if (!ext.active) { + // WASM tools / MCP servers + const activeLabel = document.createElement('span'); + activeLabel.className = 'ext-active-label'; + activeLabel.textContent = ext.active ? 'Active' : 'Installed'; + actions.appendChild(activeLabel); + + // MCP servers may be installed but inactive — show Activate button + if (ext.kind === 'mcp_server' && !ext.active) { const activateBtn = document.createElement('button'); activateBtn.className = 'btn-ext activate'; activateBtn.textContent = 'Activate'; activateBtn.addEventListener('click', () => activateExtension(ext.name)); actions.appendChild(activateBtn); - } else { - const activeLabel = document.createElement('span'); - activeLabel.className = 'ext-active-label'; - activeLabel.textContent = 'Active'; - actions.appendChild(activeLabel); } - if (ext.needs_setup) { + if (ext.needs_setup || ext.has_auth) { const configBtn = document.createElement('button'); configBtn.className = 'btn-ext configure'; - configBtn.textContent = ext.authenticated ? 'Reconfigure' : 'Setup'; + configBtn.textContent = ext.authenticated ? 'Reconfigure' : 'Configure'; configBtn.addEventListener('click', () => showConfigureModal(ext.name)); actions.appendChild(configBtn); } @@ -2008,6 +2020,11 @@ function activateExtension(name) { apiFetch('/api/extensions/' + encodeURIComponent(name) + '/activate', { method: 'POST' }) .then((res) => { if (res.success) { + // Even on success, the tool may need OAuth (e.g., WASM loaded but no token yet) + if (res.auth_url) { + showToast('Opening authentication for ' + name, 'info'); + window.open(res.auth_url, '_blank', 'width=600,height=700'); + } loadExtensions(); return; } @@ -2158,7 +2175,11 @@ function submitConfigureModal(name, fields) { .then((res) => { closeConfigureModal(); if (res.success) { - if (res.activated) { + if (res.auth_url) { + // OAuth flow started — open consent popup + showToast('Opening OAuth authorization for ' + name, 'info'); + window.open(res.auth_url, '_blank', 'width=600,height=700'); + } else if (res.activated) { showToast('Configured and activated ' + name, 'success'); } else { showToast(res.message || 'Configuration saved but activation failed', 'warning'); diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index a01aed3a..41aad382 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -388,6 +388,9 @@ pub struct ExtensionInfo { /// Whether this extension has configurable secrets (setup schema). #[serde(default)] pub needs_setup: bool, + /// Whether this extension has an auth configuration (OAuth or manual token). + #[serde(default)] + pub has_auth: bool, /// WASM channel activation status: "installed", "configured", "active", "failed". #[serde(skip_serializing_if = "Option::is_none")] pub activation_status: Option, diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index 7a4586b9..8f8cd3a7 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -17,11 +17,17 @@ //! - **Runtime**: Users can set GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET //! env vars, which take priority over built-in defaults. +use std::collections::HashMap; use std::time::Duration; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use rand::RngCore; +use sha2::{Digest, Sha256}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::net::TcpListener; +use crate::secrets::{CreateSecretParams, SecretsStore}; + // ── Built-in credentials ──────────────────────────────────────────────── pub struct OAuthCredentials { @@ -121,6 +127,9 @@ pub enum OAuthCallbackError { #[error("Timed out waiting for authorization")] Timeout, + #[error("CSRF state mismatch: expected {expected}, got {actual}")] + StateMismatch { expected: String, actual: String }, + #[error("IO error: {0}")] Io(String), } @@ -177,16 +186,22 @@ pub async fn bind_callback_listener() -> Result /// extracts the value of `param_name` (e.g., "code" or "token"), and shows a branded /// landing page using `display_name` (e.g., "Google", "Notion", "NEAR AI"). /// +/// When `expected_state` is `Some`, the callback's `state` query parameter is validated +/// against it to prevent CSRF attacks. If the state doesn't match, the callback is +/// rejected with an error page. +/// /// Times out after 5 minutes. pub async fn wait_for_callback( listener: TcpListener, path_prefix: &str, param_name: &str, display_name: &str, + expected_state: Option<&str>, ) -> Result { let path_prefix = path_prefix.to_string(); let param_name = param_name.to_string(); let display_name = display_name.to_string(); + let expected_state = expected_state.map(String::from); tokio::time::timeout(Duration::from_secs(300), async move { loop { @@ -221,17 +236,29 @@ pub async fn wait_for_callback( return Err(OAuthCallbackError::Denied); } - // Look for the target parameter - for param in query.split('&') { - let parts: Vec<&str> = param.splitn(2, '=').collect(); - if parts.len() == 2 && parts[0] == param_name { - let value = urlencoding::decode(parts[1]) - .unwrap_or_else(|_| parts[1].into()) - .into_owned(); + // Parse all query params into a map for validation + let params: HashMap<&str, String> = query + .split('&') + .filter_map(|p| { + let mut parts = p.splitn(2, '='); + let key = parts.next()?; + let val = parts.next().unwrap_or(""); + Some(( + key, + urlencoding::decode(val) + .unwrap_or_else(|_| val.into()) + .into_owned(), + )) + }) + .collect(); - let html = landing_html(&display_name, true); + // Validate CSRF state parameter + if let Some(ref expected) = expected_state { + let actual = params.get("state").cloned().unwrap_or_default(); + if actual != *expected { + let html = landing_html(&display_name, false); let response = format!( - "HTTP/1.1 200 OK\r\n\ + "HTTP/1.1 403 Forbidden\r\n\ Content-Type: text/html; charset=utf-8\r\n\ Connection: close\r\n\ \r\n\ @@ -239,11 +266,29 @@ pub async fn wait_for_callback( html ); let _ = socket.write_all(response.as_bytes()).await; - let _ = socket.shutdown().await; - - return Ok(value); + return Err(OAuthCallbackError::StateMismatch { + expected: expected.clone(), + actual, + }); } } + + // Look for the target parameter + if let Some(value) = params.get(param_name.as_str()) { + let html = landing_html(&display_name, true); + let response = format!( + "HTTP/1.1 200 OK\r\n\ + Content-Type: text/html; charset=utf-8\r\n\ + Connection: close\r\n\ + \r\n\ + {}", + html + ); + let _ = socket.write_all(response.as_bytes()).await; + let _ = socket.shutdown().await; + + return Ok(value.clone()); + } } // Not the callback we're looking for @@ -271,7 +316,288 @@ fn html_escape(s: &str) -> String { out } -/// HTML landing page shown in the browser after an OAuth redirect. +// ── Shared OAuth flow steps ───────────────────────────────────────── + +/// Response from the OAuth token exchange. +pub struct OAuthTokenResponse { + pub access_token: String, + pub refresh_token: Option, + pub expires_in: Option, +} + +/// Result of building an OAuth 2.0 authorization URL. +pub struct OAuthUrlResult { + /// The full authorization URL to redirect the user to. + pub url: String, + /// PKCE code verifier (must be sent with the token exchange request). + pub code_verifier: Option, + /// Random state parameter for CSRF protection (must be validated in callback). + pub state: String, +} + +/// Build an OAuth 2.0 authorization URL with optional PKCE and CSRF state. +/// +/// Returns an `OAuthUrlResult` containing the authorization URL, optional PKCE +/// code verifier, and a random `state` parameter for CSRF protection. The caller +/// must validate the `state` value in the callback before exchanging the code. +pub fn build_oauth_url( + authorization_url: &str, + client_id: &str, + redirect_uri: &str, + scopes: &[String], + use_pkce: bool, + extra_params: &HashMap, +) -> OAuthUrlResult { + // Generate PKCE verifier and challenge + let (code_verifier, code_challenge) = if use_pkce { + let mut verifier_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut verifier_bytes); + let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); + + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize()); + + (Some(verifier), Some(challenge)) + } else { + (None, None) + }; + + // Generate random state for CSRF protection + let mut state_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut state_bytes); + let state = URL_SAFE_NO_PAD.encode(state_bytes); + + // Build authorization URL + let mut auth_url = format!( + "{}?client_id={}&response_type=code&redirect_uri={}&state={}", + authorization_url, + urlencoding::encode(client_id), + urlencoding::encode(redirect_uri), + urlencoding::encode(&state), + ); + + if !scopes.is_empty() { + auth_url.push_str(&format!( + "&scope={}", + urlencoding::encode(&scopes.join(" ")) + )); + } + + if let Some(ref challenge) = code_challenge { + auth_url.push_str(&format!( + "&code_challenge={}&code_challenge_method=S256", + challenge + )); + } + + for (key, value) in extra_params { + auth_url.push_str(&format!( + "&{}={}", + urlencoding::encode(key), + urlencoding::encode(value) + )); + } + + OAuthUrlResult { + url: auth_url, + code_verifier, + state, + } +} + +/// Exchange an OAuth authorization code for tokens. +/// +/// POSTs to `token_url` with the authorization code and optional PKCE verifier. +/// If `client_secret` is provided, uses HTTP Basic auth; otherwise includes +/// `client_id` in the form body (for public clients). +pub async fn exchange_oauth_code( + token_url: &str, + client_id: &str, + client_secret: Option<&str>, + code: &str, + redirect_uri: &str, + code_verifier: Option<&str>, + access_token_field: &str, +) -> Result { + let client = reqwest::Client::new(); + let mut token_params = vec![ + ("grant_type", "authorization_code".to_string()), + ("code", code.to_string()), + ("redirect_uri", redirect_uri.to_string()), + ]; + + if let Some(verifier) = code_verifier { + token_params.push(("code_verifier", verifier.to_string())); + } + + let mut request = client.post(token_url); + + if let Some(secret) = client_secret { + request = request.basic_auth(client_id, Some(secret)); + } else { + token_params.push(("client_id", client_id.to_string())); + } + + let token_response = request + .form(&token_params) + .send() + .await + .map_err(|e| OAuthCallbackError::Io(format!("Token exchange request failed: {}", e)))?; + + if !token_response.status().is_success() { + let status = token_response.status(); + let body = token_response.text().await.unwrap_or_default(); + return Err(OAuthCallbackError::Io(format!( + "Token exchange failed: {} - {}", + status, body + ))); + } + + let token_data: serde_json::Value = token_response + .json() + .await + .map_err(|e| OAuthCallbackError::Io(format!("Failed to parse token response: {}", e)))?; + + let access_token = token_data + .get(access_token_field) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + // Log only the field names present, not values (which may contain tokens) + let fields: Vec<&str> = token_data + .as_object() + .map(|o| o.keys().map(|k| k.as_str()).collect()) + .unwrap_or_default(); + OAuthCallbackError::Io(format!( + "No '{}' field in token response (fields present: {:?})", + access_token_field, fields + )) + })? + .to_string(); + + let refresh_token = token_data + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(String::from); + let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64()); + + Ok(OAuthTokenResponse { + access_token, + refresh_token, + expires_in, + }) +} + +/// Store OAuth tokens (access + refresh) in the secrets store. +/// +/// Also stores the granted scopes as `{secret_name}_scopes` so that scope +/// expansion can be detected on subsequent activations. +#[allow(clippy::too_many_arguments)] +pub async fn store_oauth_tokens( + store: &(dyn SecretsStore + Send + Sync), + user_id: &str, + secret_name: &str, + provider: Option<&str>, + access_token: &str, + refresh_token: Option<&str>, + expires_in: Option, + scopes: &[String], +) -> Result<(), OAuthCallbackError> { + let mut params = CreateSecretParams::new(secret_name, access_token); + + if let Some(prov) = provider { + params = params.with_provider(prov); + } + + if let Some(secs) = expires_in { + let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64); + params = params.with_expiry(expires_at); + } + + store + .create(user_id, params) + .await + .map_err(|e| OAuthCallbackError::Io(format!("Failed to save token: {}", e)))?; + + // Store refresh token separately (no expiry, it's long-lived) + if let Some(rt) = refresh_token { + let refresh_name = format!("{}_refresh_token", secret_name); + let mut refresh_params = CreateSecretParams::new(&refresh_name, rt); + if let Some(prov) = provider { + refresh_params = refresh_params.with_provider(prov); + } + store + .create(user_id, refresh_params) + .await + .map_err(|e| OAuthCallbackError::Io(format!("Failed to save refresh token: {}", e)))?; + } + + // Store granted scopes for scope expansion detection + if !scopes.is_empty() { + let scopes_name = format!("{}_scopes", secret_name); + let scopes_value = scopes.join(" "); + let scopes_params = CreateSecretParams::new(&scopes_name, &scopes_value); + // Best-effort: scope tracking failure shouldn't block auth + let _ = store.create(user_id, scopes_params).await; + } + + Ok(()) +} + +/// Validate an OAuth token against a tool's validation endpoint. +/// +/// Sends a request to the configured endpoint with the token as a Bearer header. +/// Returns `Ok(())` if the response status matches the expected success status, +/// or an error with details if validation fails (wrong account, expired token, etc.). +pub async fn validate_oauth_token( + token: &str, + validation: &crate::tools::wasm::ValidationEndpointSchema, +) -> Result<(), OAuthCallbackError> { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?; + + let request = match validation.method.to_uppercase().as_str() { + "POST" => client.post(&validation.url), + _ => client.get(&validation.url), + }; + + let mut request = request.header("Authorization", format!("Bearer {}", token)); + + // Add custom headers from the validation schema (e.g., Notion-Version) + for (key, value) in &validation.headers { + request = request.header(key, value); + } + + let response = request + .send() + .await + .map_err(|e| OAuthCallbackError::Io(format!("Validation request failed: {}", e)))?; + + if response.status().as_u16() == validation.success_status { + Ok(()) + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let truncated: String = if body.len() > 200 { + let mut end = 200; + while end > 0 && !body.is_char_boundary(end) { + end -= 1; + } + format!("{}...", &body[..end]) + } else { + body + }; + Err(OAuthCallbackError::Io(format!( + "Token validation failed: HTTP {} (expected {}): {}", + status, validation.success_status, truncated + ))) + } +} + +// ── Landing pages ─────────────────────────────────────────────────── + pub fn landing_html(provider_name: &str, success: bool) -> String { let safe_name = html_escape(provider_name); let (icon, heading, subtitle, accent) = if success { @@ -512,4 +838,105 @@ mod tests { assert!(html.contains("#ef4444")); // red accent assert!(!html.contains("Connected")); } + + #[test] + fn test_build_oauth_url_basic() { + use std::collections::HashMap; + + use crate::cli::oauth_defaults::build_oauth_url; + + let result = build_oauth_url( + "https://accounts.google.com/o/oauth2/auth", + "my-client-id", + "http://localhost:9876/callback", + &["openid".to_string(), "email".to_string()], + false, + &HashMap::new(), + ); + + assert!( + result + .url + .starts_with("https://accounts.google.com/o/oauth2/auth?") + ); + assert!(result.url.contains("client_id=my-client-id")); + assert!(result.url.contains("response_type=code")); + assert!(result.url.contains("redirect_uri=")); + assert!(result.url.contains("scope=openid%20email")); + assert!(result.url.contains("state=")); + assert!(result.code_verifier.is_none()); + assert!(!result.state.is_empty()); + } + + #[test] + fn test_build_oauth_url_with_pkce() { + use std::collections::HashMap; + + use crate::cli::oauth_defaults::build_oauth_url; + + let result = build_oauth_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &[], + true, + &HashMap::new(), + ); + + assert!(result.url.contains("code_challenge=")); + assert!(result.url.contains("code_challenge_method=S256")); + assert!(result.code_verifier.is_some()); + let verifier = result.code_verifier.unwrap(); + assert!(!verifier.is_empty()); + } + + #[test] + fn test_build_oauth_url_with_extra_params() { + use std::collections::HashMap; + + use crate::cli::oauth_defaults::build_oauth_url; + + let mut extra = HashMap::new(); + extra.insert("access_type".to_string(), "offline".to_string()); + extra.insert("prompt".to_string(), "consent".to_string()); + + let result = build_oauth_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &["read".to_string()], + false, + &extra, + ); + + assert!(result.url.contains("access_type=offline")); + assert!(result.url.contains("prompt=consent")); + } + + #[test] + fn test_build_oauth_url_state_is_unique() { + use std::collections::HashMap; + + use crate::cli::oauth_defaults::build_oauth_url; + + let result1 = build_oauth_url( + "https://auth.example.com/authorize", + "client", + "http://localhost:9876/callback", + &[], + false, + &HashMap::new(), + ); + let result2 = build_oauth_url( + "https://auth.example.com/authorize", + "client", + "http://localhost:9876/callback", + &[], + false, + &HashMap::new(), + ); + + // State should be different each time (random) + assert_ne!(result1.state, result2.state); + } } diff --git a/src/cli/tool.rs b/src/cli/tool.rs index f099599e..1721541c 100644 --- a/src/cli/tool.rs +++ b/src/cli/tool.rs @@ -782,11 +782,7 @@ async fn auth_tool_oauth( auth: &crate::tools::wasm::AuthCapabilitySchema, oauth: &crate::tools::wasm::OAuthConfigSchema, ) -> anyhow::Result<()> { - use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; - use rand::RngCore; - use sha2::{Digest, Sha256}; - - use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT}; + use crate::cli::oauth_defaults; let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name); @@ -827,142 +823,69 @@ async fn auth_tool_oauth( println!(); let listener = oauth_defaults::bind_callback_listener().await?; - let redirect_uri = format!("http://localhost:{}/callback", OAUTH_CALLBACK_PORT); + let redirect_uri = format!("{}/callback", oauth_defaults::callback_url()); - // Generate PKCE verifier and challenge - let (code_verifier, code_challenge) = if oauth.use_pkce { - let mut verifier_bytes = [0u8; 32]; - rand::thread_rng().fill_bytes(&mut verifier_bytes); - let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); - - let mut hasher = Sha256::new(); - hasher.update(verifier.as_bytes()); - let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize()); - - (Some(verifier), Some(challenge)) - } else { - (None, None) - }; - - // Build authorization URL - let mut auth_url = format!( - "{}?client_id={}&response_type=code&redirect_uri={}", - oauth.authorization_url, - urlencoding::encode(&client_id), - urlencoding::encode(&redirect_uri) + // Build authorization URL with PKCE and CSRF state + let oauth_result = oauth_defaults::build_oauth_url( + &oauth.authorization_url, + &client_id, + &redirect_uri, + &oauth.scopes, + oauth.use_pkce, + &oauth.extra_params, ); - - if !oauth.scopes.is_empty() { - auth_url.push_str(&format!( - "&scope={}", - urlencoding::encode(&oauth.scopes.join(" ")) - )); - } - - if let Some(ref challenge) = code_challenge { - auth_url.push_str(&format!( - "&code_challenge={}&code_challenge_method=S256", - challenge - )); - } - - // Add extra params - for (key, value) in &oauth.extra_params { - auth_url.push_str(&format!( - "&{}={}", - urlencoding::encode(key), - urlencoding::encode(value) - )); - } + let code_verifier = oauth_result.code_verifier; println!(" Opening browser for {} login...", display_name); println!(); - if let Err(e) = open::that(&auth_url) { + if let Err(e) = open::that(&oauth_result.url) { println!(" Could not open browser: {}", e); println!(" Please open this URL manually:"); - println!(" {}", auth_url); + println!(" {}", oauth_result.url); } println!(" Waiting for authorization..."); - let code = - oauth_defaults::wait_for_callback(listener, "/callback", "code", display_name).await?; + let code = oauth_defaults::wait_for_callback( + listener, + "/callback", + "code", + display_name, + Some(&oauth_result.state), + ) + .await?; println!(); println!(" Exchanging code for token..."); // Exchange code for token - let client = reqwest::Client::new(); - let mut token_params = vec![ - ("grant_type", "authorization_code".to_string()), - ("code", code), - ("redirect_uri", redirect_uri), - ]; - - if let Some(ref verifier) = code_verifier { - token_params.push(("code_verifier", verifier.to_string())); - } - - // Build token request - let mut request = client.post(&oauth.token_url); - - // Use Basic auth if client_secret is provided, otherwise include client_id in body - if let Some(ref secret) = client_secret { - request = request.basic_auth(&client_id, Some(secret)); - } else { - token_params.push(("client_id", client_id)); - } - - let token_response = request.form(&token_params).send().await?; - - if !token_response.status().is_success() { - let status = token_response.status(); - let body = token_response.text().await.unwrap_or_default(); - return Err(anyhow::anyhow!( - "Token exchange failed: {} - {}", - status, - body - )); - } - - let token_data: serde_json::Value = token_response.json().await?; - let access_token = token_data - .get(&oauth.access_token_field) - .and_then(|v| v.as_str()) - .ok_or_else(|| { - anyhow::anyhow!( - "No {} in token response: {:?}", - oauth.access_token_field, - token_data - ) - })?; - - let refresh_token = token_data.get("refresh_token").and_then(|v| v.as_str()); - let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64()); - - // Save the token (with refresh token and expiry if provided) - save_token( - store, - user_id, - auth, - access_token, - refresh_token, - expires_in, + let token_response = oauth_defaults::exchange_oauth_code( + &oauth.token_url, + &client_id, + client_secret.as_deref(), + &code, + &redirect_uri, + code_verifier.as_deref(), + &oauth.access_token_field, ) .await?; - // Extract any additional info for display - let workspace_name = token_data - .get("workspace_name") - .and_then(|v| v.as_str()) - .or_else(|| token_data.get("team_name").and_then(|v| v.as_str())); + // Save tokens (access + refresh + scopes) + oauth_defaults::store_oauth_tokens( + store, + user_id, + &auth.secret_name, + auth.provider.as_deref(), + &token_response.access_token, + token_response.refresh_token.as_deref(), + token_response.expires_in, + &oauth.scopes, + ) + .await?; println!(); println!(" ✓ {} connected!", display_name); - if let Some(workspace) = workspace_name { - println!(" Workspace: {}", workspace); - } println!(); println!(" The tool can now access the API."); println!(); @@ -1107,46 +1030,15 @@ async fn validate_token( validation: &crate::tools::wasm::ValidationEndpointSchema, _secret_name: &str, ) -> anyhow::Result<()> { - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(10)) - .build()?; - - // Build request based on method - let request = match validation.method.to_uppercase().as_str() { - "GET" => client.get(&validation.url), - "POST" => client.post(&validation.url), - _ => client.get(&validation.url), - }; - - // Add authorization header (assume Bearer for now, could be extended) - let response = request - .header("Authorization", format!("Bearer {}", token)) - .header("Notion-Version", "2022-06-28") // Notion-specific, but harmless for others - .send() - .await?; - - if response.status().as_u16() == validation.success_status { - Ok(()) - } else { - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - Err(anyhow::anyhow!( - "HTTP {} (expected {}): {}", - status, - validation.success_status, - if body.len() > 100 { - format!("{}...", &body[..100]) - } else { - body - } - )) - } + crate::cli::oauth_defaults::validate_oauth_token(token, validation) + .await + .map_err(|e| anyhow::anyhow!("{}", e)) } /// Save token to secrets store. /// -/// Optionally stores a refresh token (as `{secret_name}_refresh_token`) and -/// sets `expires_at` on the access token so the runtime can auto-refresh. +/// Delegates to the shared `store_oauth_tokens` for OAuth tokens, or stores +/// directly for manual/env-var tokens (no scopes or refresh token). async fn save_token( store: &(dyn SecretsStore + Send + Sync), user_id: &str, @@ -1155,36 +1047,18 @@ async fn save_token( refresh_token: Option<&str>, expires_in: Option, ) -> anyhow::Result<()> { - let mut params = CreateSecretParams::new(&auth.secret_name, token); - - if let Some(ref provider) = auth.provider { - params = params.with_provider(provider); - } - - if let Some(secs) = expires_in { - let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64); - params = params.with_expiry(expires_at); - } - - store - .create(user_id, params) - .await - .map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?; - - // Store refresh token separately (no expiry, it's long-lived) - if let Some(rt) = refresh_token { - let refresh_name = format!("{}_refresh_token", auth.secret_name); - let mut refresh_params = CreateSecretParams::new(&refresh_name, rt); - if let Some(ref provider) = auth.provider { - refresh_params = refresh_params.with_provider(provider); - } - store - .create(user_id, refresh_params) - .await - .map_err(|e| anyhow::anyhow!("Failed to save refresh token: {}", e))?; - } - - Ok(()) + crate::cli::oauth_defaults::store_oauth_tokens( + store, + user_id, + &auth.secret_name, + auth.provider.as_deref(), + token, + refresh_token, + expires_in, + &[], // No scopes for manual/env-var tokens + ) + .await + .map_err(|e| anyhow::anyhow!("{}", e)) } /// Print success message. diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 50e4d9ac..9fafaee2 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -38,6 +38,9 @@ struct PendingAuth { _name: String, _kind: ExtensionKind, created_at: std::time::Instant, + /// Background task listening for the OAuth callback. + /// Aborted when a new auth flow starts for the same extension. + task_handle: Option>, } /// Runtime infrastructure needed for hot-activating WASM channels. @@ -58,6 +61,8 @@ pub struct SetupResult { pub message: String, /// Whether the channel was successfully activated after saving secrets. pub activated: bool, + /// OAuth authorization URL for the UI to open (if OAuth flow was started). + pub auth_url: Option, } /// Central manager for extension lifecycle operations. @@ -385,6 +390,7 @@ impl ExtensionManager { active, tools, needs_setup: false, + has_auth: false, installed: true, activation_error: None, }); @@ -411,6 +417,11 @@ impl ExtensionManager { .await .map(|e| e.display_name); let (authenticated, needs_setup) = self.check_tool_auth_status(&name).await; + let has_auth = self + .load_tool_capabilities(&name) + .await + .and_then(|c| c.auth) + .is_some(); extensions.push(InstalledExtension { name: name.clone(), kind: ExtensionKind::WasmTool, @@ -421,6 +432,7 @@ impl ExtensionManager { active, tools: if active { vec![name] } else { Vec::new() }, needs_setup, + has_auth, installed: true, activation_error: None, }); @@ -460,6 +472,7 @@ impl ExtensionManager { active, tools: Vec::new(), needs_setup, + has_auth: false, installed: true, activation_error, }); @@ -497,6 +510,7 @@ impl ExtensionManager { active: false, tools: Vec::new(), needs_setup: false, + has_auth: false, installed: false, activation_error: None, }); @@ -1338,6 +1352,7 @@ impl ExtensionManager { _name: name.to_string(), _kind: ExtensionKind::McpServer, created_at: std::time::Instant::now(), + task_handle: None, }, ); @@ -1424,23 +1439,45 @@ impl ExtensionManager { }); } - // Check if already authenticated - if self + // Check if already authenticated (with scope expansion detection) + let token_exists = self .secrets .exists(&self.user_id, &auth.secret_name) .await - .unwrap_or(false) - { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + .unwrap_or(false); + + if token_exists { + // If this tool has OAuth config, check whether new scopes are needed + let needs_reauth = if let Some(ref oauth) = auth.oauth { + let merged = self + .collect_shared_scopes(&auth.secret_name, &oauth.scopes) + .await; + let needs = self.needs_scope_expansion(&auth.secret_name, &merged).await; + tracing::debug!( + tool = name, + secret_name = %auth.secret_name, + merged_scopes = ?merged, + needs_reauth = needs, + "Scope expansion check" + ); + needs + } else { + false + }; + + if !needs_reauth { + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "authenticated".to_string(), + }); + } + // Fall through to OAuth branch for scope expansion } // If a token was provided, store it @@ -1464,6 +1501,62 @@ impl ExtensionManager { }); } + // OAuth flow: if the tool has OAuth config, start the browser-based flow. + // But only if credentials are available — if the tool has setup secrets + // for client_id/secret that aren't configured yet, return needs_setup. + if let Some(ref oauth) = auth.oauth { + let (setup_client_id_entry, setup_client_secret_entry) = + self.find_setup_credential_names(name).await; + + // Check all required (non-optional) setup credentials before starting + // OAuth, to avoid starting a flow that will fail during token exchange + // due to missing credentials. + let mut needs_setup = false; + if let Some((ref id_name, optional)) = setup_client_id_entry + && !optional + && !self + .secrets + .exists(&self.user_id, id_name) + .await + .unwrap_or(false) + { + needs_setup = true; + } + if !needs_setup + && let Some((ref secret_name, optional)) = setup_client_secret_entry + && !optional + && !self + .secrets + .exists(&self.user_id, secret_name) + .await + .unwrap_or(false) + { + needs_setup = true; + } + + if needs_setup { + let display = auth.display_name.as_deref().unwrap_or(name); + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: Some(format!( + "Configure OAuth credentials for {} in the Setup tab.", + display + )), + setup_url: auth.setup_url.clone(), + awaiting_token: false, + status: "needs_setup".to_string(), + }); + } + + return self + .start_wasm_oauth(name, &auth, oauth) + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string())); + } + // Return instructions for manual token entry let display = auth.display_name.unwrap_or_else(|| name.to_string()); let instructions = auth @@ -1534,6 +1627,353 @@ impl ExtensionManager { crate::tools::wasm::CapabilitiesFile::from_bytes(&cap_bytes).ok() } + /// Collect merged OAuth scopes from all installed tools sharing the same secret_name. + /// + /// When multiple tools share an OAuth provider (e.g., google-calendar and google-drive + /// both use `google_oauth_token`), we request all their scopes in a single OAuth flow + /// so one login covers everything. + async fn collect_shared_scopes( + &self, + secret_name: &str, + base_scopes: &[String], + ) -> Vec { + let mut all_scopes: std::collections::BTreeSet = + base_scopes.iter().cloned().collect(); + + if let Ok(tools) = discover_tools(&self.wasm_tools_dir).await { + for tool_name in tools.keys() { + if let Some(cap) = self.load_tool_capabilities(tool_name).await + && let Some(auth) = &cap.auth + && auth.secret_name == secret_name + && let Some(oauth) = &auth.oauth + { + all_scopes.extend(oauth.scopes.iter().cloned()); + } + } + } + + all_scopes.into_iter().collect() + } + + /// Check whether the stored scopes are insufficient for the merged scopes. + async fn needs_scope_expansion(&self, secret_name: &str, merged_scopes: &[String]) -> bool { + if merged_scopes.is_empty() { + return false; + } + + let scopes_key = format!("{}_scopes", secret_name); + let stored_scopes: std::collections::HashSet = + match self.secrets.get_decrypted(&self.user_id, &scopes_key).await { + Ok(secret) => { + let scopes: std::collections::HashSet = secret + .expose() + .split_whitespace() + .map(String::from) + .collect(); + tracing::debug!( + secret_name, + stored_scopes = ?scopes, + "Loaded stored scopes for expansion check" + ); + scopes + } + Err(_) => { + // No stored scopes record — this is a legacy token created before + // scope tracking. Force re-auth to ensure all required scopes are granted. + tracing::debug!( + secret_name, + "No stored scopes record, forcing re-auth for legacy token" + ); + return true; + } + }; + + // Check if any merged scope is missing from stored scopes + merged_scopes + .iter() + .any(|scope| !stored_scopes.contains(scope)) + } + + /// Find the setup secret names for OAuth client_id and client_secret. + /// + /// Scans `setup.required_secrets` for names containing "client_id" and "client_secret". + /// Returns `(Option<(name, optional)>, Option<(name, optional)>)`. + async fn find_setup_credential_names( + &self, + tool_name: &str, + ) -> (Option<(String, bool)>, Option<(String, bool)>) { + let Some(cap) = self.load_tool_capabilities(tool_name).await else { + return (None, None); + }; + let Some(setup) = &cap.setup else { + return (None, None); + }; + + let mut client_id_entry = None; + let mut client_secret_entry = None; + for secret in &setup.required_secrets { + let lower = secret.name.to_lowercase(); + if lower.ends_with("client_id") || lower == "client_id" { + client_id_entry = Some((secret.name.clone(), secret.optional)); + } else if lower.ends_with("client_secret") || lower == "client_secret" { + client_secret_entry = Some((secret.name.clone(), secret.optional)); + } + } + (client_id_entry, client_secret_entry) + } + + /// Resolve an OAuth credential value via: secrets store → inline → env var → builtin. + /// + /// For web gateway users, the secrets store is checked first because client_id/secret + /// may have been entered via the Setup tab (stored as setup secrets). + async fn resolve_oauth_credential( + &self, + inline_value: &Option, + env_var_name: &Option, + builtin_value: Option<&str>, + setup_secret_name: Option<&str>, + ) -> Option { + // 1. Check secrets store (entered via Setup tab) + if let Some(secret_name) = setup_secret_name + && let Ok(secret) = self.secrets.get_decrypted(&self.user_id, secret_name).await + { + let val = secret.expose(); + if !val.is_empty() { + return Some(val.to_string()); + } + } + + // 2. Inline value from capabilities.json + if let Some(val) = inline_value { + return Some(val.clone()); + } + + // 3. Runtime environment variable + if let Some(env) = env_var_name + && let Ok(val) = std::env::var(env) + { + return Some(val); + } + + // 4. Built-in defaults + builtin_value.map(String::from) + } + + /// Start the OAuth browser flow for a WASM tool. + /// + /// Binds a callback listener, builds the authorization URL, spawns a background + /// task to wait for the callback and exchange the code, then returns the auth URL + /// immediately so the web UI can open it. + async fn start_wasm_oauth( + &self, + name: &str, + auth: &crate::tools::wasm::AuthCapabilitySchema, + oauth: &crate::tools::wasm::OAuthConfigSchema, + ) -> Result { + use crate::cli::oauth_defaults; + + let builtin = oauth_defaults::builtin_credentials(&auth.secret_name); + + // Find setup secret names for client_id and client_secret from capabilities. + // These are the actual names used in the Setup tab (e.g., "google_oauth_client_id"), + // which may differ from "{secret_name}_client_id". + let (setup_client_id_entry, setup_client_secret_entry) = + self.find_setup_credential_names(name).await; + let setup_client_id_name = setup_client_id_entry.map(|(n, _)| n); + let setup_client_secret_name = setup_client_secret_entry.map(|(n, _)| n); + + // Resolve client_id: setup secrets → inline → env var → builtin + let client_id = self + .resolve_oauth_credential( + &oauth.client_id, + &oauth.client_id_env, + builtin.as_ref().map(|c| c.client_id), + setup_client_id_name.as_deref(), + ) + .await + .ok_or_else(|| { + let env_name = oauth + .client_id_env + .as_deref() + .unwrap_or("the client_id env var"); + let mut msg = format!( + "OAuth client_id not configured for '{}'. \ + Enter it in the Setup tab or set {} env var", + name, env_name + ); + // Only mention the Google-specific build flag for Google providers + if auth.secret_name.to_lowercase().contains("google") { + msg.push_str(", or build with IRONCLAW_GOOGLE_CLIENT_ID"); + } + msg.push('.'); + msg + })?; + + // Resolve client_secret (optional for PKCE-only flows) + let client_secret = self + .resolve_oauth_credential( + &oauth.client_secret, + &oauth.client_secret_env, + builtin.as_ref().map(|c| c.client_secret), + setup_client_secret_name.as_deref(), + ) + .await; + + // Cancel any existing pending auth for this tool (frees port 9876) + { + let mut pending = self.pending_auth.write().await; + if let Some(old) = pending.remove(name) + && let Some(handle) = old.task_handle + { + handle.abort(); + } + } + + // Bind callback listener + let listener = oauth_defaults::bind_callback_listener() + .await + .map_err(|e| format!("Failed to start OAuth callback listener: {}", e))?; + + let redirect_uri = format!("{}/callback", oauth_defaults::callback_url()); + + // Merge scopes from all tools sharing this provider + let merged_scopes = self + .collect_shared_scopes(&auth.secret_name, &oauth.scopes) + .await; + + // Build authorization URL with CSRF state + let oauth_result = oauth_defaults::build_oauth_url( + &oauth.authorization_url, + &client_id, + &redirect_uri, + &merged_scopes, + oauth.use_pkce, + &oauth.extra_params, + ); + let auth_url = oauth_result.url.clone(); + let code_verifier = oauth_result.code_verifier; + let expected_state = oauth_result.state; + + // Spawn background task: wait for callback → exchange code → validate → store tokens + let display_name = auth + .display_name + .clone() + .unwrap_or_else(|| name.to_string()); + let token_url = oauth.token_url.clone(); + let access_token_field = oauth.access_token_field.clone(); + let secret_name = auth.secret_name.clone(); + let provider = auth.provider.clone(); + let validation_endpoint = auth.validation_endpoint.clone(); + let user_id = self.user_id.clone(); + let secrets = Arc::clone(&self.secrets); + let sse_sender = self.sse_sender.read().await.clone(); + let ext_name = name.to_string(); + + let task_handle = tokio::spawn(async move { + let result: Result<(), String> = async { + let code = oauth_defaults::wait_for_callback( + listener, + "/callback", + "code", + &display_name, + Some(&expected_state), + ) + .await + .map_err(|e| e.to_string())?; + + let token_response = oauth_defaults::exchange_oauth_code( + &token_url, + &client_id, + client_secret.as_deref(), + &code, + &redirect_uri, + code_verifier.as_deref(), + &access_token_field, + ) + .await + .map_err(|e| e.to_string())?; + + // Validate the token before storing (catches wrong account, etc.) + if let Some(ref validation) = validation_endpoint { + oauth_defaults::validate_oauth_token(&token_response.access_token, validation) + .await + .map_err(|e| e.to_string())?; + } + + oauth_defaults::store_oauth_tokens( + secrets.as_ref(), + &user_id, + &secret_name, + provider.as_deref(), + &token_response.access_token, + token_response.refresh_token.as_deref(), + token_response.expires_in, + &merged_scopes, + ) + .await + .map_err(|e| e.to_string())?; + + Ok(()) + } + .await; + + // Broadcast SSE event + let (success, message) = match result { + Ok(()) => (true, format!("{} authenticated successfully", display_name)), + Err(ref e) => ( + false, + format!("{} authentication failed: {}", display_name, e), + ), + }; + + match &result { + Ok(()) => { + tracing::info!( + tool = %ext_name, + "OAuth completed successfully" + ); + } + Err(e) => { + tracing::warn!( + tool = %ext_name, + error = %e, + "WASM tool OAuth failed" + ); + } + } + + if let Some(ref sender) = sse_sender { + let _ = sender.send(crate::channels::web::types::SseEvent::AuthCompleted { + extension_name: ext_name, + success, + message, + }); + } + }); + + // Store pending auth with task handle + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::WasmTool, + created_at: std::time::Instant::now(), + task_handle: Some(task_handle), + }, + ); + + Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: Some(auth_url), + callback_type: Some("local".to_string()), + instructions: None, + setup_url: None, + awaiting_token: false, + status: "awaiting_authorization".to_string(), + }) + } + /// Check whether a WASM tool's required setup secrets are provided. /// /// Returns `(authenticated, needs_setup)` — same semantics as `check_channel_auth_status`. @@ -2261,7 +2701,16 @@ impl ExtensionManager { async fn cleanup_expired_auths(&self) { let mut pending = self.pending_auth.write().await; - pending.retain(|_, auth| auth.created_at.elapsed() < std::time::Duration::from_secs(300)); + pending.retain(|_, auth| { + let expired = auth.created_at.elapsed() >= std::time::Duration::from_secs(300); + if expired { + // Abort the background listener task to free port 9876 + if let Some(ref handle) = auth.task_handle { + handle.abort(); + } + } + !expired + }); } /// Get the setup schema for an extension (secret fields and their status). @@ -2478,16 +2927,55 @@ impl ExtensionManager { } } - // For tools, save and attempt auto-activation + // For tools, save and attempt auto-activation, then check auth. if kind == ExtensionKind::WasmTool { match self.activate_wasm_tool(name).await { Ok(result) => { - return Ok(SetupResult { - message: format!( + // Delete existing OAuth token so auth() starts a fresh flow. + // Done AFTER activation succeeds to avoid losing tokens on failure. + // This covers Reconfigure: user wants to re-auth (switch account, update creds). + if let Some(cap) = self.load_tool_capabilities(name).await + && let Some(ref auth_cfg) = cap.auth + && auth_cfg.oauth.is_some() + { + let _ = self + .secrets + .delete(&self.user_id, &auth_cfg.secret_name) + .await; + let _ = self + .secrets + .delete(&self.user_id, &format!("{}_scopes", auth_cfg.secret_name)) + .await; + let _ = self + .secrets + .delete( + &self.user_id, + &format!("{}_refresh_token", auth_cfg.secret_name), + ) + .await; + } + + // Check if auth is needed (OAuth or manual token). + // This is safe to call here — cancel-and-retry prevents port conflicts. + let mut auth_url = None; + if let Ok(auth_result) = self.auth(name, None).await { + auth_url = auth_result.auth_url; + } + let message = if auth_url.is_some() { + format!( + "Configuration saved and tool '{}' activated. Complete OAuth in your browser.", + name + ) + } else { + format!( "Configuration saved and tool '{}' activated. {}", name, result.message - ), + ) + }; + return Ok(SetupResult { + message, activated: true, + auth_url, }); } Err(e) => { @@ -2499,6 +2987,7 @@ impl ExtensionManager { return Ok(SetupResult { message: format!("Configuration saved for '{}'.", name), activated: false, + auth_url: None, }); } } @@ -2515,6 +3004,7 @@ impl ExtensionManager { name, result.message ), activated: true, + auth_url: None, }) } Err(e) => { @@ -2536,6 +3026,7 @@ impl ExtensionManager { name, e ), activated: false, + auth_url: None, }) } } diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index c0d45c90..353b6ff9 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -204,6 +204,9 @@ pub struct InstalledExtension { /// Whether this extension has a setup schema (required_secrets) that can be configured. #[serde(default)] pub needs_setup: bool, + /// Whether this extension has an auth configuration (OAuth or manual token). + #[serde(default)] + pub has_auth: bool, /// Whether this extension is installed locally (false = available in registry but not installed). #[serde(default = "default_true")] pub installed: bool, diff --git a/src/llm/session.rs b/src/llm/session.rs index 2dedfe56..7a410ef4 100644 --- a/src/llm/session.rs +++ b/src/llm/session.rs @@ -347,7 +347,7 @@ impl SessionManager { // The NEAR AI API redirects to: {frontend_callback}/auth/callback?token=X&... let session_token = - oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI") + oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI", None) .await .map_err(|e| LlmError::SessionRenewalFailed { provider: "nearai".to_string(), diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs index 98a6a62b..bd7b203c 100644 --- a/src/tools/mcp/auth.rs +++ b/src/tools/mcp/auth.rs @@ -531,7 +531,7 @@ pub async fn wait_for_authorization_callback( listener: TcpListener, server_name: &str, ) -> Result { - oauth_defaults::wait_for_callback(listener, "/callback", "code", server_name) + oauth_defaults::wait_for_callback(listener, "/callback", "code", server_name, None) .await .map_err(|e| match e { oauth_defaults::OAuthCallbackError::Denied => AuthError::AuthorizationDenied, @@ -539,6 +539,9 @@ pub async fn wait_for_authorization_callback( oauth_defaults::OAuthCallbackError::PortInUse(_, msg) => { AuthError::Http(format!("Port error: {}", msg)) } + oauth_defaults::OAuthCallbackError::StateMismatch { .. } => { + AuthError::Http("CSRF state mismatch in OAuth callback".to_string()) + } oauth_defaults::OAuthCallbackError::Io(msg) => AuthError::Http(msg), }) } diff --git a/src/tools/wasm/capabilities_schema.rs b/src/tools/wasm/capabilities_schema.rs index 50d8f338..97561df6 100644 --- a/src/tools/wasm/capabilities_schema.rs +++ b/src/tools/wasm/capabilities_schema.rs @@ -512,6 +512,11 @@ pub struct ValidationEndpointSchema { /// Expected HTTP status code for success (defaults to 200). #[serde(default = "default_success_status")] pub success_status: u16, + + /// Additional headers to send with the validation request. + /// Used for service-specific requirements (e.g., Notion-Version for Notion API). + #[serde(default)] + pub headers: HashMap, } fn default_method() -> String {