mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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:
+400
-24
@@ -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(¶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<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
@@ -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()),
|
||||
|
||||
@@ -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
@@ -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;
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user