Fix hosted OAuth refresh via proxy (#1602)

* Fix hosted OAuth refresh via proxy

* Address OAuth refresh review feedback

* Address new OAuth refresh review comments

* Address additional OAuth refresh review feedback

* Harden proxy exchange redirects
This commit is contained in:
Henry Park
2026-03-24 13:51:30 -07:00
committed by GitHub
parent f3da30a454
commit dcb2d89e3a
8 changed files with 1407 additions and 99 deletions
+400 -24
View File
@@ -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<String>,
builtin: Option<&OAuthCredentials>,
exchange_proxy_configured: bool,
) -> Option<String> {
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<String, String>,
}
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<OAuthTokenResponse, OAuthCallbackError> {
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<OAuthTokenResponse, OAuthCallbackError> {
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(&params)
.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<String>,
form: HashMap<String, String>,
}
#[derive(Clone)]
struct MockProxyState {
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
exchange_redirect_target: String,
refresh_redirect_target: String,
}
struct MockProxyServer {
addr: SocketAddr,
requests: Arc<Mutex<Vec<RecordedProxyRequest>>>,
shutdown_tx: Option<oneshot::Sender<()>>,
server_task: Option<tokio::task::JoinHandle<()>>,
}
impl MockProxyServer {
async fn start() -> Self {
async fn exchange_handler(
State(state): State<MockProxyState>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
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<MockProxyState>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
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<MockProxyState>) -> Redirect {
Redirect::temporary(&state.exchange_redirect_target)
}
async fn refresh_redirect_handler(State(state): State<MockProxyState>) -> 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<RecordedProxyRequest> {
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() {
+12 -23
View File
@@ -53,22 +53,6 @@ struct HostedOAuthFlowStart {
flow: crate::cli::oauth_defaults::PendingOAuthFlow,
}
fn hosted_proxy_client_secret(
client_secret: &Option<String>,
builtin: Option<&crate::cli::oauth_defaults::OAuthCredentials>,
exchange_proxy_configured: bool,
) -> Option<String> {
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()),
+141
View File
@@ -418,6 +418,7 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
let oauth = auth.oauth.as_ref()?;
let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name);
let exchange_proxy_url = crate::cli::oauth_defaults::exchange_proxy_url();
let client_id = oauth
.client_id
@@ -440,11 +441,21 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
.and_then(|env| std::env::var(env).ok())
})
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
let client_secret = crate::cli::oauth_defaults::hosted_proxy_client_secret(
&client_secret,
builtin.as_ref(),
exchange_proxy_url.is_some(),
);
let gateway_token = crate::config::helpers::env_or_override("GATEWAY_AUTH_TOKEN")
.map(|token| token.trim().to_string())
.filter(|token| !token.is_empty());
Some(OAuthRefreshConfig {
token_url: oauth.token_url.clone(),
client_id,
client_secret,
exchange_proxy_url,
gateway_token,
secret_name: auth.secret_name.clone(),
provider: auth.provider.clone(),
})
@@ -711,9 +722,44 @@ mod tests {
use tempfile::TempDir;
use crate::config::helpers::lock_env;
use crate::testing::credentials::{TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET};
use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools};
/// Restores a test-scoped env var override on drop.
struct EnvVarGuard {
key: String,
previous: Option<String>,
}
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"));
}
// ---------------------------------------------------------------
+443 -38
View File
@@ -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<String>,
/// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080").
pub exchange_proxy_url: Option<String>,
/// Gateway auth token for authenticating with the hosted OAuth proxy.
pub gateway_token: Option<String>,
/// 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<DecryptedSecret> {
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<String>,
form: HashMap<String, String>,
}
struct MockProxyServer {
addr: SocketAddr,
requests: Arc<AsyncMutex<Vec<RecordedProxyRequest>>>,
shutdown_tx: Option<oneshot::Sender<()>>,
server_task: Option<tokio::task::JoinHandle<()>>,
}
impl MockProxyServer {
async fn start() -> Self {
async fn refresh_handler(
State(requests): State<Arc<AsyncMutex<Vec<RecordedProxyRequest>>>>,
headers: HeaderMap,
Form(form): Form<HashMap<String, String>>,
) -> Json<serde_json::Value> {
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<RecordedProxyRequest> {
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;
+9
View File
@@ -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
+114 -14
View File
@@ -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,
+61
View File
@@ -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)
+227
View File
@@ -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