mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Merge pull request #1687 from nearai/staging-promote/1d577782-23620402154
chore: promote staging to main (2026-03-26 22:07 UTC)
This commit is contained in:
+434
-2
@@ -836,10 +836,10 @@ async fn oauth_callback_handler(
|
||||
|
||||
let result: Result<(), String> = async {
|
||||
let token_response = if let Some(proxy_url) = &exchange_proxy_url {
|
||||
let gateway_token = flow.gateway_token.as_deref().unwrap_or_default();
|
||||
let oauth_proxy_auth_token = flow.oauth_proxy_auth_token().unwrap_or_default();
|
||||
oauth_defaults::exchange_via_proxy(oauth_defaults::ProxyTokenExchangeRequest {
|
||||
proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_url: &flow.token_url,
|
||||
client_id: &flow.client_id,
|
||||
client_secret: flow.client_secret.as_deref(),
|
||||
@@ -3057,6 +3057,160 @@ mod tests {
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct RecordedOauthProxyRequest {
|
||||
authorization: Option<String>,
|
||||
form: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockOauthProxyState {
|
||||
requests: Arc<tokio::sync::Mutex<Vec<RecordedOauthProxyRequest>>>,
|
||||
}
|
||||
|
||||
struct MockOauthProxyServer {
|
||||
addr: std::net::SocketAddr,
|
||||
requests: Arc<tokio::sync::Mutex<Vec<RecordedOauthProxyRequest>>>,
|
||||
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
server_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockOauthProxyServer {
|
||||
async fn start() -> Self {
|
||||
async fn exchange_handler(
|
||||
State(state): State<MockOauthProxyState>,
|
||||
headers: axum::http::HeaderMap,
|
||||
axum::Form(form): axum::Form<std::collections::HashMap<String, String>>,
|
||||
) -> Json<serde_json::Value> {
|
||||
state.requests.lock().await.push(RecordedOauthProxyRequest {
|
||||
authorization: headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string),
|
||||
form,
|
||||
});
|
||||
Json(serde_json::json!({
|
||||
"access_token": "proxy-access-token",
|
||||
"refresh_token": "proxy-refresh-token",
|
||||
"expires_in": 7200
|
||||
}))
|
||||
}
|
||||
|
||||
let requests = Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind mock oauth proxy");
|
||||
let addr = listener.local_addr().expect("mock oauth proxy addr");
|
||||
let app = Router::new()
|
||||
.route("/oauth/exchange", post(exchange_handler))
|
||||
.with_state(MockOauthProxyState {
|
||||
requests: Arc::clone(&requests),
|
||||
});
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::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<RecordedOauthProxyRequest> {
|
||||
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 MockOauthProxyServer {
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: 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.original {
|
||||
std::env::set_var(self.key, value);
|
||||
} else {
|
||||
std::env::remove_var(self.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard {
|
||||
let original = std::env::var(key).ok();
|
||||
// SAFETY: Tests use lock_env() to serialize environment access.
|
||||
unsafe {
|
||||
if let Some(value) = value {
|
||||
std::env::set_var(key, value);
|
||||
} else {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
EnvVarGuard { key, original }
|
||||
}
|
||||
|
||||
fn fresh_pending_oauth_flow(
|
||||
secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync>,
|
||||
sse_manager: Option<Arc<SseManager>>,
|
||||
oauth_proxy_auth_token: Option<String>,
|
||||
) -> crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
extension_name: "test_tool".to_string(),
|
||||
display_name: "Test Tool".to_string(),
|
||||
token_url: "https://example.com/token".to_string(),
|
||||
client_id: "client123".to_string(),
|
||||
client_secret: None,
|
||||
redirect_uri: "https://example.com/oauth/callback".to_string(),
|
||||
code_verifier: Some("test-code-verifier".to_string()),
|
||||
access_token_field: "access_token".to_string(),
|
||||
secret_name: "test_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
validation_endpoint: None,
|
||||
scopes: vec!["email".to_string()],
|
||||
user_id: "test".to_string(),
|
||||
secrets,
|
||||
sse_manager,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extensions_setup_submit_returns_failure_when_not_activated() {
|
||||
use axum::body::Body;
|
||||
@@ -3714,6 +3868,284 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_accepts_versioned_hosted_state_without_instance_name() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
|
||||
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
|
||||
TEST_GATEWAY_CRYPTO_KEY.to_string(),
|
||||
))
|
||||
.expect("crypto"),
|
||||
)));
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
|
||||
|
||||
let Some(created_at) = expired_flow_created_at() else {
|
||||
eprintln!(
|
||||
"Skipping versioned OAuth state without instance test: monotonic uptime below expiry window"
|
||||
);
|
||||
return;
|
||||
};
|
||||
let flow = crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
extension_name: "test_tool".to_string(),
|
||||
display_name: "Test Tool".to_string(),
|
||||
token_url: "https://example.com/token".to_string(),
|
||||
client_id: "client123".to_string(),
|
||||
client_secret: None,
|
||||
redirect_uri: "https://example.com/oauth/callback".to_string(),
|
||||
code_verifier: None,
|
||||
access_token_field: "access_token".to_string(),
|
||||
secret_name: "test_token".to_string(),
|
||||
provider: None,
|
||||
validation_endpoint: None,
|
||||
scopes: vec![],
|
||||
user_id: "test".to_string(),
|
||||
secrets,
|
||||
sse_manager: None,
|
||||
gateway_token: None,
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at,
|
||||
};
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None);
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Authorization Failed"));
|
||||
assert!(
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.read()
|
||||
.await
|
||||
.get("test_nonce")
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_happy_path_with_gateway_token_fallback() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let proxy = MockOauthProxyServer::start().await;
|
||||
// Keep the process-wide env locked for the full callback so the handler
|
||||
// sees a stable proxy URL/token configuration throughout the test.
|
||||
let _env_guard = crate::config::helpers::lock_env();
|
||||
let _exchange_url_guard =
|
||||
set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url()));
|
||||
let _proxy_auth_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
|
||||
let secrets = test_secrets_store();
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets));
|
||||
let sse_mgr = Arc::new(SseManager::new());
|
||||
let mut receiver = sse_mgr.sender().subscribe();
|
||||
let flow = fresh_pending_oauth_flow(
|
||||
Arc::clone(&secrets),
|
||||
Some(Arc::clone(&sse_mgr)),
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
);
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", Some("myinstance"));
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Test Tool Connected"));
|
||||
|
||||
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("code").map(String::as_str),
|
||||
Some("fake_code")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("test-code-verifier")
|
||||
);
|
||||
|
||||
let access_token = secrets
|
||||
.get_decrypted("test", "test_token")
|
||||
.await
|
||||
.expect("access token stored");
|
||||
assert_eq!(access_token.expose(), "proxy-access-token");
|
||||
|
||||
let refresh_token = secrets
|
||||
.get_decrypted("test", "test_token_refresh_token")
|
||||
.await
|
||||
.expect("refresh token stored");
|
||||
assert_eq!(refresh_token.expose(), "proxy-refresh-token");
|
||||
|
||||
match receiver.recv().await.expect("auth_completed event").event {
|
||||
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(extension_name, "test_tool");
|
||||
assert!(success, "OAuth callback should broadcast success");
|
||||
}
|
||||
event => panic!("expected AuthCompleted event, got {event:?}"),
|
||||
}
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_happy_path_with_dedicated_proxy_auth_token() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let proxy = MockOauthProxyServer::start().await;
|
||||
// Keep the process-wide env locked for the full callback so the handler
|
||||
// sees a stable proxy URL/token configuration throughout the test.
|
||||
let _env_guard = crate::config::helpers::lock_env();
|
||||
let _exchange_url_guard =
|
||||
set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url()));
|
||||
let _proxy_auth_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-oauth-proxy-secret"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
|
||||
let secrets = test_secrets_store();
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets));
|
||||
let sse_mgr = Arc::new(SseManager::new());
|
||||
let mut receiver = sse_mgr.sender().subscribe();
|
||||
let flow = fresh_pending_oauth_flow(
|
||||
Arc::clone(&secrets),
|
||||
Some(Arc::clone(&sse_mgr)),
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
);
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None);
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Test Tool Connected"));
|
||||
|
||||
let requests = proxy.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer shared-oauth-proxy-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code").map(String::as_str),
|
||||
Some("fake_code")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("test-code-verifier")
|
||||
);
|
||||
|
||||
let access_token = secrets
|
||||
.get_decrypted("test", "test_token")
|
||||
.await
|
||||
.expect("access token stored");
|
||||
assert_eq!(access_token.expose(), "proxy-access-token");
|
||||
|
||||
let refresh_token = secrets
|
||||
.get_decrypted("test", "test_token_refresh_token")
|
||||
.await
|
||||
.expect("refresh token stored");
|
||||
assert_eq!(refresh_token.expose(), "proxy-refresh-token");
|
||||
|
||||
match receiver.recv().await.expect("auth_completed event").event {
|
||||
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(extension_name, "test_tool");
|
||||
assert!(success, "OAuth callback should broadcast success");
|
||||
}
|
||||
event => panic!("expected AuthCompleted event, got {event:?}"),
|
||||
}
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
// --- Slack relay OAuth CSRF tests ---
|
||||
|
||||
fn test_relay_oauth_router(state: Arc<GatewayState>) -> Router {
|
||||
|
||||
+184
-5
@@ -473,7 +473,8 @@ pub struct PendingOAuthFlow {
|
||||
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
/// SSE broadcast manager for notifying the web UI.
|
||||
pub sse_manager: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: Option<String>,
|
||||
/// Additional form params for the token exchange request.
|
||||
/// Used for provider-specific requirements such as RFC 8707 `resource`.
|
||||
@@ -496,6 +497,12 @@ impl std::fmt::Debug for PendingOAuthFlow {
|
||||
}
|
||||
}
|
||||
|
||||
impl PendingOAuthFlow {
|
||||
pub fn oauth_proxy_auth_token(&self) -> Option<&str> {
|
||||
self.gateway_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter.
|
||||
pub type PendingOAuthRegistry = Arc<RwLock<HashMap<String, PendingOAuthFlow>>>;
|
||||
|
||||
@@ -529,6 +536,22 @@ pub fn exchange_proxy_url() -> Option<String> {
|
||||
.filter(|url| !url.is_empty())
|
||||
}
|
||||
|
||||
/// Returns the configured OAuth proxy auth token, if any.
|
||||
///
|
||||
/// New hosted infra can inject a dedicated shared proxy secret via
|
||||
/// `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`. Existing hosted instances continue to
|
||||
/// work by falling back to `GATEWAY_AUTH_TOKEN`.
|
||||
pub fn oauth_proxy_auth_token() -> Option<String> {
|
||||
fn normalized_env_value(key: &str) -> Option<String> {
|
||||
crate::config::helpers::env_or_override(key)
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
normalized_env_value("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN")
|
||||
.or_else(|| normalized_env_value("GATEWAY_AUTH_TOKEN"))
|
||||
}
|
||||
|
||||
/// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout).
|
||||
pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300);
|
||||
|
||||
@@ -674,6 +697,8 @@ pub fn strip_instance_prefix(state: &str) -> &str {
|
||||
|
||||
pub struct ProxyTokenExchangeRequest<'a> {
|
||||
pub proxy_url: &'a str,
|
||||
/// OAuth proxy auth token.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: &'a str,
|
||||
pub token_url: &'a str,
|
||||
pub client_id: &'a str,
|
||||
@@ -687,6 +712,8 @@ pub struct ProxyTokenExchangeRequest<'a> {
|
||||
|
||||
pub struct ProxyRefreshTokenRequest<'a> {
|
||||
pub proxy_url: &'a str,
|
||||
/// OAuth proxy auth token.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: &'a str,
|
||||
pub token_url: &'a str,
|
||||
pub client_id: &'a str,
|
||||
@@ -729,7 +756,7 @@ fn oauth_token_response_from_json(
|
||||
|
||||
/// Exchange an OAuth authorization code via the platform's token exchange proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
/// Authenticated via an OAuth proxy auth token (Bearer header). The caller may
|
||||
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||
/// the provider requires it.
|
||||
///
|
||||
@@ -741,7 +768,7 @@ pub async fn exchange_via_proxy(
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
if request.gateway_token.is_empty() {
|
||||
return Err(OAuthCallbackError::Io(
|
||||
"Gateway auth token is required for proxy token exchange".to_string(),
|
||||
"OAuth proxy auth token is required for proxy token exchange".to_string(),
|
||||
));
|
||||
}
|
||||
let exchange_url = format!("{}/oauth/exchange", request.proxy_url.trim_end_matches('/'));
|
||||
@@ -796,7 +823,7 @@ pub async fn exchange_via_proxy(
|
||||
|
||||
/// Refresh an OAuth access token via the platform's token refresh proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
/// Authenticated via an OAuth proxy 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(
|
||||
@@ -804,7 +831,7 @@ pub async fn refresh_token_via_proxy(
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
if request.gateway_token.is_empty() {
|
||||
return Err(OAuthCallbackError::Io(
|
||||
"Gateway auth token is required for proxy token refresh".to_string(),
|
||||
"OAuth proxy auth token is required for proxy token refresh".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
@@ -1010,6 +1037,37 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
if let Some(ref value) = self.original {
|
||||
std::env::set_var(self.key, value);
|
||||
} else {
|
||||
std::env::remove_var(self.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard {
|
||||
let original = std::env::var(key).ok();
|
||||
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
if let Some(value) = value {
|
||||
std::env::set_var(key, value);
|
||||
} else {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
EnvVarGuard { key, original }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hosted_proxy_client_secret_suppresses_builtin_secret() {
|
||||
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
|
||||
@@ -1030,6 +1088,79 @@ mod tests {
|
||||
assert_eq!(result, client_secret);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_exchange_via_proxy_sends_auth_and_form() {
|
||||
let server = MockProxyServer::start().await;
|
||||
let mut extra_token_params = HashMap::new();
|
||||
extra_token_params.insert("resource".to_string(), "https://mcp.notion.com".to_string());
|
||||
|
||||
let response = super::exchange_via_proxy(super::ProxyTokenExchangeRequest {
|
||||
proxy_url: &server.base_url(),
|
||||
gateway_token: "shared-oauth-proxy-secret",
|
||||
code: "auth-code-123",
|
||||
redirect_uri: "https://oauth.example.com/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: &extra_token_params,
|
||||
})
|
||||
.await
|
||||
.expect("proxy exchange 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 shared-oauth-proxy-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code").map(String::as_str),
|
||||
Some("auth-code-123")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("redirect_uri").map(String::as_str),
|
||||
Some("https://oauth.example.com/oauth/callback")
|
||||
);
|
||||
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("access_token_field")
|
||||
.map(String::as_str),
|
||||
Some("access_token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("code-verifier-123")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("resource").map(String::as_str),
|
||||
Some("https://mcp.notion.com")
|
||||
);
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_token_via_proxy_sends_auth_and_form() {
|
||||
let server = MockProxyServer::start().await;
|
||||
@@ -1535,6 +1666,54 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_prefers_dedicated_env() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-proxy-secret"),
|
||||
);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("shared-proxy-secret")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_falls_back_to_gateway_token() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("gateway-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_whitespace_dedicated_env_falls_back_to_gateway_token() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", Some(" "));
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("gateway-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_returns_none_when_unset() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
|
||||
assert_eq!(crate::cli::oauth_defaults::oauth_proxy_auth_token(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_instance_prefix_with_colon() {
|
||||
use crate::cli::oauth_defaults::strip_instance_prefix;
|
||||
|
||||
@@ -192,6 +192,9 @@ pub struct JobContext {
|
||||
/// but subsequent tools (e.g., `json`) may need the full output. This
|
||||
/// stash stores the complete, unsanitized output so tools can reference
|
||||
/// previous results by ID via `$tool_call_id` parameter syntax.
|
||||
///
|
||||
/// Also used for cross-tool implicit state (keys prefixed with `__`) such
|
||||
/// as `__routine_last_name` for fallback recovery in routine tool chains.
|
||||
#[serde(skip)]
|
||||
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
||||
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
||||
|
||||
@@ -403,9 +403,10 @@ pub struct ExtensionManager {
|
||||
/// when running in gateway mode, consumed by the web gateway's
|
||||
/// `/oauth/callback` handler.
|
||||
pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry,
|
||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||
/// Read once at construction from `GATEWAY_AUTH_TOKEN` env var.
|
||||
gateway_token: Option<String>,
|
||||
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||
/// Resolved once at construction from `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`,
|
||||
/// then `GATEWAY_AUTH_TOKEN` as a backward-compatible fallback.
|
||||
oauth_proxy_auth_token: Option<String>,
|
||||
/// Relay config captured at startup. Used by `auth_channel_relay` and
|
||||
/// `activate_channel_relay` instead of re-reading env vars.
|
||||
relay_config: Option<crate::config::RelayConfig>,
|
||||
@@ -535,7 +536,7 @@ impl ExtensionManager {
|
||||
activation_errors: RwLock::new(HashMap::new()),
|
||||
sse_manager: RwLock::new(None),
|
||||
pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(),
|
||||
gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(),
|
||||
oauth_proxy_auth_token: crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
relay_config: crate::config::RelayConfig::from_env(),
|
||||
relay_event_tx: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
relay_signing_secret_cache: Arc::new(std::sync::Mutex::new(None)),
|
||||
@@ -2788,7 +2789,7 @@ impl ExtensionManager {
|
||||
user_id: user_id.to_string(),
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_manager: self.sse_manager.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||
token_exchange_extra_params,
|
||||
client_id_secret_name: if server.oauth.is_none() {
|
||||
Some(server.client_id_secret_name())
|
||||
@@ -3305,7 +3306,7 @@ impl ExtensionManager {
|
||||
user_id: user_id.to_string(),
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_manager: self.sse_manager.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now(),
|
||||
|
||||
@@ -650,6 +650,23 @@ pub(crate) fn routine_update_parameters_schema() -> Value {
|
||||
})
|
||||
}
|
||||
|
||||
const ROUTINE_LAST_NAME_STASH_KEY: &str = "__routine_last_name";
|
||||
|
||||
async fn stash_last_routine_name(ctx: &JobContext, name: &str) {
|
||||
ctx.tool_output_stash
|
||||
.write()
|
||||
.await
|
||||
.insert(ROUTINE_LAST_NAME_STASH_KEY.to_string(), name.to_string());
|
||||
}
|
||||
|
||||
async fn restore_last_routine_name(ctx: &JobContext) -> Option<String> {
|
||||
ctx.tool_output_stash
|
||||
.read()
|
||||
.await
|
||||
.get(ROUTINE_LAST_NAME_STASH_KEY)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn nested_object<'a>(params: &'a Value, field: &str) -> Option<&'a Map<String, Value>> {
|
||||
params.get(field).and_then(Value::as_object)
|
||||
}
|
||||
@@ -1093,6 +1110,7 @@ impl Tool for RoutineCreateTool {
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
let normalized = parse_routine_create_request(¶ms)?;
|
||||
stash_last_routine_name(ctx, &normalized.name).await;
|
||||
let trigger = build_routine_trigger(&normalized.trigger);
|
||||
let action =
|
||||
build_routine_action(&normalized.name, &normalized.prompt, &normalized.execution);
|
||||
@@ -1274,6 +1292,7 @@ impl Tool for RoutineUpdateTool {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let name = require_str(¶ms, "name")?;
|
||||
stash_last_routine_name(ctx, name).await;
|
||||
|
||||
let mut routine = self
|
||||
.store
|
||||
@@ -1411,11 +1430,24 @@ impl Tool for RoutineDeleteTool {
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let name = require_str(¶ms, "name")?;
|
||||
let name = if let Some(name) = params.get("name").and_then(|v| v.as_str()) {
|
||||
if name.trim().is_empty() {
|
||||
return Err(ToolError::InvalidParameters(
|
||||
"'name' parameter cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
name.to_string()
|
||||
} else {
|
||||
restore_last_routine_name(ctx).await.ok_or_else(|| {
|
||||
ToolError::InvalidParameters(
|
||||
"missing 'name' parameter and no previous routine target to infer".to_string(),
|
||||
)
|
||||
})?
|
||||
};
|
||||
|
||||
let routine = self
|
||||
.store
|
||||
.get_routine_by_name(&ctx.user_id, name)
|
||||
.get_routine_by_name(&ctx.user_id, &name)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
|
||||
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
|
||||
@@ -1430,7 +1462,7 @@ impl Tool for RoutineDeleteTool {
|
||||
self.engine.refresh_event_cache().await;
|
||||
|
||||
let result = serde_json::json!({
|
||||
"name": name,
|
||||
"name": &name,
|
||||
"deleted": deleted,
|
||||
});
|
||||
|
||||
|
||||
+19
-1
@@ -117,6 +117,11 @@ impl McpClient {
|
||||
/// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`.
|
||||
///
|
||||
/// Returns an error if the config uses a non-HTTP transport.
|
||||
///
|
||||
/// **Note:** The session manager is NOT wired into the transport. For
|
||||
/// production use, prefer `create_client_from_config()` which constructs
|
||||
/// the transport with session tracking.
|
||||
#[cfg(test)]
|
||||
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
||||
if !matches!(
|
||||
config.effective_transport(),
|
||||
@@ -214,7 +219,14 @@ impl McpClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a session manager for Streamable HTTP session tracking.
|
||||
/// Attach a session manager to the **client** only.
|
||||
///
|
||||
/// **Warning:** This does NOT wire the session manager into the underlying
|
||||
/// `HttpMcpTransport`, so the transport will not capture `Mcp-Session-Id`
|
||||
/// from responses. For production use, construct the transport with
|
||||
/// `HttpMcpTransport::with_session_manager()` and pass it to
|
||||
/// `new_with_transport()` instead. See `create_client_from_config()`.
|
||||
#[cfg(test)]
|
||||
pub fn with_session_manager(mut self, session_manager: Arc<McpSessionManager>) -> Self {
|
||||
self.session_manager = Some(session_manager);
|
||||
self
|
||||
@@ -235,6 +247,12 @@ impl McpClient {
|
||||
self.session_manager.is_some()
|
||||
}
|
||||
|
||||
/// Get the underlying transport (test-only).
|
||||
#[cfg(test)]
|
||||
pub(crate) fn transport(&self) -> &Arc<dyn McpTransport> {
|
||||
&self.transport
|
||||
}
|
||||
|
||||
/// Get the next request ID.
|
||||
fn next_request_id(&self) -> u64 {
|
||||
self.next_id.fetch_add(1, Ordering::SeqCst)
|
||||
|
||||
+101
-16
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
||||
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::mcp::config::{EffectiveTransport, McpServerConfig};
|
||||
use crate::tools::mcp::http_transport::HttpMcpTransport;
|
||||
use crate::tools::mcp::{McpClient, McpProcessManager, McpSessionManager, McpTransport};
|
||||
|
||||
/// Error returned when MCP client creation fails.
|
||||
@@ -78,33 +79,37 @@ pub async fn create_client_from_config(
|
||||
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
||||
}
|
||||
EffectiveTransport::Http => {
|
||||
// Authenticated (OAuth) path: tokens exist or server requires auth.
|
||||
if let Some(ref secrets) = secrets {
|
||||
let has_tokens =
|
||||
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
||||
|
||||
if has_tokens || server.requires_auth() {
|
||||
Ok(McpClient::new_authenticated(
|
||||
return Ok(McpClient::new_authenticated(
|
||||
server,
|
||||
Arc::clone(session_manager),
|
||||
Arc::clone(secrets),
|
||||
user_id,
|
||||
))
|
||||
} else {
|
||||
Ok(McpClient::new_with_config(server)
|
||||
.map_err(|e| McpFactoryError::InvalidConfig {
|
||||
name: server_name.clone(),
|
||||
reason: e.to_string(),
|
||||
})?
|
||||
.with_session_manager(Arc::clone(session_manager)))
|
||||
));
|
||||
}
|
||||
} else {
|
||||
Ok(McpClient::new_with_config(server)
|
||||
.map_err(|e| McpFactoryError::InvalidConfig {
|
||||
name: server_name,
|
||||
reason: e.to_string(),
|
||||
})?
|
||||
.with_session_manager(Arc::clone(session_manager)))
|
||||
}
|
||||
|
||||
// Non-OAuth HTTP: wire the session manager into the *transport* so
|
||||
// it captures `Mcp-Session-Id` from responses. Passing it only to
|
||||
// the client (via `with_session_manager`) is not enough — the
|
||||
// transport must know about it to read/write the header.
|
||||
let transport = Arc::new(
|
||||
HttpMcpTransport::new(server.url.clone(), server.name.clone())
|
||||
.with_session_manager(Arc::clone(session_manager)),
|
||||
);
|
||||
Ok(McpClient::new_with_transport(
|
||||
server.name.clone(),
|
||||
transport,
|
||||
Some(Arc::clone(session_manager)),
|
||||
secrets,
|
||||
user_id,
|
||||
Some(server),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -134,4 +139,84 @@ mod tests {
|
||||
"non-OAuth HTTP clients must carry a session manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression test: the factory must wire the session manager into the
|
||||
/// *transport*, not just the client. Otherwise the transport never
|
||||
/// captures `Mcp-Session-Id` from responses and subsequent requests
|
||||
/// lack the header, causing the server to reject them.
|
||||
#[tokio::test]
|
||||
async fn test_factory_non_oauth_http_transport_captures_session_id() {
|
||||
use axum::http::header::HeaderName;
|
||||
use axum::{Router, http::StatusCode, response::IntoResponse, routing::post};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const SESSION_ID: &str = "test-session-abc123";
|
||||
|
||||
async fn session_echo() -> impl IntoResponse {
|
||||
let body = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"result": {}
|
||||
})
|
||||
.to_string();
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(
|
||||
HeaderName::from_static("mcp-session-id"),
|
||||
SESSION_ID.to_string(),
|
||||
)],
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
let app = Router::new().route("/", post(session_echo));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let url = format!("http://127.0.0.1:{}", addr.port());
|
||||
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
let server = McpServerConfig::new("session-test", &url);
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
let process_manager = Arc::new(McpProcessManager::new());
|
||||
|
||||
let client = create_client_from_config(
|
||||
server,
|
||||
&session_manager,
|
||||
&process_manager,
|
||||
None,
|
||||
"test-user",
|
||||
)
|
||||
.await
|
||||
.expect("factory should succeed for HTTP config");
|
||||
|
||||
// Pre-create a session entry so that update_session_id has something to update.
|
||||
// In production, the MCP initialize handshake calls get_or_create before responses arrive.
|
||||
session_manager.get_or_create("session-test", &url).await;
|
||||
|
||||
// Send a request through the client's transport to trigger session capture.
|
||||
use crate::tools::mcp::protocol::McpRequest;
|
||||
let request = McpRequest {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: Some(1),
|
||||
method: "test".to_string(),
|
||||
params: Some(serde_json::json!({})),
|
||||
};
|
||||
let headers = std::collections::HashMap::new();
|
||||
client
|
||||
.transport()
|
||||
.send(&request, &headers)
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
// Verify the session manager captured the session ID from the response.
|
||||
let captured = session_manager.get_session_id("session-test").await;
|
||||
assert_eq!(
|
||||
captured.as_deref(),
|
||||
Some(SESSION_ID),
|
||||
"transport must capture Mcp-Session-Id into session manager"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -494,6 +494,34 @@ mod tests {
|
||||
assert_eq!(echoed["authorization"], "Bearer oauth-token");
|
||||
}
|
||||
|
||||
/// Regression test for #1436: 202 Accepted responses for notifications
|
||||
/// were parsed as JSON, causing "Failed to parse MCP response" errors
|
||||
/// that broke the MCP session handshake.
|
||||
#[tokio::test]
|
||||
async fn test_wire_202_accepted_for_notification() {
|
||||
use axum::{Router, http::StatusCode, routing::post};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
async fn accept_notification() -> StatusCode {
|
||||
StatusCode::ACCEPTED
|
||||
}
|
||||
|
||||
let app = Router::new().route("/", post(accept_notification));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let url = format!("http://127.0.0.1:{}", addr.port());
|
||||
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
let transport = HttpMcpTransport::new(&url, "test-202");
|
||||
let request = McpRequest::initialized_notification();
|
||||
let response = transport.send(&request, &HashMap::new()).await.unwrap();
|
||||
assert!(response.result.is_none());
|
||||
assert!(response.error.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wire_custom_auth_preserved_when_no_per_request_auth() {
|
||||
let (url, _handle) = spawn_echo_server().await;
|
||||
|
||||
@@ -446,16 +446,14 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
||||
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());
|
||||
let oauth_proxy_auth_token = crate::cli::oauth_defaults::oauth_proxy_auth_token();
|
||||
|
||||
Some(OAuthRefreshConfig {
|
||||
token_url: oauth.token_url.clone(),
|
||||
client_id,
|
||||
client_secret,
|
||||
exchange_proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
secret_name: auth.secret_name.clone(),
|
||||
provider: auth.provider.clone(),
|
||||
})
|
||||
@@ -891,6 +889,11 @@ 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);
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
@@ -982,6 +985,7 @@ mod tests {
|
||||
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);
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
|
||||
// google_oauth_token should fall back to built-in credentials
|
||||
let caps = CapabilitiesFile {
|
||||
@@ -1021,6 +1025,7 @@ mod tests {
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
|
||||
@@ -1061,6 +1066,7 @@ mod tests {
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
let _client_secret_guard =
|
||||
@@ -1095,6 +1101,47 @@ mod tests {
|
||||
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_oauth_refresh_config_hosted_proxy_prefers_dedicated_proxy_auth_token() {
|
||||
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 _oauth_proxy_token_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-oauth-proxy-secret"),
|
||||
);
|
||||
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.gateway_token.as_deref(),
|
||||
Some("shared-oauth-proxy-secret")
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
// Security regression tests
|
||||
// ---------------------------------------------------------------
|
||||
|
||||
@@ -62,7 +62,8 @@ pub struct OAuthRefreshConfig {
|
||||
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.
|
||||
/// OAuth proxy auth token for authenticating with the hosted OAuth proxy.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
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`.
|
||||
@@ -71,6 +72,12 @@ pub struct OAuthRefreshConfig {
|
||||
pub provider: Option<String>,
|
||||
}
|
||||
|
||||
impl OAuthRefreshConfig {
|
||||
fn oauth_proxy_auth_token(&self) -> Option<&str> {
|
||||
self.gateway_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-resolved credential for host-based injection.
|
||||
///
|
||||
/// Built before each WASM execution by decrypting secrets from the store.
|
||||
@@ -1218,9 +1225,9 @@ async fn refresh_oauth_token(
|
||||
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 {
|
||||
let Some(oauth_proxy_auth_token) = config.oauth_proxy_auth_token() else {
|
||||
tracing::warn!(
|
||||
"OAuth refresh proxy is configured, but no gateway auth token is available"
|
||||
"OAuth refresh proxy is configured, but no OAuth proxy auth token is available"
|
||||
);
|
||||
return false;
|
||||
};
|
||||
@@ -1235,7 +1242,7 @@ async fn refresh_oauth_token(
|
||||
let token_response = match oauth_defaults::refresh_token_via_proxy(
|
||||
oauth_defaults::ProxyRefreshTokenRequest {
|
||||
proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_url: &config.token_url,
|
||||
client_id: &config.client_id,
|
||||
client_secret: config.client_secret.as_deref(),
|
||||
@@ -2704,7 +2711,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() {
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_oauth_proxy_auth_token()
|
||||
{
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
|
||||
@@ -205,7 +205,44 @@ mod tests {
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Test 5: routine_manual_create_defaults_to_tools_enabled
|
||||
// Test 5: routine_update_fail_delete_fallback
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn routine_update_fail_delete_fallback() {
|
||||
let trace = LlmTrace::from_file(concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json"
|
||||
))
|
||||
.expect("failed to load routine_update_fail_delete_fallback.json");
|
||||
|
||||
let rig = TestRigBuilder::new()
|
||||
.with_trace(trace.clone())
|
||||
.with_auto_approve_tools(true)
|
||||
.build()
|
||||
.await;
|
||||
|
||||
rig.send_message("Try converting a routine trigger, then recover by deleting it")
|
||||
.await;
|
||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
||||
|
||||
rig.verify_trace_expects(&trace, &responses);
|
||||
|
||||
let completed = rig.tool_calls_completed();
|
||||
assert!(
|
||||
completed.iter().any(|(n, ok)| n == "routine_update" && !ok),
|
||||
"routine_update should fail in this regression path: {completed:?}"
|
||||
);
|
||||
assert!(
|
||||
completed.iter().any(|(n, ok)| n == "routine_delete" && *ok),
|
||||
"routine_delete should recover successfully via preserved routine identity: {completed:?}"
|
||||
);
|
||||
|
||||
rig.shutdown();
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Test 6: routine_manual_create_defaults_to_tools_enabled
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
@@ -246,7 +283,7 @@ mod tests {
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Test 6: routine_manual_create_explicit_no_tools
|
||||
// Test 7: routine_manual_create_explicit_no_tools
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
@@ -287,7 +324,7 @@ mod tests {
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Test 7: routine_history
|
||||
// Test 8: routine_history
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
{
|
||||
"model_name": "test-routine-update-fail-delete-fallback",
|
||||
"expects": {
|
||||
"tools_used": ["routine_create", "routine_update", "routine_delete"],
|
||||
"tool_results_contain": {
|
||||
"routine_update": "Cannot update schedule or timezone on a non-cron routine.",
|
||||
"routine_delete": "temp-routine"
|
||||
},
|
||||
"min_responses": 1
|
||||
},
|
||||
"steps": [
|
||||
{
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_rc_fallback",
|
||||
"name": "routine_create",
|
||||
"arguments": {
|
||||
"name": "temp-routine",
|
||||
"trigger_type": "manual",
|
||||
"prompt": "Temporary routine for fallback test."
|
||||
}
|
||||
}
|
||||
],
|
||||
"input_tokens": 120,
|
||||
"output_tokens": 40
|
||||
}
|
||||
},
|
||||
{
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_ru_fallback",
|
||||
"name": "routine_update",
|
||||
"arguments": {
|
||||
"name": "temp-routine",
|
||||
"schedule": "0 */10 * * * *"
|
||||
}
|
||||
}
|
||||
],
|
||||
"input_tokens": 200,
|
||||
"output_tokens": 30
|
||||
}
|
||||
},
|
||||
{
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_rd_fallback",
|
||||
"name": "routine_delete",
|
||||
"arguments": {}
|
||||
}
|
||||
],
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 20
|
||||
}
|
||||
},
|
||||
{
|
||||
"response": {
|
||||
"type": "text",
|
||||
"content": "I recovered from the failed update and cleaned up the original routine.",
|
||||
"input_tokens": 380,
|
||||
"output_tokens": 25
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user