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:
Illia Polosukhin
2026-02-02 23:22:52 -08:00
co-authored by Claude Opus 4.5
parent 45bbfa026d
commit 32bfd24154
21 changed files with 5115 additions and 66 deletions
+357
View File
@@ -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());
}
}
+423
View File
@@ -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());
}
}
+429
View File
@@ -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
View File
@@ -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
View File
@@ -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,
};
+422
View File
@@ -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());
}
}
+616
View File
@@ -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(&params.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,
&params.user_id,
&params.name,
&params.version,
&params.description,
&params.wasm_binary,
&binary_hash,
&params.parameters_schema,
&params.source_url,
&params.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());
}
}
+8 -3
View File
@@ -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());
}
}