AI Agent 工具沙箱：用 WasmEdge 隔离 LLM 生成代码，冷启动 3ms 的 VPS 实践

# 1. 安装 WasmEdge（约 40MB，含 AOT 编译器）
curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/master/utils/install.sh \
  | bash -s -- -p /opt/wasmedge --version 0.14.1
echo 'export PATH=/opt/wasmedge/bin:$PATH' >> /etc/profile.d/wasmedge.sh

wasmedge --version
# wasmedge version 0.14.1-rc.1

# 2. 装 Rust wasm32-wasi target（只需一次）
apt install -y build-essential pkg-config
curl -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal
source ~/.cargo/env
rustup target add wasm32-wasi

// tools/log_grep/src/main.rs
use std::io::{self, Read, Write};
use regex::Regex;

fn main() {
    let mut input = String::new();
    io::stdin().read_to_string(&mut input).unwrap();
    // 网关以 JSON 传入：{"path":"/logs/error.log","pattern":"timeout","limit":20}
    let v: serde_json::Value = serde_json::from_str(&input).unwrap();
    let path = v["path"].as_str().unwrap_or("");
    let pat  = v["pattern"].as_str().unwrap_or("");
    let limit = v["limit"].as_u64().unwrap_or(20) as usize;

    // 双保险：路径必须在 /logs 下且不含 ..（即便网关被绕过）
    if !path.starts_with("/logs/") || path.contains("..") {
        println!("{{\"error\":\"path rejected\"}}");
        return;
    }
    let re = match Regex::new(pat) {
        Ok(r) => r,
        Err(e) => { println!("{{\"error\":\"bad regex: {e}\"}}"); return; }
    };
    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(e) => { println!("{{\"error\":\"read failed: {e}\"}}"); return; }
    };
    let hits: Vec<String> = content.lines()
        .filter(|l| re.is_match(l))
        .take(limit)
        .map(|l| l.chars().take(500).collect())   // 截断，防超大行打爆上下文
        .collect();
    println!("{}", serde_json::json!({"count": hits.len(), "lines": hits}));
}

cd tools/log_grep && cargo build --release --target wasm32-wasi
ls -lh target/wasm32-wasi/release/log_grep.wasm
# -rwxr-xr-x  2.6M  log_grep.wasm

# 3. 用 AOT 预编译，把冷启动从 ~40ms 降到个位数 ms
wasmedge compile log_grep.wasm log_grep_aot.wasm
# 编译后约 3MB，之后每次执行都是原生代码

# gateway.py
import json, re, sqlite3, subprocess, hashlib, time, os

WASMEDGE = "/opt/wasmedge/bin/wasmedge"
SANDBOX  = "/srv/agent/tools"

# 能力表：工具 -> (wasm 文件, 预开只读目录, 参数 schema)
ALLOWLIST = {
    "log_grep": {
        "bin": "log_grep_aot.wasm",
        "preopen": ["/logs:/var/log/nginx:ro"],
        "schema": {"path": r"^/logs/[\w./-]+$", "pattern": r"^.{1,200}$"},
        "limit": {"memory_mb": 128, "timeout_s": 10, "max_args": 3},
    },
    "csv_stats": {
        "bin": "csv_stats_aot.wasm",
        "preopen": ["/data:/srv/agent/data:ro"],
        "schema": {"path": r"^/data/[\w./-]+$"},
        "limit": {"memory_mb": 256, "timeout_s": 20, "max_args": 1},
    },
}

def audit(tool, args, code, out_hash, latency_ms):
    con = sqlite3.connect("/srv/agent/audit.db")
    con.execute("""CREATE TABLE IF NOT EXISTS calls(
        ts REAL, tool TEXT, args TEXT, exit_code INT, out_sha256 TEXT, latency_ms INT)""")
    con.execute("INSERT INTO calls VALUES (?,?,?,?,?,?)",
                (time.time(), tool, json.dumps(args), code, out_hash, latency_ms))
    con.commit()

def check(tool, args):
    spec = ALLOWLIST.get(tool)
    if not spec:
        raise ValueError(f"tool '{tool}' not allowed")
    if set(args) - set(spec["schema"]):
        raise ValueError("unexpected arg keys")
    for k, rx in spec["schema"].items():
        if k in args and not re.match(rx, str(args[k])):
            raise ValueError(f"arg '{k}' rejected by schema")
    return spec

def run_tool(tool, args):
    spec = check(tool, args)
    lim = spec["limit"]
    argv = [WASMEDGE, "run"]
    for p in spec["preopen"]:
        argv += ["--dir", p]
    argv += ["--memory-limit", f"{lim['memory_mb']}MB", os.path.join(SANDBOX, spec["bin"])]

    t0 = time.time()
    try:
        r = subprocess.run(argv, input=json.dumps(args), capture_output=True,
                           timeout=lim["timeout_s"], text=True)
        code, out = r.returncode, r.stdout
        if r.returncode != 0 and r.stderr:
            out = json.dumps({"error": r.stderr[:400]})
    except subprocess.TimeoutExpired:
        code, out = 124, json.dumps({"error": "timeout"})
    ms = int((time.time() - t0) * 1000)
    audit(tool, args, code, hashlib.sha256(out.encode()).hexdigest()[:16], ms)
    return {"ok": code == 0, "result": json.loads(out) if out else None, "latency_ms": ms}

# agent_tools.py —— 暴露给本地 Ollama 的工具描述（OpenAI tools 格式）
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "log_grep",
            "description": "在 Nginx 日志中按正则检索行，返回最多 20 条。只能访问日志目录。",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "形如 /logs/error.log"},
                    "pattern": {"type": "string", "description": "RE2 正则表达式"},
                },
                "required": ["path", "pattern"],
            },
        },
    },
]

# 在 Agent 主循环里
from gateway import run_tool

def dispatch(call):
    name = call.function.name
    args = json.loads(call.function.arguments or "{}")
    r = run_tool(name, args)          # 任何异常都在这里变成结构化错误
    return json.dumps(r, ensure_ascii=False)[:8000]

# /etc/systemd/system/agent-gateway.service
[Unit]
Description=AI Agent Tool Gateway (WasmEdge sandbox)
After=network.target ollama.service

[Service]
Type=simple
User=agent
WorkingDirectory=/srv/agent
ExecStart=/usr/bin/python3 /srv/agent/gateway_api.py
Restart=on-failure
RestartSec=3

# 网关自身也最小化：
NoNewPrivileges=yes
PrivateTmp=yes
PrivateDevices=yes
ProtectSystem=strict
ProtectHome=yes
ReadOnlyPaths=/srv/agent/tools /var/log/nginx
ReadWritePaths=/srv/agent/audit.db /srv/agent/run
MemoryMax=1G
CPUQuota=150%
RestrictAddressFamilies=AF_INET AF_INET6
SystemCallFilter=@system-service
SystemCallErrorNumber=EPERM

[Install]
WantedBy=multi-user.target