mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
refactor: consolidate tool approval into single param-aware method (#274)
* refactor: consolidate tool approval into single param-aware method Replace the two confusing approval methods (requires_approval() and requires_approval_for()) with a single requires_approval(&self, params) returning a 3-variant ApprovalRequirement enum (Never, UnlessAutoApproved, Always). This enables param-aware approval decisions: HTTP calls without auth headers now skip approval entirely, while authenticated requests always require it. Shell tool merges its destructive-command detection into the same method. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: add credential injection to built-in HTTP tool Wire the WASM credential injection system into the built-in HTTP tool so credentials are auto-injected at the boundary (zero-exposure model). - Add SharedCredentialRegistry: thread-safe, append-only registry of credential mappings populated by WASM tools at registration time - Add credential_detect module with broad auth detection for headers (12 exact + 5 substring matches), header values (7 auth scheme prefixes), and URL query params (17 exact + 5 substring matches) - HttpTool now accepts optional credential registry + secrets store, auto-injects matching credentials in execute(), and uses broader auth detection in requires_approval() - ToolRegistry passes credential registry to HttpTool at startup and populates it when WASM tools register - Remove old hardcoded AUTH_HEADER_NAMES / has_auth_headers in favor of the new params_contain_manual_credentials() Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR #274 review comments (query param injection, lock poisoning, visibility) - Fix injected query params not being sent on outbound HTTP requests by also calling .query() on the RequestBuilder alongside parsed_url mutation - Recover from poisoned RwLock in SharedCredentialRegistry instead of silently ignoring failures, with tracing::warn for visibility - Narrow inject_credential and host_matches_pattern to pub(crate) to avoid committing to them as stable public API Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
542268fde5
commit
2cdd1acb1e
@@ -286,8 +286,8 @@ impl Tool for <Name>Tool {
|
||||
false // Set true if tool processes external data
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
false // Set true if tool is destructive or contacts external services
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> crate::tools::tool::ApprovalRequirement {
|
||||
crate::tools::tool::ApprovalRequirement::Never // Set to UnlessAutoApproved or Always as needed
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
+26
-29
@@ -284,32 +284,8 @@ impl Agent {
|
||||
for (idx, original_tc) in tool_calls.iter().enumerate() {
|
||||
let mut tc = original_tc.clone();
|
||||
|
||||
// Check if tool requires approval (skipped when auto_approve_tools is set)
|
||||
if !self.config.auto_approve_tools
|
||||
&& let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
let mut is_auto_approved = {
|
||||
let sess = session.lock().await;
|
||||
sess.is_tool_auto_approved(&tc.name)
|
||||
};
|
||||
|
||||
// Override auto-approval for destructive parameters
|
||||
if is_auto_approved && tool.requires_approval_for(&tc.arguments) {
|
||||
tracing::info!(
|
||||
tool = %tc.name,
|
||||
"Parameters require explicit approval despite auto-approve"
|
||||
);
|
||||
is_auto_approved = false;
|
||||
}
|
||||
|
||||
if !is_auto_approved {
|
||||
approval_needed = Some((idx, tc, tool));
|
||||
break; // remaining tools are deferred
|
||||
}
|
||||
}
|
||||
|
||||
// Hook: BeforeToolCall
|
||||
// Hook: BeforeToolCall (runs before approval so hooks can
|
||||
// modify parameters — approval is checked on final params)
|
||||
let event = crate::hooks::HookEvent::ToolCall {
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
@@ -352,6 +328,27 @@ impl Agent {
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Check if tool requires approval on the final (post-hook)
|
||||
// parameters. Skipped when auto_approve_tools is set.
|
||||
if !self.config.auto_approve_tools
|
||||
&& let Some(tool) = self.tools().get(&tc.name).await
|
||||
{
|
||||
use crate::tools::ApprovalRequirement;
|
||||
let needs_approval = match tool.requires_approval(&tc.arguments) {
|
||||
ApprovalRequirement::Never => false,
|
||||
ApprovalRequirement::UnlessAutoApproved => {
|
||||
let sess = session.lock().await;
|
||||
!sess.is_tool_auto_approved(&tc.name)
|
||||
}
|
||||
ApprovalRequirement::Always => true,
|
||||
};
|
||||
|
||||
if needs_approval {
|
||||
approval_needed = Some((idx, tc, tool));
|
||||
break; // remaining tools are deferred
|
||||
}
|
||||
}
|
||||
|
||||
let preflight_idx = preflight.len();
|
||||
preflight.push((tc.clone(), PreflightOutcome::Runnable));
|
||||
runnable.push((preflight_idx, tc));
|
||||
@@ -910,9 +907,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shell_destructive_command_requires_approval_for() {
|
||||
// ShellTool::requires_approval_for should detect destructive commands.
|
||||
// This exercises the same code path used inline in run_agentic_loop.
|
||||
fn test_shell_destructive_command_requires_explicit_approval() {
|
||||
// requires_explicit_approval() detects destructive commands that
|
||||
// should return ApprovalRequirement::Always from ShellTool.
|
||||
use crate::tools::builtin::shell::requires_explicit_approval;
|
||||
|
||||
let destructive_cmds = [
|
||||
|
||||
@@ -357,7 +357,7 @@ impl Scheduler {
|
||||
.into());
|
||||
}
|
||||
|
||||
if tool.requires_approval() {
|
||||
if tool.requires_approval(¶ms).is_required() {
|
||||
return Err(crate::error::ToolError::AuthRequired {
|
||||
name: tool_name.to_string(),
|
||||
}
|
||||
|
||||
+9
-10
@@ -746,19 +746,18 @@ impl Agent {
|
||||
)> = None;
|
||||
|
||||
for (idx, tc) in deferred_tool_calls.iter().enumerate() {
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
let is_auto_approved = {
|
||||
let sess = session.lock().await;
|
||||
let mut approved = sess.is_tool_auto_approved(&tc.name);
|
||||
if approved && tool.requires_approval_for(&tc.arguments) {
|
||||
approved = false;
|
||||
if let Some(tool) = self.tools().get(&tc.name).await {
|
||||
use crate::tools::ApprovalRequirement;
|
||||
let needs_approval = match tool.requires_approval(&tc.arguments) {
|
||||
ApprovalRequirement::Never => false,
|
||||
ApprovalRequirement::UnlessAutoApproved => {
|
||||
let sess = session.lock().await;
|
||||
!sess.is_tool_auto_approved(&tc.name)
|
||||
}
|
||||
approved
|
||||
ApprovalRequirement::Always => true,
|
||||
};
|
||||
|
||||
if !is_auto_approved {
|
||||
if needs_approval {
|
||||
approval_needed = Some((idx, tc.clone(), tool));
|
||||
break; // remaining tools stay deferred
|
||||
}
|
||||
|
||||
+1
-1
@@ -432,7 +432,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
})?;
|
||||
|
||||
// Tools requiring approval are blocked in autonomous jobs
|
||||
if tool.requires_approval() {
|
||||
if tool.requires_approval(params).is_required() {
|
||||
return Err(crate::error::ToolError::AuthRequired {
|
||||
name: tool_name.to_string(),
|
||||
}
|
||||
|
||||
+9
-2
@@ -724,8 +724,15 @@ async fn main() -> anyhow::Result<()> {
|
||||
let safety = Arc::new(SafetyLayer::new(&config.safety));
|
||||
tracing::info!("Safety layer initialized");
|
||||
|
||||
// Initialize tool registry
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
// Initialize tool registry with credential injection support
|
||||
let credential_registry = Arc::new(ironclaw::tools::wasm::SharedCredentialRegistry::new());
|
||||
let tools = if let Some(ref ss) = secrets_store {
|
||||
Arc::new(
|
||||
ToolRegistry::new().with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)),
|
||||
)
|
||||
} else {
|
||||
Arc::new(ToolRegistry::new())
|
||||
};
|
||||
tools.register_builtin_tools();
|
||||
|
||||
// Create embeddings provider if configured
|
||||
|
||||
@@ -0,0 +1,381 @@
|
||||
//! Broad detection of manually-provided credentials in HTTP request parameters.
|
||||
//!
|
||||
//! Used by the built-in HTTP tool to decide whether approval is needed when
|
||||
//! the LLM provides auth data directly in headers or URL query parameters.
|
||||
|
||||
/// Check whether HTTP request parameters contain manually-provided credentials.
|
||||
///
|
||||
/// Inspects headers (name/value), URL query parameters, and URL userinfo
|
||||
/// for patterns that indicate authentication data.
|
||||
pub fn params_contain_manual_credentials(params: &serde_json::Value) -> bool {
|
||||
headers_contain_credentials(params)
|
||||
|| url_contains_credential_params(params)
|
||||
|| url_contains_userinfo(params)
|
||||
}
|
||||
|
||||
/// Header names that are exact matches for credential-carrying headers (case-insensitive).
|
||||
const AUTH_HEADER_EXACT: &[&str] = &[
|
||||
"authorization",
|
||||
"proxy-authorization",
|
||||
"cookie",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"x-auth-token",
|
||||
"x-token",
|
||||
"x-access-token",
|
||||
"x-session-token",
|
||||
"x-csrf-token",
|
||||
"x-secret",
|
||||
"x-api-secret",
|
||||
];
|
||||
|
||||
/// Substrings in header names that suggest credentials (case-insensitive).
|
||||
/// Note: "key" is excluded to avoid false positives like "X-Idempotency-Key".
|
||||
const AUTH_HEADER_SUBSTRINGS: &[&str] = &["auth", "token", "secret", "credential", "password"];
|
||||
|
||||
/// Value prefixes that indicate auth schemes (case-insensitive).
|
||||
const AUTH_VALUE_PREFIXES: &[&str] = &[
|
||||
"bearer ",
|
||||
"basic ",
|
||||
"token ",
|
||||
"digest ",
|
||||
"hoba ",
|
||||
"mutual ",
|
||||
"aws4-hmac-sha256 ",
|
||||
];
|
||||
|
||||
/// URL query parameter names that are exact matches for credentials (case-insensitive).
|
||||
const AUTH_QUERY_EXACT: &[&str] = &[
|
||||
"api_key",
|
||||
"apikey",
|
||||
"api-key",
|
||||
"access_token",
|
||||
"token",
|
||||
"key",
|
||||
"secret",
|
||||
"password",
|
||||
"auth",
|
||||
"auth_token",
|
||||
"session_token",
|
||||
"client_secret",
|
||||
"client_id",
|
||||
"app_key",
|
||||
"app_secret",
|
||||
"sig",
|
||||
"signature",
|
||||
];
|
||||
|
||||
/// Substrings in query parameter names that suggest credentials (case-insensitive).
|
||||
const AUTH_QUERY_SUBSTRINGS: &[&str] = &["token", "secret", "auth", "password", "credential"];
|
||||
|
||||
fn header_name_is_credential(name: &str) -> bool {
|
||||
let lower = name.to_lowercase();
|
||||
|
||||
if AUTH_HEADER_EXACT.contains(&lower.as_str()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
AUTH_HEADER_SUBSTRINGS.iter().any(|sub| lower.contains(sub))
|
||||
}
|
||||
|
||||
fn header_value_is_credential(value: &str) -> bool {
|
||||
let lower = value.to_lowercase();
|
||||
AUTH_VALUE_PREFIXES.iter().any(|pfx| lower.starts_with(pfx))
|
||||
}
|
||||
|
||||
fn headers_contain_credentials(params: &serde_json::Value) -> bool {
|
||||
match params.get("headers") {
|
||||
Some(serde_json::Value::Object(map)) => map.iter().any(|(k, v)| {
|
||||
header_name_is_credential(k) || v.as_str().is_some_and(header_value_is_credential)
|
||||
}),
|
||||
Some(serde_json::Value::Array(items)) => items.iter().any(|item| {
|
||||
let name_match = item
|
||||
.get("name")
|
||||
.and_then(|n| n.as_str())
|
||||
.is_some_and(header_name_is_credential);
|
||||
let value_match = item
|
||||
.get("value")
|
||||
.and_then(|v| v.as_str())
|
||||
.is_some_and(header_value_is_credential);
|
||||
name_match || value_match
|
||||
}),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn query_param_is_credential(name: &str) -> bool {
|
||||
let lower = name.to_lowercase();
|
||||
|
||||
if AUTH_QUERY_EXACT.contains(&lower.as_str()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
AUTH_QUERY_SUBSTRINGS.iter().any(|sub| lower.contains(sub))
|
||||
}
|
||||
|
||||
fn url_contains_credential_params(params: &serde_json::Value) -> bool {
|
||||
let url_str = match params.get("url").and_then(|u| u.as_str()) {
|
||||
Some(u) => u,
|
||||
None => return false,
|
||||
};
|
||||
|
||||
let parsed = match url::Url::parse(url_str) {
|
||||
Ok(u) => u,
|
||||
Err(_) => return false,
|
||||
};
|
||||
|
||||
parsed
|
||||
.query_pairs()
|
||||
.any(|(name, _)| query_param_is_credential(&name))
|
||||
}
|
||||
|
||||
/// Detect credentials embedded in URL userinfo (e.g., `https://user:pass@host/`).
|
||||
fn url_contains_userinfo(params: &serde_json::Value) -> bool {
|
||||
let url_str = match params.get("url").and_then(|u| u.as_str()) {
|
||||
Some(u) => u,
|
||||
None => return false,
|
||||
};
|
||||
|
||||
let parsed = match url::Url::parse(url_str) {
|
||||
Ok(u) => u,
|
||||
Err(_) => return false,
|
||||
};
|
||||
|
||||
// Non-empty username or password in the URL indicates embedded credentials
|
||||
!parsed.username().is_empty() || parsed.password().is_some()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── Header name exact match ────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_authorization_header_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com",
|
||||
"headers": {"Authorization": "Bearer token123"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_exact_header_names() {
|
||||
for name in AUTH_HEADER_EXACT {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {name.to_string(): "some_value"}
|
||||
});
|
||||
assert!(
|
||||
params_contain_manual_credentials(¶ms),
|
||||
"Header '{}' should be detected",
|
||||
name
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_header_name_case_insensitive() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"AUTHORIZATION": "value"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── Header name substring match ────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_header_substring_auth() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"X-Custom-Auth-Header": "value"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_header_substring_token() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"X-My-Token": "value"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── Header value prefix match ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_bearer_value_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"X-Custom": "Bearer sk-abc123"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_basic_value_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"X-Custom": "Basic dXNlcjpwYXNz"}
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── Array-format headers ───────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_array_format_header_name() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": [{"name": "Authorization", "value": "Bearer token"}]
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_array_format_header_value_prefix() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": [{"name": "X-Custom", "value": "Token abc123"}]
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── URL query parameter detection ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_url_api_key_param() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data?api_key=abc123"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_access_token_param() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data?access_token=xyz"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_query_substring_match() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data?my_auth_code=xyz"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_query_case_insensitive() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data?API_KEY=abc"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── False positive checks ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_idempotency_key_not_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "POST",
|
||||
"url": "https://api.example.com",
|
||||
"headers": {"X-Idempotency-Key": "uuid-1234"}
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_type_not_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"Content-Type": "application/json", "Accept": "text/html"}
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_headers_no_query() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com/path"
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_query_params() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/search?q=hello&page=1&limit=10"
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_headers() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {}
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_url_returns_false() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "not a url"
|
||||
});
|
||||
assert!(!params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
// ── URL userinfo detection ─────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_url_userinfo_with_password_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://user:[email protected]/data"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_userinfo_username_only_detected() {
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://[email protected]/data"
|
||||
});
|
||||
assert!(params_contain_manual_credentials(¶ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_without_userinfo_not_detected_by_userinfo_check() {
|
||||
// This specifically tests that url_contains_userinfo returns false
|
||||
// for a normal URL (the broader function may still detect query params).
|
||||
assert!(!url_contains_userinfo(&serde_json::json!({
|
||||
"url": "https://api.example.com/data"
|
||||
})));
|
||||
}
|
||||
}
|
||||
@@ -7,11 +7,13 @@
|
||||
//! - Enforcing safety policies
|
||||
//! - Detecting secret leakage in outputs
|
||||
|
||||
mod credential_detect;
|
||||
mod leak_detector;
|
||||
mod policy;
|
||||
mod sanitizer;
|
||||
mod validator;
|
||||
|
||||
pub use credential_detect::params_contain_manual_credentials;
|
||||
pub use leak_detector::{
|
||||
LeakAction, LeakDetectionError, LeakDetector, LeakMatch, LeakPattern, LeakScanResult,
|
||||
LeakSeverity,
|
||||
|
||||
@@ -45,7 +45,7 @@ use crate::llm::{
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Requirement specification for building software.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -1019,8 +1019,8 @@ impl Tool for BuildSoftwareTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Building software should require approval
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ use async_trait::async_trait;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::extensions::{ExtensionKind, ExtensionManager};
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||
|
||||
// ── tool_search ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -158,8 +158,8 @@ impl Tool for ToolInstallTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -253,8 +253,8 @@ impl Tool for ToolAuthTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -478,8 +478,8 @@ impl Tool for ToolRemoveTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -500,11 +500,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tool_install_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ToolInstallTool {
|
||||
manager: test_manager_stub(),
|
||||
};
|
||||
assert_eq!(tool.name(), "tool_install");
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("name").is_some());
|
||||
assert!(schema["properties"].get("url").is_some());
|
||||
@@ -512,11 +516,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tool_auth_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ToolAuthTool {
|
||||
manager: test_manager_stub(),
|
||||
};
|
||||
assert_eq!(tool.name(), "tool_auth");
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("name").is_some());
|
||||
// token param must NOT be in schema (security: tokens never go through LLM)
|
||||
@@ -528,31 +536,43 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tool_activate_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ToolActivateTool {
|
||||
manager: test_manager_stub(),
|
||||
};
|
||||
assert_eq!(tool.name(), "tool_activate");
|
||||
assert!(!tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::Never
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_list_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ToolListTool {
|
||||
manager: test_manager_stub(),
|
||||
};
|
||||
assert_eq!(tool.name(), "tool_list");
|
||||
assert!(!tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::Never
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("kind").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_remove_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ToolRemoveTool {
|
||||
manager: test_manager_stub(),
|
||||
};
|
||||
assert_eq!(tool.name(), "tool_remove");
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
/// Create a stub manager for schema tests (these don't call execute).
|
||||
|
||||
@@ -11,7 +11,9 @@ use async_trait::async_trait;
|
||||
use tokio::fs;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::{Tool, ToolDomain, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::tool::{
|
||||
ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str,
|
||||
};
|
||||
use crate::workspace::paths as ws_paths;
|
||||
|
||||
/// Well-known workspace filenames that must go through memory_write, not write_file.
|
||||
@@ -265,8 +267,8 @@ impl Tool for ReadFileTool {
|
||||
true // File content could contain anything
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Reading local files should require approval
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn domain(&self) -> ToolDomain {
|
||||
@@ -372,8 +374,8 @@ impl Tool for WriteFileTool {
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // File writes should require approval
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
@@ -488,8 +490,8 @@ impl Tool for ListDirTool {
|
||||
false // Directory listings are safe
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Directory listings can leak filesystem structure
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn domain(&self) -> ToolDomain {
|
||||
@@ -697,8 +699,8 @@ impl Tool for ApplyPatchTool {
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // File edits should require approval
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
|
||||
+310
-6
@@ -2,6 +2,7 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::{IpAddr, ToSocketAddrs};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -10,7 +11,9 @@ use reqwest::Client;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::safety::LeakDetector;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::wasm::{InjectedCredentials, SharedCredentialRegistry, inject_credential};
|
||||
|
||||
/// Maximum response body size (5 MB).
|
||||
///
|
||||
@@ -22,6 +25,8 @@ const MAX_RESPONSE_SIZE: usize = 5 * 1024 * 1024;
|
||||
/// Tool for making HTTP requests.
|
||||
pub struct HttpTool {
|
||||
client: Client,
|
||||
credential_registry: Option<Arc<SharedCredentialRegistry>>,
|
||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl HttpTool {
|
||||
@@ -33,7 +38,22 @@ impl HttpTool {
|
||||
.build()
|
||||
.expect("Failed to create HTTP client");
|
||||
|
||||
Self { client }
|
||||
Self {
|
||||
client,
|
||||
credential_registry: None,
|
||||
secrets_store: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a credential registry and secrets store for auto-injection.
|
||||
pub fn with_credentials(
|
||||
mut self,
|
||||
registry: Arc<SharedCredentialRegistry>,
|
||||
secrets_store: Arc<dyn SecretsStore + Send + Sync>,
|
||||
) -> Self {
|
||||
self.credential_registry = Some(registry);
|
||||
self.secrets_store = Some(secrets_store);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
@@ -146,6 +166,15 @@ fn parse_headers_param(
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract host from URL in params (for approval checks).
|
||||
fn extract_host_from_params(params: &serde_json::Value) -> Option<String> {
|
||||
params
|
||||
.get("url")
|
||||
.and_then(|u| u.as_str())
|
||||
.and_then(|u| reqwest::Url::parse(u).ok())
|
||||
.and_then(|u| u.host_str().map(|h| h.to_string()))
|
||||
}
|
||||
|
||||
impl Default for HttpTool {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
@@ -211,10 +240,10 @@ impl Tool for HttpTool {
|
||||
let method = require_str(¶ms, "method")?;
|
||||
|
||||
let url = require_str(¶ms, "url")?;
|
||||
let parsed_url = validate_url(url)?;
|
||||
let mut parsed_url = validate_url(url)?;
|
||||
|
||||
// Parse headers
|
||||
let headers_vec = parse_headers_param(params.get("headers"))?;
|
||||
let mut headers_vec = parse_headers_param(params.get("headers"))?;
|
||||
|
||||
// Build request
|
||||
let mut request = match method.to_uppercase().as_str() {
|
||||
@@ -261,6 +290,41 @@ impl Tool for HttpTool {
|
||||
None
|
||||
};
|
||||
|
||||
// Credential injection from shared registry
|
||||
if let (Some(registry), Some(store)) = (
|
||||
self.credential_registry.as_ref(),
|
||||
self.secrets_store.as_ref(),
|
||||
) {
|
||||
let host = parsed_url.host_str().unwrap_or("");
|
||||
let matched: Vec<crate::secrets::CredentialMapping> = registry.find_for_host(host);
|
||||
for mapping in &matched {
|
||||
match store
|
||||
.get_decrypted(&_ctx.user_id, &mapping.secret_name)
|
||||
.await
|
||||
{
|
||||
Ok(secret) => {
|
||||
let mut injected = InjectedCredentials::empty();
|
||||
inject_credential(&mut injected, &mapping.location, &secret);
|
||||
for (name, value) in &injected.headers {
|
||||
request = request.header(name.as_str(), value.as_str());
|
||||
headers_vec.push((name.clone(), value.clone()));
|
||||
}
|
||||
for (name, value) in &injected.query_params {
|
||||
parsed_url.query_pairs_mut().append_pair(name, value);
|
||||
request = request.query(&[(name.as_str(), value.as_str())]);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
secret = %mapping.secret_name,
|
||||
error = %e,
|
||||
"Failed to inject credential for HTTP tool"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Leak detection on outbound request (url/headers/body)
|
||||
let detector = LeakDetector::new();
|
||||
detector
|
||||
@@ -352,8 +416,20 @@ impl Tool for HttpTool {
|
||||
true // External data always needs sanitization
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // HTTP requests go to external services, require user approval
|
||||
fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement {
|
||||
// 1. Manual auth headers/query params in LLM params
|
||||
if crate::safety::params_contain_manual_credentials(params) {
|
||||
return ApprovalRequirement::Always;
|
||||
}
|
||||
// 2. Target host has credential mappings (will be auto-injected)
|
||||
if let Some(ref registry) = self.credential_registry
|
||||
&& let Some(host) = extract_host_from_params(params)
|
||||
&& registry.has_credentials_for_host(&host)
|
||||
{
|
||||
return ApprovalRequirement::Always;
|
||||
}
|
||||
// Default: outbound HTTP still needs approval unless auto-approved
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -467,4 +543,232 @@ mod tests {
|
||||
"body schema must include a type for OpenAI-compatible tool validation"
|
||||
);
|
||||
}
|
||||
|
||||
// ── Approval requirement tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_no_auth_headers_returns_unless_auto_approved() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data"
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_header_object_format_returns_always() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data",
|
||||
"headers": {"Authorization": "Bearer token123"}
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_header_array_format_returns_always() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data",
|
||||
"headers": [{"name": "Authorization", "value": "Bearer token123"}]
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_header_case_insensitive() {
|
||||
let tool = HttpTool::new();
|
||||
|
||||
// Object format with mixed case
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"AUTHORIZATION": "Bearer x"}
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
|
||||
// Array format with mixed case
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": [{"name": "X-Api-Key", "value": "key123"}]
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_auth_header_names_detected() {
|
||||
let tool = HttpTool::new();
|
||||
for header_name in [
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"cookie",
|
||||
"proxy-authorization",
|
||||
"x-auth-token",
|
||||
"api-key",
|
||||
"x-token",
|
||||
"x-access-token",
|
||||
"x-session-token",
|
||||
"x-csrf-token",
|
||||
"x-secret",
|
||||
"x-api-secret",
|
||||
] {
|
||||
let mut headers = serde_json::Map::new();
|
||||
headers.insert(header_name.to_string(), serde_json::json!("value"));
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": headers
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::Always,
|
||||
"Header '{}' should trigger Always approval",
|
||||
header_name
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_non_auth_headers_return_unless_auto_approved() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"Content-Type": "application/json", "Accept": "text/html"}
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_headers_return_unless_auto_approved() {
|
||||
let tool = HttpTool::new();
|
||||
|
||||
// Empty object
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {}
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
|
||||
// Empty array
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": []
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
// ── Credential registry approval tests ─────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_host_with_credential_mapping_returns_always() {
|
||||
use crate::secrets::CredentialMapping;
|
||||
use crate::tools::wasm::SharedCredentialRegistry;
|
||||
|
||||
let registry = Arc::new(SharedCredentialRegistry::new());
|
||||
registry.add_mappings(vec![CredentialMapping::bearer(
|
||||
"openai_key",
|
||||
"api.openai.com",
|
||||
)]);
|
||||
|
||||
let tool = HttpTool::new().with_credentials(
|
||||
registry,
|
||||
// secrets_store is not used in requires_approval, just needs to be present
|
||||
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
|
||||
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
|
||||
"0123456789abcdef0123456789abcdef".to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
))),
|
||||
);
|
||||
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.openai.com/v1/models"
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_host_without_credential_mapping_returns_unless_auto_approved() {
|
||||
use crate::tools::wasm::SharedCredentialRegistry;
|
||||
|
||||
let registry = Arc::new(SharedCredentialRegistry::new());
|
||||
// Empty registry - no credential mappings
|
||||
|
||||
let tool = HttpTool::new().with_credentials(
|
||||
registry,
|
||||
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
|
||||
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
|
||||
"0123456789abcdef0123456789abcdef".to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
))),
|
||||
);
|
||||
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data"
|
||||
});
|
||||
assert_eq!(
|
||||
tool.requires_approval(¶ms),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_query_param_credential_returns_always() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://api.example.com/data?api_key=secret123"
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bearer_value_in_custom_header_returns_always() {
|
||||
let tool = HttpTool::new();
|
||||
let params = serde_json::json!({
|
||||
"method": "GET",
|
||||
"url": "https://example.com",
|
||||
"headers": {"X-Custom": "Bearer sk-test123"}
|
||||
});
|
||||
assert_eq!(tool.requires_approval(¶ms), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_host_from_params_valid() {
|
||||
let params = serde_json::json!({
|
||||
"url": "https://api.example.com/path"
|
||||
});
|
||||
assert_eq!(
|
||||
extract_host_from_params(¶ms),
|
||||
Some("api.example.com".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_host_from_params_missing_url() {
|
||||
let params = serde_json::json!({"method": "GET"});
|
||||
assert_eq!(extract_host_from_params(¶ms), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,7 +22,7 @@ use crate::history::SandboxJobRecord;
|
||||
use crate::orchestrator::auth::CredentialGrant;
|
||||
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||
|
||||
/// Resolve a job ID from a full UUID or a short prefix (like git short SHAs).
|
||||
///
|
||||
@@ -1005,8 +1005,8 @@ impl Tool for CancelJobTool {
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Canceling a job should require approval
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
@@ -1268,8 +1268,8 @@ impl Tool for JobPromptTool {
|
||||
Ok(ToolOutput::success(result, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
@@ -1611,10 +1611,14 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_prompt_tool_requires_approval() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let queue: PromptQueue =
|
||||
Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new()));
|
||||
let tool = test_prompt_tool(queue);
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
+35
-22
@@ -55,7 +55,9 @@ use tokio::process::Command;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::sandbox::{SandboxManager, SandboxPolicy};
|
||||
use crate::tools::tool::{Tool, ToolDomain, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::tool::{
|
||||
ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str,
|
||||
};
|
||||
|
||||
/// Maximum output size before truncation (64KB).
|
||||
const MAX_OUTPUT_SIZE: usize = 64 * 1024;
|
||||
@@ -696,11 +698,7 @@ impl Tool for ShellTool {
|
||||
Ok(ToolOutput::success(result, duration))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true // Shell commands should require approval
|
||||
}
|
||||
|
||||
fn requires_approval_for(&self, params: &serde_json::Value) -> bool {
|
||||
fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement {
|
||||
let cmd = params
|
||||
.get("command")
|
||||
.and_then(|c| c.as_str().map(String::from))
|
||||
@@ -714,10 +712,10 @@ impl Tool for ShellTool {
|
||||
if let Some(ref cmd) = cmd
|
||||
&& requires_explicit_approval(cmd)
|
||||
{
|
||||
return true;
|
||||
return ApprovalRequirement::Always;
|
||||
}
|
||||
|
||||
false
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
@@ -861,31 +859,46 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_approval_for_destructive_command() {
|
||||
fn test_requires_approval_destructive_command() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ShellTool::new();
|
||||
// Destructive commands must return true even though shell already
|
||||
// requires base approval -- the distinction matters for auto-approve override.
|
||||
assert!(tool.requires_approval_for(&serde_json::json!({"command": "rm -rf /tmp"})));
|
||||
assert!(tool.requires_approval_for(
|
||||
&serde_json::json!({"command": "git push --force origin main"})
|
||||
));
|
||||
assert!(tool.requires_approval_for(&serde_json::json!({"command": "DROP TABLE users;"})));
|
||||
// Destructive commands must return Always to bypass auto-approve.
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"command": "rm -rf /tmp"})),
|
||||
ApprovalRequirement::Always
|
||||
);
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"command": "git push --force origin main"})),
|
||||
ApprovalRequirement::Always
|
||||
);
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"command": "DROP TABLE users;"})),
|
||||
ApprovalRequirement::Always
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_approval_for_safe_command() {
|
||||
fn test_requires_approval_safe_command() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ShellTool::new();
|
||||
// Safe commands should not override auto-approval; only destructive ones do.
|
||||
assert!(!tool.requires_approval_for(&serde_json::json!({"command": "cargo build"})));
|
||||
assert!(!tool.requires_approval_for(&serde_json::json!({"command": "echo hello"})));
|
||||
// Safe commands return UnlessAutoApproved (can be auto-approved).
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"command": "cargo build"})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"command": "echo hello"})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_approval_for_string_encoded_args() {
|
||||
fn test_requires_approval_string_encoded_args() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = ShellTool::new();
|
||||
// When arguments are string-encoded JSON (rare LLM behavior).
|
||||
let args = serde_json::Value::String(r#"{"command": "rm -rf /tmp/stuff"}"#.to_string());
|
||||
assert!(tool.requires_approval_for(&args));
|
||||
assert_eq!(tool.requires_approval(&args), ApprovalRequirement::Always);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -10,7 +10,7 @@ use async_trait::async_trait;
|
||||
use crate::context::JobContext;
|
||||
use crate::skills::catalog::SkillCatalog;
|
||||
use crate::skills::registry::SkillRegistry;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str};
|
||||
|
||||
// ── skill_list ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -356,8 +356,8 @@ impl Tool for SkillInstallTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -553,8 +553,8 @@ impl Tool for SkillRemoveTool {
|
||||
Ok(ToolOutput::success(output, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
true
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
|
||||
@@ -575,27 +575,39 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_skill_list_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = SkillListTool::new(test_registry());
|
||||
assert_eq!(tool.name(), "skill_list");
|
||||
assert!(!tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::Never
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema.get("properties").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_skill_search_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = SkillSearchTool::new(test_registry(), test_catalog());
|
||||
assert_eq!(tool.name(), "skill_search");
|
||||
assert!(!tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::Never
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("query").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_skill_install_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = SkillInstallTool::new(test_registry(), test_catalog());
|
||||
assert_eq!(tool.name(), "skill_install");
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("name").is_some());
|
||||
assert!(schema["properties"].get("url").is_some());
|
||||
@@ -604,9 +616,13 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_skill_remove_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
let tool = SkillRemoveTool::new(test_registry());
|
||||
assert_eq!(tool.name(), "skill_remove");
|
||||
assert!(tool.requires_approval());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"].get("name").is_some());
|
||||
}
|
||||
|
||||
@@ -18,7 +18,7 @@ use crate::tools::mcp::protocol::{
|
||||
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
|
||||
};
|
||||
use crate::tools::mcp::session::McpSessionManager;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
/// MCP client for communicating with MCP servers.
|
||||
///
|
||||
@@ -538,9 +538,13 @@ impl Tool for McpToolWrapper {
|
||||
true // MCP tools are external, always sanitize
|
||||
}
|
||||
|
||||
fn requires_approval(&self) -> bool {
|
||||
// Check the destructive_hint annotation from the MCP server
|
||||
self.tool.requires_approval()
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
// Delegate to the MCP protocol type's own requires_approval() bool method
|
||||
if self.tool.requires_approval() {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
} else {
|
||||
ApprovalRequirement::Never
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -21,4 +21,4 @@ pub use builder::{
|
||||
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
||||
};
|
||||
pub use registry::ToolRegistry;
|
||||
pub use tool::{Tool, ToolDomain, ToolError, ToolOutput};
|
||||
pub use tool::{ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput};
|
||||
|
||||
+51
-3
@@ -24,8 +24,8 @@ use crate::tools::builtin::{
|
||||
};
|
||||
use crate::tools::tool::{Tool, ToolDomain};
|
||||
use crate::tools::wasm::{
|
||||
Capabilities, OAuthRefreshConfig, ResourceLimits, WasmError, WasmStorageError, WasmToolRuntime,
|
||||
WasmToolStore, WasmToolWrapper,
|
||||
Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError,
|
||||
WasmStorageError, WasmToolRuntime, WasmToolStore, WasmToolWrapper,
|
||||
};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
@@ -73,6 +73,10 @@ pub struct ToolRegistry {
|
||||
tools: RwLock<HashMap<String, Arc<dyn Tool>>>,
|
||||
/// Tracks which names were registered as built-in (protected from shadowing).
|
||||
builtin_names: RwLock<std::collections::HashSet<String>>,
|
||||
/// Shared credential registry populated by WASM tools, consumed by HTTP tool.
|
||||
credential_registry: Option<Arc<SharedCredentialRegistry>>,
|
||||
/// Secrets store for credential injection (shared with HTTP tool).
|
||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
@@ -81,9 +85,27 @@ impl ToolRegistry {
|
||||
Self {
|
||||
tools: RwLock::new(HashMap::new()),
|
||||
builtin_names: RwLock::new(std::collections::HashSet::new()),
|
||||
credential_registry: None,
|
||||
secrets_store: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a registry with credential injection support.
|
||||
pub fn with_credentials(
|
||||
mut self,
|
||||
credential_registry: Arc<SharedCredentialRegistry>,
|
||||
secrets_store: Arc<dyn SecretsStore + Send + Sync>,
|
||||
) -> Self {
|
||||
self.credential_registry = Some(credential_registry);
|
||||
self.secrets_store = Some(secrets_store);
|
||||
self
|
||||
}
|
||||
|
||||
/// Get a reference to the shared credential registry.
|
||||
pub fn credential_registry(&self) -> Option<&Arc<SharedCredentialRegistry>> {
|
||||
self.credential_registry.as_ref()
|
||||
}
|
||||
|
||||
/// Register a tool. Rejects dynamic tools that try to shadow a built-in name.
|
||||
pub async fn register(&self, tool: Arc<dyn Tool>) {
|
||||
let name = tool.name().to_string();
|
||||
@@ -176,7 +198,12 @@ impl ToolRegistry {
|
||||
self.register_sync(Arc::new(EchoTool));
|
||||
self.register_sync(Arc::new(TimeTool));
|
||||
self.register_sync(Arc::new(JsonTool));
|
||||
self.register_sync(Arc::new(HttpTool::new()));
|
||||
|
||||
let mut http = HttpTool::new();
|
||||
if let (Some(cr), Some(ss)) = (&self.credential_registry, &self.secrets_store) {
|
||||
http = http.with_credentials(Arc::clone(cr), Arc::clone(ss));
|
||||
}
|
||||
self.register_sync(Arc::new(http));
|
||||
|
||||
tracing::info!("Registered {} built-in tools", self.count());
|
||||
}
|
||||
@@ -419,6 +446,14 @@ impl ToolRegistry {
|
||||
.prepare(reg.name, reg.wasm_bytes, reg.limits)
|
||||
.await?;
|
||||
|
||||
// Extract credential mappings before capabilities are moved into the wrapper
|
||||
let credential_mappings: Vec<crate::secrets::CredentialMapping> = reg
|
||||
.capabilities
|
||||
.http
|
||||
.as_ref()
|
||||
.map(|http| http.credentials.values().cloned().collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
// Create the wrapper
|
||||
let mut wrapper = WasmToolWrapper::new(Arc::clone(reg.runtime), prepared, reg.capabilities);
|
||||
|
||||
@@ -439,6 +474,19 @@ impl ToolRegistry {
|
||||
// Register the tool
|
||||
self.register(Arc::new(wrapper)).await;
|
||||
|
||||
// Add credential mappings to the shared registry (for HTTP tool injection)
|
||||
if let Some(cr) = &self.credential_registry
|
||||
&& !credential_mappings.is_empty()
|
||||
{
|
||||
let count = credential_mappings.len();
|
||||
cr.add_mappings(credential_mappings);
|
||||
tracing::debug!(
|
||||
name = reg.name,
|
||||
credential_count = count,
|
||||
"Added credential mappings from WASM tool"
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!(name = reg.name, "Registered WASM tool");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+35
-27
@@ -9,6 +9,25 @@ use thiserror::Error;
|
||||
|
||||
use crate::context::JobContext;
|
||||
|
||||
/// How much approval a specific tool invocation requires.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ApprovalRequirement {
|
||||
/// No approval needed.
|
||||
Never,
|
||||
/// Needs approval, but session auto-approve can bypass.
|
||||
UnlessAutoApproved,
|
||||
/// Always needs explicit approval (even if auto-approved).
|
||||
Always,
|
||||
}
|
||||
|
||||
impl ApprovalRequirement {
|
||||
/// Whether this invocation requires approval in contexts where
|
||||
/// auto-approve is irrelevant (e.g. autonomous worker/scheduler).
|
||||
pub fn is_required(&self) -> bool {
|
||||
!matches!(self, Self::Never)
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a tool should execute: orchestrator process or inside a container.
|
||||
///
|
||||
/// Orchestrator tools run in the main agent process (memory access, job mgmt, etc).
|
||||
@@ -160,31 +179,14 @@ pub trait Tool: Send + Sync {
|
||||
true
|
||||
}
|
||||
|
||||
/// Whether this tool requires explicit user approval before execution.
|
||||
/// Whether this tool invocation requires user approval.
|
||||
///
|
||||
/// Returns false by default since most tools run in a sandboxed/virtualized
|
||||
/// environment. Only tools that make external network calls or perform
|
||||
/// destructive operations should return true.
|
||||
///
|
||||
/// When true, the agent will prompt the user for confirmation before
|
||||
/// executing this tool.
|
||||
fn requires_approval(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether this specific invocation should override auto-approval.
|
||||
///
|
||||
/// This method is called after checking `requires_approval()` and finding that
|
||||
/// the tool is auto-approved for this session. Return `true` to force approval
|
||||
/// for this specific invocation despite auto-approval (for example, for
|
||||
/// destructive operations like `rm -rf` or `git push --force`).
|
||||
///
|
||||
/// Return `false` to allow auto-approval to proceed normally.
|
||||
///
|
||||
/// The default returns `false`. Override only if you need parameter-aware
|
||||
/// approval gating.
|
||||
fn requires_approval_for(&self, _params: &serde_json::Value) -> bool {
|
||||
false
|
||||
/// Returns `Never` by default (most tools run in a sandboxed environment).
|
||||
/// Override to return `UnlessAutoApproved` for tools that need approval
|
||||
/// but can be session-auto-approved, or `Always` for invocations that
|
||||
/// must always prompt (e.g. destructive shell commands, HTTP with auth).
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::Never
|
||||
}
|
||||
|
||||
/// Maximum time this tool is allowed to run before the caller kills it.
|
||||
@@ -347,9 +349,15 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_approval_for_default() {
|
||||
fn test_requires_approval_default() {
|
||||
let tool = EchoTool;
|
||||
// Default requires_approval_for() returns false, allowing auto-approval.
|
||||
assert!(!tool.requires_approval_for(&serde_json::json!({"message": "hi"})));
|
||||
// Default requires_approval() returns Never.
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({"message": "hi"})),
|
||||
ApprovalRequirement::Never
|
||||
);
|
||||
assert!(!ApprovalRequirement::Never.is_required());
|
||||
assert!(ApprovalRequirement::UnlessAutoApproved.is_required());
|
||||
assert!(ApprovalRequirement::Always.is_required());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,6 +23,7 @@
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use crate::secrets::{
|
||||
CredentialLocation, CredentialMapping, DecryptedSecret, SecretError, SecretsStore,
|
||||
@@ -59,6 +60,88 @@ impl From<SecretError> for InjectionError {
|
||||
}
|
||||
}
|
||||
|
||||
/// Thread-safe, append-only registry of credential mappings from all installed tools.
|
||||
///
|
||||
/// Aggregates credential mappings from WASM tools so the built-in HTTP tool can
|
||||
/// auto-inject credentials for matching hosts. Uses `std::sync::RwLock` so
|
||||
/// `requires_approval` (sync) can query it without async.
|
||||
pub struct SharedCredentialRegistry {
|
||||
mappings: RwLock<Vec<CredentialMapping>>,
|
||||
}
|
||||
|
||||
impl SharedCredentialRegistry {
|
||||
/// Create an empty registry.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
mappings: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add credential mappings (called when WASM tools register).
|
||||
pub fn add_mappings(&self, mappings: impl IntoIterator<Item = CredentialMapping>) {
|
||||
match self.mappings.write() {
|
||||
Ok(mut guard) => {
|
||||
guard.extend(mappings);
|
||||
}
|
||||
Err(poisoned) => {
|
||||
tracing::warn!(
|
||||
"SharedCredentialRegistry RwLock poisoned during add_mappings; recovering"
|
||||
);
|
||||
let mut guard = poisoned.into_inner();
|
||||
guard.extend(mappings);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if any credential mapping matches this host (sync, for requires_approval).
|
||||
pub fn has_credentials_for_host(&self, host: &str) -> bool {
|
||||
let guard = match self.mappings.read() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!(
|
||||
"SharedCredentialRegistry RwLock poisoned during has_credentials_for_host; recovering"
|
||||
);
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
guard.iter().any(|mapping| {
|
||||
mapping
|
||||
.host_patterns
|
||||
.iter()
|
||||
.any(|pattern| host_matches_pattern(host, pattern))
|
||||
})
|
||||
}
|
||||
|
||||
/// Get all credential mappings matching a host (for injection).
|
||||
pub fn find_for_host(&self, host: &str) -> Vec<CredentialMapping> {
|
||||
let guard = match self.mappings.read() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!(
|
||||
"SharedCredentialRegistry RwLock poisoned during find_for_host; recovering"
|
||||
);
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
guard
|
||||
.iter()
|
||||
.filter(|mapping| {
|
||||
mapping
|
||||
.host_patterns
|
||||
.iter()
|
||||
.any(|pattern| host_matches_pattern(host, pattern))
|
||||
})
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SharedCredentialRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of credential injection.
|
||||
#[derive(Debug)]
|
||||
pub struct InjectedCredentials {
|
||||
@@ -431,4 +514,78 @@ mod tests {
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// ── SharedCredentialRegistry tests ─────────────────────────────────
|
||||
|
||||
use crate::tools::wasm::credential_injector::SharedCredentialRegistry;
|
||||
|
||||
#[test]
|
||||
fn test_shared_registry_empty() {
|
||||
let registry = SharedCredentialRegistry::new();
|
||||
assert!(!registry.has_credentials_for_host("api.example.com"));
|
||||
assert!(registry.find_for_host("api.example.com").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shared_registry_add_and_find() {
|
||||
let registry = SharedCredentialRegistry::new();
|
||||
registry.add_mappings(vec![
|
||||
CredentialMapping::bearer("openai_key", "api.openai.com"),
|
||||
CredentialMapping::header("github_token", "X-GitHub-Token", "*.github.com"),
|
||||
]);
|
||||
|
||||
assert!(registry.has_credentials_for_host("api.openai.com"));
|
||||
assert!(!registry.has_credentials_for_host("api.anthropic.com"));
|
||||
|
||||
let found = registry.find_for_host("api.openai.com");
|
||||
assert_eq!(found.len(), 1);
|
||||
assert_eq!(found[0].secret_name, "openai_key");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shared_registry_wildcard_host() {
|
||||
let registry = SharedCredentialRegistry::new();
|
||||
registry.add_mappings(vec![CredentialMapping::bearer("gh_token", "*.github.com")]);
|
||||
|
||||
assert!(registry.has_credentials_for_host("api.github.com"));
|
||||
assert!(registry.has_credentials_for_host("uploads.github.com"));
|
||||
assert!(!registry.has_credentials_for_host("github.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shared_registry_multiple_adds() {
|
||||
let registry = SharedCredentialRegistry::new();
|
||||
registry.add_mappings(vec![CredentialMapping::bearer("key1", "api.example.com")]);
|
||||
registry.add_mappings(vec![CredentialMapping::bearer("key2", "api.example.com")]);
|
||||
|
||||
let found = registry.find_for_host("api.example.com");
|
||||
assert_eq!(found.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shared_registry_thread_safety() {
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
|
||||
let registry = Arc::new(SharedCredentialRegistry::new());
|
||||
|
||||
let handles: Vec<_> = (0..4)
|
||||
.map(|i| {
|
||||
let r = Arc::clone(®istry);
|
||||
thread::spawn(move || {
|
||||
r.add_mappings(vec![CredentialMapping::bearer(
|
||||
format!("key_{}", i),
|
||||
"api.example.com",
|
||||
)]);
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
for h in handles {
|
||||
h.join().unwrap();
|
||||
}
|
||||
|
||||
let found = registry.find_for_host("api.example.com");
|
||||
assert_eq!(found.len(), 4);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -104,7 +104,10 @@ pub use capabilities::{
|
||||
|
||||
// Security components (V2)
|
||||
pub use allowlist::{AllowlistResult, AllowlistValidator, DenyReason};
|
||||
pub use credential_injector::{CredentialInjector, InjectedCredentials, InjectionError};
|
||||
pub(crate) use credential_injector::inject_credential;
|
||||
pub use credential_injector::{
|
||||
CredentialInjector, InjectedCredentials, InjectionError, SharedCredentialRegistry,
|
||||
};
|
||||
pub use rate_limiter::{LimitType, RateLimitError, RateLimitResult, RateLimiter};
|
||||
|
||||
// Storage (V2)
|
||||
|
||||
Reference in New Issue
Block a user