mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
Add WASM sandbox secure API extension
Extends the WASM sandbox with HTTP API capabilities, secrets management, tool aliasing, and leak detection. Key security principle: WASM never sees credentials, injection happens at host boundary. New modules: - secrets: AES-256-GCM encrypted storage with HKDF key derivation - leak_detector: Aho-Corasick + regex pattern matching for secret exfiltration - capabilities: Extended capability system (HTTP, ToolInvoke, Secrets) - allowlist: HTTP endpoint validation with glob patterns - credential_injector: Host-boundary credential injection - rate_limiter: Sliding window per-tool rate limiting - storage: WASM binary storage with BLAKE3 integrity verification Leak detection happens at two points: 1. Before HTTP request (prevents exfiltration via URL/headers/body) 2. After response (prevents exposure in outputs returned to WASM) Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.5
parent
45bbfa026d
commit
32bfd24154
@@ -0,0 +1,357 @@
|
||||
//! HTTP endpoint allowlist validation.
|
||||
//!
|
||||
//! Validates that HTTP requests from WASM tools only go to allowed endpoints.
|
||||
//! This is the first line of defense against unauthorized API access.
|
||||
//!
|
||||
//! # Validation Flow
|
||||
//!
|
||||
//! ```text
|
||||
//! WASM HTTP request ──► Parse URL ──► Check allowlist ──► Allow/Deny
|
||||
//! │ │
|
||||
//! │ ├─► Host match?
|
||||
//! │ ├─► Path prefix match?
|
||||
//! │ └─► Method allowed?
|
||||
//! │
|
||||
//! └─► Validate URL format
|
||||
//! ```
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use crate::tools::wasm::capabilities::EndpointPattern;
|
||||
|
||||
/// Result of allowlist validation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AllowlistResult {
|
||||
/// Request is allowed.
|
||||
Allowed,
|
||||
/// Request is denied with reason.
|
||||
Denied(DenyReason),
|
||||
}
|
||||
|
||||
impl AllowlistResult {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, AllowlistResult::Allowed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Reason why a request was denied.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DenyReason {
|
||||
/// URL could not be parsed.
|
||||
InvalidUrl(String),
|
||||
/// Host is not in the allowlist.
|
||||
HostNotAllowed(String),
|
||||
/// Path does not match any allowed prefix.
|
||||
PathNotAllowed { host: String, path: String },
|
||||
/// HTTP method is not allowed for this endpoint.
|
||||
MethodNotAllowed { method: String, host: String },
|
||||
/// Allowlist is empty (no endpoints configured).
|
||||
EmptyAllowlist,
|
||||
/// URL scheme is not HTTPS.
|
||||
InsecureScheme(String),
|
||||
}
|
||||
|
||||
impl fmt::Display for DenyReason {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
DenyReason::InvalidUrl(url) => write!(f, "Invalid URL: {}", url),
|
||||
DenyReason::HostNotAllowed(host) => write!(f, "Host not in allowlist: {}", host),
|
||||
DenyReason::PathNotAllowed { host, path } => {
|
||||
write!(f, "Path not allowed for host {}: {}", host, path)
|
||||
}
|
||||
DenyReason::MethodNotAllowed { method, host } => {
|
||||
write!(f, "Method {} not allowed for host {}", method, host)
|
||||
}
|
||||
DenyReason::EmptyAllowlist => write!(f, "No endpoints in allowlist"),
|
||||
DenyReason::InsecureScheme(scheme) => {
|
||||
write!(f, "Insecure scheme: {} (only HTTPS allowed)", scheme)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validates HTTP requests against an allowlist.
|
||||
pub struct AllowlistValidator {
|
||||
patterns: Vec<EndpointPattern>,
|
||||
/// Whether to require HTTPS (default: true).
|
||||
require_https: bool,
|
||||
}
|
||||
|
||||
impl AllowlistValidator {
|
||||
/// Create a new validator with the given patterns.
|
||||
pub fn new(patterns: Vec<EndpointPattern>) -> Self {
|
||||
Self {
|
||||
patterns,
|
||||
require_https: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Allow HTTP (insecure) requests. Use with caution.
|
||||
pub fn allow_http(mut self) -> Self {
|
||||
self.require_https = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// Check if a request is allowed.
|
||||
pub fn validate(&self, url: &str, method: &str) -> AllowlistResult {
|
||||
// Check for empty allowlist
|
||||
if self.patterns.is_empty() {
|
||||
return AllowlistResult::Denied(DenyReason::EmptyAllowlist);
|
||||
}
|
||||
|
||||
// Parse the URL
|
||||
let parsed = match parse_url(url) {
|
||||
Ok(p) => p,
|
||||
Err(e) => return AllowlistResult::Denied(DenyReason::InvalidUrl(e)),
|
||||
};
|
||||
|
||||
// Check HTTPS requirement
|
||||
if self.require_https && parsed.scheme != "https" {
|
||||
return AllowlistResult::Denied(DenyReason::InsecureScheme(parsed.scheme.clone()));
|
||||
}
|
||||
|
||||
// Find a matching pattern
|
||||
for pattern in &self.patterns {
|
||||
if pattern.matches(&parsed.host, &parsed.path, method) {
|
||||
return AllowlistResult::Allowed;
|
||||
}
|
||||
}
|
||||
|
||||
// No pattern matched, figure out why for better error messages
|
||||
let host_matches: Vec<_> = self
|
||||
.patterns
|
||||
.iter()
|
||||
.filter(|p| p.host_matches(&parsed.host))
|
||||
.collect();
|
||||
|
||||
if host_matches.is_empty() {
|
||||
AllowlistResult::Denied(DenyReason::HostNotAllowed(parsed.host))
|
||||
} else {
|
||||
// Host matches but path/method doesn't
|
||||
let path_matches: Vec<_> = host_matches
|
||||
.iter()
|
||||
.filter(|p| {
|
||||
p.path_prefix.is_none()
|
||||
|| parsed
|
||||
.path
|
||||
.starts_with(p.path_prefix.as_deref().unwrap_or(""))
|
||||
})
|
||||
.collect();
|
||||
|
||||
if path_matches.is_empty() {
|
||||
AllowlistResult::Denied(DenyReason::PathNotAllowed {
|
||||
host: parsed.host,
|
||||
path: parsed.path,
|
||||
})
|
||||
} else {
|
||||
AllowlistResult::Denied(DenyReason::MethodNotAllowed {
|
||||
method: method.to_string(),
|
||||
host: parsed.host,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if any pattern would allow this host.
|
||||
pub fn host_allowed(&self, host: &str) -> bool {
|
||||
self.patterns.iter().any(|p| p.host_matches(host))
|
||||
}
|
||||
|
||||
/// Get all allowed hosts (for debugging/logging).
|
||||
pub fn allowed_hosts(&self) -> Vec<&str> {
|
||||
self.patterns.iter().map(|p| p.host.as_str()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Parsed URL components.
|
||||
struct ParsedUrl {
|
||||
scheme: String,
|
||||
host: String,
|
||||
path: String,
|
||||
}
|
||||
|
||||
/// Simple URL parser (avoids pulling in a full URL crate).
|
||||
fn parse_url(url: &str) -> Result<ParsedUrl, String> {
|
||||
// Find scheme
|
||||
let (scheme, rest) = url
|
||||
.split_once("://")
|
||||
.ok_or_else(|| "Missing scheme (expected http:// or https://)".to_string())?;
|
||||
|
||||
let scheme = scheme.to_lowercase();
|
||||
if scheme != "http" && scheme != "https" {
|
||||
return Err(format!("Unsupported scheme: {}", scheme));
|
||||
}
|
||||
|
||||
// Split host from path
|
||||
let (host_and_port, path) = match rest.find('/') {
|
||||
Some(idx) => (&rest[..idx], &rest[idx..]),
|
||||
None => (rest, "/"),
|
||||
};
|
||||
|
||||
// Remove port from host
|
||||
let host = match host_and_port.rfind(':') {
|
||||
Some(idx) => {
|
||||
// Make sure this isn't an IPv6 address
|
||||
if host_and_port.starts_with('[') {
|
||||
// IPv6: [::1]:8080 or [::1]
|
||||
if let Some(bracket_idx) = host_and_port.find(']') {
|
||||
// Extract the IPv6 address without brackets
|
||||
&host_and_port[1..bracket_idx]
|
||||
} else {
|
||||
return Err("Invalid IPv6 address".to_string());
|
||||
}
|
||||
} else {
|
||||
&host_and_port[..idx]
|
||||
}
|
||||
}
|
||||
None => host_and_port,
|
||||
};
|
||||
|
||||
// Validate host
|
||||
if host.is_empty() {
|
||||
return Err("Empty host".to_string());
|
||||
}
|
||||
|
||||
Ok(ParsedUrl {
|
||||
scheme,
|
||||
host: host.to_lowercase(),
|
||||
path: path.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::allowlist::{AllowlistValidator, DenyReason};
|
||||
use crate::tools::wasm::capabilities::EndpointPattern;
|
||||
|
||||
fn validator_with_patterns() -> AllowlistValidator {
|
||||
AllowlistValidator::new(vec![
|
||||
EndpointPattern::host("api.openai.com").with_path_prefix("/v1/"),
|
||||
EndpointPattern::host("api.anthropic.com")
|
||||
.with_path_prefix("/v1/messages")
|
||||
.with_methods(vec!["POST".to_string()]),
|
||||
EndpointPattern::host("*.example.com"),
|
||||
])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allowed_request() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("https://api.openai.com/v1/chat/completions", "POST");
|
||||
assert!(result.is_allowed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denied_wrong_host() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("https://evil.com/steal/data", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::HostNotAllowed(_)));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denied_wrong_path() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("https://api.openai.com/v2/different", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::PathNotAllowed { .. }));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_denied_wrong_method() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
// Anthropic endpoint only allows POST
|
||||
let result = validator.validate("https://api.anthropic.com/v1/messages", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::MethodNotAllowed { .. }));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_host() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("https://api.example.com/anything", "GET");
|
||||
assert!(result.is_allowed());
|
||||
|
||||
let result = validator.validate("https://sub.api.example.com/anything", "GET");
|
||||
assert!(result.is_allowed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_require_https() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("http://api.openai.com/v1/chat", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::InsecureScheme(_)));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allow_http() {
|
||||
let validator = validator_with_patterns().allow_http();
|
||||
|
||||
let result = validator.validate("http://api.example.com/test", "GET");
|
||||
assert!(result.is_allowed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_allowlist() {
|
||||
let validator = AllowlistValidator::new(vec![]);
|
||||
|
||||
let result = validator.validate("https://anything.com/", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::EmptyAllowlist));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_url() {
|
||||
let validator = validator_with_patterns();
|
||||
|
||||
let result = validator.validate("not-a-url", "GET");
|
||||
assert!(!result.is_allowed());
|
||||
|
||||
if let super::AllowlistResult::Denied(reason) = result {
|
||||
assert!(matches!(reason, DenyReason::InvalidUrl(_)));
|
||||
} else {
|
||||
panic!("Expected denied");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_url_with_port() {
|
||||
let validator =
|
||||
AllowlistValidator::new(vec![EndpointPattern::host("localhost")]).allow_http();
|
||||
|
||||
let result = validator.validate("http://localhost:8080/api", "GET");
|
||||
assert!(result.is_allowed());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,423 @@
|
||||
//! Extended capabilities for WASM sandbox.
|
||||
//!
|
||||
//! Defines the capability system that controls what a WASM tool can do.
|
||||
//! All capabilities are opt-in; tools have NO access by default.
|
||||
//!
|
||||
//! # Capability Types
|
||||
//!
|
||||
//! - **Workspace**: Read files from the agent's workspace
|
||||
//! - **HTTP**: Make HTTP requests to allowlisted endpoints
|
||||
//! - **ToolInvoke**: Call other tools via aliases
|
||||
//! - **Secrets**: Check if secrets exist (never read values)
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::secrets::CredentialMapping;
|
||||
|
||||
/// All capabilities that can be granted to a WASM tool.
|
||||
///
|
||||
/// By default, all capabilities are `None` (disabled).
|
||||
/// Each must be explicitly granted.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct Capabilities {
|
||||
/// Read files from workspace.
|
||||
pub workspace_read: Option<WorkspaceCapability>,
|
||||
/// Make HTTP requests.
|
||||
pub http: Option<HttpCapability>,
|
||||
/// Invoke other tools.
|
||||
pub tool_invoke: Option<ToolInvokeCapability>,
|
||||
/// Check if secrets exist.
|
||||
pub secrets: Option<SecretsCapability>,
|
||||
}
|
||||
|
||||
impl Capabilities {
|
||||
/// Create capabilities with no permissions.
|
||||
pub fn none() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Enable workspace read with the given allowed prefixes.
|
||||
pub fn with_workspace_read(mut self, prefixes: Vec<String>) -> Self {
|
||||
self.workspace_read = Some(WorkspaceCapability {
|
||||
allowed_prefixes: prefixes,
|
||||
reader: None,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable HTTP requests with the given configuration.
|
||||
pub fn with_http(mut self, http: HttpCapability) -> Self {
|
||||
self.http = Some(http);
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable tool invocation with the given aliases.
|
||||
pub fn with_tool_invoke(mut self, aliases: HashMap<String, String>) -> Self {
|
||||
self.tool_invoke = Some(ToolInvokeCapability {
|
||||
aliases,
|
||||
rate_limit: RateLimitConfig::default(),
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable secret existence checks.
|
||||
pub fn with_secrets(mut self, allowed: Vec<String>) -> Self {
|
||||
self.secrets = Some(SecretsCapability {
|
||||
allowed_names: allowed,
|
||||
});
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Workspace read capability configuration.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct WorkspaceCapability {
|
||||
/// Allowed path prefixes (e.g., ["context/", "daily/"]).
|
||||
/// Empty means all paths allowed (within safety constraints).
|
||||
pub allowed_prefixes: Vec<String>,
|
||||
/// Function to actually read from workspace.
|
||||
/// This is injected by the runtime to avoid coupling to workspace impl.
|
||||
pub reader: Option<Arc<dyn WorkspaceReader>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for WorkspaceCapability {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("WorkspaceCapability")
|
||||
.field("allowed_prefixes", &self.allowed_prefixes)
|
||||
.field("reader", &self.reader.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for reading from workspace (allows mocking in tests).
|
||||
pub trait WorkspaceReader: Send + Sync {
|
||||
fn read(&self, path: &str) -> Option<String>;
|
||||
}
|
||||
|
||||
/// HTTP request capability configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HttpCapability {
|
||||
/// Allowed endpoint patterns.
|
||||
pub allowlist: Vec<EndpointPattern>,
|
||||
/// Credential mappings (secret name -> injection location).
|
||||
pub credentials: HashMap<String, CredentialMapping>,
|
||||
/// Rate limiting configuration.
|
||||
pub rate_limit: RateLimitConfig,
|
||||
/// Maximum request body size in bytes.
|
||||
pub max_request_bytes: usize,
|
||||
/// Maximum response body size in bytes.
|
||||
pub max_response_bytes: usize,
|
||||
/// Request timeout.
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for HttpCapability {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
allowlist: Vec::new(),
|
||||
credentials: HashMap::new(),
|
||||
rate_limit: RateLimitConfig::default(),
|
||||
max_request_bytes: 1024 * 1024, // 1 MB
|
||||
max_response_bytes: 10 * 1024 * 1024, // 10 MB
|
||||
timeout: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpCapability {
|
||||
/// Create a new HTTP capability with an allowlist.
|
||||
pub fn new(allowlist: Vec<EndpointPattern>) -> Self {
|
||||
Self {
|
||||
allowlist,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a credential mapping.
|
||||
pub fn with_credential(mut self, name: impl Into<String>, mapping: CredentialMapping) -> Self {
|
||||
self.credentials.insert(name.into(), mapping);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set rate limiting.
|
||||
pub fn with_rate_limit(mut self, rate_limit: RateLimitConfig) -> Self {
|
||||
self.rate_limit = rate_limit;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set request timeout.
|
||||
pub fn with_timeout(mut self, timeout: Duration) -> Self {
|
||||
self.timeout = timeout;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set max request body size.
|
||||
pub fn with_max_request_bytes(mut self, bytes: usize) -> Self {
|
||||
self.max_request_bytes = bytes;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set max response body size.
|
||||
pub fn with_max_response_bytes(mut self, bytes: usize) -> Self {
|
||||
self.max_response_bytes = bytes;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Pattern for matching allowed HTTP endpoints.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EndpointPattern {
|
||||
/// Hostname pattern (e.g., "api.example.com", "*.example.com").
|
||||
pub host: String,
|
||||
/// Path prefix (e.g., "/v1/", "/api/").
|
||||
pub path_prefix: Option<String>,
|
||||
/// Allowed HTTP methods (empty = all methods allowed).
|
||||
pub methods: Vec<String>,
|
||||
}
|
||||
|
||||
impl EndpointPattern {
|
||||
/// Create a pattern for a specific host.
|
||||
pub fn host(host: impl Into<String>) -> Self {
|
||||
Self {
|
||||
host: host.into(),
|
||||
path_prefix: None,
|
||||
methods: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a path prefix constraint.
|
||||
pub fn with_path_prefix(mut self, prefix: impl Into<String>) -> Self {
|
||||
self.path_prefix = Some(prefix.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Restrict to specific HTTP methods.
|
||||
pub fn with_methods(mut self, methods: Vec<String>) -> Self {
|
||||
self.methods = methods;
|
||||
self
|
||||
}
|
||||
|
||||
/// Check if this pattern matches a URL and method.
|
||||
pub fn matches(&self, url_host: &str, url_path: &str, method: &str) -> bool {
|
||||
// Check host
|
||||
if !self.host_matches(url_host) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check path prefix
|
||||
if let Some(ref prefix) = self.path_prefix {
|
||||
if !url_path.starts_with(prefix) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Check method
|
||||
if !self.methods.is_empty() {
|
||||
let method_upper = method.to_uppercase();
|
||||
if !self
|
||||
.methods
|
||||
.iter()
|
||||
.any(|m| m.to_uppercase() == method_upper)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// Check if host pattern matches (public for allowlist validation).
|
||||
pub fn host_matches(&self, url_host: &str) -> bool {
|
||||
if self.host == url_host {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Support wildcard: *.example.com matches sub.example.com
|
||||
if let Some(suffix) = self.host.strip_prefix("*.") {
|
||||
if url_host.ends_with(suffix) && url_host.len() > suffix.len() {
|
||||
// Ensure there's a dot before the suffix (or it's the whole thing)
|
||||
let prefix = &url_host[..url_host.len() - suffix.len()];
|
||||
if prefix.ends_with('.') || prefix.is_empty() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool invocation capability.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ToolInvokeCapability {
|
||||
/// Mapping from alias to real tool name.
|
||||
/// WASM calls tools by alias, never by real name.
|
||||
pub aliases: HashMap<String, String>,
|
||||
/// Rate limiting for tool calls.
|
||||
pub rate_limit: RateLimitConfig,
|
||||
}
|
||||
|
||||
impl ToolInvokeCapability {
|
||||
/// Create with a set of aliases.
|
||||
pub fn new(aliases: HashMap<String, String>) -> Self {
|
||||
Self {
|
||||
aliases,
|
||||
rate_limit: RateLimitConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve an alias to a real tool name.
|
||||
pub fn resolve_alias(&self, alias: &str) -> Option<&str> {
|
||||
self.aliases.get(alias).map(|s| s.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
/// Secrets capability (existence check only).
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SecretsCapability {
|
||||
/// Secret names this tool can check existence of.
|
||||
/// Supports glob: "openai_*" matches "openai_key", "openai_org".
|
||||
pub allowed_names: Vec<String>,
|
||||
}
|
||||
|
||||
impl SecretsCapability {
|
||||
/// Check if a secret name is allowed.
|
||||
pub fn is_allowed(&self, name: &str) -> bool {
|
||||
for pattern in &self.allowed_names {
|
||||
if pattern == name {
|
||||
return true;
|
||||
}
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
if name.starts_with(prefix) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Rate limiting configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RateLimitConfig {
|
||||
/// Maximum requests per minute.
|
||||
pub requests_per_minute: u32,
|
||||
/// Maximum requests per hour.
|
||||
pub requests_per_hour: u32,
|
||||
}
|
||||
|
||||
impl Default for RateLimitConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
requests_per_minute: 60,
|
||||
requests_per_hour: 1000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RateLimitConfig {
|
||||
/// Create a restrictive rate limit.
|
||||
pub fn restrictive() -> Self {
|
||||
Self {
|
||||
requests_per_minute: 10,
|
||||
requests_per_hour: 100,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a permissive rate limit.
|
||||
pub fn permissive() -> Self {
|
||||
Self {
|
||||
requests_per_minute: 120,
|
||||
requests_per_hour: 5000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::capabilities::{Capabilities, EndpointPattern, SecretsCapability};
|
||||
|
||||
#[test]
|
||||
fn test_capabilities_default_is_none() {
|
||||
let caps = Capabilities::default();
|
||||
assert!(caps.workspace_read.is_none());
|
||||
assert!(caps.http.is_none());
|
||||
assert!(caps.tool_invoke.is_none());
|
||||
assert!(caps.secrets.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_endpoint_pattern_exact_host() {
|
||||
let pattern = EndpointPattern::host("api.example.com");
|
||||
|
||||
assert!(pattern.matches("api.example.com", "/", "GET"));
|
||||
assert!(!pattern.matches("other.example.com", "/", "GET"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_endpoint_pattern_wildcard_host() {
|
||||
let pattern = EndpointPattern::host("*.example.com");
|
||||
|
||||
assert!(pattern.matches("api.example.com", "/", "GET"));
|
||||
assert!(pattern.matches("sub.api.example.com", "/", "GET"));
|
||||
assert!(!pattern.matches("example.com", "/", "GET"));
|
||||
assert!(!pattern.matches("notexample.com", "/", "GET"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_endpoint_pattern_path_prefix() {
|
||||
let pattern = EndpointPattern::host("api.example.com").with_path_prefix("/v1/");
|
||||
|
||||
assert!(pattern.matches("api.example.com", "/v1/users", "GET"));
|
||||
assert!(pattern.matches("api.example.com", "/v1/", "GET"));
|
||||
assert!(!pattern.matches("api.example.com", "/v2/users", "GET"));
|
||||
assert!(!pattern.matches("api.example.com", "/", "GET"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_endpoint_pattern_methods() {
|
||||
let pattern = EndpointPattern::host("api.example.com")
|
||||
.with_methods(vec!["GET".to_string(), "POST".to_string()]);
|
||||
|
||||
assert!(pattern.matches("api.example.com", "/", "GET"));
|
||||
assert!(pattern.matches("api.example.com", "/", "get")); // case insensitive
|
||||
assert!(pattern.matches("api.example.com", "/", "POST"));
|
||||
assert!(!pattern.matches("api.example.com", "/", "DELETE"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_secrets_capability_exact_match() {
|
||||
let cap = SecretsCapability {
|
||||
allowed_names: vec!["openai_key".to_string()],
|
||||
};
|
||||
|
||||
assert!(cap.is_allowed("openai_key"));
|
||||
assert!(!cap.is_allowed("anthropic_key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_secrets_capability_glob() {
|
||||
let cap = SecretsCapability {
|
||||
allowed_names: vec!["openai_*".to_string()],
|
||||
};
|
||||
|
||||
assert!(cap.is_allowed("openai_key"));
|
||||
assert!(cap.is_allowed("openai_org"));
|
||||
assert!(!cap.is_allowed("anthropic_key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_capabilities_builder() {
|
||||
let caps = Capabilities::none()
|
||||
.with_workspace_read(vec!["context/".to_string()])
|
||||
.with_secrets(vec!["test_*".to_string()]);
|
||||
|
||||
assert!(caps.workspace_read.is_some());
|
||||
assert!(caps.secrets.is_some());
|
||||
assert!(caps.http.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,429 @@
|
||||
//! Credential injection for WASM HTTP requests.
|
||||
//!
|
||||
//! Injects secrets into HTTP requests at the host boundary.
|
||||
//! WASM tools NEVER see the actual credential values.
|
||||
//!
|
||||
//! # Injection Flow
|
||||
//!
|
||||
//! ```text
|
||||
//! WASM requests HTTP ──► Host receives request ──► Match credentials by host
|
||||
//! │
|
||||
//! ┌───────────────────┘
|
||||
//! ▼
|
||||
//! Decrypt secret from store
|
||||
//! │
|
||||
//! ▼
|
||||
//! Inject into request:
|
||||
//! ├─► Authorization header (Bearer/Basic)
|
||||
//! ├─► Custom header (X-API-Key, etc.)
|
||||
//! └─► Query parameter
|
||||
//! │
|
||||
//! ▼
|
||||
//! Execute HTTP request
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::secrets::{
|
||||
CredentialLocation, CredentialMapping, DecryptedSecret, SecretError, SecretsStore,
|
||||
};
|
||||
|
||||
/// Error during credential injection.
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum InjectionError {
|
||||
#[error("Secret not found: {0}")]
|
||||
SecretNotFound(String),
|
||||
|
||||
#[error("Secret access denied: {0}")]
|
||||
AccessDenied(String),
|
||||
|
||||
#[error("Secret has expired: {0}")]
|
||||
SecretExpired(String),
|
||||
|
||||
#[error("Decryption failed: {0}")]
|
||||
DecryptionFailed(String),
|
||||
|
||||
#[error("No matching credential for host: {0}")]
|
||||
NoMatchingCredential(String),
|
||||
}
|
||||
|
||||
impl From<SecretError> for InjectionError {
|
||||
fn from(e: SecretError) -> Self {
|
||||
match e {
|
||||
SecretError::NotFound(name) => InjectionError::SecretNotFound(name),
|
||||
SecretError::Expired => InjectionError::SecretExpired("unknown".to_string()),
|
||||
SecretError::AccessDenied => InjectionError::AccessDenied("unknown".to_string()),
|
||||
SecretError::DecryptionFailed(msg) => InjectionError::DecryptionFailed(msg),
|
||||
_ => InjectionError::DecryptionFailed(e.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of credential injection.
|
||||
#[derive(Debug)]
|
||||
pub struct InjectedCredentials {
|
||||
/// Headers to add to the request.
|
||||
pub headers: HashMap<String, String>,
|
||||
/// Query parameters to add.
|
||||
pub query_params: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl InjectedCredentials {
|
||||
pub fn empty() -> Self {
|
||||
Self {
|
||||
headers: HashMap::new(),
|
||||
query_params: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.headers.is_empty() && self.query_params.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Injects credentials into HTTP requests.
|
||||
pub struct CredentialInjector {
|
||||
mappings: HashMap<String, CredentialMapping>,
|
||||
allowed_secrets: Vec<String>,
|
||||
}
|
||||
|
||||
impl CredentialInjector {
|
||||
/// Create a new injector with the given mappings.
|
||||
pub fn new(mappings: HashMap<String, CredentialMapping>, allowed_secrets: Vec<String>) -> Self {
|
||||
Self {
|
||||
mappings,
|
||||
allowed_secrets,
|
||||
}
|
||||
}
|
||||
|
||||
/// Find credentials that should be injected for a given host.
|
||||
pub fn find_credentials_for_host(&self, host: &str) -> Vec<&CredentialMapping> {
|
||||
self.mappings
|
||||
.values()
|
||||
.filter(|mapping| {
|
||||
mapping
|
||||
.host_patterns
|
||||
.iter()
|
||||
.any(|pattern| host_matches_pattern(host, pattern))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Inject credentials for an HTTP request.
|
||||
///
|
||||
/// Returns the headers and query params to add to the request.
|
||||
pub async fn inject(
|
||||
&self,
|
||||
user_id: &str,
|
||||
host: &str,
|
||||
store: &dyn SecretsStore,
|
||||
) -> Result<InjectedCredentials, InjectionError> {
|
||||
let matching_mappings = self.find_credentials_for_host(host);
|
||||
|
||||
if matching_mappings.is_empty() {
|
||||
// No credentials needed for this host
|
||||
return Ok(InjectedCredentials::empty());
|
||||
}
|
||||
|
||||
let mut result = InjectedCredentials::empty();
|
||||
|
||||
for mapping in matching_mappings {
|
||||
// Check if secret is in allowed list
|
||||
if !self.is_secret_allowed(&mapping.secret_name) {
|
||||
return Err(InjectionError::AccessDenied(mapping.secret_name.clone()));
|
||||
}
|
||||
|
||||
// Get the decrypted secret
|
||||
let secret = store
|
||||
.get_decrypted(user_id, &mapping.secret_name)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
SecretError::NotFound(name) => InjectionError::SecretNotFound(name),
|
||||
SecretError::Expired => {
|
||||
InjectionError::SecretExpired(mapping.secret_name.clone())
|
||||
}
|
||||
_ => InjectionError::DecryptionFailed(e.to_string()),
|
||||
})?;
|
||||
|
||||
// Inject based on location
|
||||
inject_credential(&mut result, &mapping.location, &secret);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Check if a secret name is in the allowed list.
|
||||
fn is_secret_allowed(&self, name: &str) -> bool {
|
||||
for pattern in &self.allowed_secrets {
|
||||
if pattern == name {
|
||||
return true;
|
||||
}
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
if name.starts_with(prefix) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Inject a single credential into the result.
|
||||
fn inject_credential(
|
||||
result: &mut InjectedCredentials,
|
||||
location: &CredentialLocation,
|
||||
secret: &DecryptedSecret,
|
||||
) {
|
||||
match location {
|
||||
CredentialLocation::AuthorizationBearer => {
|
||||
result.headers.insert(
|
||||
"Authorization".to_string(),
|
||||
format!("Bearer {}", secret.expose()),
|
||||
);
|
||||
}
|
||||
CredentialLocation::AuthorizationBasic { username } => {
|
||||
let credentials = format!("{}:{}", username, secret.expose());
|
||||
let encoded = base64_encode(credentials.as_bytes());
|
||||
result
|
||||
.headers
|
||||
.insert("Authorization".to_string(), format!("Basic {}", encoded));
|
||||
}
|
||||
CredentialLocation::Header { name, prefix } => {
|
||||
let value = match prefix {
|
||||
Some(p) => format!("{}{}", p, secret.expose()),
|
||||
None => secret.expose().to_string(),
|
||||
};
|
||||
result.headers.insert(name.clone(), value);
|
||||
}
|
||||
CredentialLocation::QueryParam { name } => {
|
||||
result
|
||||
.query_params
|
||||
.insert(name.clone(), secret.expose().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a host matches a pattern (supports wildcards).
|
||||
fn host_matches_pattern(host: &str, pattern: &str) -> bool {
|
||||
if pattern == host {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Support wildcard: *.example.com matches sub.example.com
|
||||
if let Some(suffix) = pattern.strip_prefix("*.") {
|
||||
if host.ends_with(suffix) && host.len() > suffix.len() {
|
||||
let prefix = &host[..host.len() - suffix.len()];
|
||||
if prefix.ends_with('.') || prefix.is_empty() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Simple base64 encoding (avoids extra dependency).
|
||||
fn base64_encode(input: &[u8]) -> String {
|
||||
const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
|
||||
let mut result = String::new();
|
||||
let mut i = 0;
|
||||
|
||||
while i < input.len() {
|
||||
let b0 = input[i];
|
||||
let b1 = if i + 1 < input.len() { input[i + 1] } else { 0 };
|
||||
let b2 = if i + 2 < input.len() { input[i + 2] } else { 0 };
|
||||
|
||||
result.push(ALPHABET[(b0 >> 2) as usize] as char);
|
||||
result.push(ALPHABET[(((b0 & 0x03) << 4) | (b1 >> 4)) as usize] as char);
|
||||
|
||||
if i + 1 < input.len() {
|
||||
result.push(ALPHABET[(((b1 & 0x0f) << 2) | (b2 >> 6)) as usize] as char);
|
||||
} else {
|
||||
result.push('=');
|
||||
}
|
||||
|
||||
if i + 2 < input.len() {
|
||||
result.push(ALPHABET[(b2 & 0x3f) as usize] as char);
|
||||
} else {
|
||||
result.push('=');
|
||||
}
|
||||
|
||||
i += 3;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use secrecy::SecretString;
|
||||
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, InMemorySecretsStore,
|
||||
SecretsCrypto, SecretsStore,
|
||||
};
|
||||
use crate::tools::wasm::credential_injector::{
|
||||
CredentialInjector, base64_encode, host_matches_pattern,
|
||||
};
|
||||
|
||||
fn test_store() -> InMemorySecretsStore {
|
||||
let key = "0123456789abcdef0123456789abcdef";
|
||||
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||
InMemorySecretsStore::new(crypto)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_host_matches_exact() {
|
||||
assert!(host_matches_pattern("api.openai.com", "api.openai.com"));
|
||||
assert!(!host_matches_pattern("api.openai.com", "other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_host_matches_wildcard() {
|
||||
assert!(host_matches_pattern("api.example.com", "*.example.com"));
|
||||
assert!(host_matches_pattern("sub.api.example.com", "*.example.com"));
|
||||
assert!(!host_matches_pattern("example.com", "*.example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_base64_encode() {
|
||||
assert_eq!(base64_encode(b"hello"), "aGVsbG8=");
|
||||
assert_eq!(base64_encode(b"user:pass"), "dXNlcjpwYXNz");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inject_bearer() {
|
||||
let store = test_store();
|
||||
store
|
||||
.create("user1", CreateSecretParams::new("openai_key", "sk-test123"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut mappings = HashMap::new();
|
||||
mappings.insert(
|
||||
"openai".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "openai_key".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["api.openai.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let injector = CredentialInjector::new(mappings, vec!["openai_key".to_string()]);
|
||||
let result = injector
|
||||
.inject("user1", "api.openai.com", &store)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.headers.get("Authorization"),
|
||||
Some(&"Bearer sk-test123".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inject_custom_header() {
|
||||
let store = test_store();
|
||||
store
|
||||
.create("user1", CreateSecretParams::new("api_key", "secret123"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut mappings = HashMap::new();
|
||||
mappings.insert(
|
||||
"custom".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "api_key".to_string(),
|
||||
location: CredentialLocation::Header {
|
||||
name: "X-API-Key".to_string(),
|
||||
prefix: None,
|
||||
},
|
||||
host_patterns: vec!["*.example.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let injector = CredentialInjector::new(mappings, vec!["api_key".to_string()]);
|
||||
let result = injector
|
||||
.inject("user1", "api.example.com", &store)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.headers.get("X-API-Key"),
|
||||
Some(&"secret123".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_inject_basic_auth() {
|
||||
let store = test_store();
|
||||
store
|
||||
.create("user1", CreateSecretParams::new("password", "mypassword"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut mappings = HashMap::new();
|
||||
mappings.insert(
|
||||
"basic".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "password".to_string(),
|
||||
location: CredentialLocation::AuthorizationBasic {
|
||||
username: "myuser".to_string(),
|
||||
},
|
||||
host_patterns: vec!["api.service.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
let injector = CredentialInjector::new(mappings, vec!["password".to_string()]);
|
||||
let result = injector
|
||||
.inject("user1", "api.service.com", &store)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// myuser:mypassword base64 encoded
|
||||
let expected = format!("Basic {}", base64_encode(b"myuser:mypassword"));
|
||||
assert_eq!(result.headers.get("Authorization"), Some(&expected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_no_credentials_for_host() {
|
||||
let store = test_store();
|
||||
|
||||
let injector = CredentialInjector::new(HashMap::new(), vec![]);
|
||||
let result = injector
|
||||
.inject("user1", "unknown.com", &store)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_access_denied_for_secret() {
|
||||
let store = test_store();
|
||||
store
|
||||
.create("user1", CreateSecretParams::new("secret_key", "value"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut mappings = HashMap::new();
|
||||
mappings.insert(
|
||||
"test".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "secret_key".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["api.test.com".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
// Empty allowed list = nothing allowed
|
||||
let injector = CredentialInjector::new(mappings, vec![]);
|
||||
let result = injector.inject("user1", "api.test.com", &store).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
+253
-42
@@ -2,10 +2,29 @@
|
||||
//!
|
||||
//! Implements a minimal, security-focused host API following VMLogic patterns
|
||||
//! from NEAR blockchain. The principle is: deny by default, grant minimal capabilities.
|
||||
//!
|
||||
//! # Extended API (V2)
|
||||
//!
|
||||
//! In addition to the basic log/time/workspace functions, the host now provides:
|
||||
//!
|
||||
//! - **http_request**: Make HTTP requests to allowlisted endpoints with credential injection
|
||||
//! - **tool_invoke**: Call other tools via aliases
|
||||
//! - **secret_exists**: Check if a secret exists (never read values)
|
||||
//!
|
||||
//! # Security Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! WASM Tool ──▶ Host Function ──▶ Allowlist ──▶ Credential ──▶ Execute
|
||||
//! (untrusted) (boundary) Validator Injector Request
|
||||
//! │
|
||||
//! ▼
|
||||
//! ◀────── Leak Detector ◀────── Response
|
||||
//! (sanitized, no secrets)
|
||||
//! ```
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::error::WasmError;
|
||||
|
||||
/// Maximum log entries per execution (prevents log spam attacks).
|
||||
@@ -44,46 +63,10 @@ pub struct LogEntry {
|
||||
pub timestamp_millis: u64,
|
||||
}
|
||||
|
||||
/// Capabilities that can be granted to a WASM tool.
|
||||
///
|
||||
/// By default, tools have NO capabilities. Each must be explicitly granted.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct Capabilities {
|
||||
/// If Some, tool can read from workspace at these paths.
|
||||
/// Empty vec means workspace access granted but no paths allowed yet.
|
||||
/// None means workspace access completely disabled.
|
||||
pub workspace_read: Option<WorkspaceCapability>,
|
||||
}
|
||||
|
||||
/// Workspace read capability configuration.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct WorkspaceCapability {
|
||||
/// Allowed path prefixes (e.g., ["context/", "daily/"]).
|
||||
/// Empty means all paths allowed (within safety constraints).
|
||||
pub allowed_prefixes: Vec<String>,
|
||||
/// Function to actually read from workspace.
|
||||
/// This is injected by the runtime to avoid coupling to workspace impl.
|
||||
pub reader: Option<Arc<dyn WorkspaceReader>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for WorkspaceCapability {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("WorkspaceCapability")
|
||||
.field("allowed_prefixes", &self.allowed_prefixes)
|
||||
.field("reader", &self.reader.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for reading from workspace (allows mocking in tests).
|
||||
pub trait WorkspaceReader: Send + Sync {
|
||||
fn read(&self, path: &str) -> Option<String>;
|
||||
}
|
||||
|
||||
/// Host state maintained during WASM execution.
|
||||
///
|
||||
/// This is the "VMLogic" equivalent, it tracks all side effects and enforces limits.
|
||||
#[derive(Debug)]
|
||||
/// Extended in V2 to support HTTP requests, tool invocation, and secret checks.
|
||||
pub struct HostState {
|
||||
/// Collected log entries.
|
||||
logs: Vec<LogEntry>,
|
||||
@@ -93,6 +76,25 @@ pub struct HostState {
|
||||
capabilities: Capabilities,
|
||||
/// Count of log entries dropped due to rate limiting.
|
||||
logs_dropped: usize,
|
||||
/// User ID for secret/credential lookups.
|
||||
user_id: Option<String>,
|
||||
/// HTTP request count for rate limiting within this execution.
|
||||
http_request_count: u32,
|
||||
/// Tool invoke count for rate limiting within this execution.
|
||||
tool_invoke_count: u32,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HostState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("HostState")
|
||||
.field("logs_count", &self.logs.len())
|
||||
.field("logging_enabled", &self.logging_enabled)
|
||||
.field("logs_dropped", &self.logs_dropped)
|
||||
.field("user_id", &self.user_id)
|
||||
.field("http_request_count", &self.http_request_count)
|
||||
.field("tool_invoke_count", &self.tool_invoke_count)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl HostState {
|
||||
@@ -103,6 +105,22 @@ impl HostState {
|
||||
logging_enabled: true,
|
||||
capabilities,
|
||||
logs_dropped: 0,
|
||||
user_id: None,
|
||||
http_request_count: 0,
|
||||
tool_invoke_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new host state with user context.
|
||||
pub fn new_with_user(capabilities: Capabilities, user_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
logs: Vec::new(),
|
||||
logging_enabled: true,
|
||||
capabilities,
|
||||
logs_dropped: 0,
|
||||
user_id: Some(user_id.into()),
|
||||
http_request_count: 0,
|
||||
tool_invoke_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -111,6 +129,16 @@ impl HostState {
|
||||
Self::new(Capabilities::default())
|
||||
}
|
||||
|
||||
/// Get the user ID if set.
|
||||
pub fn user_id(&self) -> Option<&str> {
|
||||
self.user_id.as_deref()
|
||||
}
|
||||
|
||||
/// Get the capabilities.
|
||||
pub fn capabilities(&self) -> &Capabilities {
|
||||
&self.capabilities
|
||||
}
|
||||
|
||||
/// Log a message from WASM.
|
||||
///
|
||||
/// Returns Ok(()) if logged, Err if rate limited or too long.
|
||||
@@ -204,6 +232,114 @@ impl HostState {
|
||||
pub fn logs_dropped(&self) -> usize {
|
||||
self.logs_dropped
|
||||
}
|
||||
|
||||
/// Check if a secret exists (does not expose value).
|
||||
///
|
||||
/// Returns false if:
|
||||
/// - Secrets capability not granted
|
||||
/// - Secret name not in allowed list
|
||||
/// - User ID not set
|
||||
pub fn secret_exists(&self, name: &str) -> bool {
|
||||
let capability = match &self.capabilities.secrets {
|
||||
Some(cap) => cap,
|
||||
None => return false,
|
||||
};
|
||||
|
||||
// Check if name is allowed
|
||||
capability.is_allowed(name)
|
||||
}
|
||||
|
||||
/// Check if HTTP capability is available for a given URL and method.
|
||||
///
|
||||
/// Returns an error message if not allowed.
|
||||
pub fn check_http_allowed(&self, url: &str, method: &str) -> Result<(), String> {
|
||||
let capability = self
|
||||
.capabilities
|
||||
.http
|
||||
.as_ref()
|
||||
.ok_or_else(|| "HTTP capability not granted".to_string())?;
|
||||
|
||||
// Use the allowlist validator
|
||||
use crate::tools::wasm::allowlist::AllowlistValidator;
|
||||
|
||||
let validator = AllowlistValidator::new(capability.allowlist.clone());
|
||||
let result = validator.validate(url, method);
|
||||
|
||||
if result.is_allowed() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("HTTP request not allowed: {:?}", result))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if tool invocation is allowed for an alias.
|
||||
///
|
||||
/// Returns the real tool name if allowed, error otherwise.
|
||||
pub fn check_tool_invoke_allowed(&self, alias: &str) -> Result<String, String> {
|
||||
let capability = self
|
||||
.capabilities
|
||||
.tool_invoke
|
||||
.as_ref()
|
||||
.ok_or_else(|| "Tool invocation capability not granted".to_string())?;
|
||||
|
||||
capability
|
||||
.resolve_alias(alias)
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| format!("Unknown tool alias: {}", alias))
|
||||
}
|
||||
|
||||
/// Increment HTTP request counter and check rate limit.
|
||||
///
|
||||
/// Returns error if rate limit exceeded.
|
||||
pub fn record_http_request(&mut self) -> Result<(), String> {
|
||||
// Verify HTTP capability exists
|
||||
let _capability = self
|
||||
.capabilities
|
||||
.http
|
||||
.as_ref()
|
||||
.ok_or_else(|| "HTTP capability not granted".to_string())?;
|
||||
|
||||
self.http_request_count += 1;
|
||||
|
||||
// Simple per-execution rate limit (additional to global rate limiter)
|
||||
// This prevents a single execution from making too many requests
|
||||
const MAX_REQUESTS_PER_EXECUTION: u32 = 50;
|
||||
if self.http_request_count > MAX_REQUESTS_PER_EXECUTION {
|
||||
return Err(format!(
|
||||
"Too many HTTP requests in single execution (max {})",
|
||||
MAX_REQUESTS_PER_EXECUTION
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Increment tool invoke counter and check rate limit.
|
||||
///
|
||||
/// Returns error if rate limit exceeded.
|
||||
pub fn record_tool_invoke(&mut self) -> Result<(), String> {
|
||||
self.tool_invoke_count += 1;
|
||||
|
||||
const MAX_INVOKES_PER_EXECUTION: u32 = 20;
|
||||
if self.tool_invoke_count > MAX_INVOKES_PER_EXECUTION {
|
||||
return Err(format!(
|
||||
"Too many tool invocations in single execution (max {})",
|
||||
MAX_INVOKES_PER_EXECUTION
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get HTTP request count for this execution.
|
||||
pub fn http_request_count(&self) -> u32 {
|
||||
self.http_request_count
|
||||
}
|
||||
|
||||
/// Get tool invoke count for this execution.
|
||||
pub fn tool_invoke_count(&self) -> u32 {
|
||||
self.tool_invoke_count
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate a workspace path for security.
|
||||
@@ -243,12 +379,15 @@ fn validate_workspace_path(path: &str) -> Result<(), WasmError> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::host::{
|
||||
Capabilities, HostState, LogLevel, MAX_LOG_ENTRIES, MAX_LOG_MESSAGE_BYTES,
|
||||
WorkspaceCapability, WorkspaceReader, validate_workspace_path,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::tools::wasm::capabilities::{
|
||||
Capabilities, SecretsCapability, WorkspaceCapability, WorkspaceReader,
|
||||
};
|
||||
use crate::tools::wasm::host::{
|
||||
HostState, LogLevel, MAX_LOG_ENTRIES, MAX_LOG_MESSAGE_BYTES, validate_workspace_path,
|
||||
};
|
||||
|
||||
struct MockReader {
|
||||
content: String,
|
||||
}
|
||||
@@ -330,6 +469,7 @@ mod tests {
|
||||
allowed_prefixes: vec![],
|
||||
reader: Some(reader),
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let state = HostState::new(capabilities);
|
||||
@@ -348,6 +488,7 @@ mod tests {
|
||||
allowed_prefixes: vec!["context/".to_string()],
|
||||
reader: Some(reader),
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let state = HostState::new(capabilities);
|
||||
@@ -392,4 +533,74 @@ mod tests {
|
||||
assert!(validate_workspace_path("projects/alpha/notes.md").is_ok());
|
||||
assert!(validate_workspace_path("MEMORY.md").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_secret_exists_no_capability() {
|
||||
let state = HostState::minimal();
|
||||
assert!(!state.secret_exists("any_secret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_secret_exists_with_capability() {
|
||||
let capabilities = Capabilities {
|
||||
secrets: Some(SecretsCapability {
|
||||
allowed_names: vec!["openai_*".to_string(), "exact_name".to_string()],
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let state = HostState::new(capabilities);
|
||||
|
||||
// Glob match
|
||||
assert!(state.secret_exists("openai_key"));
|
||||
assert!(state.secret_exists("openai_org"));
|
||||
|
||||
// Exact match
|
||||
assert!(state.secret_exists("exact_name"));
|
||||
|
||||
// Not allowed
|
||||
assert!(!state.secret_exists("stripe_key"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_http_request_rate_limit() {
|
||||
// Create state with HTTP capability enabled
|
||||
let capabilities = Capabilities {
|
||||
http: Some(crate::tools::wasm::capabilities::HttpCapability::default()),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
|
||||
// Should allow up to 50 requests
|
||||
for _ in 0..50 {
|
||||
assert!(state.record_http_request().is_ok());
|
||||
}
|
||||
|
||||
// 51st should fail
|
||||
assert!(state.record_http_request().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_invoke_rate_limit() {
|
||||
// Create state with tool invoke capability enabled
|
||||
let capabilities = Capabilities {
|
||||
tool_invoke: Some(crate::tools::wasm::capabilities::ToolInvokeCapability::default()),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
|
||||
// Should allow up to 20 invocations
|
||||
for _ in 0..20 {
|
||||
assert!(state.record_tool_invoke().is_ok());
|
||||
}
|
||||
|
||||
// 21st should fail
|
||||
assert!(state.record_tool_invoke().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_with_user() {
|
||||
let state = HostState::new_with_user(Capabilities::default(), "user123");
|
||||
assert_eq!(state.user_id(), Some("user123"));
|
||||
}
|
||||
}
|
||||
|
||||
+50
-20
@@ -10,26 +10,23 @@
|
||||
//!
|
||||
//! - **Memory limits**: Memory growth is bounded via ResourceLimiter.
|
||||
//!
|
||||
//! - **Minimal host API**: Only log, time, and optional workspace read.
|
||||
//! - **Extended host API (V2)**: log, time, workspace, HTTP, tool invoke, secrets
|
||||
//!
|
||||
//! - **Capability-based security**: Features are opt-in via Capabilities.
|
||||
//!
|
||||
//! # Architecture
|
||||
//! # Architecture (V2)
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────────────────────────────┐
|
||||
//! │ Tool Registration │
|
||||
//! │ WASM bytes → Validate → Compile (AOT) → PreparedModule (cached) │
|
||||
//! └─────────────────────────────────────────────────────────────────────┘
|
||||
//! │
|
||||
//! ▼
|
||||
//! ┌─────────────────────────────────────────────────────────────────────┐
|
||||
//! │ Tool Execution │
|
||||
//! │ JSON params → WasmToolWrapper → Fresh Instance → Execute → Result │
|
||||
//! │ ↓ ↓ │
|
||||
//! │ ResourceLimiter HostState │
|
||||
//! │ (memory, fuel) (log, time, workspace) │
|
||||
//! └─────────────────────────────────────────────────────────────────────┘
|
||||
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||
//! │ WASM Tool Execution │
|
||||
//! │ │
|
||||
//! │ WASM Tool ──▶ Host Function ──▶ Allowlist ──▶ Credential ──▶ Execute │
|
||||
//! │ (untrusted) (boundary) Validator Injector Request │
|
||||
//! │ │ │
|
||||
//! │ ▼ │
|
||||
//! │ ◀────── Leak Detector ◀────── Response │
|
||||
//! │ (sanitized, no secrets) │
|
||||
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! # Security Constraints
|
||||
@@ -40,17 +37,22 @@
|
||||
//! | Memory exhaustion | ResourceLimiter, 10MB default |
|
||||
//! | Infinite loops | Epoch interruption + tokio timeout |
|
||||
//! | Filesystem access | No WASI FS, only host workspace_read |
|
||||
//! | Network access | No network host functions |
|
||||
//! | Network access | Allowlisted endpoints only |
|
||||
//! | Credential exposure | Injection at host boundary only |
|
||||
//! | Secret exfiltration | Leak detector scans all outputs |
|
||||
//! | Log spam | Max 1000 entries, 4KB per message |
|
||||
//! | Path traversal | Validate paths (no `..`, no `/` prefix) |
|
||||
//! | Trap recovery | Discard instance, never reuse |
|
||||
//! | Side channels | Fresh instance per execution |
|
||||
//! | Rate abuse | Per-tool rate limiting |
|
||||
//! | WASM tampering | BLAKE3 hash verification on load |
|
||||
//! | Direct tool access | Tool aliasing (indirection layer) |
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use near_agent::tools::wasm::{WasmToolRuntime, WasmRuntimeConfig, WasmToolWrapper};
|
||||
//! use near_agent::tools::wasm::host::Capabilities;
|
||||
//! use near_agent::tools::wasm::Capabilities;
|
||||
//! use std::sync::Arc;
|
||||
//!
|
||||
//! // Create runtime
|
||||
@@ -60,24 +62,52 @@
|
||||
//! let wasm_bytes = std::fs::read("my_tool.wasm")?;
|
||||
//! let prepared = runtime.prepare("my_tool", &wasm_bytes, None).await?;
|
||||
//!
|
||||
//! // Create wrapper with minimal capabilities
|
||||
//! let tool = WasmToolWrapper::new(runtime, prepared, Capabilities::default());
|
||||
//! // Create wrapper with HTTP capability
|
||||
//! let capabilities = Capabilities::none()
|
||||
//! .with_http(HttpCapability::new(vec![
|
||||
//! EndpointPattern::host("api.openai.com").with_path_prefix("/v1/"),
|
||||
//! ]));
|
||||
//! let tool = WasmToolWrapper::new(runtime, prepared, capabilities);
|
||||
//!
|
||||
//! // Execute (implements Tool trait)
|
||||
//! let output = tool.execute(serde_json::json!({"input": "test"}), &ctx).await?;
|
||||
//! ```
|
||||
|
||||
mod allowlist;
|
||||
mod capabilities;
|
||||
mod credential_injector;
|
||||
mod error;
|
||||
mod host;
|
||||
mod limits;
|
||||
mod rate_limiter;
|
||||
mod runtime;
|
||||
mod storage;
|
||||
mod wrapper;
|
||||
|
||||
// Core types
|
||||
pub use error::{TrapCode, TrapInfo, WasmError};
|
||||
pub use host::{Capabilities, HostState, LogEntry, LogLevel, WorkspaceCapability, WorkspaceReader};
|
||||
pub use host::{HostState, LogEntry, LogLevel};
|
||||
pub use limits::{
|
||||
DEFAULT_FUEL_LIMIT, DEFAULT_MEMORY_LIMIT, DEFAULT_TIMEOUT, FuelConfig, ResourceLimits,
|
||||
WasmResourceLimiter,
|
||||
};
|
||||
pub use runtime::{PreparedModule, WasmRuntimeConfig, WasmToolRuntime};
|
||||
pub use wrapper::WasmToolWrapper;
|
||||
|
||||
// Capabilities (V2)
|
||||
pub use capabilities::{
|
||||
Capabilities, EndpointPattern, HttpCapability, RateLimitConfig, SecretsCapability,
|
||||
ToolInvokeCapability, WorkspaceCapability, WorkspaceReader,
|
||||
};
|
||||
|
||||
// Security components (V2)
|
||||
pub use allowlist::{AllowlistResult, AllowlistValidator, DenyReason};
|
||||
pub use credential_injector::{CredentialInjector, InjectedCredentials, InjectionError};
|
||||
pub use rate_limiter::{LimitType, RateLimitError, RateLimitResult, RateLimiter};
|
||||
|
||||
// Storage (V2)
|
||||
pub use storage::{
|
||||
PostgresWasmToolStore, StoreToolParams, StoredCapabilities, StoredWasmTool,
|
||||
StoredWasmToolWithBinary, ToolStatus, TrustLevel, WasmStorageError, WasmToolStore,
|
||||
compute_binary_hash, verify_binary_integrity,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,422 @@
|
||||
//! Rate limiting for WASM tool operations.
|
||||
//!
|
||||
//! Provides per-tool rate limiting for HTTP requests and tool invocations.
|
||||
//! Uses a sliding window algorithm for smooth rate enforcement.
|
||||
//!
|
||||
//! # Rate Limit Algorithm
|
||||
//!
|
||||
//! Uses a simplified sliding window counter:
|
||||
//! - Track request counts for current minute and hour windows
|
||||
//! - Reset counters when window expires
|
||||
//! - Increment counter and check against limits
|
||||
//!
|
||||
//! # Persistence
|
||||
//!
|
||||
//! Rate limit state can be persisted to PostgreSQL for cross-process
|
||||
//! rate limiting (useful for distributed deployments).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::tools::wasm::capabilities::RateLimitConfig;
|
||||
|
||||
/// Result of a rate limit check.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum RateLimitResult {
|
||||
/// Request is allowed.
|
||||
Allowed {
|
||||
/// Remaining requests in the current minute.
|
||||
remaining_minute: u32,
|
||||
/// Remaining requests in the current hour.
|
||||
remaining_hour: u32,
|
||||
},
|
||||
/// Request is rate limited.
|
||||
Limited {
|
||||
/// When the rate limit will reset.
|
||||
retry_after: Duration,
|
||||
/// Which limit was exceeded.
|
||||
limit_type: LimitType,
|
||||
},
|
||||
}
|
||||
|
||||
impl RateLimitResult {
|
||||
pub fn is_allowed(&self) -> bool {
|
||||
matches!(self, RateLimitResult::Allowed { .. })
|
||||
}
|
||||
}
|
||||
|
||||
/// Which rate limit was exceeded.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LimitType {
|
||||
PerMinute,
|
||||
PerHour,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for LimitType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
LimitType::PerMinute => write!(f, "per-minute"),
|
||||
LimitType::PerHour => write!(f, "per-hour"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// State for a single rate limit window.
|
||||
#[derive(Debug, Clone)]
|
||||
struct WindowState {
|
||||
window_start: Instant,
|
||||
count: u32,
|
||||
}
|
||||
|
||||
impl WindowState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
window_start: Instant::now(),
|
||||
count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the window has expired and reset if needed.
|
||||
fn maybe_reset(&mut self, window_duration: Duration) {
|
||||
if self.window_start.elapsed() >= window_duration {
|
||||
self.window_start = Instant::now();
|
||||
self.count = 0;
|
||||
}
|
||||
}
|
||||
|
||||
/// Time until window resets.
|
||||
fn time_until_reset(&self, window_duration: Duration) -> Duration {
|
||||
let elapsed = self.window_start.elapsed();
|
||||
if elapsed >= window_duration {
|
||||
Duration::ZERO
|
||||
} else {
|
||||
window_duration - elapsed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Rate limit state for a single tool.
|
||||
#[derive(Debug)]
|
||||
struct ToolRateLimitState {
|
||||
minute_window: WindowState,
|
||||
hour_window: WindowState,
|
||||
}
|
||||
|
||||
impl ToolRateLimitState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
minute_window: WindowState::new(),
|
||||
hour_window: WindowState::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// In-memory rate limiter for WASM tools.
|
||||
pub struct RateLimiter {
|
||||
/// State per (user_id, tool_name).
|
||||
state: RwLock<HashMap<(String, String), ToolRateLimitState>>,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
/// Create a new rate limiter.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a request is allowed and record it if so.
|
||||
pub async fn check_and_record(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tool_name: &str,
|
||||
config: &RateLimitConfig,
|
||||
) -> RateLimitResult {
|
||||
let key = (user_id.to_string(), tool_name.to_string());
|
||||
|
||||
let mut state = self.state.write().await;
|
||||
let tool_state = state.entry(key).or_insert_with(ToolRateLimitState::new);
|
||||
|
||||
// Reset windows if expired
|
||||
tool_state
|
||||
.minute_window
|
||||
.maybe_reset(Duration::from_secs(60));
|
||||
tool_state
|
||||
.hour_window
|
||||
.maybe_reset(Duration::from_secs(3600));
|
||||
|
||||
// Check minute limit
|
||||
if tool_state.minute_window.count >= config.requests_per_minute {
|
||||
return RateLimitResult::Limited {
|
||||
retry_after: tool_state
|
||||
.minute_window
|
||||
.time_until_reset(Duration::from_secs(60)),
|
||||
limit_type: LimitType::PerMinute,
|
||||
};
|
||||
}
|
||||
|
||||
// Check hour limit
|
||||
if tool_state.hour_window.count >= config.requests_per_hour {
|
||||
return RateLimitResult::Limited {
|
||||
retry_after: tool_state
|
||||
.hour_window
|
||||
.time_until_reset(Duration::from_secs(3600)),
|
||||
limit_type: LimitType::PerHour,
|
||||
};
|
||||
}
|
||||
|
||||
// Record the request
|
||||
tool_state.minute_window.count += 1;
|
||||
tool_state.hour_window.count += 1;
|
||||
|
||||
RateLimitResult::Allowed {
|
||||
remaining_minute: config.requests_per_minute - tool_state.minute_window.count,
|
||||
remaining_hour: config.requests_per_hour - tool_state.hour_window.count,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check without recording (for preview/estimation).
|
||||
pub async fn check(
|
||||
&self,
|
||||
user_id: &str,
|
||||
tool_name: &str,
|
||||
config: &RateLimitConfig,
|
||||
) -> RateLimitResult {
|
||||
let key = (user_id.to_string(), tool_name.to_string());
|
||||
|
||||
let mut state = self.state.write().await;
|
||||
let tool_state = state.entry(key).or_insert_with(ToolRateLimitState::new);
|
||||
|
||||
// Reset windows if expired
|
||||
tool_state
|
||||
.minute_window
|
||||
.maybe_reset(Duration::from_secs(60));
|
||||
tool_state
|
||||
.hour_window
|
||||
.maybe_reset(Duration::from_secs(3600));
|
||||
|
||||
// Check minute limit
|
||||
if tool_state.minute_window.count >= config.requests_per_minute {
|
||||
return RateLimitResult::Limited {
|
||||
retry_after: tool_state
|
||||
.minute_window
|
||||
.time_until_reset(Duration::from_secs(60)),
|
||||
limit_type: LimitType::PerMinute,
|
||||
};
|
||||
}
|
||||
|
||||
// Check hour limit
|
||||
if tool_state.hour_window.count >= config.requests_per_hour {
|
||||
return RateLimitResult::Limited {
|
||||
retry_after: tool_state
|
||||
.hour_window
|
||||
.time_until_reset(Duration::from_secs(3600)),
|
||||
limit_type: LimitType::PerHour,
|
||||
};
|
||||
}
|
||||
|
||||
RateLimitResult::Allowed {
|
||||
remaining_minute: config.requests_per_minute - tool_state.minute_window.count,
|
||||
remaining_hour: config.requests_per_hour - tool_state.hour_window.count,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current usage for a tool.
|
||||
pub async fn get_usage(&self, user_id: &str, tool_name: &str) -> Option<(u32, u32)> {
|
||||
let key = (user_id.to_string(), tool_name.to_string());
|
||||
let state = self.state.read().await;
|
||||
|
||||
state
|
||||
.get(&key)
|
||||
.map(|s| (s.minute_window.count, s.hour_window.count))
|
||||
}
|
||||
|
||||
/// Clear rate limit state for a tool (for testing or manual reset).
|
||||
pub async fn clear(&self, user_id: &str, tool_name: &str) {
|
||||
let key = (user_id.to_string(), tool_name.to_string());
|
||||
self.state.write().await.remove(&key);
|
||||
}
|
||||
|
||||
/// Clear all rate limit state.
|
||||
pub async fn clear_all(&self) {
|
||||
self.state.write().await.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RateLimiter {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Error when rate limited.
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
#[error("Rate limited ({limit_type}), retry after {retry_after:?}")]
|
||||
pub struct RateLimitError {
|
||||
pub retry_after: Duration,
|
||||
pub limit_type: LimitType,
|
||||
}
|
||||
|
||||
impl From<RateLimitResult> for Result<(), RateLimitError> {
|
||||
fn from(result: RateLimitResult) -> Self {
|
||||
match result {
|
||||
RateLimitResult::Allowed { .. } => Ok(()),
|
||||
RateLimitResult::Limited {
|
||||
retry_after,
|
||||
limit_type,
|
||||
} => Err(RateLimitError {
|
||||
retry_after,
|
||||
limit_type,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::capabilities::RateLimitConfig;
|
||||
use crate::tools::wasm::rate_limiter::{LimitType, RateLimitResult, RateLimiter};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_allowed_within_limits() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 10,
|
||||
requests_per_hour: 100,
|
||||
};
|
||||
|
||||
let result = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
match result {
|
||||
RateLimitResult::Allowed {
|
||||
remaining_minute,
|
||||
remaining_hour,
|
||||
} => {
|
||||
assert_eq!(remaining_minute, 9);
|
||||
assert_eq!(remaining_hour, 99);
|
||||
}
|
||||
_ => panic!("Expected allowed"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_minute_limit_exceeded() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 2,
|
||||
requests_per_hour: 100,
|
||||
};
|
||||
|
||||
// Use up the minute limit
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
// Third request should be limited
|
||||
let result = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
match result {
|
||||
RateLimitResult::Limited {
|
||||
limit_type,
|
||||
retry_after,
|
||||
} => {
|
||||
assert_eq!(limit_type, LimitType::PerMinute);
|
||||
assert!(retry_after.as_secs() <= 60);
|
||||
}
|
||||
_ => panic!("Expected limited"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hour_limit_exceeded() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 100, // High minute limit
|
||||
requests_per_hour: 2, // Low hour limit
|
||||
};
|
||||
|
||||
// Use up the hour limit
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
// Third request should be limited
|
||||
let result = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
match result {
|
||||
RateLimitResult::Limited { limit_type, .. } => {
|
||||
assert_eq!(limit_type, LimitType::PerHour);
|
||||
}
|
||||
_ => panic!("Expected limited"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_user_isolation() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 1,
|
||||
requests_per_hour: 10,
|
||||
};
|
||||
|
||||
// User1 uses their limit
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
let result1 = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
// User2 should still have their limit
|
||||
let result2 = limiter.check_and_record("user2", "tool1", &config).await;
|
||||
|
||||
assert!(!result1.is_allowed());
|
||||
assert!(result2.is_allowed());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tool_isolation() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 1,
|
||||
requests_per_hour: 10,
|
||||
};
|
||||
|
||||
// Tool1 uses its limit
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
let result1 = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
// Tool2 should still have its limit
|
||||
let result2 = limiter.check_and_record("user1", "tool2", &config).await;
|
||||
|
||||
assert!(!result1.is_allowed());
|
||||
assert!(result2.is_allowed());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_usage() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig::default();
|
||||
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
|
||||
let usage = limiter.get_usage("user1", "tool1").await;
|
||||
assert_eq!(usage, Some((3, 3)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clear() {
|
||||
let limiter = RateLimiter::new();
|
||||
let config = RateLimitConfig {
|
||||
requests_per_minute: 1,
|
||||
requests_per_hour: 10,
|
||||
};
|
||||
|
||||
limiter.check_and_record("user1", "tool1", &config).await;
|
||||
let result1 = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
assert!(!result1.is_allowed());
|
||||
|
||||
limiter.clear("user1", "tool1").await;
|
||||
|
||||
let result2 = limiter.check_and_record("user1", "tool1", &config).await;
|
||||
assert!(result2.is_allowed());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,616 @@
|
||||
//! WASM binary storage with integrity verification.
|
||||
//!
|
||||
//! Stores compiled WASM tools in PostgreSQL with BLAKE3 hash verification.
|
||||
//! On load, the hash is verified to detect tampering.
|
||||
//!
|
||||
//! # Storage Flow
|
||||
//!
|
||||
//! ```text
|
||||
//! WASM bytes ──► BLAKE3 hash ──► Store in PostgreSQL
|
||||
//! │ (binary + hash)
|
||||
//! │
|
||||
//! └──► Later: Load ──► Verify hash ──► Return bytes
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use deadpool_postgres::Pool;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::tools::wasm::capabilities::{
|
||||
Capabilities, EndpointPattern, HttpCapability, RateLimitConfig, SecretsCapability,
|
||||
ToolInvokeCapability,
|
||||
};
|
||||
|
||||
/// Trust level for a WASM tool.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TrustLevel {
|
||||
/// Built-in system tool (highest trust).
|
||||
System,
|
||||
/// Audited and verified tool.
|
||||
Verified,
|
||||
/// User-uploaded tool (untrusted).
|
||||
User,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TrustLevel {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
TrustLevel::System => write!(f, "system"),
|
||||
TrustLevel::Verified => write!(f, "verified"),
|
||||
TrustLevel::User => write!(f, "user"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for TrustLevel {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"system" => Ok(TrustLevel::System),
|
||||
"verified" => Ok(TrustLevel::Verified),
|
||||
"user" => Ok(TrustLevel::User),
|
||||
_ => Err(format!("Unknown trust level: {}", s)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Status of a WASM tool.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolStatus {
|
||||
/// Tool is active and can be used.
|
||||
Active,
|
||||
/// Tool is disabled (manually or due to errors).
|
||||
Disabled,
|
||||
/// Tool is quarantined (suspected malicious).
|
||||
Quarantined,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ToolStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
ToolStatus::Active => write!(f, "active"),
|
||||
ToolStatus::Disabled => write!(f, "disabled"),
|
||||
ToolStatus::Quarantined => write!(f, "quarantined"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::str::FromStr for ToolStatus {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"active" => Ok(ToolStatus::Active),
|
||||
"disabled" => Ok(ToolStatus::Disabled),
|
||||
"quarantined" => Ok(ToolStatus::Quarantined),
|
||||
_ => Err(format!("Unknown status: {}", s)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A stored WASM tool.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoredWasmTool {
|
||||
pub id: Uuid,
|
||||
pub user_id: String,
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
pub description: String,
|
||||
pub parameters_schema: serde_json::Value,
|
||||
pub source_url: Option<String>,
|
||||
pub trust_level: TrustLevel,
|
||||
pub status: ToolStatus,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Full tool data including binary (not returned by default for efficiency).
|
||||
#[derive(Debug)]
|
||||
pub struct StoredWasmToolWithBinary {
|
||||
pub tool: StoredWasmTool,
|
||||
pub wasm_binary: Vec<u8>,
|
||||
pub binary_hash: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Capabilities stored in the database.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoredCapabilities {
|
||||
pub id: Uuid,
|
||||
pub wasm_tool_id: Uuid,
|
||||
pub http_allowlist: Vec<EndpointPattern>,
|
||||
pub allowed_secrets: Vec<String>,
|
||||
pub tool_aliases: HashMap<String, String>,
|
||||
pub requests_per_minute: u32,
|
||||
pub requests_per_hour: u32,
|
||||
pub max_request_body_bytes: i64,
|
||||
pub max_response_body_bytes: i64,
|
||||
pub workspace_read_prefixes: Vec<String>,
|
||||
pub http_timeout_secs: i32,
|
||||
}
|
||||
|
||||
impl StoredCapabilities {
|
||||
/// Convert to runtime Capabilities struct.
|
||||
pub fn to_capabilities(&self) -> Capabilities {
|
||||
let mut caps = Capabilities::default();
|
||||
|
||||
// Workspace read
|
||||
if !self.workspace_read_prefixes.is_empty() {
|
||||
caps = caps.with_workspace_read(self.workspace_read_prefixes.clone());
|
||||
}
|
||||
|
||||
// HTTP capability
|
||||
if !self.http_allowlist.is_empty() {
|
||||
caps.http = Some(HttpCapability {
|
||||
allowlist: self.http_allowlist.clone(),
|
||||
credentials: HashMap::new(), // Loaded separately
|
||||
rate_limit: RateLimitConfig {
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
requests_per_hour: self.requests_per_hour,
|
||||
},
|
||||
max_request_bytes: self.max_request_body_bytes as usize,
|
||||
max_response_bytes: self.max_response_body_bytes as usize,
|
||||
timeout: std::time::Duration::from_secs(self.http_timeout_secs as u64),
|
||||
});
|
||||
}
|
||||
|
||||
// Tool invoke capability
|
||||
if !self.tool_aliases.is_empty() {
|
||||
caps.tool_invoke = Some(ToolInvokeCapability {
|
||||
aliases: self.tool_aliases.clone(),
|
||||
rate_limit: RateLimitConfig {
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
requests_per_hour: self.requests_per_hour,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Secrets capability
|
||||
if !self.allowed_secrets.is_empty() {
|
||||
caps.secrets = Some(SecretsCapability {
|
||||
allowed_names: self.allowed_secrets.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
caps
|
||||
}
|
||||
}
|
||||
|
||||
/// Error from WASM storage operations.
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum WasmStorageError {
|
||||
#[error("Tool not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Tool is disabled")]
|
||||
Disabled,
|
||||
|
||||
#[error("Tool is quarantined")]
|
||||
Quarantined,
|
||||
|
||||
#[error("Binary integrity check failed: hash mismatch")]
|
||||
IntegrityCheckFailed,
|
||||
|
||||
#[error("Database error: {0}")]
|
||||
Database(String),
|
||||
|
||||
#[error("Invalid data: {0}")]
|
||||
InvalidData(String),
|
||||
}
|
||||
|
||||
/// Trait for WASM tool storage.
|
||||
#[async_trait]
|
||||
pub trait WasmToolStore: Send + Sync {
|
||||
/// Store a new WASM tool.
|
||||
async fn store(&self, params: StoreToolParams) -> Result<StoredWasmTool, WasmStorageError>;
|
||||
|
||||
/// Get tool metadata (without binary).
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<StoredWasmTool, WasmStorageError>;
|
||||
|
||||
/// Get tool with binary (verifies integrity).
|
||||
async fn get_with_binary(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<StoredWasmToolWithBinary, WasmStorageError>;
|
||||
|
||||
/// Get tool capabilities.
|
||||
async fn get_capabilities(
|
||||
&self,
|
||||
tool_id: Uuid,
|
||||
) -> Result<Option<StoredCapabilities>, WasmStorageError>;
|
||||
|
||||
/// List all tools for a user.
|
||||
async fn list(&self, user_id: &str) -> Result<Vec<StoredWasmTool>, WasmStorageError>;
|
||||
|
||||
/// Update tool status.
|
||||
async fn update_status(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
status: ToolStatus,
|
||||
) -> Result<(), WasmStorageError>;
|
||||
|
||||
/// Delete a tool.
|
||||
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, WasmStorageError>;
|
||||
}
|
||||
|
||||
/// Parameters for storing a new tool.
|
||||
pub struct StoreToolParams {
|
||||
pub user_id: String,
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
pub description: String,
|
||||
pub wasm_binary: Vec<u8>,
|
||||
pub parameters_schema: serde_json::Value,
|
||||
pub source_url: Option<String>,
|
||||
pub trust_level: TrustLevel,
|
||||
}
|
||||
|
||||
/// Compute BLAKE3 hash of WASM binary.
|
||||
pub fn compute_binary_hash(binary: &[u8]) -> Vec<u8> {
|
||||
let hash = blake3::hash(binary);
|
||||
hash.as_bytes().to_vec()
|
||||
}
|
||||
|
||||
/// Verify binary integrity against stored hash.
|
||||
pub fn verify_binary_integrity(binary: &[u8], expected_hash: &[u8]) -> bool {
|
||||
let actual_hash = compute_binary_hash(binary);
|
||||
actual_hash == expected_hash
|
||||
}
|
||||
|
||||
/// PostgreSQL implementation of WasmToolStore.
|
||||
pub struct PostgresWasmToolStore {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
impl PostgresWasmToolStore {
|
||||
pub fn new(pool: Pool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WasmToolStore for PostgresWasmToolStore {
|
||||
async fn store(&self, params: StoreToolParams) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let binary_hash = compute_binary_hash(¶ms.wasm_binary);
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
|
||||
let row = client
|
||||
.query_one(
|
||||
r#"
|
||||
INSERT INTO wasm_tools (
|
||||
id, user_id, name, version, description, wasm_binary, binary_hash,
|
||||
parameters_schema, source_url, trust_level, status, created_at, updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, 'active', $11, $11)
|
||||
ON CONFLICT (user_id, name, version) DO UPDATE SET
|
||||
description = EXCLUDED.description,
|
||||
wasm_binary = EXCLUDED.wasm_binary,
|
||||
binary_hash = EXCLUDED.binary_hash,
|
||||
parameters_schema = EXCLUDED.parameters_schema,
|
||||
source_url = EXCLUDED.source_url,
|
||||
updated_at = NOW()
|
||||
RETURNING id, user_id, name, version, description, parameters_schema,
|
||||
source_url, trust_level, status, created_at, updated_at
|
||||
"#,
|
||||
&[
|
||||
&id,
|
||||
¶ms.user_id,
|
||||
¶ms.name,
|
||||
¶ms.version,
|
||||
¶ms.description,
|
||||
¶ms.wasm_binary,
|
||||
&binary_hash,
|
||||
¶ms.parameters_schema,
|
||||
¶ms.source_url,
|
||||
¶ms.trust_level.to_string(),
|
||||
&now,
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
row_to_tool(&row)
|
||||
}
|
||||
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let row = client
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, parameters_schema,
|
||||
source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = $1 AND name = $2 AND status = 'active'
|
||||
ORDER BY version DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
&[&user_id, &name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match row {
|
||||
Some(r) => {
|
||||
let tool = row_to_tool(&r)?;
|
||||
match tool.status {
|
||||
ToolStatus::Active => Ok(tool),
|
||||
ToolStatus::Disabled => Err(WasmStorageError::Disabled),
|
||||
ToolStatus::Quarantined => Err(WasmStorageError::Quarantined),
|
||||
}
|
||||
}
|
||||
None => Err(WasmStorageError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_with_binary(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<StoredWasmToolWithBinary, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let row = client
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, wasm_binary, binary_hash,
|
||||
parameters_schema, source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = $1 AND name = $2 AND status = 'active'
|
||||
ORDER BY version DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
&[&user_id, &name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match row {
|
||||
Some(r) => {
|
||||
let wasm_binary: Vec<u8> = r.get("wasm_binary");
|
||||
let binary_hash: Vec<u8> = r.get("binary_hash");
|
||||
|
||||
// Verify integrity
|
||||
if !verify_binary_integrity(&wasm_binary, &binary_hash) {
|
||||
tracing::error!(
|
||||
user_id = user_id,
|
||||
name = name,
|
||||
"WASM binary integrity check failed"
|
||||
);
|
||||
return Err(WasmStorageError::IntegrityCheckFailed);
|
||||
}
|
||||
|
||||
let tool = row_to_tool(&r)?;
|
||||
|
||||
match tool.status {
|
||||
ToolStatus::Active => Ok(StoredWasmToolWithBinary {
|
||||
tool,
|
||||
wasm_binary,
|
||||
binary_hash,
|
||||
}),
|
||||
ToolStatus::Disabled => Err(WasmStorageError::Disabled),
|
||||
ToolStatus::Quarantined => Err(WasmStorageError::Quarantined),
|
||||
}
|
||||
}
|
||||
None => Err(WasmStorageError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_capabilities(
|
||||
&self,
|
||||
tool_id: Uuid,
|
||||
) -> Result<Option<StoredCapabilities>, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let row = client
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT id, wasm_tool_id, http_allowlist, allowed_secrets, tool_aliases,
|
||||
requests_per_minute, requests_per_hour, max_request_body_bytes,
|
||||
max_response_body_bytes, workspace_read_prefixes, http_timeout_secs
|
||||
FROM tool_capabilities
|
||||
WHERE wasm_tool_id = $1
|
||||
"#,
|
||||
&[&tool_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match row {
|
||||
Some(r) => {
|
||||
let http_allowlist_json: serde_json::Value = r.get("http_allowlist");
|
||||
let tool_aliases_json: serde_json::Value = r.get("tool_aliases");
|
||||
|
||||
let http_allowlist: Vec<EndpointPattern> =
|
||||
serde_json::from_value(http_allowlist_json).unwrap_or_default();
|
||||
let tool_aliases: HashMap<String, String> =
|
||||
serde_json::from_value(tool_aliases_json).unwrap_or_default();
|
||||
|
||||
Ok(Some(StoredCapabilities {
|
||||
id: r.get("id"),
|
||||
wasm_tool_id: r.get("wasm_tool_id"),
|
||||
http_allowlist,
|
||||
allowed_secrets: r.get("allowed_secrets"),
|
||||
tool_aliases,
|
||||
requests_per_minute: r.get::<_, i32>("requests_per_minute") as u32,
|
||||
requests_per_hour: r.get::<_, i32>("requests_per_hour") as u32,
|
||||
max_request_body_bytes: r.get("max_request_body_bytes"),
|
||||
max_response_body_bytes: r.get("max_response_body_bytes"),
|
||||
workspace_read_prefixes: r.get("workspace_read_prefixes"),
|
||||
http_timeout_secs: r.get("http_timeout_secs"),
|
||||
}))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list(&self, user_id: &str) -> Result<Vec<StoredWasmTool>, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let rows = client
|
||||
.query(
|
||||
r#"
|
||||
SELECT DISTINCT ON (name) id, user_id, name, version, description,
|
||||
parameters_schema, source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = $1
|
||||
ORDER BY name, version DESC
|
||||
"#,
|
||||
&[&user_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
rows.into_iter().map(|r| row_to_tool(&r)).collect()
|
||||
}
|
||||
|
||||
async fn update_status(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
status: ToolStatus,
|
||||
) -> Result<(), WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let result = client
|
||||
.execute(
|
||||
"UPDATE wasm_tools SET status = $1, updated_at = NOW() WHERE user_id = $2 AND name = $3",
|
||||
&[&status.to_string(), &user_id, &name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
if result == 0 {
|
||||
return Err(WasmStorageError::NotFound(name.to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, WasmStorageError> {
|
||||
let client = self
|
||||
.pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let result = client
|
||||
.execute(
|
||||
"DELETE FROM wasm_tools WHERE user_id = $1 AND name = $2",
|
||||
&[&user_id, &name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
Ok(result > 0)
|
||||
}
|
||||
}
|
||||
|
||||
fn row_to_tool(row: &tokio_postgres::Row) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let trust_level_str: String = row.get("trust_level");
|
||||
let status_str: String = row.get("status");
|
||||
|
||||
Ok(StoredWasmTool {
|
||||
id: row.get("id"),
|
||||
user_id: row.get("user_id"),
|
||||
name: row.get("name"),
|
||||
version: row.get("version"),
|
||||
description: row.get("description"),
|
||||
parameters_schema: row.get("parameters_schema"),
|
||||
source_url: row.get("source_url"),
|
||||
trust_level: trust_level_str
|
||||
.parse()
|
||||
.map_err(WasmStorageError::InvalidData)?,
|
||||
status: status_str
|
||||
.parse()
|
||||
.map_err(WasmStorageError::InvalidData)?,
|
||||
created_at: row.get("created_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::storage::{
|
||||
ToolStatus, TrustLevel, compute_binary_hash, verify_binary_integrity,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_compute_hash() {
|
||||
let binary = b"(module)";
|
||||
let hash = compute_binary_hash(binary);
|
||||
assert_eq!(hash.len(), 32); // BLAKE3 produces 32-byte hash
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_integrity_success() {
|
||||
let binary = b"test wasm binary content";
|
||||
let hash = compute_binary_hash(binary);
|
||||
assert!(verify_binary_integrity(binary, &hash));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_integrity_failure() {
|
||||
let binary = b"test wasm binary content";
|
||||
let hash = compute_binary_hash(binary);
|
||||
let tampered = b"tampered wasm binary content";
|
||||
assert!(!verify_binary_integrity(tampered, &hash));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trust_level_parse() {
|
||||
assert_eq!("system".parse::<TrustLevel>().unwrap(), TrustLevel::System);
|
||||
assert_eq!(
|
||||
"verified".parse::<TrustLevel>().unwrap(),
|
||||
TrustLevel::Verified
|
||||
);
|
||||
assert_eq!("user".parse::<TrustLevel>().unwrap(), TrustLevel::User);
|
||||
assert!("invalid".parse::<TrustLevel>().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_status_parse() {
|
||||
assert_eq!("active".parse::<ToolStatus>().unwrap(), ToolStatus::Active);
|
||||
assert_eq!(
|
||||
"disabled".parse::<ToolStatus>().unwrap(),
|
||||
ToolStatus::Disabled
|
||||
);
|
||||
assert_eq!(
|
||||
"quarantined".parse::<ToolStatus>().unwrap(),
|
||||
ToolStatus::Quarantined
|
||||
);
|
||||
assert!("invalid".parse::<ToolStatus>().is_err());
|
||||
}
|
||||
}
|
||||
@@ -12,8 +12,9 @@ use wasmtime::component::{Component, Linker, Val};
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::error::WasmError;
|
||||
use crate::tools::wasm::host::{Capabilities, HostState, LogLevel};
|
||||
use crate::tools::wasm::host::{HostState, LogLevel};
|
||||
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
||||
use crate::tools::wasm::runtime::{PreparedModule, WasmToolRuntime};
|
||||
|
||||
@@ -380,10 +381,11 @@ impl std::fmt::Debug for WasmToolWrapper {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::host::Capabilities;
|
||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||
|
||||
#[test]
|
||||
fn test_wrapper_creation() {
|
||||
// This test verifies the runtime can be created
|
||||
@@ -399,5 +401,8 @@ mod tests {
|
||||
fn test_capabilities_default() {
|
||||
let caps = Capabilities::default();
|
||||
assert!(caps.workspace_read.is_none());
|
||||
assert!(caps.http.is_none());
|
||||
assert!(caps.tool_invoke.is_none());
|
||||
assert!(caps.secrets.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user