diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index 3b57872f..e9001909 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -62,6 +62,30 @@ pub fn builtin_client_id_override_env(secret_name: &str) -> Option<&'static str> } } +/// Suppress the baked-in desktop OAuth client secret when a hosted proxy is configured. +/// +/// In hosted deployments, IronClaw may resolve the platform Google client ID from +/// environment variables while still falling back to the baked-in desktop secret. +/// That client_id/client_secret mismatch breaks Google token exchange and refresh. +/// +/// When the proxy is configured, the platform will inject the correct server-side +/// secret for matching platform credentials, so the baked-in secret must be omitted. +pub fn hosted_proxy_client_secret( + client_secret: &Option, + builtin: Option<&OAuthCredentials>, + exchange_proxy_configured: bool, +) -> Option { + if !exchange_proxy_configured { + return client_secret.clone(); + } + + let builtin_secret = builtin.map(|credentials| credentials.client_secret); + match (client_secret, builtin_secret) { + (Some(resolved), Some(baked_in)) if resolved == baked_in => None, + _ => client_secret.clone(), + } +} + // ── Shared callback server ────────────────────────────────────────────── // Core OAuth callback infrastructure is defined in `crate::llm::oauth_helpers` @@ -661,6 +685,48 @@ pub struct ProxyTokenExchangeRequest<'a> { pub extra_token_params: &'a HashMap, } +pub struct ProxyRefreshTokenRequest<'a> { + pub proxy_url: &'a str, + pub gateway_token: &'a str, + pub token_url: &'a str, + pub client_id: &'a str, + pub client_secret: Option<&'a str>, + pub refresh_token: &'a str, + pub provider: Option<&'a str>, +} + +fn oauth_token_response_from_json( + token_data: serde_json::Value, + access_token_field: &str, +) -> Result { + let access_token = token_data + .get(access_token_field) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + 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 proxy 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, + }) +} + /// Exchange an OAuth authorization code via the platform's token exchange proxy. /// /// Authenticated via the gateway auth token (Bearer header). The caller may @@ -682,6 +748,7 @@ pub async fn exchange_via_proxy( let client = reqwest::Client::builder() .timeout(Duration::from_secs(60)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?; let mut params = vec![ @@ -724,41 +791,350 @@ pub async fn exchange_via_proxy( .json() .await .map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?; + oauth_token_response_from_json(token_data, request.access_token_field) +} - let access_token = token_data - .get(request.access_token_field) - .and_then(|v| v.as_str()) - .ok_or_else(|| { - 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 proxy response (fields present: {:?})", - request.access_token_field, fields - )) - })? - .to_string(); +/// Refresh an OAuth access token via the platform's token refresh proxy. +/// +/// Authenticated via the gateway auth token (Bearer header). The caller may +/// either rely on proxy-side secret lookup or forward a `client_secret` when +/// the provider requires it. +pub async fn refresh_token_via_proxy( + request: ProxyRefreshTokenRequest<'_>, +) -> Result { + if request.gateway_token.is_empty() { + return Err(OAuthCallbackError::Io( + "Gateway auth token is required for proxy token refresh".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()); + let refresh_url = format!("{}/oauth/refresh", request.proxy_url.trim_end_matches('/')); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(15)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?; - Ok(OAuthTokenResponse { - access_token, - refresh_token, - expires_in, - }) + let mut params = vec![ + ("refresh_token", request.refresh_token.to_string()), + ("token_url", request.token_url.to_string()), + ("client_id", request.client_id.to_string()), + ]; + if let Some(secret) = request.client_secret { + params.push(("client_secret", secret.to_string())); + } + if let Some(provider) = request.provider { + params.push(("provider", provider.to_string())); + } + + let response = client + .post(&refresh_url) + .bearer_auth(request.gateway_token) + .form(¶ms) + .send() + .await + .map_err(|e| { + OAuthCallbackError::Io(format!("Token refresh proxy request failed: {}", e)) + })?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(OAuthCallbackError::Io(format!( + "Token refresh proxy failed: {} - {}", + status, body + ))); + } + + let token_data: serde_json::Value = response + .json() + .await + .map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?; + + oauth_token_response_from_json(token_data, "access_token") } #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::net::SocketAddr; + use std::sync::Arc; + + use axum::extract::{Form, State}; + use axum::http::HeaderMap; + use axum::response::Redirect; + use axum::routing::post; + use axum::{Json, Router}; + use serde_json::json; + use tokio::net::TcpListener; + use tokio::sync::{Mutex, oneshot}; + use crate::cli::oauth_defaults::{ builtin_credentials, callback_host, callback_url, is_loopback_host, landing_html, }; use crate::config::helpers::lock_env; + use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET}; + + #[derive(Clone, Debug, PartialEq, Eq)] + struct RecordedProxyRequest { + authorization: Option, + form: HashMap, + } + + #[derive(Clone)] + struct MockProxyState { + requests: Arc>>, + exchange_redirect_target: String, + refresh_redirect_target: String, + } + + struct MockProxyServer { + addr: SocketAddr, + requests: Arc>>, + shutdown_tx: Option>, + server_task: Option>, + } + + impl MockProxyServer { + async fn start() -> Self { + async fn exchange_handler( + State(state): State, + headers: HeaderMap, + Form(form): Form>, + ) -> Json { + state.requests.lock().await.push(RecordedProxyRequest { + authorization: headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + form, + }); + Json(json!({ + "access_token": "proxy-access-token", + "refresh_token": "proxy-refresh-token", + "expires_in": 7200 + })) + } + + async fn refresh_handler( + State(state): State, + headers: HeaderMap, + Form(form): Form>, + ) -> Json { + state.requests.lock().await.push(RecordedProxyRequest { + authorization: headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + form, + }); + Json(json!({ + "access_token": "proxy-access-token", + "refresh_token": "proxy-refresh-token", + "expires_in": 7200 + })) + } + + async fn exchange_redirect_handler(State(state): State) -> Redirect { + Redirect::temporary(&state.exchange_redirect_target) + } + + async fn refresh_redirect_handler(State(state): State) -> Redirect { + Redirect::temporary(&state.refresh_redirect_target) + } + + let requests = Arc::new(Mutex::new(Vec::new())); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind mock proxy"); + let addr = listener.local_addr().expect("read mock proxy addr"); + let exchange_redirect_target = format!("http://{addr}/oauth/exchange"); + let refresh_redirect_target = format!("http://{addr}/oauth/refresh"); + let app = Router::new() + .route("/oauth/exchange", post(exchange_handler)) + .route("/oauth/refresh", post(refresh_handler)) + .route("/redirect/oauth/exchange", post(exchange_redirect_handler)) + .route("/redirect/oauth/refresh", post(refresh_redirect_handler)) + .with_state(MockProxyState { + requests: Arc::clone(&requests), + exchange_redirect_target, + refresh_redirect_target, + }); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let server_task = tokio::spawn(async move { + let _ = axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = shutdown_rx.await; + }) + .await; + }); + + Self { + addr, + requests, + shutdown_tx: Some(shutdown_tx), + server_task: Some(server_task), + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.addr) + } + + fn redirecting_base_url(&self) -> String { + format!("{}/redirect", self.base_url()) + } + + async fn requests(&self) -> Vec { + self.requests.lock().await.clone() + } + + async fn shutdown(mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + let _ = task.await; + } + } + } + + impl Drop for MockProxyServer { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + task.abort(); + } + } + } + + #[test] + fn test_hosted_proxy_client_secret_suppresses_builtin_secret() { + let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds"); + let client_secret = Some(builtin.client_secret.to_string()); + + let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true); + + assert_eq!(result, None); + } + + #[test] + fn test_hosted_proxy_client_secret_preserves_explicit_secret() { + let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds"); + let client_secret = Some("hosted-server-secret".to_string()); + + let result = super::hosted_proxy_client_secret(&client_secret, Some(&builtin), true); + + assert_eq!(result, client_secret); + } + + #[tokio::test] + async fn test_refresh_token_via_proxy_sends_auth_and_form() { + let server = MockProxyServer::start().await; + + let response = super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest { + proxy_url: &server.base_url(), + gateway_token: "gateway-test-token", + token_url: "https://oauth2.googleapis.com/token", + client_id: TEST_OAUTH_CLIENT_ID, + client_secret: Some(TEST_OAUTH_CLIENT_SECRET), + refresh_token: "refresh-token-123", + provider: Some("google"), + }) + .await + .expect("proxy refresh succeeds"); + + assert_eq!(response.access_token, "proxy-access-token"); + assert_eq!( + response.refresh_token.as_deref(), + Some("proxy-refresh-token") + ); + assert_eq!(response.expires_in, Some(7200)); + + let requests = server.requests().await; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].authorization.as_deref(), + Some("Bearer gateway-test-token") + ); + assert_eq!( + requests[0].form.get("token_url").map(String::as_str), + Some("https://oauth2.googleapis.com/token") + ); + assert_eq!( + requests[0].form.get("client_id").map(String::as_str), + Some(TEST_OAUTH_CLIENT_ID) + ); + assert_eq!( + requests[0].form.get("client_secret").map(String::as_str), + Some(TEST_OAUTH_CLIENT_SECRET) + ); + assert_eq!( + requests[0].form.get("refresh_token").map(String::as_str), + Some("refresh-token-123") + ); + assert_eq!( + requests[0].form.get("provider").map(String::as_str), + Some("google") + ); + + server.shutdown().await; + } + + #[tokio::test] + async fn test_exchange_via_proxy_does_not_follow_redirects() { + let server = MockProxyServer::start().await; + + let error = match super::exchange_via_proxy(super::ProxyTokenExchangeRequest { + proxy_url: &server.redirecting_base_url(), + gateway_token: "gateway-test-token", + code: "auth-code-123", + redirect_uri: "http://localhost:3000/oauth/callback", + token_url: "https://oauth2.googleapis.com/token", + client_id: TEST_OAUTH_CLIENT_ID, + client_secret: Some(TEST_OAUTH_CLIENT_SECRET), + access_token_field: "access_token", + code_verifier: Some("code-verifier-123"), + extra_token_params: &HashMap::new(), + }) + .await + { + Ok(_) => panic!("redirected proxy exchange should fail"), + Err(error) => error, + }; + + assert!(error.to_string().contains("307")); + assert!(server.requests().await.is_empty()); + + server.shutdown().await; + } + + #[tokio::test] + async fn test_refresh_token_via_proxy_does_not_follow_redirects() { + let server = MockProxyServer::start().await; + + let error = match super::refresh_token_via_proxy(super::ProxyRefreshTokenRequest { + proxy_url: &server.redirecting_base_url(), + gateway_token: "gateway-test-token", + token_url: "https://oauth2.googleapis.com/token", + client_id: TEST_OAUTH_CLIENT_ID, + client_secret: Some(TEST_OAUTH_CLIENT_SECRET), + refresh_token: "refresh-token-123", + provider: Some("google"), + }) + .await + { + Ok(_) => panic!("redirected proxy refresh should fail"), + Err(error) => error, + }; + + assert!(error.to_string().contains("307")); + assert!(server.requests().await.is_empty()); + + server.shutdown().await; + } #[test] fn test_is_loopback_host() { diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 39654305..0f308352 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -53,22 +53,6 @@ struct HostedOAuthFlowStart { flow: crate::cli::oauth_defaults::PendingOAuthFlow, } -fn hosted_proxy_client_secret( - client_secret: &Option, - builtin: Option<&crate::cli::oauth_defaults::OAuthCredentials>, - exchange_proxy_configured: bool, -) -> Option { - if !exchange_proxy_configured { - return client_secret.clone(); - } - - let builtin_secret = builtin.map(|credentials| credentials.client_secret); - match (client_secret, builtin_secret) { - (Some(resolved), Some(baked_in)) if resolved == baked_in => None, - _ => client_secret.clone(), - } -} - fn normalize_oauth_callback_path(path: &str) -> String { let trimmed_path = path.trim_end_matches('/'); if trimmed_path.is_empty() { @@ -3199,7 +3183,7 @@ impl ExtensionManager { // apps. Sending the desktop secret would cause a client_id/secret // mismatch because the container's GOOGLE_OAUTH_CLIENT_ID is the web // app, not the desktop app. - let proxy_client_secret = hosted_proxy_client_secret( + let proxy_client_secret = oauth_defaults::hosted_proxy_client_secret( &client_secret, builtin.as_ref(), oauth_defaults::exchange_proxy_url().is_some(), @@ -5714,7 +5698,7 @@ mod tests { use crate::extensions::manager::{ ChannelRuntimeState, FallbackDecision, TelegramBindingData, TelegramBindingResult, TelegramOwnerBindingState, build_wasm_channel_runtime_config_updates, - combine_install_errors, fallback_decision, hosted_proxy_client_secret, infer_kind_from_url, + combine_install_errors, fallback_decision, infer_kind_from_url, normalize_hosted_callback_url, send_telegram_text_message, telegram_message_matches_verification_code, }; @@ -7966,7 +7950,8 @@ mod tests { let builtin_ref = builtin.as_ref(); let secret = Some(builtin_ref.unwrap().client_secret.to_string()); - let result = hosted_proxy_client_secret(&secret, builtin_ref, true); + let result = + crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, true); assert_eq!( result, None, "built-in desktop secret must be suppressed when the exchange proxy is configured" @@ -7978,7 +7963,8 @@ mod tests { let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token"); let secret = Some("user-entered-custom-secret".to_string()); - let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true); + let result = + crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true); assert_eq!( result, Some("user-entered-custom-secret".to_string()), @@ -7992,7 +7978,8 @@ mod tests { let builtin_ref = builtin.as_ref(); let secret = Some(builtin_ref.unwrap().client_secret.to_string()); - let result = hosted_proxy_client_secret(&secret, builtin_ref, false); + let result = + crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin_ref, false); assert_eq!( result, secret, "built-in secret must be kept when the callback will exchange directly" @@ -8003,7 +7990,8 @@ mod tests { fn test_proxy_client_secret_none_stays_none() { let builtin = crate::cli::oauth_defaults::builtin_credentials("google_oauth_token"); - let result = hosted_proxy_client_secret(&None, builtin.as_ref(), true); + let result = + crate::cli::oauth_defaults::hosted_proxy_client_secret(&None, builtin.as_ref(), true); assert_eq!( result, None, "None secret stays None even when the exchange proxy is configured" @@ -8017,7 +8005,8 @@ mod tests { assert!(builtin.is_none()); let secret = Some("dcr-secret".to_string()); - let result = hosted_proxy_client_secret(&secret, builtin.as_ref(), true); + let result = + crate::cli::oauth_defaults::hosted_proxy_client_secret(&secret, builtin.as_ref(), true); assert_eq!( result, Some("dcr-secret".to_string()), diff --git a/src/tools/wasm/loader.rs b/src/tools/wasm/loader.rs index b50fc717..2a7ed040 100644 --- a/src/tools/wasm/loader.rs +++ b/src/tools/wasm/loader.rs @@ -418,6 +418,7 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option Option, + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: Tests use lock_env() to serialize environment access. + unsafe { + if let Some(ref value) = self.previous { + std::env::set_var(&self.key, value); + } else { + std::env::remove_var(&self.key); + } + } + } + } + + fn set_env_var(key: &str, value: Option<&str>) -> EnvVarGuard { + let previous = std::env::var(key).ok(); + // SAFETY: Tests use lock_env() to serialize environment access. + unsafe { + match value { + Some(value) => std::env::set_var(key, value), + None => std::env::remove_var(key), + } + } + EnvVarGuard { + key: key.to_string(), + previous, + } + } + #[test] fn wit_version_compat_none_is_ok() { // Pre-versioning extensions (no wit_version declared) should always pass @@ -871,6 +917,8 @@ mod tests { config.client_secret, Some(TEST_OAUTH_CLIENT_SECRET.to_string()) ); + assert_eq!(config.exchange_proxy_url, None); + assert_eq!(config.gateway_token, None); assert_eq!(config.secret_name, "google_oauth_token"); assert_eq!(config.provider, Some("google".to_string())); } @@ -931,6 +979,10 @@ mod tests { AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, }; + let _guard = lock_env(); + let _proxy_guard = set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", None); + let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None); + // google_oauth_token should fall back to built-in credentials let caps = CapabilitiesFile { auth: Some(AuthCapabilitySchema { @@ -952,6 +1004,95 @@ mod tests { let config = config.unwrap(); assert!(!config.client_id.is_empty()); assert!(config.client_secret.is_some()); + assert_eq!(config.exchange_proxy_url, None); + assert_eq!(config.gateway_token, None); + } + + #[test] + fn test_resolve_oauth_refresh_config_hosted_proxy_populates_env_and_suppresses_builtin_secret() + { + use crate::tools::wasm::capabilities_schema::{ + AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, + }; + + let _guard = lock_env(); + let _proxy_guard = set_env_var( + "IRONCLAW_OAUTH_EXCHANGE_URL", + Some("https://compose-api.example.com"), + ); + let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token")); + let _client_id_guard = + set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id")); + + let caps = CapabilitiesFile { + auth: Some(AuthCapabilitySchema { + secret_name: "google_oauth_token".to_string(), + provider: Some("google".to_string()), + oauth: Some(OAuthConfigSchema { + authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(), + token_url: "https://oauth2.googleapis.com/token".to_string(), + client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + + let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config"); + assert_eq!(config.client_id, "hosted-google-client-id"); + assert_eq!(config.client_secret, None); + assert_eq!( + config.exchange_proxy_url.as_deref(), + Some("https://compose-api.example.com") + ); + assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token")); + } + + #[test] + fn test_resolve_oauth_refresh_config_hosted_proxy_preserves_explicit_secret() { + use crate::tools::wasm::capabilities_schema::{ + AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, + }; + + let _guard = lock_env(); + let _proxy_guard = set_env_var( + "IRONCLAW_OAUTH_EXCHANGE_URL", + Some("https://compose-api.example.com"), + ); + let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token")); + let _client_id_guard = + set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id")); + let _client_secret_guard = + set_env_var("GOOGLE_OAUTH_CLIENT_SECRET", Some("hosted-server-secret")); + + let caps = CapabilitiesFile { + auth: Some(AuthCapabilitySchema { + secret_name: "google_oauth_token".to_string(), + provider: Some("google".to_string()), + oauth: Some(OAuthConfigSchema { + authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(), + token_url: "https://oauth2.googleapis.com/token".to_string(), + client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()), + client_secret_env: Some("GOOGLE_OAUTH_CLIENT_SECRET".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + + let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config"); + assert_eq!(config.client_id, "hosted-google-client-id"); + assert_eq!( + config.client_secret.as_deref(), + Some("hosted-server-secret") + ); + assert_eq!( + config.exchange_proxy_url.as_deref(), + Some("https://compose-api.example.com") + ); + assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token")); } // --------------------------------------------------------------- diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index 33fcedb9..05508e97 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -19,7 +19,7 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView}; use crate::context::JobContext; use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor}; use crate::safety::LeakDetector; -use crate::secrets::SecretsStore; +use crate::secrets::{DecryptedSecret, SecretsStore}; use crate::tools::tool::{Tool, ToolError, ToolOutput}; use crate::tools::wasm::capabilities::Capabilities; use crate::tools::wasm::credential_injector::{ @@ -44,6 +44,7 @@ wasmtime::component::bindgen!({ }); // Alias the export interface types for convenience. +use crate::cli::oauth_defaults; use exports::near::agent::tool as wit_tool; /// Configuration needed to refresh an expired OAuth access token. @@ -59,6 +60,10 @@ pub struct OAuthRefreshConfig { pub client_id: String, /// OAuth client_secret (optional, some providers use PKCE without a secret). pub client_secret: Option, + /// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080"). + pub exchange_proxy_url: Option, + /// Gateway auth token for authenticating with the hosted OAuth proxy. + pub gateway_token: Option, /// Secret name of the access token (e.g., "google_oauth_token"). /// The refresh token lives at `{secret_name}_refresh_token`. pub secret_name: String, @@ -1210,6 +1215,53 @@ async fn refresh_oauth_token( user_id: &str, config: &OAuthRefreshConfig, ) -> bool { + let refresh_name = format!("{}_refresh_token", config.secret_name); + + if let Some(proxy_url) = config.exchange_proxy_url.as_deref() { + let Some(gateway_token) = config.gateway_token.as_deref() else { + tracing::warn!( + "OAuth refresh proxy is configured, but no gateway auth token is available" + ); + return false; + }; + + // In hosted mode, the configured exchange proxy owns the outbound token + // refresh and validation policy for the provider token_url. Direct-mode + // HTTPS/private-IP checks remain in place for self-hosted refreshes below. + let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await { + Some(secret) => secret, + None => return false, + }; + let token_response = match oauth_defaults::refresh_token_via_proxy( + oauth_defaults::ProxyRefreshTokenRequest { + proxy_url, + gateway_token, + token_url: &config.token_url, + client_id: &config.client_id, + client_secret: config.client_secret.as_deref(), + refresh_token: refresh_secret.expose(), + provider: config.provider.as_deref(), + }, + ) + .await + { + Ok(response) => response, + Err(error) => { + tracing::warn!(error = %error, "OAuth token refresh via proxy failed"); + return false; + } + }; + + return persist_refreshed_oauth_tokens( + store, + user_id, + config, + &refresh_name, + token_response, + ) + .await; + } + // SSRF defense: token_url comes from the tool's capabilities file. if !config.token_url.starts_with("https://") { tracing::warn!( @@ -1227,19 +1279,6 @@ async fn refresh_oauth_token( return false; } - let refresh_name = format!("{}_refresh_token", config.secret_name); - let refresh_secret = match store.get_decrypted(user_id, &refresh_name).await { - Ok(s) => s, - Err(e) => { - tracing::debug!( - secret_name = %refresh_name, - error = %e, - "No refresh token available, skipping token refresh" - ); - return false; - } - }; - let client = match reqwest::Client::builder() .timeout(Duration::from_secs(15)) .redirect(reqwest::redirect::Policy::none()) @@ -1252,6 +1291,10 @@ async fn refresh_oauth_token( } }; + let refresh_secret = match load_oauth_refresh_secret(store, user_id, &refresh_name).await { + Some(secret) => secret, + None => return false, + }; let mut params = vec![ ("grant_type", "refresh_token".to_string()), ("refresh_token", refresh_secret.expose().to_string()), @@ -1287,22 +1330,55 @@ async fn refresh_oauth_token( return false; } }; - - let new_access_token = match token_data.get("access_token").and_then(|v| v.as_str()) { - Some(t) => t, + let token_response = match token_data.get("access_token").and_then(|v| v.as_str()) { + Some(access_token) => oauth_defaults::OAuthTokenResponse { + access_token: access_token.to_string(), + refresh_token: token_data + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(str::to_string), + expires_in: token_data.get("expires_in").and_then(|v| v.as_u64()), + }, None => { tracing::warn!("Token refresh response missing access_token field"); return false; } }; - // Store the new access token with expiry + persist_refreshed_oauth_tokens(store, user_id, config, &refresh_name, token_response).await +} + +async fn load_oauth_refresh_secret( + store: &(dyn SecretsStore + Send + Sync), + user_id: &str, + refresh_name: &str, +) -> Option { + match store.get_decrypted(user_id, refresh_name).await { + Ok(secret) => Some(secret), + Err(error) => { + tracing::debug!( + secret_name = %refresh_name, + error = %error, + "No refresh token available, skipping token refresh" + ); + None + } + } +} + +async fn persist_refreshed_oauth_tokens( + store: &(dyn SecretsStore + Send + Sync), + user_id: &str, + config: &OAuthRefreshConfig, + refresh_name: &str, + token_response: oauth_defaults::OAuthTokenResponse, +) -> bool { let mut access_params = - crate::secrets::CreateSecretParams::new(&config.secret_name, new_access_token); + crate::secrets::CreateSecretParams::new(&config.secret_name, &token_response.access_token); if let Some(ref provider) = config.provider { access_params = access_params.with_provider(provider); } - if let Some(expires_in) = token_data.get("expires_in").and_then(|v| v.as_u64()) { + if let Some(expires_in) = token_response.expires_in { let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in as i64); access_params = access_params.with_expiry(expires_at); } @@ -1312,10 +1388,8 @@ async fn refresh_oauth_token( return false; } - // Store rotated refresh token if the provider sent a new one - if let Some(new_refresh) = token_data.get("refresh_token").and_then(|v| v.as_str()) { - let mut refresh_params = - crate::secrets::CreateSecretParams::new(&refresh_name, new_refresh); + if let Some(new_refresh) = token_response.refresh_token.as_deref() { + let mut refresh_params = crate::secrets::CreateSecretParams::new(refresh_name, new_refresh); if let Some(ref provider) = config.provider { refresh_params = refresh_params.with_provider(provider); } @@ -1664,9 +1738,18 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::net::SocketAddr; use std::sync::{Arc, Mutex}; use async_trait::async_trait; + use axum::extract::{Form, State}; + use axum::http::HeaderMap; + use axum::routing::post; + use axum::{Json, Router}; + use serde_json::json; + use tokio::net::TcpListener; + use tokio::sync::{Mutex as AsyncMutex, oneshot}; use uuid::Uuid; use crate::context::JobContext; @@ -1756,6 +1839,95 @@ mod tests { } } + #[derive(Clone, Debug, PartialEq, Eq)] + struct RecordedProxyRequest { + authorization: Option, + form: HashMap, + } + + struct MockProxyServer { + addr: SocketAddr, + requests: Arc>>, + shutdown_tx: Option>, + server_task: Option>, + } + + impl MockProxyServer { + async fn start() -> Self { + async fn refresh_handler( + State(requests): State>>>, + headers: HeaderMap, + Form(form): Form>, + ) -> Json { + requests.lock().await.push(RecordedProxyRequest { + authorization: headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + form, + }); + Json(json!({ + "access_token": "mock-refreshed-access-token", + "refresh_token": "mock-rotated-refresh-token", + "expires_in": 3600 + })) + } + + let requests = Arc::new(AsyncMutex::new(Vec::new())); + let app = Router::new() + .route("/oauth/refresh", post(refresh_handler)) + .with_state(Arc::clone(&requests)); + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind mock proxy"); + let addr = listener.local_addr().expect("read mock proxy addr"); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let server_task = tokio::spawn(async move { + let _ = axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = shutdown_rx.await; + }) + .await; + }); + + Self { + addr, + requests, + shutdown_tx: Some(shutdown_tx), + server_task: Some(server_task), + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.addr) + } + + async fn requests(&self) -> Vec { + self.requests.lock().await.clone() + } + + async fn shutdown(mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + let _ = task.await; + } + } + } + + impl Drop for MockProxyServer { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + task.abort(); + } + } + } + #[test] fn test_wrapper_creation() { // This test verifies the runtime can be created @@ -2094,8 +2266,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_bearer() { - use std::collections::HashMap; - use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, }; @@ -2141,8 +2311,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_owner_scope_bearer() { - use std::collections::HashMap; - use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, }; @@ -2188,8 +2356,6 @@ mod tests { #[tokio::test] async fn test_execute_resolves_host_credentials_from_owner_scope_context() { - use std::collections::HashMap; - use crate::secrets::{CredentialLocation, CredentialMapping}; use crate::tools::wasm::capabilities::HttpCapability; @@ -2239,8 +2405,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_missing_secret() { - use std::collections::HashMap; - use crate::secrets::{CredentialLocation, CredentialMapping}; use crate::tools::wasm::capabilities::HttpCapability; use crate::tools::wasm::wrapper::resolve_host_credentials; @@ -2272,8 +2436,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_skips_refresh_when_not_expired() { - use std::collections::HashMap; - use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, }; @@ -2315,6 +2477,8 @@ mod tests { token_url: "https://oauth2.googleapis.com/token".to_string(), client_id: TEST_OAUTH_CLIENT_ID.to_string(), client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()), + exchange_proxy_url: None, + gateway_token: None, secret_name: "google_oauth_token".to_string(), provider: Some("google".to_string()), }; @@ -2331,8 +2495,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_skips_refresh_no_config() { - use std::collections::HashMap; - use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, }; @@ -2376,8 +2538,6 @@ mod tests { #[tokio::test] async fn test_resolve_host_credentials_skips_refresh_no_expires_at() { - use std::collections::HashMap; - use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, }; @@ -2417,6 +2577,8 @@ mod tests { token_url: "https://oauth2.googleapis.com/token".to_string(), client_id: TEST_OAUTH_CLIENT_ID.to_string(), client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()), + exchange_proxy_url: None, + gateway_token: None, secret_name: "google_oauth_token".to_string(), provider: Some("google".to_string()), }; @@ -2431,6 +2593,249 @@ mod tests { ); } + #[tokio::test] + async fn test_resolve_host_credentials_refreshes_via_proxy_without_direct_token_url_validation() + { + use crate::secrets::{ + CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, + }; + use crate::tools::wasm::capabilities::HttpCapability; + use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials}; + + let proxy = MockProxyServer::start().await; + let store = test_secrets_store(); + + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token", "expired-access-token") + .with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)), + ) + .await + .unwrap(); + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"), + ) + .await + .unwrap(); + + let mut credentials = HashMap::new(); + credentials.insert( + "google_oauth_token".to_string(), + CredentialMapping { + secret_name: "google_oauth_token".to_string(), + location: CredentialLocation::AuthorizationBearer, + host_patterns: vec!["www.googleapis.com".to_string()], + }, + ); + + let caps = Capabilities { + http: Some(HttpCapability { + credentials, + ..Default::default() + }), + ..Default::default() + }; + + let oauth_config = OAuthRefreshConfig { + token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(), + client_id: "hosted-google-client-id".to_string(), + client_secret: None, + exchange_proxy_url: Some(proxy.base_url()), + gateway_token: Some("gateway-test-token".to_string()), + secret_name: "google_oauth_token".to_string(), + provider: Some("google".to_string()), + }; + + let resolved = + resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await; + assert_eq!(resolved.len(), 1); + assert_eq!( + resolved[0].headers.get("Authorization"), + Some(&"Bearer mock-refreshed-access-token".to_string()) + ); + + let access_secret = store.get("user1", "google_oauth_token").await.unwrap(); + assert!( + access_secret + .expires_at + .expect("refreshed access token expiry") + > chrono::Utc::now() + ); + let access_value = store + .get_decrypted("user1", "google_oauth_token") + .await + .unwrap(); + assert_eq!(access_value.expose(), "mock-refreshed-access-token"); + + let refresh_value = store + .get_decrypted("user1", "google_oauth_token_refresh_token") + .await + .unwrap(); + assert_eq!(refresh_value.expose(), "mock-rotated-refresh-token"); + + let requests = proxy.requests().await; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].authorization.as_deref(), + Some("Bearer gateway-test-token") + ); + assert_eq!( + requests[0].form.get("client_id").map(String::as_str), + Some("hosted-google-client-id") + ); + assert_eq!( + requests[0].form.get("token_url").map(String::as_str), + Some("http://127.0.0.1:9/provider-token-endpoint") + ); + assert_eq!( + requests[0].form.get("refresh_token").map(String::as_str), + Some("stored-refresh-token") + ); + assert_eq!( + requests[0].form.get("provider").map(String::as_str), + Some("google") + ); + assert!(!requests[0].form.contains_key("client_secret")); + + proxy.shutdown().await; + } + + #[tokio::test] + async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() { + use crate::secrets::{ + CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, + }; + use crate::tools::wasm::capabilities::HttpCapability; + use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials}; + + let store = RecordingSecretsStore::new(); + + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token", "expired-access-token") + .with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)), + ) + .await + .unwrap(); + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"), + ) + .await + .unwrap(); + + let mut credentials = HashMap::new(); + credentials.insert( + "google_oauth_token".to_string(), + CredentialMapping { + secret_name: "google_oauth_token".to_string(), + location: CredentialLocation::AuthorizationBearer, + host_patterns: vec!["www.googleapis.com".to_string()], + }, + ); + + let caps = Capabilities { + http: Some(HttpCapability { + credentials, + ..Default::default() + }), + ..Default::default() + }; + + let oauth_config = OAuthRefreshConfig { + token_url: "https://oauth2.googleapis.com/token".to_string(), + client_id: "hosted-google-client-id".to_string(), + client_secret: None, + exchange_proxy_url: Some("https://compose-api.example.com".to_string()), + gateway_token: None, + secret_name: "google_oauth_token".to_string(), + provider: Some("google".to_string()), + }; + + let resolved = + resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await; + assert!(resolved.is_empty()); + + let lookups = store.decrypted_lookups(); + assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string()))); + assert!(!lookups.contains(&( + "user1".to_string(), + "google_oauth_token_refresh_token".to_string(), + ))); + } + + #[tokio::test] + async fn test_resolve_host_credentials_skips_refresh_token_lookup_for_invalid_direct_token_url() + { + use crate::secrets::{ + CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, + }; + use crate::tools::wasm::capabilities::HttpCapability; + use crate::tools::wasm::wrapper::{OAuthRefreshConfig, resolve_host_credentials}; + + let store = RecordingSecretsStore::new(); + + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token", "expired-access-token") + .with_expiry(chrono::Utc::now() - chrono::Duration::hours(1)), + ) + .await + .unwrap(); + store + .create( + "user1", + CreateSecretParams::new("google_oauth_token_refresh_token", "stored-refresh-token"), + ) + .await + .unwrap(); + + let mut credentials = HashMap::new(); + credentials.insert( + "google_oauth_token".to_string(), + CredentialMapping { + secret_name: "google_oauth_token".to_string(), + location: CredentialLocation::AuthorizationBearer, + host_patterns: vec!["www.googleapis.com".to_string()], + }, + ); + + let caps = Capabilities { + http: Some(HttpCapability { + credentials, + ..Default::default() + }), + ..Default::default() + }; + + let oauth_config = OAuthRefreshConfig { + token_url: "http://127.0.0.1:9/provider-token-endpoint".to_string(), + client_id: TEST_OAUTH_CLIENT_ID.to_string(), + client_secret: Some(TEST_OAUTH_CLIENT_SECRET.to_string()), + exchange_proxy_url: None, + gateway_token: None, + secret_name: "google_oauth_token".to_string(), + provider: Some("google".to_string()), + }; + + let resolved = + resolve_host_credentials(&caps, Some(&store), "user1", Some(&oauth_config)).await; + assert!(resolved.is_empty()); + + let lookups = store.decrypted_lookups(); + assert!(lookups.contains(&("user1".to_string(), "google_oauth_token".to_string()))); + assert!(!lookups.contains(&( + "user1".to_string(), + "google_oauth_token_refresh_token".to_string(), + ))); + } + #[test] fn test_is_private_ip_v4() { use std::net::IpAddr; diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0cf5e6dc..46b7b752 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -53,6 +53,7 @@ HEADED=1 pytest scenarios/ | `test_skills.py` | Skills tab UI visibility, ClawHub search (skipped if registry unreachable), install + remove lifecycle | | `test_sse_reconnect.py` | SSE reconnects after programmatic `eventSource.close()` + `connectSSE()`; history is reloaded after reconnect | | `test_tool_approval.py` | Approval card appears, buttons disable on approve/deny, parameters toggle via `page.evaluate("showApproval(...)")`; the waiting-approval regression uses a real HTTP tool call | +| `test_oauth_refresh.py` | Hosted Gmail OAuth regression: complete setup via `/oauth/callback`, expire the stored access token in libSQL, trigger a real `gmail` tool call through `/api/chat/send`, and verify refresh goes through the mock `/oauth/refresh` proxy without forwarding `client_secret` | ## `helpers.py` @@ -75,6 +76,7 @@ All fixtures are defined in `tests/e2e/conftest.py`. Running `pytest scenarios/` | `ironclaw_binary` | Checks `target/debug/ironclaw`; if absent, runs `cargo build --no-default-features --features libsql` (timeout 600s). | | `mock_llm_server` | Starts `mock_llm.py --port 0`, reads the assigned port from stdout, waits for `/v1/models` to return 200. Yields the base URL. | | `ironclaw_server` | Starts the ironclaw binary with a minimal env (see below), waits for `/api/health` (timeout 60s). Yields the base URL. On teardown sends **SIGINT** (not SIGTERM) so the tokio ctrl_c handler triggers a graceful shutdown and LLVM coverage data is flushed. | +| `hosted_oauth_refresh_server` | Starts a second ironclaw instance with a dedicated libSQL DB and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id`, while still pointing `IRONCLAW_OAUTH_EXCHANGE_URL` at `mock_llm.py`. Yields a dict with `base_url`, `db_path`, `gateway_user_id`, and `mock_llm_url` for the hosted refresh regression scenario. | | `browser` | Launches a single Chromium instance (headless by default; set `HEADED=1` for headed). Shared across all tests. | ### Function-scoped fixtures @@ -100,6 +102,8 @@ EMBEDDING_ENABLED=false, SKILLS_ENABLED=true ONBOARD_COMPLETED=true # prevents setup wizard ``` +The `hosted_oauth_refresh_server` fixture uses the same baseline, but with its own DB/home tempdirs and `GOOGLE_OAUTH_CLIENT_ID=hosted-google-client-id` so hosted OAuth flows exercise proxy credential injection instead of the baked-in desktop Google app. + The binary is also started with `--no-onboard`. Coverage env vars (`CARGO_LLVM_COV*`, `LLVM_*`, `CARGO_ENCODED_RUSTFLAGS`, `CARGO_INCREMENTAL`) are forwarded from the outer environment when present. ## Mock LLM (`mock_llm.py`) @@ -113,6 +117,11 @@ python mock_llm.py --port 0 It serves `POST /v1/chat/completions` (streaming + non-streaming) and `GET /v1/models`. Responses are pattern-matched from `CANNED_RESPONSES` against the last user message. Unmatched messages return `"I understand your request."`. The model name reported is always `"mock-model"`. +It also hosts OAuth test endpoints: +- `POST /oauth/exchange` for hosted auth-code exchange +- `POST /oauth/refresh` for hosted refresh-token exchange +- `GET /__mock/oauth/state` and `POST /__mock/oauth/reset` so HTTP E2E scenarios can assert exact proxy payloads and reset counters between setup and refresh assertions + To add a new canned response: ```python # In mock_llm.py diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 06c7da03..1496f93f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -113,6 +113,15 @@ def _reserve_loopback_sockets(count: int) -> list[socket.socket]: raise +def _forward_coverage_env(env: dict[str, str]) -> None: + """Forward cargo-llvm-cov env vars into child processes when present.""" + cov_env_prefixes = ("CARGO_LLVM_COV", "LLVM_") + cov_env_extras = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL") + for key, val in os.environ.items(): + if key.startswith(cov_env_prefixes) or key in cov_env_extras: + env[key] = val + + @pytest.fixture(scope="session") def ironclaw_binary(): """Ensure ironclaw binary is built. Returns the binary path.""" @@ -264,14 +273,7 @@ async def ironclaw_server( "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, } - # Forward LLVM coverage instrumentation env vars when present - # (allows cargo-llvm-cov to collect profraw data from E2E runs). - # Use prefix matching to stay resilient to cargo-llvm-cov changes. - COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_") - COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL") - for key, val in os.environ.items(): - if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS: - env[key] = val + _forward_coverage_env(env) proc = await asyncio.create_subprocess_exec( ironclaw_binary, "--no-onboard", stdin=asyncio.subprocess.DEVNULL, @@ -310,6 +312,109 @@ async def ironclaw_server( proc.kill() +@pytest.fixture(scope="session") +async def hosted_oauth_refresh_server( + ironclaw_binary, + mock_llm_server, + wasm_tools_dir, +): + """Start a hosted-mode ironclaw instance for OAuth refresh regression tests.""" + reserved = _reserve_loopback_sockets(2) + db_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-db-") + home_tmpdir = tempfile.TemporaryDirectory(prefix="ironclaw-e2e-hosted-oauth-home-") + + try: + gateway_port = reserved[0].getsockname()[1] + http_port = reserved[1].getsockname()[1] + for sock in reserved: + if sock.fileno() != -1: + sock.close() + + db_path = os.path.join(db_tmpdir.name, "hosted-oauth-refresh.db") + home_dir = home_tmpdir.name + env = { + "PATH": os.environ.get("PATH", "/usr/bin:/bin"), + "HOME": home_dir, + "IRONCLAW_BASE_DIR": os.path.join(home_dir, ".ironclaw"), + "RUST_LOG": "ironclaw=info", + "RUST_BACKTRACE": "1", + "IRONCLAW_OWNER_ID": OWNER_SCOPE_ID, + "GATEWAY_ENABLED": "true", + "GATEWAY_HOST": "127.0.0.1", + "GATEWAY_PORT": str(gateway_port), + "GATEWAY_AUTH_TOKEN": AUTH_TOKEN, + "GATEWAY_USER_ID": OWNER_SCOPE_ID, + "HTTP_HOST": "127.0.0.1", + "HTTP_PORT": str(http_port), + "HTTP_WEBHOOK_SECRET": HTTP_WEBHOOK_SECRET, + "CLI_ENABLED": "false", + "LLM_BACKEND": "openai_compatible", + "LLM_BASE_URL": mock_llm_server, + "LLM_MODEL": "mock-model", + "DATABASE_BACKEND": "libsql", + "LIBSQL_PATH": db_path, + "SECRETS_MASTER_KEY": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "SANDBOX_ENABLED": "false", + "SKILLS_ENABLED": "true", + "ROUTINES_ENABLED": "true", + "HEARTBEAT_ENABLED": "false", + "EMBEDDING_ENABLED": "false", + "WASM_ENABLED": "true", + "WASM_TOOLS_DIR": wasm_tools_dir, + "WASM_CHANNELS_DIR": _WASM_CHANNELS_TMPDIR.name, + "ONBOARD_COMPLETED": "true", + "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", + "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, + "GOOGLE_OAUTH_CLIENT_ID": "hosted-google-client-id", + } + _forward_coverage_env(env) + + proc = await asyncio.create_subprocess_exec( + ironclaw_binary, "--no-onboard", + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + ) + base_url = f"http://127.0.0.1:{gateway_port}" + try: + await wait_for_ready(f"{base_url}/api/health", timeout=60) + yield { + "base_url": base_url, + "db_path": db_path, + "gateway_user_id": OWNER_SCOPE_ID, + "mock_llm_url": mock_llm_server, + } + except TimeoutError: + returncode = proc.returncode + stderr_bytes = b"" + if proc.stderr: + try: + stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2) + except (asyncio.TimeoutError, Exception): + pass + stderr_text = stderr_bytes.decode("utf-8", errors="replace") + if proc.returncode is None: + proc.kill() + pytest.fail( + f"hosted oauth refresh server failed to start on port {gateway_port} " + f"(returncode={returncode}).\nstderr:\n{stderr_text}" + ) + finally: + if proc.returncode is None: + proc.send_signal(signal.SIGINT) + try: + await asyncio.wait_for(proc.wait(), timeout=10) + except asyncio.TimeoutError: + proc.kill() + finally: + for sock in reserved: + if sock.fileno() != -1: + sock.close() + db_tmpdir.cleanup() + home_tmpdir.cleanup() + + @pytest.fixture(scope="session") async def http_channel_server(ironclaw_server, server_ports): """HTTP webhook channel base URL.""" @@ -362,12 +467,7 @@ async def http_channel_server_without_secret( "IRONCLAW_OAUTH_CALLBACK_URL": "https://oauth.test.example/oauth/callback", "IRONCLAW_OAUTH_EXCHANGE_URL": mock_llm_server, } - # Forward LLVM coverage instrumentation env vars when present - COV_ENV_PREFIXES = ("CARGO_LLVM_COV", "LLVM_") - COV_ENV_EXTRAS = ("CARGO_ENCODED_RUSTFLAGS", "CARGO_INCREMENTAL") - for key, val in os.environ.items(): - if key.startswith(COV_ENV_PREFIXES) or key in COV_ENV_EXTRAS: - env[key] = val + _forward_coverage_env(env) proc = await asyncio.create_subprocess_exec( ironclaw_binary, "--no-onboard", stdin=asyncio.subprocess.DEVNULL, diff --git a/tests/e2e/mock_llm.py b/tests/e2e/mock_llm.py index 359c22d5..1147662c 100644 --- a/tests/e2e/mock_llm.py +++ b/tests/e2e/mock_llm.py @@ -34,6 +34,15 @@ TOOL_CALL_PATTERNS = [ "body": {"label": m.group("label")}, }, ), + ( + re.compile(r"check gmail unread|gmail unread", re.IGNORECASE), + "gmail", + lambda _: { + "action": "list_messages", + "query": "is:unread", + "max_results": 1, + }, + ), (re.compile(r"what time|current time", re.IGNORECASE), "time", lambda _: {"operation": "now"}), ( re.compile( @@ -91,6 +100,15 @@ TOOL_CALL_PATTERNS = [ ] +def _new_oauth_state() -> dict: + return { + "exchange_count": 0, + "refresh_count": 0, + "last_exchange": None, + "last_refresh": None, + } + + def _last_user_content(messages: list[dict]) -> str: for msg in reversed(messages): if msg.get("role") == "user": @@ -272,6 +290,12 @@ async def oauth_exchange(request: web.Request) -> web.Response: specific token params such as RFC 8707 `resource` are forwarded here. """ data = await request.post() + oauth_state = request.app["oauth_state"] + oauth_state["exchange_count"] += 1 + oauth_state["last_exchange"] = { + "authorization": request.headers.get("Authorization"), + "form": dict(data), + } code = data.get("code", "") access_token_field = data.get("access_token_field", "access_token") @@ -290,6 +314,39 @@ async def oauth_exchange(request: web.Request) -> web.Response: }) +async def oauth_refresh(request: web.Request) -> web.Response: + """Mock OAuth token refresh proxy for hosted refresh E2E tests.""" + data = await request.post() + oauth_state = request.app["oauth_state"] + oauth_state["refresh_count"] += 1 + oauth_state["last_refresh"] = { + "authorization": request.headers.get("Authorization"), + "form": dict(data), + } + + if request.headers.get("Authorization") != "Bearer e2e-test-token": + return web.json_response({"error": "invalid_gateway_auth"}, status=401) + if data.get("client_id") != "hosted-google-client-id": + return web.json_response({"error": "invalid_client_id"}, status=400) + if "client_secret" in data: + return web.json_response({"error": "unexpected_client_secret"}, status=400) + + return web.json_response({ + "access_token": "mock-refreshed-access-token", + "refresh_token": "mock-rotated-refresh-token", + "expires_in": 3600, + }) + + +async def oauth_state_handler(request: web.Request) -> web.Response: + return web.json_response(request.app["oauth_state"]) + + +async def oauth_reset(request: web.Request) -> web.Response: + request.app["oauth_state"] = _new_oauth_state() + return web.json_response({"ok": True}) + + async def models(_request: web.Request) -> web.Response: return web.json_response({ "object": "list", @@ -424,12 +481,16 @@ def main(): parser.add_argument("--port", type=int, default=0) args = parser.parse_args() app = web.Application() + app["oauth_state"] = _new_oauth_state() # Register both /v1/ and non-/v1/ paths (rig-core omits the /v1/ prefix) app.router.add_post("/v1/chat/completions", chat_completions) app.router.add_post("/chat/completions", chat_completions) app.router.add_get("/v1/models", models) app.router.add_get("/models", models) app.router.add_post("/oauth/exchange", oauth_exchange) + app.router.add_post("/oauth/refresh", oauth_refresh) + app.router.add_get("/__mock/oauth/state", oauth_state_handler) + app.router.add_post("/__mock/oauth/reset", oauth_reset) # Mock MCP server endpoints app.router.add_post("/mcp", mcp_endpoint) app.router.add_post("/mcp-400", mcp_endpoint_400) diff --git a/tests/e2e/scenarios/test_oauth_refresh.py b/tests/e2e/scenarios/test_oauth_refresh.py new file mode 100644 index 00000000..50871f7f --- /dev/null +++ b/tests/e2e/scenarios/test_oauth_refresh.py @@ -0,0 +1,227 @@ +"""Hosted OAuth refresh HTTP regression test. + +Runs a real ironclaw binary in hosted mode, expires a stored Gmail access +token in the libSQL database, triggers a real gmail tool call through the +chat API, and verifies that refresh uses the hosted proxy endpoint. +""" + +import asyncio +import sqlite3 +from datetime import datetime, timezone +from urllib.parse import parse_qs, urlparse + +import httpx + +from helpers import api_get, api_post + + +def _extract_state(auth_url: str) -> str: + parsed = urlparse(auth_url) + state = parse_qs(parsed.query).get("state", [None])[0] + assert state, f"auth_url should include state: {auth_url}" + return state + + +def _parse_timestamp(value: str | None) -> datetime | None: + if value is None: + return None + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def _expire_access_token(db_path: str, user_id: str, secret_name: str) -> None: + with sqlite3.connect(db_path) as conn: + cursor = conn.execute( + """ + UPDATE secrets + SET expires_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-1 hour') + WHERE user_id = ?1 AND name = ?2 + """, + (user_id, secret_name), + ) + conn.commit() + assert cursor.rowcount == 1, f"Expected one secret row for {user_id}/{secret_name}" + + +def _find_secret_row( + db_path: str, + secret_name: str, +) -> tuple[str, str | None, str | None]: + with sqlite3.connect(db_path) as conn: + row = conn.execute( + """ + SELECT user_id, expires_at, updated_at + FROM secrets + WHERE name = ?1 + ORDER BY updated_at DESC + LIMIT 1 + """, + (secret_name,), + ).fetchone() + assert row is not None, f"Missing secret row for {secret_name}" + return row[0], row[1], row[2] + + +async def _get_extension(base_url: str, name: str) -> dict | None: + response = await api_get(base_url, "/api/extensions", timeout=15) + response.raise_for_status() + for extension in response.json().get("extensions", []): + if extension["name"] == name: + return extension + return None + + +async def _reset_mock_oauth_state(mock_base_url: str) -> None: + async with httpx.AsyncClient() as client: + response = await client.post(f"{mock_base_url}/__mock/oauth/reset", timeout=10) + response.raise_for_status() + + +async def _get_mock_oauth_state(mock_base_url: str) -> dict: + async with httpx.AsyncClient() as client: + response = await client.get(f"{mock_base_url}/__mock/oauth/state", timeout=10) + response.raise_for_status() + return response.json() + + +async def _approve_pending_request(base_url: str, thread_id: str, request_id: str) -> None: + response = await api_post( + base_url, + "/api/chat/approval", + json={"request_id": request_id, "action": "approve", "thread_id": thread_id}, + timeout=15, + ) + assert response.status_code == 202, ( + f"Approval submission failed: {response.status_code} {response.text[:400]}" + ) + + +async def _wait_for_gmail_tool_call(base_url: str, thread_id: str, timeout: float = 30.0) -> dict: + approved_request_ids = set() + for _ in range(int(timeout * 2)): + response = await api_get( + base_url, + f"/api/chat/history?thread_id={thread_id}", + timeout=15, + ) + response.raise_for_status() + history = response.json() + + pending = history.get("pending_approval") + if pending and pending["request_id"] not in approved_request_ids: + await _approve_pending_request(base_url, thread_id, pending["request_id"]) + approved_request_ids.add(pending["request_id"]) + + for turn in history.get("turns", []): + for tool_call in turn.get("tool_calls", []): + if tool_call.get("name") == "gmail": + return history + + await asyncio.sleep(0.5) + + raise AssertionError(f"Timed out waiting for gmail tool call in thread {thread_id}") + + +async def _wait_for_refresh_request(mock_base_url: str, timeout: float = 20.0) -> dict: + for _ in range(int(timeout * 2)): + state = await _get_mock_oauth_state(mock_base_url) + if state.get("refresh_count") == 1: + return state + await asyncio.sleep(0.5) + raise AssertionError("Timed out waiting for exactly one OAuth refresh request") + + +async def test_hosted_gmail_oauth_refresh_uses_proxy(hosted_oauth_refresh_server): + server = hosted_oauth_refresh_server["base_url"] + db_path = hosted_oauth_refresh_server["db_path"] + mock_base_url = hosted_oauth_refresh_server["mock_llm_url"] + + install_response = await api_post( + server, + "/api/extensions/install", + json={"name": "gmail"}, + timeout=180, + ) + assert install_response.status_code == 200, install_response.text + assert install_response.json().get("success") is True + + setup_response = await api_post( + server, + "/api/extensions/gmail/setup", + json={"secrets": {}}, + timeout=30, + ) + assert setup_response.status_code == 200, setup_response.text + setup_data = setup_response.json() + assert setup_data.get("success") is True, setup_data + auth_url = setup_data.get("auth_url") + assert auth_url, setup_data + auth_params = parse_qs(urlparse(auth_url).query) + assert auth_params.get("client_id") == ["hosted-google-client-id"] + + async with httpx.AsyncClient() as client: + callback_response = await client.get( + f"{server}/oauth/callback", + params={"code": "mock_auth_code", "state": _extract_state(auth_url)}, + timeout=30, + follow_redirects=True, + ) + + assert callback_response.status_code == 200, callback_response.text[:400] + callback_body = callback_response.text.lower() + assert "connected" in callback_body or "success" in callback_body + + gmail = await _get_extension(server, "gmail") + assert gmail is not None, "gmail should be installed" + assert gmail["authenticated"] is True, gmail + assert "gmail" in gmail.get("tools", []), gmail + + await _reset_mock_oauth_state(mock_base_url) + + stored_user_id, expires_before, updated_before = _find_secret_row( + db_path, "google_oauth_token" + ) + assert _parse_timestamp(expires_before) is not None + assert _parse_timestamp(updated_before) is not None + + await asyncio.sleep(0.1) + _expire_access_token(db_path, stored_user_id, "google_oauth_token") + + thread_response = await api_post(server, "/api/chat/thread/new", timeout=15) + assert thread_response.status_code == 200, thread_response.text + thread_id = thread_response.json()["id"] + + send_response = await api_post( + server, + "/api/chat/send", + json={"content": "check gmail unread", "thread_id": thread_id}, + timeout=30, + ) + assert send_response.status_code == 202, send_response.text + + history = await _wait_for_gmail_tool_call(server, thread_id) + assert any( + tool_call.get("name") == "gmail" + for turn in history.get("turns", []) + for tool_call in turn.get("tool_calls", []) + ), history + + oauth_state = await _wait_for_refresh_request(mock_base_url) + assert oauth_state["refresh_count"] == 1, oauth_state + last_refresh = oauth_state["last_refresh"] + assert last_refresh is not None, oauth_state + assert last_refresh["authorization"] == "Bearer e2e-test-token" + assert last_refresh["form"]["client_id"] == "hosted-google-client-id" + assert "client_secret" not in last_refresh["form"], last_refresh + + refreshed_user_id, expires_after, updated_after = _find_secret_row( + db_path, "google_oauth_token" + ) + assert refreshed_user_id == stored_user_id + expires_after_dt = _parse_timestamp(expires_after) + updated_after_dt = _parse_timestamp(updated_after) + updated_before_dt = _parse_timestamp(updated_before) + assert expires_after_dt is not None + assert updated_after_dt is not None + assert updated_before_dt is not None + assert expires_after_dt > datetime.now(timezone.utc) + assert updated_after_dt > updated_before_dt