From 29b7dc3bbb4a20be375cb03df63b7b93ecefdfd9 Mon Sep 17 00:00:00 2001 From: jocay Date: Thu, 27 Aug 2026 14:19:10 +0800 Subject: [PATCH] =?UTF-8?q?bolt:=20=E6=9E=81=E9=80=9F=E5=8D=95=E6=96=87?= =?UTF-8?q?=E4=BB=B6=20MCP=20=E5=91=BD=E4=BB=A4=E6=89=A7=E8=A1=8C=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 手写 JSON-RPC 2.0 + MCP (2025-06-18) 协议实现,仅依赖 serde_json + libc - 双传输:stdio(本地 AI 客户端)与 Streamable HTTP(远程连接) - 工具:run_command(超时/取消/进度通知/进程树杀除)与 shell 持久会话四件套 - 安全:命令白黑名单、Bearer 认证、Origin 校验 - LTO fat + codegen-units=1,ping 吞吐约 1.4 万 msg/s - 8 个单元测试 + stdio/HTTP 端到端集成测试(31 项) --- .gitignore | 1 + Cargo.lock | 112 +++ Cargo.toml | 18 + README.md | 106 +++ docs/bolt_http_test.py | 153 ++++ docs/bolt_stdio_test.py | 173 ++++ src/main.rs | 1906 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 2469 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 README.md create mode 100644 docs/bolt_http_test.py create mode 100644 docs/bolt_stdio_test.py create mode 100644 src/main.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..821b536 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,112 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "bolt" +version = "0.1.0" +dependencies = [ + "libc", + "serde_json", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "syn" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..a085096 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "bolt" +version = "0.1.0" +edition = "2021" +description = "bolt — 极速单文件 MCP 命令执行服务器 (stdio / Streamable HTTP)" +license = "MIT" + +[dependencies] +serde_json = "1" + +[target.'cfg(unix)'.dependencies] +libc = "0.2" + +[profile.release] +opt-level = 3 +lto = "fat" +codegen-units = 1 +strip = true diff --git a/README.md b/README.md new file mode 100644 index 0000000..41fc0ab --- /dev/null +++ b/README.md @@ -0,0 +1,106 @@ +# bolt — 极速单文件 MCP 命令执行服务器 + +让 AI(Claude Code / Claude Desktop / Cursor 等 MCP 客户端)通过 **MCP 协议**连接后,在服务器上执行 shell 命令。 + +- **单文件**:全部代码在 `src/main.rs`(约 1900 行,其中 400 行为单元测试),仅依赖 `serde_json`(Unix 下加 `libc`) +- **零框架**:手写 JSON-RPC 2.0 + MCP 协议 + Streamable HTTP 服务器,无 async runtime、无重量级依赖 +- **双传输**:`stdio`(本地 AI 客户端默认)与 `Streamable HTTP`(远程连接) + +## 快速开始 + +```bash +cargo build --release # 产物: target/release/bolt (~1MB) +``` + +**本地模式(stdio)**:由 MCP 客户端拉起本进程 + +```bash +# Claude Code +claude mcp add --transport stdio bolt -- /abs/path/to/bolt + +# 或通用 MCP 客户端配置 +# {"mcpServers": {"bolt": {"command": "/abs/path/to/bolt"}}} +``` + +**远程模式(Streamable HTTP)**:AI 通过网络连接 + +```bash +bolt --http 127.0.0.1:8080 # 仅本机 +bolt --http 0.0.0.0:8080 --token mysecret # 对外 + Bearer 认证 +bolt --http 0.0.0.0:8080 --token mysecret --allow "git,docker,curl" + +# Claude Code 连接远程 +claude mcp add --transport http bolt http://127.0.0.1:8080/mcp +``` + +## 工具 + +| 工具 | 说明 | +|---|---| +| `run_command` | 执行 shell 命令:完整 shell 语法(管道/重定向/`&&`)、`cwd`、`env`、`timeout_ms`(超时自动杀进程树)、`max_output_bytes`(输出截断防膨胀)、`shell=false` 直接执行二进制;客户端可传 `_meta.progressToken` 接收进度通知,支持 `notifications/cancelled` 取消 | +| `shell_open` | 开启持久化交互式 shell 会话(`cd`、环境变量等状态跨次保留),可指定 `shell`(sh/bash/zsh/fish)与 `cwd` | +| `shell_write` | 向会话写入输入(自动补换行) | +| `shell_read` | 读取会话新输出,可阻塞等待(`timeout_ms`,0 = 等到会话退出) | +| `shell_close` | 关闭并杀掉会话进程树 | + +`run_command` 返回:`exit_code` / `signal` / `timed_out` / `cancelled` / `duration_ms` / `stdout` / `stderr` / 截断标志等(同时提供 `content` 与 `structuredContent`)。 + +## 配置 + +CLI 参数(环境变量 `BOLT_*` 为回退): + +| 参数 | 默认 | 说明 | +|---|---|---| +| `--http ` | 无(stdio) | 启用 Streamable HTTP 监听地址,如 `127.0.0.1:8080` | +| `--token ` | 无 | HTTP 要求 `Authorization: Bearer ` | +| `--allow ` | 空 | 命令白名单(命令首词前缀匹配;非空时仅允许列表内命令) | +| `--deny ` | 空 | 命令黑名单(前缀匹配) | +| `--timeout-ms ` | 30000 | 命令默认超时(工具参数可覆盖,0 = 无限) | +| `--cwd ` | 当前 | 命令默认工作目录 | +| `--max-output ` | 1M | 单流输出捕获上限(支持 K/M/G) | +| `--allow-origin ` | localhost/127.0.0.1 | 额外允许的 HTTP Origin(防 DNS rebinding) | + +```bash +export BOLT_TOKEN=xxx BOLT_ALLOW="git,docker" BOLT_TIMEOUT_MS=60000 +bolt --http 127.0.0.1:8080 +``` + +## 协议实现(2025-06-18 规范) + +- `initialize`(协议版本协商,支持 2024-11-05 / 2025-03-26 / 2025-06-18)、`ping`、`tools/list`、`tools/call`、`shutdown` +- `notifications/initialized`、`notifications/cancelled`、`notifications/exit`、`notifications/progress`(进度上报) +- 工具执行错误以 `isError: true` 返回;协议错误返回标准 JSON-RPC 错误码(-32601/-32602/-32700 等) +- Streamable HTTP:POST 请求(通知 → 202;请求 → JSON 或 SSE 响应)、GET 长连接 SSE 流(服务器通知广播 + 10s 心跳)、CORS、Origin 校验、Bearer 认证、无状态(不要求会话 ID) + +## 性能 + +- 单线程事件循环 + 复用缓冲区,无锁热路径;输出计数用原子操作,进度推送不占主循环 +- 子进程独占进程组(`setpgid`),超时/取消按进程组 `SIGKILL`,连带杀掉整棵进程树,无僵尸 +- 实测(本机,含 Python 客户端逐字节 I/O 开销): + - ping 往返吞吐 ≈ **14,700 msg/s** + - 并发 `echo` 命令平均 **4.7 ms/次**(进程 spawn + 捕获输出) +- 构建配置:`lto = "fat"`、`codegen-units = 1`、`strip = true`,二进制约 0.7MB + +## 测试 + +```bash +cargo test --release # 8 个单元测试(协议分发/策略/HTTP 解析/执行引擎/会话) +``` + +集成测试脚本在 `docs/`(stdio 12 项、HTTP 19 项,覆盖超时/取消/进程树/进度/并发/auth/Origin): + +```bash +python3 /tmp/bolt_stdio_test.py # stdio 端到端 +python3 /tmp/bolt_http_test.py # HTTP 端到端 +``` + +## 安全说明 + +- 默认只接受本机连接(`127.0.0.1`),对外公开需显式 `0.0.0.0` + `--token` +- 该服务器等于远程 root shell 入口,务必配置白名单(`--allow`)与认证后再暴露公网 +- Unix 下命令在独立进程组运行,取消/超时会连带杀掉子进程 + +## 平台 + +- Linux / macOS(完整功能:进程组杀、信号处理) +- Windows(基础功能:`cmd /C`、`taskkill /T` 杀进程树;无信号优雅退出) \ No newline at end of file diff --git a/docs/bolt_http_test.py b/docs/bolt_http_test.py new file mode 100644 index 0000000..6b8e478 --- /dev/null +++ b/docs/bolt_http_test.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""bolt HTTP 模式端到端测试""" +import http.client, json, subprocess, threading, time + +BOLT = "/root/bolt/target/release/bolt" +PORT = 18080 +passed = failed = 0 +def check(name, cond, detail=""): + global passed, failed + if cond: + passed += 1 + print(f" PASS {name}") + else: + failed += 1 + print(f" FAIL {name} {detail}") + +proc = subprocess.Popen([BOLT, "--http", f"127.0.0.1:{PORT}"], stderr=subprocess.PIPE) +time.sleep(0.5) + +def post(obj, auth=None, origin=None, accept="application/json", path="/mcp"): + c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=30) + headers = {"Content-Type": "application/json", "Accept": accept} + if auth: headers["Authorization"] = auth + if origin: headers["Origin"] = origin + body = json.dumps(obj) + c.request("POST", path, body=body, headers=headers) + r = c.getresponse() + data = r.read() + ct = r.getheader("Content-Type", "") + c.close() + return r.status, ct, data + +# 1. initialize(JSON 响应) +st, ct, data = post({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"}}}) +m = json.loads(data) +check("init json", st == 200 and "application/json" in ct and m["result"]["serverInfo"]["name"] == "bolt", f"{st} {ct} {data[:100]}") + +# 2. SSE 响应(accept 含 text/event-stream) +st, ct, data = post({"jsonrpc": "2.0", "id": 2, "method": "ping"}, accept="application/json, text/event-stream") +sseline = data.split(b"data: ")[1].split(b"\r\n")[0].strip() +m = json.loads(sseline) +check("ping sse", st == 200 and "text/event-stream" in ct and m.get("result") == {}, f"{st} {ct} {data[:150]}") + +# 3. 通知 → 202 +st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/initialized"}) +check("notif 202", st == 202 and len(data) == 0, f"{st} {data[:50]}") + +# 4. tools/call +st, ct, data = post({"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo HTTP_OK && exit 5"}}}) +m = json.loads(data) +sc = m["result"]["structuredContent"] +check("run over http", sc["exit_code"] == 5 and sc["stdout"].strip() == "HTTP_OK", str(sc)) + +# 5. GET SSE 流:开流收心跳,并发命令进度通知广播 +got_event = threading.Event() +got_heartbeat = threading.Event() +msgs = [] +def sse_get(): + c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=35) + c.request("GET", "/mcp", headers={"Accept": "text/event-stream"}) + r = c.getresponse() + assert r.status == 200 and "text/event-stream" in r.getheader("Content-Type") + buf = b"" + end = time.time() + 30 + while time.time() < end: + chunk = r.read(1) + if not chunk: break + buf += chunk + if buf.endswith(b"\r\n\r\n"): + msgs.append(buf.decode(errors="replace")) + buf = b"" + if msgs and (": ping" in msgs[-1]): + got_heartbeat.set() + c.close() +t = threading.Thread(target=sse_get) +t.start() +time.sleep(0.7) +# 发一个带进度 token 的命令(约 1.2s) +st, ct, data = post({"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo h-$x; sleep 0.3; done"}, "_meta": {"progressToken": "h1"}}}) +m = json.loads(data) +check("long cmd done", m["result"]["structuredContent"]["stdout"].count("h-") == 4, str(m)[:100]) +saw_progress = any("notifications/progress" in x and "h1" in x for x in msgs) +check("sse got progress broadcast", saw_progress, str(msgs)[:200]) +t.join(timeout=25) +saw_hb = got_heartbeat.is_set() +check("sse heartbeat", saw_hb, str(msgs)[:200]) + +# 6. 取消:连接 A 长命令,连接 B 发 cancelled +st, ct, data = post({"jsonrpc": "2.0", "id": 5, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo nothing"}}}) # 预热(连接池无关,每次新连接) + +def long_cmd(): + return post({"jsonrpc": "2.0", "id": 6, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "sleep 30"}}}) +res = {} +t2 = threading.Thread(target=lambda: res.update(r=long_cmd())) +t2.start() +time.sleep(0.5) +st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 6}}) +check("cancel notif 202", st == 202) +t2.join(timeout=10) +m = json.loads(res["r"][2]) +check("http cancel works", m["result"]["structuredContent"]["cancelled"] is True, str(m)[:150]) + +# 7. 并发两个命令 +def run(cmd, rid): + return post({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": cmd}}}) +results = {} +ts = [threading.Thread(target=lambda c=cmd, i=i: results.update({i: run(c, i)})) for i, cmd in [(10, "sleep 0.5; echo A"), (11, "sleep 0.1; echo B")]] +for t in ts: t.start() +for t in ts: t.join(timeout=10) +check("http concurrent", json.loads(results[10][2])["result"]["structuredContent"]["stdout"].strip() == "A" and json.loads(results[11][2])["result"]["structuredContent"]["stdout"].strip() == "B") + +# 8. 错误路径 +st, _, _ = post({"jsonrpc": "2.0", "id": 99, "method": "nope"}) +check("unknown method", st == 200) +c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5) +c.request("POST", "/mcp", body=b"this is {not json", headers={"Content-Type": "application/json", "Accept": "application/json"}) +r = c.getresponse(); bd = r.read(); c.close() +check("bad json -> 400", r.status == 400, str(r.status)) +st, _, data = post({"jsonrpc": "2.0", "id": 99, "method": "tools/call", "params": {"name": "nope", "arguments": {}}}) +check("unknown tool -32602", json.loads(data)["error"]["code"] == -32602) +c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5) +c.request("GET", "/other", headers={"Accept": "text/event-stream"}) +r = c.getresponse(); r.read(); c.close() +check("404 other path", r.status == 404) +c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5) +c.request("PUT", "/mcp") +r = c.getresponse(); r.read(); c.close() +check("405 method", r.status == 405) +c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5) +c.request("OPTIONS", "/mcp") +r = c.getresponse(); r.read(); c.close() +check("options cors", r.status == 204 and r.getheader("Access-Control-Allow-Origin") == "*") + +# 9. origin 校验 +st, ct, data = post({"jsonrpc": "2.0", "id": 50, "method": "ping"}, origin="https://evil.example.com") +check("evil origin 403", st == 403, str(st)) +st, ct, data = post({"jsonrpc": "2.0", "id": 51, "method": "ping"}, origin="http://localhost:4173") +check("localhost origin ok", st == 200, str(st)) + +# 10. GET 无 Accept 头也返回 SSE(宽松) +c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5) +c.request("GET", "/mcp") +r = c.getresponse() +c.close() +check("get sse default", r.status == 200 and "text/event-stream" in r.getheader("Content-Type")) + +proc.kill() +err = proc.stderr.read().decode() +print(f"\n=== HTTP: {passed} passed, {failed} failed ===") +if err.strip(): + print("bolt stderr:", err[:400]) +sys_exit = 1 if failed else 0 +import sys; sys.exit(sys_exit) \ No newline at end of file diff --git a/docs/bolt_stdio_test.py b/docs/bolt_stdio_test.py new file mode 100644 index 0000000..03bd2a4 --- /dev/null +++ b/docs/bolt_stdio_test.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""bolt stdio 模式端到端协议测试驱动""" +import faulthandler, json, os, subprocess, sys, threading, time + + +BOLT = "/root/bolt/target/release/bolt" +proc = subprocess.Popen([BOLT], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0) +sent = {} +next_id = [0] + +def send(method, params=None): + next_id[0] += 1 + msg = {"jsonrpc": "2.0", "id": next_id[0], "method": method} + if params is not None: + msg["params"] = params + proc.stdin.write((json.dumps(msg) + "\n").encode()) + proc.stdin.flush() + return next_id[0] + +def send_notif(method, params=None): + msg = {"jsonrpc": "2.0", "method": method} + if params is not None: + msg["params"] = params + proc.stdin.write((json.dumps(msg) + "\n").encode()) + proc.stdin.flush() + +def recv(timeout=30): + import select + end = time.time() + timeout + line = b"" + while True: + r, _, _ = select.select([proc.stdout], [], [], max(0.0, end - time.time())) + if not r: + raise TimeoutError(f"timeout, got so far: {line[:200]}") + ch = proc.stdout.read(1) + if not ch: + raise TimeoutError("EOF") + if ch == b"\n": + return json.loads(line) + line += ch + +received = {} # 乱序到达的响应缓存(id -> msg) +def wait_id(i, timeout=30): + if i in received: + return received.pop(i) + deadline = time.time() + timeout + pend.setdefault("extra", []).append(f"awaiting {i}") + while time.time() < deadline: + try: + m = recv(timeout=min(3, deadline - time.time())) + except TimeoutError as e: + raise TimeoutError(f"no response for id {i}: {e}; extras: {json.dumps(pend.get('extra', []))[:800]}") + if m.get("id") == i: + return m + if m.get("id") is not None: + received[m["id"]] = m + else: + pend.setdefault("extra", []).append(m) + raise TimeoutError(f"no response for id {i}, extras: {json.dumps(pend.get('extra', []))[:800]}") + +pend = {} +passed = failed = 0 +def check(name, cond, detail=""): + global passed, failed + if cond: + passed += 1 + print(f" PASS {name}") + else: + failed += 1 + print(f" FAIL {name} {detail}") + +# 1. initialize +i = send("initialize", {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "test", "version": "1"}}) +r = wait_id(i) +check("initialize", r["result"]["protocolVersion"] == "2025-06-18" and r["result"]["serverInfo"]["name"] == "bolt", str(r)) +send_notif("notifications/initialized") + +# 2. tools/list +i = send("tools/list") +r = wait_id(i) +names = [t["name"] for t in r["result"]["tools"]] +check("tools/list", names == ["run_command", "shell_open", "shell_write", "shell_read", "shell_close"], str(names)) + +# 3. run_command 基本 +i = send("tools/call", {"name": "run_command", "arguments": {"command": "echo hello; printf 'oops' >&2; exit 7"}}) +r = wait_id(i) +sc = r["result"]["structuredContent"] +check("run basic", sc["exit_code"] == 7 and sc["stdout"] == "hello\n" and sc["stderr"] == "oops", str(sc)) + +# 4. 超时 +t0 = time.time() +i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30", "timeout_ms": 400}}) +r = wait_id(i) +dt = time.time() - t0 +sc = r["result"]["structuredContent"] +check("timeout kill", sc["timed_out"] and not sc["cancelled"] and dt < 5, f"{sc} dt={dt:.1f}") + +# 5. 取消:发一个长命令,再发 cancelled +i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30 && echo done"}}) +time.sleep(0.4) +send_notif("notifications/cancelled", {"requestId": i}) +r = wait_id(i) +sc = r["result"]["structuredContent"] +check("cancel kill", sc["cancelled"] and sc["exit_code"] is None, str(sc)) + +# 6. 进程组连带杀:sh -c 'sleep 30 & sleep 30' 的子进程也应变僵尸 +i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30 & sleep 30", "timeout_ms": 400}}) +r = wait_id(i) +check("kill whole tree", r["result"]["structuredContent"]["timed_out"]) + +# 7. 进度通知(_meta.progressToken) +i = send("tools/call", {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo tick-$x; sleep 0.3; done"}, "_meta": {"progressToken": "tok-1"}}) +got_progress = False +deadline = time.time() + 10 +while time.time() < deadline: + m = recv(timeout=deadline - time.time()) + if m.get("method") == "notifications/progress" and m.get("params", {}).get("progressToken") == "tok-1": + got_progress = True + break + if m.get("id") == i: + break +check("progress notification", got_progress) +# 排空管道:打印残留 +import select +import os +while True: + r, _, _ = select.select([proc.stdout], [], [], 0) + if not r: break + try: + m = recv(timeout=2) + print(" LEFT:", m.get("id"), m.get("method")) + except Exception: + break + +# 8. 并发:同时两个命令 +i1 = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 0.6; echo ONE"}}) +i2 = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 0.2; echo TWO"}}) +r1, r2 = wait_id(i1), wait_id(i2) +check("concurrent", r1["result"]["structuredContent"]["stdout"].strip() == "ONE" and r2["result"]["structuredContent"]["stdout"].strip() == "TWO") + +# 9. shell 会话 +i = send("tools/call", {"name": "shell_open", "arguments": {}}) +r = wait_id(i) +sid = r["result"]["session_id"] +check("shell open", isinstance(sid, int)) +send("tools/call", {"name": "shell_write", "arguments": {"session_id": sid, "input": "cd /tmp && echo IN_TMP; echo DONE_$$; exit"}}) +# 输出先于退出标志到达,循环读取直到会话退出 +all_out = "" +for _ in range(10): + i = send("tools/call", {"name": "shell_read", "arguments": {"session_id": sid, "timeout_ms": 3000}}) + r = wait_id(i) + res = r["result"] + all_out += res["output"] + if res["exited"]: + break +check("shell read output", "IN_TMP" in all_out and res["exited"], f"out={all_out!r}") +i = send("tools/call", {"name": "shell_close", "arguments": {"session_id": sid}}) +r = wait_id(i) +check("shell close", r["result"]["closed"] is True) + +# 10. stdio 关闭 → 进程退出 +send_notif("notifications/exit") +try: + proc.wait(timeout=3) + check("graceful exit", True) +except subprocess.TimeoutExpired: + check("graceful exit", False) + +err = proc.stderr.read().decode() +print(f"\n=== {passed} passed, {failed} failed ===") +if err.strip(): + print("stderr:", err[:500]) +sys.exit(1 if failed else 0) \ No newline at end of file diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..c4cc9d8 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,1906 @@ +//! bolt — 极速单文件 MCP 命令执行服务器 +//! +//! AI(Claude Desktop / Claude Code / Cursor 等 MCP 客户端)连接后即可在服务器上执行命令。 +//! +//! 特性: +//! * 零框架、纯手写 JSON-RPC 2.0 + MCP 协议(仅依赖 serde_json,Unix 下加 libc) +//! * 双传输:stdio(本地 AI 客户端默认方式)与 Streamable HTTP(远程连接) +//! * 工具: +//! - run_command 执行 shell 命令(支持管道/重定向/超时/环境变量/输出截断/进度通知) +//! - shell_open 开启持久化交互式 shell 会话(状态保留,可连续对话) +//! - shell_write 向会话写入输入 +//! - shell_read 读取会话输出(可阻塞等待) +//! - shell_close 关闭会话 +//! * 取消:客户端发送 notifications/cancelled 可杀掉进行中的命令(Unix 下按进程组杀,连带子进程) +//! * 安全:命令 allow/deny 前缀白黑名单、HTTP Bearer Token、Origin 校验、默认本机绑定 +//! * 性能:单线程事件循环 + 复用缓冲区、原子计数、无锁热路径、LTO fat + codegen-units=1 +//! +//! 用法: +//! bolt # stdio 模式(默认,MCP 客户端拉起本进程) +//! bolt --http 127.0.0.1:8080 # Streamable HTTP 模式,远程 AI 客户端连接 +//! bolt --http 0.0.0.0:8080 --token mysecret --allow "git,docker,curl" +//! +//! 环境变量(CLI 参数优先):BOLT_TOKEN / BOLT_ALLOW / BOLT_DENY / BOLT_TIMEOUT_MS / +//! BOLT_CWD / BOLT_MAX_OUTPUT / BOLT_ALLOW_ORIGIN +//! +//! 客户端配置示例: +//! Claude Code(本地): claude mcp add --transport stdio bolt -- /abs/path/bolt +//! Claude Code(远程): claude mcp add --transport http bolt http://host:8080/mcp +//! 通用 MCP 客户端: {"mcpServers": {"bolt": {"command": "/abs/path/bolt"}}} + +use std::collections::HashMap; +use std::env; +use std::io::{self, BufRead, BufReader, BufWriter, Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::path::Path; +use std::process::{Command, ExitStatus, Stdio}; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::{mpsc, Arc, Condvar, Mutex, OnceLock}; +use std::thread; +use std::time::{Duration, Instant}; + +#[cfg(unix)] +use std::os::unix::process::CommandExt; + +use serde_json::{json, Value}; + +// --------------------------------------------------------------------------- +// 常量 +// --------------------------------------------------------------------------- + +const VERSION: &str = env!("CARGO_PKG_VERSION"); +const PROTOCOL_VERSIONS: [&str; 3] = ["2024-11-05", "2025-03-26", "2025-06-18"]; +const LATEST_PROTOCOL: &str = "2025-06-18"; + +const MAX_LINE_BYTES: usize = 16 * 1024 * 1024; // stdio 单条消息上限 +const MAX_HTTP_HEAD: usize = 64 * 1024; // HTTP 请求头上限 +const MAX_HTTP_BODY: usize = 4 * 1024 * 1024; // HTTP body 上限 +const DEFAULT_TIMEOUT_MS: u64 = 30_000; // 命令默认超时 +const MAX_TIMEOUT_MS: u64 = 6 * 3600 * 1000; // 命令超时上限(6h) +const DEFAULT_MAX_OUTPUT: usize = 1024 * 1024; // 单流默认捕获上限 +const MAX_MAX_OUTPUT: usize = 64 * 1024 * 1024; +const SESSION_BUF_CAP: usize = 16 * 1024 * 1024; // 会话环形缓冲上限 +const SESSION_WRITE_TIMEOUT: Duration = Duration::from_secs(10); // 会话写入超时 +const SSE_HEARTBEAT: Duration = Duration::from_secs(10); // SSE 保活间隔 + +// --------------------------------------------------------------------------- +// 配置 +// --------------------------------------------------------------------------- + +enum Mode { + Stdio, + Http(String), +} + +struct Config { + mode: Mode, + token: Option, + allow: Vec, + deny: Vec, + timeout_ms: u64, + cwd: Option, + max_output: usize, + allow_origins: Vec, +} + +fn split_list(s: &str) -> Vec { + s.split(',') + .map(str::trim) + .filter(|x| !x.is_empty()) + .map(str::to_string) + .collect() +} + +/// 解析 "1024" / "64K" / "1M" / "2G" 形式的字节数 +fn parse_size(s: &str) -> Option { + let s = s.trim(); + let (num, mult) = match s.chars().last() { + Some('k') | Some('K') => (&s[..s.len() - 1], 1 << 10), + Some('m') | Some('M') => (&s[..s.len() - 1], 1 << 20), + Some('g') | Some('G') => (&s[..s.len() - 1], 1 << 30), + _ => (s, 1), + }; + num.parse::().ok().map(|n| n.saturating_mul(mult)) +} + +fn parse_args() -> Option { + let mut mode = Mode::Stdio; + let mut token: Option = None; + let mut allow: Vec = Vec::new(); + let mut deny: Vec = Vec::new(); + let mut timeout_ms = DEFAULT_TIMEOUT_MS; + let mut cwd: Option = None; + let mut max_output = DEFAULT_MAX_OUTPUT; + let mut allow_origins: Vec = Vec::new(); + + let args: Vec = env::args().skip(1).collect(); + let mut i = 0; + while i < args.len() { + let a = args[i].as_str(); + let next = |i: &mut usize| -> Option { + *i += 1; + args.get(*i).cloned() + }; + match a { + "-h" | "--help" => return None, + "--http" => mode = Mode::Http(next(&mut i)?), + "--token" => token = next(&mut i), + "--allow" => allow = split_list(&next(&mut i)?), + "--deny" => deny = split_list(&next(&mut i)?), + "--timeout-ms" => timeout_ms = next(&mut i)?.parse().ok()?, + "--cwd" => cwd = next(&mut i), + "--max-output" => max_output = parse_size(&next(&mut i)?)?, + "--allow-origin" => allow_origins.push(next(&mut i)?), + _ => { + eprintln!("[bolt] 未知参数: {a}"); + return None; + } + } + i += 1; + } + // 环境变量回退(CLI 优先) + if token.is_none() { + token = env::var("BOLT_TOKEN").ok().filter(|s| !s.is_empty()); + } + if allow.is_empty() { + allow = split_list(&env::var("BOLT_ALLOW").unwrap_or_default()); + } + if deny.is_empty() { + deny = split_list(&env::var("BOLT_DENY").unwrap_or_default()); + } + if let Ok(v) = env::var("BOLT_TIMEOUT_MS") { + if let Ok(n) = v.parse() { + timeout_ms = n; + } + } + if cwd.is_none() { + cwd = env::var("BOLT_CWD").ok().filter(|s| !s.is_empty()); + } + if let Ok(v) = env::var("BOLT_MAX_OUTPUT") { + if let Some(n) = parse_size(&v) { + max_output = n; + } + } + if allow_origins.is_empty() { + allow_origins = split_list(&env::var("BOLT_ALLOW_ORIGIN").unwrap_or_default()); + } + Some(Config { + mode, + token, + allow, + deny, + timeout_ms, + cwd, + max_output, + allow_origins, + }) +} + +fn print_help() { + println!( + r#"bolt {VERSION} — 极速单文件 MCP 命令执行服务器 +AI 通过 MCP(stdio 或 Streamable HTTP)连接后即可在服务器上执行命令。 + +用法: + bolt [选项] + +选项: + --http 启用 Streamable HTTP 传输并监听(如 127.0.0.1:8080 或 :8080) + 缺省为 stdio 模式(由 MCP 客户端拉起本进程) + --token HTTP 模式要求 Authorization: Bearer + --allow 命令白名单(命令首词前缀匹配,非空时仅允许列出的命令) + --deny 命令黑名单(前缀匹配) + --timeout-ms 命令默认超时毫秒(0 = 无限,工具参数可覆盖;默认 30000) + --cwd 命令默认工作目录 + --max-output 单流输出捕获上限,支持 K/M/G 后缀(默认 1M) + --allow-origin 额外允许的 HTTP Origin(默认仅 localhost/127.0.0.1) + -h, --help 显示帮助 + +环境变量(CLI 优先):BOLT_TOKEN BOLT_ALLOW BOLT_DENY BOLT_TIMEOUT_MS BOLT_CWD BOLT_MAX_OUTPUT BOLT_ALLOW_ORIGIN + +客户端配置: + Claude Code 本地: claude mcp add --transport stdio bolt -- /abs/path/bolt + Claude Code 远程: claude mcp add --transport http bolt http://host:8080/mcp + 通用 MCP 客户端: {{"mcpServers": {{"bolt": {{"command": "/abs/path/bolt"}}}}}} +"# + ); +} + +// --------------------------------------------------------------------------- +// 出站通道:stdio 单写者锁 / HTTP 广播到 SSE 流 +// --------------------------------------------------------------------------- + +struct Hub { + map: Mutex>>, + next: AtomicU64, +} + +impl Hub { + fn add(&self, tx: mpsc::Sender) -> u64 { + let id = self.next.fetch_add(1, Ordering::Relaxed) + 1; + self.map.lock().unwrap().insert(id, tx); + id + } + fn remove(&self, id: u64) { + self.map.lock().unwrap().remove(&id); + } + /// 广播(发送失败即剔除死连接;mpsc 无界通道,不阻塞) + fn broadcast(&self, msg: &str) { + let mut map = self.map.lock().unwrap(); + map.retain(|_, tx| tx.send(msg.to_string()).is_ok()); + } +} + +#[derive(Clone)] +enum Sink { + Stdio(Arc>), + Http(Arc), +} + +impl Sink { + #[inline] + fn send(&self, msg: &str) { + match self { + Sink::Stdio(o) => { + if let Ok(mut g) = o.lock() { + let _ = g.write_all(msg.as_bytes()); + let _ = g.write_all(b"\n"); + let _ = g.flush(); + } + } + Sink::Http(h) => h.broadcast(msg), + } + } +} + +// --------------------------------------------------------------------------- +// 应用状态 +// --------------------------------------------------------------------------- + +struct App { + cfg: Config, + sink: Sink, + hub: Arc, + exit_flag: Arc, + inflight: Mutex>>, // 进行中的命令(id -> 可杀句柄) + sessions: Mutex>>, + next_session: AtomicU64, +} + +type Shared = Arc; + +impl App { + fn new(cfg: Config) -> App { + let hub = Arc::new(Hub { + map: Mutex::new(HashMap::new()), + next: AtomicU64::new(0), + }); + let sink = match &cfg.mode { + Mode::Stdio => Sink::Stdio(Arc::new(Mutex::new(io::stdout()))), + Mode::Http(_) => Sink::Http(hub.clone()), + }; + App { + cfg, + sink, + hub, + exit_flag: Arc::new(AtomicBool::new(false)), + inflight: Mutex::new(HashMap::new()), + sessions: Mutex::new(HashMap::new()), + next_session: AtomicU64::new(0), + } + } +} + +/// 进程组杀掉 pid 对应的整棵进程树 +#[cfg(unix)] +fn kill_tree(pid: u32) { + unsafe { + libc::kill(-(pid as i32), libc::SIGKILL); + } +} +#[cfg(windows)] +fn kill_tree(pid: u32) { + let _ = Command::new("taskkill") + .args(["/PID", &pid.to_string(), "/T", "/F"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); +} + +/// 可取消句柄:注册 pid 前收到取消则标记,注册后立即补杀 +struct Killable { + pid: AtomicU32, + want_kill: AtomicBool, + cancelled: AtomicBool, +} + +impl Default for Killable { + fn default() -> Self { + Killable { + pid: AtomicU32::new(0), + want_kill: AtomicBool::new(false), + cancelled: AtomicBool::new(false), + } + } +} + +impl Killable { + fn register_pid(&self, pid: u32) { + self.pid.store(pid, Ordering::SeqCst); + if self.want_kill.load(Ordering::SeqCst) { + kill_tree(pid); + } + } + fn kill(&self, mark_cancelled: bool) { + if mark_cancelled { + self.cancelled.store(true, Ordering::SeqCst); + } + self.want_kill.store(true, Ordering::SeqCst); + let p = self.pid.load(Ordering::SeqCst); + if p != 0 { + kill_tree(p); + } + } +} + +fn id_key(id: &Value) -> String { + serde_json::to_string(id).unwrap_or_default() +} + +fn register_inflight(app: &Shared, id: &Value) -> Arc { + let k = Arc::new(Killable::default()); + app.inflight.lock().unwrap().insert(id_key(id), k.clone()); + k +} + +fn unregister_inflight(app: &Shared, id: &Value, k: &Arc) { + let mut m = app.inflight.lock().unwrap(); + if m.get(&id_key(id)).map(|x| Arc::ptr_eq(x, k)).unwrap_or(false) { + m.remove(&id_key(id)); + } +} + +// --------------------------------------------------------------------------- +// JSON-RPC / MCP 辅助 +// --------------------------------------------------------------------------- + +fn rpc_error(id: &Value, code: i64, message: &str) -> Value { + json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}) +} + +const PARSE_ERROR: &str = + r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"parse error"}}"#; + +/// 工具执行错误(非协议错误):结果中 isError: true +fn tool_error(msg: &str) -> Value { + json!({"content": [{"type": "text", "text": format!("error: {msg}")}], "isError": true}) +} + +/// 命令白/黑名单检查(first_token 为命令首词) +fn policy_check(cfg: &Config, first_token: &str) -> Result<(), String> { + if !cfg.deny.is_empty() && cfg.deny.iter().any(|d| first_token.starts_with(d)) { + return Err(format!("命令 '{first_token}' 被服务器策略禁止")); + } + if !cfg.allow.is_empty() && !cfg.allow.iter().any(|a| first_token.starts_with(a)) { + return Err(format!("命令 '{first_token}' 不在服务器白名单中")); + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// 工具:run_command 执行引擎 +// --------------------------------------------------------------------------- + +struct RunReq { + command: String, + args: Vec, + shell: bool, + cwd: Option, + env: Vec<(String, String)>, + timeout_ms: u64, // 0 = 无限 + max_output: usize, +} + +fn parse_run_params(p: &Value, cfg: &Config) -> Result { + let command = p + .get("command") + .and_then(Value::as_str) + .ok_or("'command' 必须是非空字符串")? + .to_string(); + let shell = p.get("shell").and_then(Value::as_bool).unwrap_or(true); + let args: Vec = match p.get("args") { + None => Vec::new(), + Some(Value::Array(a)) => a + .iter() + .map(|x| { + x.as_str() + .map(str::to_string) + .ok_or_else(|| "'args' 必须是字符串数组".to_string()) + }) + .collect::>()?, + Some(_) => return Err("'args' 必须是字符串数组".into()), + }; + if shell && !args.is_empty() { + return Err("shell=true 时不能同时使用 'args'(shell 语法直接写在 command 里)".into()); + } + let cwd = match p.get("cwd") { + None => cfg.cwd.clone(), + Some(v) => Some(v.as_str().ok_or("'cwd' 必须是字符串")?.to_string()), + }; + let mut env_map: Vec<(String, String)> = Vec::new(); + if let Some(e) = p.get("env") { + let obj = e.as_object().ok_or("'env' 必须是对象")?; + for (k, v) in obj { + let s = v.as_str().ok_or_else(|| format!("env 的值 '{k}' 必须是字符串"))?; + env_map.push((k.clone(), s.to_string())); + } + } + let timeout_ms = match p.get("timeout_ms") { + None => cfg.timeout_ms, + Some(v) => { + let n = v.as_u64().ok_or("'timeout_ms' 必须是整数")?; + if n > MAX_TIMEOUT_MS { + return Err(format!("'timeout_ms' 不能超过 {MAX_TIMEOUT_MS}")); + } + n + } + }; + let max_output = match p.get("max_output_bytes") { + None => cfg.max_output, + Some(v) => { + let n = v.as_u64().ok_or("'max_output_bytes' 必须是整数")? as usize; + if !(1024..=MAX_MAX_OUTPUT).contains(&n) { + return Err(format!("'max_output_bytes' 必须在 1024..{MAX_MAX_OUTPUT}")); + } + n + } + }; + Ok(RunReq { + command, + args, + shell, + cwd, + env: env_map, + timeout_ms, + max_output, + }) +} + +struct ExecOutcome { + exit_code: Option, + signal: Option, + timed_out: bool, + cancelled: bool, + duration_ms: u64, + stdout: String, + stderr: String, + trunc_out: bool, + trunc_err: bool, +} + +struct OutBuf { + bytes: Vec, + truncated: bool, +} + +impl Default for OutBuf { + fn default() -> Self { + OutBuf { + bytes: Vec::new(), + truncated: false, + } + } +} + +struct OutShared { + stdout: Arc>, + stderr: Arc>, + total: Arc, +} + +/// 管道读取循环:读取到 cap 后继续排空(防止子进程写满管道阻塞),只标记截断 +fn pipe_read_loop( + mut pipe: R, + out: Arc>, + total: Arc, + cap: usize, +) { + let mut chunk = [0u8; 16 * 1024]; + loop { + match pipe.read(&mut chunk) { + Ok(0) => break, + Ok(n) => { + total.fetch_add(n as u64, Ordering::Relaxed); + let mut b = out.lock().unwrap(); + if b.bytes.len() < cap { + let take = (cap - b.bytes.len()).min(n); + b.bytes.extend_from_slice(&chunk[..take]); + if take < n { + b.truncated = true; + } + } else { + b.truncated = true; + } + } + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => break, + } + } +} + +/// 进度通知循环:每 200ms 上报已读字节数 +fn progress_loop(token: Value, total: Arc, done: Arc, sink: Sink) { + let mut last = 0u64; + while !done.load(Ordering::Relaxed) { + thread::sleep(Duration::from_millis(200)); + let t = total.load(Ordering::Relaxed); + if t != last { + last = t; + let msg = json!({ + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": token, + "progress": t, + "message": format!("已读取 {t} 字节输出"), + } + }); + sink.send(&msg.to_string()); + } + } +} + +fn exec(req: &RunReq, killable: &Killable, progress: Option, app: &Shared) -> Result { + // 构建命令 + let mut cmd = if req.shell { + #[cfg(windows)] + let mut c = Command::new("cmd"); + #[cfg(windows)] + c.arg("/C").arg(&req.command); + #[cfg(unix)] + let mut c = Command::new("sh"); + #[cfg(unix)] + c.arg("-c").arg(&req.command); + c + } else { + let mut c = Command::new(&req.command); + c.args(&req.args); + c + }; + cmd.stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + if let Some(cwd) = &req.cwd { + if !Path::new(cwd).is_dir() { + return Err(format!("工作目录不存在: {cwd}")); + } + cmd.current_dir(cwd); + } + cmd.envs(req.env.iter().map(|(k, v)| (k.as_str(), v.as_str()))); + #[cfg(unix)] + unsafe { + // 子进程放入独立进程组,便于整树杀除 + cmd.pre_exec(|| { + libc::setpgid(0, 0); + Ok(()) + }); + } + + let start = Instant::now(); + let mut child = cmd + .spawn() + .map_err(|e| format!("启动命令失败: {e}"))?; + killable.register_pid(child.id()); + + let out = Arc::new(OutShared { + stdout: Arc::new(Mutex::new(OutBuf::default())), + stderr: Arc::new(Mutex::new(OutBuf::default())), + total: Arc::new(AtomicU64::new(0)), + }); + let done = Arc::new(AtomicBool::new(false)); + + let stdout_pipe = child.stdout.take().unwrap(); + let stderr_pipe = child.stderr.take().unwrap(); + let t_out = { + let buf = out.stdout.clone(); + let total = out.total.clone(); + let cap = req.max_output; + thread::Builder::new() + .name("bolt-out".into()) + .spawn(move || pipe_read_loop(stdout_pipe, buf, total, cap)) + .unwrap() + }; + let t_err = { + let buf = out.stderr.clone(); + let total = out.total.clone(); + let cap = req.max_output; + thread::Builder::new() + .name("bolt-err".into()) + .spawn(move || pipe_read_loop(stderr_pipe, buf, total, cap)) + .unwrap() + }; + let t_prog = progress.map(|tok| { + let total = out.total.clone(); + let done = done.clone(); + let sink = app.sink.clone(); + thread::Builder::new() + .name("bolt-progress".into()) + .spawn(move || progress_loop(tok, total, done, sink)) + .unwrap() + }); + + // 收割线程独占 child.wait(),避免僵尸进程 + let (tx, rx) = mpsc::channel(); + let mut child2 = child; + let reaper = thread::Builder::new() + .name("bolt-reaper".into()) + .spawn(move || { + let st = child2.wait(); + let _ = tx.send(st); + }) + .unwrap(); + + let mut timed_out = false; + let status: Option = if req.timeout_ms == 0 { + rx.recv().ok().and_then(Result::ok) + } else { + match rx.recv_timeout(Duration::from_millis(req.timeout_ms)) { + Ok(Ok(st)) => Some(st), + Ok(Err(_)) => None, + Err(mpsc::RecvTimeoutError::Timeout) => { + timed_out = true; + killable.kill(false); // 超时杀,不计为客户端取消 + rx.recv().ok().and_then(Result::ok) + } + Err(mpsc::RecvTimeoutError::Disconnected) => None, + } + }; + + done.store(true, Ordering::Relaxed); + let _ = t_out.join(); + let _ = t_err.join(); + if let Some(t) = t_prog { + let _ = t.join(); + } + let _ = reaper.join(); + + let stdout = String::from_utf8_lossy(&out.stdout.lock().unwrap().bytes).into_owned(); + let stderr = String::from_utf8_lossy(&out.stderr.lock().unwrap().bytes).into_owned(); + let trunc_out = out.stdout.lock().unwrap().truncated; + let trunc_err = out.stderr.lock().unwrap().truncated; + + Ok(ExecOutcome { + exit_code: status.as_ref().and_then(ExitStatus::code), + signal: status.as_ref().map(exit_signal).unwrap_or(None), + timed_out, + cancelled: killable.cancelled.load(Ordering::SeqCst), + duration_ms: start.elapsed().as_millis() as u64, + stdout, + stderr, + trunc_out, + trunc_err, + }) +} + +#[cfg(unix)] +fn exit_signal(st: &ExitStatus) -> Option { + use std::os::unix::process::ExitStatusExt; + st.signal() +} +#[cfg(not(unix))] +fn exit_signal(_st: &ExitStatus) -> Option { + None +} + +fn run_result(o: &ExecOutcome, cwd: Option<&str>) -> Value { + let result = json!({ + "exit_code": o.exit_code, + "signal": o.signal, + "timed_out": o.timed_out, + "cancelled": o.cancelled, + "duration_ms": o.duration_ms, + "stdout": o.stdout, + "stderr": o.stderr, + "truncated_stdout": o.trunc_out, + "truncated_stderr": o.trunc_err, + "cwd": cwd, + }); + json!({ + "content": [{"type": "text", "text": result.to_string()}], + "structuredContent": result, + }) +} + +// --------------------------------------------------------------------------- +// 工具:shell 持久会话 +// --------------------------------------------------------------------------- + +struct Session { + pid: u32, + stdin: Mutex>, + out: Arc, +} + +struct SessionOut { + bytes: Mutex>, + off: Mutex, // 已读偏移(相对当前 Vec,裁剪后可能失效) + gen: AtomicU64, + cv: Condvar, + exited: Mutex>, + dropped: AtomicU64, +} + +fn session_read_loop(mut pipe: R, out: Arc) { + let mut chunk = [0u8; 16 * 1024]; + loop { + match pipe.read(&mut chunk) { + Ok(0) => break, + Ok(n) => { + let mut b = out.bytes.lock().unwrap(); + if b.len() >= SESSION_BUF_CAP { + let drop_n = b.len() / 2; + out.dropped.fetch_add(drop_n as u64, Ordering::Relaxed); + b.drain(..drop_n); + } + b.extend_from_slice(&chunk[..n]); + out.gen.fetch_add(1, Ordering::Relaxed); + out.cv.notify_all(); + } + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => break, + } + } +} + +fn shell_open(app: &Shared, shell: Option<&str>, cwd: Option<&str>) -> Result { + let shell = shell.unwrap_or(if cfg!(windows) { "cmd" } else { "sh" }); + policy_check(&app.cfg, shell)?; + let mut cmd = Command::new(shell); + cmd.stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + if let Some(c) = cwd { + if !Path::new(c).is_dir() { + return Err(format!("工作目录不存在: {c}")); + } + cmd.current_dir(c); + } + #[cfg(unix)] + unsafe { + cmd.pre_exec(|| { + libc::setpgid(0, 0); + Ok(()) + }); + } + let mut child = cmd + .spawn() + .map_err(|e| format!("启动 shell '{shell}' 失败: {e}"))?; + let pid = child.id(); + let stdin = child.stdin.take().unwrap(); + let stdout_pipe = child.stdout.take().unwrap(); + let stderr_pipe = child.stderr.take().unwrap(); + let out = Arc::new(SessionOut { + bytes: Mutex::new(Vec::new()), + off: Mutex::new(0), + gen: AtomicU64::new(0), + cv: Condvar::new(), + exited: Mutex::new(None), + dropped: AtomicU64::new(0), + }); + let r1 = { + let out1 = out.clone(); + thread::Builder::new() + .name("bolt-sess-out".into()) + .spawn(move || session_read_loop(stdout_pipe, out1)) + .unwrap() + }; + let r2 = { + let out2 = out.clone(); + thread::Builder::new() + .name("bolt-sess-err".into()) + .spawn(move || session_read_loop(stderr_pipe, out2)) + .unwrap() + }; + // 收割线程:先等两个 reader 结束(EOF 即输出全部读入),再置退出标志, + // 保证 shell_read 看到 exited 时所有输出都已进入缓冲区(消除丢失输出的竞态) + let out_r = out.clone(); + let _ = { + thread::Builder::new() + .name("bolt-sess-reaper".into()) + .spawn(move || { + let st = child.wait(); + let _ = r1.join(); + let _ = r2.join(); + let code = st.ok().and_then(|s| s.code()); + *out_r.exited.lock().unwrap() = code; + out_r.cv.notify_all(); + }) + .unwrap() + }; + let id = app.next_session.fetch_add(1, Ordering::Relaxed) + 1; + app.sessions.lock().unwrap().insert( + id, + Arc::new(Session { + pid, + stdin: Mutex::new(Some(stdin)), + out, + }), + ); + Ok(json!({"session_id": id, "shell": shell, "pid": pid})) +} + +fn shell_write(session: &Session, input: &str, newline: bool) -> Result { + if session.out.exited.lock().unwrap().is_some() { + return Err("会话已退出".into()); + } + let mut data = input.to_string(); + if newline && !data.is_empty() && !data.ends_with('\n') { + data.push('\n'); + } + let mut stdin_guard = session.stdin.lock().unwrap(); + let mut stdin = stdin_guard.take().ok_or("会话标准输入已关闭")?; + let (tx, rx) = mpsc::channel(); + let n = data.len(); + thread::spawn(move || { + let r = stdin.write_all(data.as_bytes()); + let _ = tx.send((r, stdin)); // 把 stdin 一起传回,便于后续复用 + }); + match rx.recv_timeout(SESSION_WRITE_TIMEOUT) { + Ok((Ok(()), stdin)) => { + *stdin_guard = Some(stdin); + Ok(json!({"ok": true, "bytes_written": n})) + } + Ok((Err(e), _)) => Err(format!("写入失败(会话可能已退出): {e}")), + Err(_) => Err("写入超时(进程未读取输入),可用 shell_close 关闭会话".into()), + } +} + +fn shell_read(session: &Session, timeout_ms: u64, max_bytes: usize) -> Value { + let out = &session.out; + let deadline = (timeout_ms != 0).then(|| Instant::now() + Duration::from_millis(timeout_ms)); + loop { + // 1) 有新输出则返回 + { + let b = out.bytes.lock().unwrap(); + let mut off = out.off.lock().unwrap(); + if *off > b.len() { + *off = 0; // 缓冲被裁剪,重置偏移 + } + if *off < b.len() { + let take = (b.len() - *off).min(max_bytes); + let s = String::from_utf8_lossy(&b[*off..*off + take]).into_owned(); + *off += take; + let truncated = *off < b.len(); + drop(off); + let exited = *out.exited.lock().unwrap(); + let dropped = out.dropped.load(Ordering::Relaxed); + let prefix = if dropped > 0 { + format!("[…已丢弃 {dropped} 字节…]\n") + } else { + String::new() + }; + return json!({ + "output": prefix + &s, + "exited": exited.is_some(), + "exit_code": exited, + "truncated": truncated, + }); + } + } + // 2) 无新输出但已退出 + if let Some(code) = *out.exited.lock().unwrap() { + return json!({"output": "", "exited": true, "exit_code": code, "truncated": false}); + } + // 3) 等待新输出或退出 + let wait = match deadline { + Some(d) => d.saturating_duration_since(Instant::now()), + None => Duration::from_secs(365 * 24 * 3600), + }; + if wait.is_zero() { + return json!({"output": "", "exited": false, "exit_code": null, "truncated": false}); + } + let gen = out.gen.load(Ordering::Relaxed); + let (guard, tres) = out.cv.wait_timeout(out.bytes.lock().unwrap(), wait).unwrap(); + drop(guard); + if out.gen.load(Ordering::Relaxed) != gen || out.exited.lock().unwrap().is_some() { + continue; // 有新事件,重新检查 + } + if tres.timed_out() { + return json!({"output": "", "exited": false, "exit_code": null, "truncated": false}); + } + } +} + +fn shell_close(app: &Shared, id: u64) -> Result { + let s = app + .sessions + .lock() + .unwrap() + .remove(&id) + .ok_or_else(|| format!("会话不存在: {id}"))?; + kill_tree(s.pid); + drop(s.stdin.lock().unwrap().take()); // 关闭 stdin,子进程看到 EOF + Ok(json!({"closed": true, "session_id": id})) +} + +// --------------------------------------------------------------------------- +// 工具列表(静态) +// --------------------------------------------------------------------------- + +static TOOLS: OnceLock = OnceLock::new(); + +fn tools_list() -> &'static Value { + TOOLS.get_or_init(|| { + json!({ + "tools": [ + { + "name": "run_command", + "description": "在服务器上执行 shell 命令并捕获输出。shell=true(默认)时支持完整 shell 语法:管道、重定向、&&、环境变量等;shell=false 时配合 args 直接执行二进制(不经 shell)。输出超限自动截断,超时自动杀进程,客户端可通过 notifications/cancelled 取消。", + "inputSchema": { + "type": "object", + "properties": { + "command": {"type": "string", "description": "要执行的命令行(shell=true 时为 shell 语法)"}, + "args": {"type": "array", "items": {"type": "string"}, "description": "shell=false 时的直接执行参数"}, + "shell": {"type": "boolean", "description": "是否经系统 shell 执行(sh -c / cmd /C)。默认 true"}, + "cwd": {"type": "string", "description": "工作目录,默认服务器配置目录"}, + "env": {"type": "object", "additionalProperties": {"type": "string"}, "description": "附加环境变量"}, + "timeout_ms": {"type": "integer", "description": "超时毫秒,0=无限,默认 30000,上限 21600000"}, + "max_output_bytes": {"type": "integer", "description": "stdout/stderr 各自捕获上限,默认 1048576"} + }, + "required": ["command"] + } + }, + { + "name": "shell_open", + "description": "启动持久化交互式 shell 会话(cwd、环境变量等状态跨次保留)。与 shell_write / shell_read 配合使用。无 PTY,交互式程序可能表现为非交互模式。", + "inputSchema": { + "type": "object", + "properties": { + "shell": {"type": "string", "description": "shell 程序,如 sh/bash/zsh/fish,默认 sh(windows 为 cmd)"}, + "cwd": {"type": "string", "description": "会话工作目录"} + } + } + }, + { + "name": "shell_write", + "description": "向会话写入一行输入(自动补换行)。", + "inputSchema": { + "type": "object", + "properties": { + "session_id": {"type": "integer", "description": "shell_open 返回的会话 ID"}, + "input": {"type": "string", "description": "要写入的内容"}, + "newline": {"type": "boolean", "description": "是否自动补换行,默认 true"} + }, + "required": ["session_id", "input"] + } + }, + { + "name": "shell_read", + "description": "读取会话自上次读取以来的新输出;若无输出则阻塞至 timeout_ms(0 = 等到会话退出)。", + "inputSchema": { + "type": "object", + "properties": { + "session_id": {"type": "integer"}, + "timeout_ms": {"type": "integer", "description": "等待毫秒,默认 5000,0=等到退出"}, + "max_bytes": {"type": "integer", "description": "本次最多返回字节数,默认 1048576"} + }, + "required": ["session_id"] + } + }, + { + "name": "shell_close", + "description": "关闭并杀掉会话进程树。", + "inputSchema": { + "type": "object", + "properties": { + "session_id": {"type": "integer"} + }, + "required": ["session_id"] + } + } + ] + }) + }) +} + +// --------------------------------------------------------------------------- +// MCP 分发 +// --------------------------------------------------------------------------- + +fn initialize_result(params: &Value) -> Value { + let client_version = params.get("protocolVersion").and_then(Value::as_str).unwrap_or(""); + let version = if PROTOCOL_VERSIONS.contains(&client_version) { + client_version.to_string() + } else { + LATEST_PROTOCOL.to_string() + }; + json!({ + "protocolVersion": version, + "capabilities": {"tools": {"listChanged": false}}, + "serverInfo": {"name": "bolt", "version": VERSION}, + "instructions": "你是连接到 bolt 命令执行服务器的 AI。你可以用 run_command 在服务器上执行 shell 命令(支持管道、重定向、超时与取消);需要保留状态(如 cd、环境变量)时,用 shell_open 开启持久会话,配合 shell_write/shell_read 交互。注意命令受服务器策略限制,输出可能被截断。" + }) +} + +/// 执行工具调用,返回工具结果(Ok)或协议错误(Err((code, msg))) +/// killable:外部可取消句柄(注册在 inflight 表中,供 notifications/cancelled 使用); +/// 为 None 时(测试或内部调用)自动创建本地句柄。 +fn exec_tool_call(app: &Shared, name: &str, args: &Value, progress: Option<&Value>, killable: Option<&Killable>) -> Result { + match name { + "run_command" => { + let req = parse_run_params(args, &app.cfg).map_err(|e| (-32602, e))?; + let token = req.command.split_whitespace().next().unwrap_or(""); + if let Err(e) = policy_check(&app.cfg, token) { + return Ok(tool_error(&e)); + } + let local_k; + let k = match killable { + Some(k) => k, + None => { + local_k = Killable::default(); + &local_k + } + }; + match exec(&req, k, progress.cloned(), app) { + Ok(o) => Ok(run_result(&o, req.cwd.as_deref())), + Err(e) => Ok(tool_error(&e)), + } + } + "shell_open" => { + let shell = args.get("shell").and_then(Value::as_str); + let cwd = args.get("cwd").and_then(Value::as_str); + match shell_open(app, shell, cwd) { + Ok(v) => Ok(v), + Err(e) => Ok(tool_error(&e)), + } + } + "shell_write" => { + let id = args + .get("session_id") + .and_then(Value::as_u64) + .ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?; + let input = args + .get("input") + .and_then(Value::as_str) + .ok_or_else(|| (-32602, "'input' 必须是字符串".to_string()))?; + let newline = args.get("newline").and_then(Value::as_bool).unwrap_or(true); + let session = match app.sessions.lock().unwrap().get(&id) { + Some(s) => s.clone(), + None => return Ok(tool_error(&format!("会话不存在: {id}"))), + }; + match shell_write(&session, input, newline) { + Ok(v) => Ok(v), + Err(e) => Ok(tool_error(&e)), + } + } + "shell_read" => { + let id = args + .get("session_id") + .and_then(Value::as_u64) + .ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?; + let timeout_ms = args.get("timeout_ms").and_then(Value::as_u64).unwrap_or(5000); + let max_bytes = args + .get("max_bytes") + .and_then(Value::as_u64) + .map(|n| n as usize) + .unwrap_or(DEFAULT_MAX_OUTPUT); + let session = match app.sessions.lock().unwrap().get(&id) { + Some(s) => s.clone(), + None => return Ok(tool_error(&format!("会话不存在: {id}"))), + }; + Ok(shell_read(&session, timeout_ms, max_bytes)) + } + "shell_close" => { + let id = args + .get("session_id") + .and_then(Value::as_u64) + .ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?; + match shell_close(app, id) { + Ok(v) => Ok(v), + Err(e) => Ok(tool_error(&e)), + } + } + _ => Err((-32602, format!("未知工具: {name}"))), + } +} + +// 处理 messages,返回需要发送的 JSON-RPC 响应(通知/异步工具调用返回 None) +fn process_message(v: &Value, app: &Shared, async_exec: bool) -> Option { + let obj = match v.as_object() { + Some(o) => o, + None => return Some(rpc_error(&Value::Null, -32600, "invalid request")), + }; + let method = match obj.get("method").and_then(Value::as_str) { + Some(m) => m, + None => return Some(rpc_error(obj.get("id").unwrap_or(&Value::Null), -32600, "invalid request")), + }; + let params = obj.get("params").cloned().unwrap_or_else(|| json!({})); + let id = obj.get("id").cloned(); + + if id.is_none() { + // 通知 + match method { + "notifications/initialized" => {} + "notifications/cancelled" => { + if let Some(rid) = params.get("requestId") { + let key = id_key(rid); + if let Some(k) = app.inflight.lock().unwrap().get(&key).cloned() { + k.kill(true); + } + } + } + "notifications/exit" => { + app.exit_flag.store(true, Ordering::SeqCst); + } + _ => {} + } + return None; + } + + let id = id.unwrap(); + let resp = match method { + "initialize" => json!({"jsonrpc": "2.0", "id": id, "result": initialize_result(¶ms)}), + "ping" => json!({"jsonrpc": "2.0", "id": id, "result": {}}), + "tools/list" => json!({"jsonrpc": "2.0", "id": id, "result": tools_list()}), + "tools/call" => { + let name = params.get("name").and_then(Value::as_str).unwrap_or(""); + let args = params.get("arguments").cloned().unwrap_or_else(|| json!({})); + let progress = params.get("_meta").and_then(|m| m.get("progressToken")).cloned(); + if async_exec { + // stdio 模式:worker 线程执行,完成后经 sink 发响应(不阻塞主循环) + let app2 = app.clone(); + let id2 = id.clone(); + let name2 = name.to_string(); + let args2 = args; + let tok2 = progress; + let k = register_inflight(app, &id); + thread::Builder::new() + .name("bolt-tool".into()) + .spawn(move || { + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + exec_tool_call(&app2, &name2, &args2, tok2.as_ref(), Some(&k)) + })); + unregister_inflight(&app2, &id2, &k); + let resp = match r { + Ok(Ok(v)) => json!({"jsonrpc": "2.0", "id": id2, "result": v}), + Ok(Err((code, msg))) => rpc_error(&id2, code, &msg), + Err(_) => rpc_error(&id2, -32603, "internal error"), + }; + app2.sink.send(&resp.to_string()); + }) + .unwrap(); + return None; + } else { + let k = register_inflight(app, &id); + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + exec_tool_call(app, name, &args, progress.as_ref(), Some(&k)) + })); + unregister_inflight(app, &id, &k); + match r { + Ok(Ok(v)) => json!({"jsonrpc": "2.0", "id": id, "result": v}), + Ok(Err((code, msg))) => rpc_error(&id, code, &msg), + Err(_) => rpc_error(&id, -32603, "internal error"), + } + } + } + "shutdown" => json!({"jsonrpc": "2.0", "id": id, "result": {}}), + _ => rpc_error(&id, -32601, &format!("method not found: {method}")), + }; + Some(resp) +} + +// --------------------------------------------------------------------------- +// stdio 传输 +// --------------------------------------------------------------------------- + +/// 读取一行(以 \n 结尾,不含换行),上限 cap 字节;返回 false 表示 EOF +fn read_line_capped(r: &mut R, cap: usize, buf: &mut Vec) -> io::Result { + loop { + let (newline_at, n) = { + let available = r.fill_buf()?; + if available.is_empty() { + return Ok(!buf.is_empty()); + } + let room = cap.saturating_sub(buf.len()); + if room == 0 { + return Ok(true); // 超限,由调用方报 parse error + } + let upto = available.len().min(room); + let nl = available[..upto].iter().position(|&b| b == b'\n'); + match nl { + Some(i) => { + buf.extend_from_slice(&available[..i]); + (true, i + 1) + } + None => { + buf.extend_from_slice(&available[..upto]); + (false, upto) + } + } + }; + r.consume(n); + if newline_at { + return Ok(true); + } + } +} + +fn run_stdio(app: Shared) -> io::Result<()> { + let stdin = io::stdin(); + let mut reader = BufReader::with_capacity(64 * 1024, stdin); + let mut buf: Vec = Vec::with_capacity(512); + loop { + // 顶部检查退出标志:exit 通知/信号处理后,即使 stdin 仍阻塞读取也能退出 + if app.exit_flag.load(Ordering::SeqCst) { + break; + } + buf.clear(); + match read_line_capped(&mut reader, MAX_LINE_BYTES, &mut buf) { + Ok(true) => {} + Ok(false) => break, // EOF:客户端关闭 stdin + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(e) => { + eprintln!("[bolt] stdin 错误: {e}"); + break; + } + } + // 去掉尾部 \r\n + let line = match buf.iter().rposition(|&b| b == b'\n') { + Some(i) => &buf[..i], + None => &buf[..], + }; + let line = match line.iter().rposition(|&b| b == b'\r') { + Some(i) if i + 1 == line.len() => &line[..i], + _ => line, + }; + if line.is_empty() { + continue; + } + let v = match serde_json::from_slice::(line) { + Ok(v) => v, + Err(_) => { + app.sink.send(PARSE_ERROR); + continue; + } + }; + if let Some(resp) = process_message(&v, &app, true) { + app.sink.send(&resp.to_string()); + } + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Streamable HTTP 传输 +// --------------------------------------------------------------------------- + +struct ReqHead { + method: String, + path: String, + conn_close: bool, + headers: HashMap, +} + +fn read_http_head(r: &mut R) -> io::Result>> { + let mut buf = Vec::with_capacity(512); + loop { + if buf.len() > MAX_HTTP_HEAD { + return Err(io::Error::new(io::ErrorKind::InvalidData, "header too large")); + } + let n = r.read_until(b'\n', &mut buf)?; + if n == 0 { + return Ok(if buf.is_empty() { None } else { Some(buf) }); + } + // 空行结束头部(\r\n\r\n 或 \n\n) + let len = buf.len(); + if len >= 2 && buf[len - 1] == b'\n' && buf[len - 2] == b'\n' { + return Ok(Some(buf)); + } + if len >= 4 && buf[len - 1] == b'\n' && buf[len - 2] == b'\r' && buf[len - 3] == b'\n' && buf[len - 4] == b'\r' { + return Ok(Some(buf)); + } + } +} + +fn parse_http_head(buf: &[u8]) -> Result { + let s = std::str::from_utf8(buf).map_err(|_| ())?; + let mut lines = s.lines(); + let req_line = lines.next().ok_or(())?; + let mut parts = req_line.split_whitespace(); + let method = parts.next().ok_or(())?.to_string(); + let raw_path = parts.next().ok_or(())?.to_string(); + let path = raw_path.split('?').next().unwrap_or("").to_string(); + let mut headers = HashMap::new(); + let mut conn_close = false; + for line in lines { + let line = line.trim_end_matches('\r'); + if line.is_empty() { + continue; + } + if let Some((k, v)) = line.split_once(':') { + let key = k.trim().to_ascii_lowercase(); + let val = v.trim().to_string(); + if key == "connection" && val.eq_ignore_ascii_case("close") { + conn_close = true; + } + headers.insert(key, val); + } + } + Ok(ReqHead { + method, + path, + conn_close, + headers, + }) +} + +fn origin_allowed(cfg: &Config, origin: &str) -> bool { + if cfg.allow_origins.iter().any(|o| origin == o || origin.starts_with(&format!("{o}:"))) { + return true; + } + let host = origin + .split("://") + .nth(1) + .unwrap_or(origin) + .split(':') + .next() + .unwrap_or(""); + host == "localhost" || host == "127.0.0.1" || host == "[::1]" +} + +fn respond_plain(w: &mut impl Write, status: &str, body: &[u8]) -> io::Result<()> { + write!( + w, + "HTTP/1.1 {status}\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: {}\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n", + body.len() + )?; + w.write_all(body)?; + w.flush() +} + +fn respond_json(w: &mut impl Write, status: &str, body: &str) -> io::Result<()> { + write!( + w, + "HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n", + body.len() + )?; + w.write_all(body.as_bytes())?; + w.flush() +} + +fn respond_options(w: &mut impl Write) -> io::Result<()> { + write!( + w, + "HTTP/1.1 204 No Content\r\nAccess-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type, Authorization, Accept\r\nAccess-Control-Max-Age: 86400\r\nConnection: keep-alive\r\n\r\n" + )?; + w.flush() +} + +fn respond_401(w: &mut impl Write) -> io::Result<()> { + write!( + w, + "HTTP/1.1 401 Unauthorized\r\nWWW-Authenticate: Bearer\r\nContent-Length: 0\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n" + )?; + w.flush() +} + +/// POST 的单事件 SSE 响应(规范:响应发送后服务器 SHOULD 关闭流) +fn write_sse_event(w: &mut impl Write, data: &str) -> io::Result<()> { + write!( + w, + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\nConnection: close\r\nAccess-Control-Allow-Origin: *\r\n\r\nevent: message\r\ndata: {data}\r\n\r\n" + )?; + w.flush() +} + +/// GET:长连接 SSE 流,接收服务器通知,15s 心跳保活 +fn sse_stream(w: &mut impl Write, app: &Shared) { + let _ = write!( + w, + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n" + ); + let _ = w.flush(); + let (tx, rx) = mpsc::channel::(); + let id = app.hub.add(tx); + loop { + match rx.recv_timeout(SSE_HEARTBEAT) { + Ok(msg) => { + let mut ev = String::with_capacity(msg.len() + 32); + ev.push_str("event: message\r\ndata: "); + ev.push_str(&msg); + ev.push_str("\r\n\r\n"); + if w.write_all(ev.as_bytes()).and_then(|_| w.flush()).is_err() { + break; + } + } + Err(mpsc::RecvTimeoutError::Timeout) => { + if w.write_all(b": ping\r\n\r\n").and_then(|_| w.flush()).is_err() { + break; + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + } + app.hub.remove(id); +} + +fn conn_loop(stream: TcpStream, app: Shared) { + let _ = stream.set_read_timeout(Some(Duration::from_secs(120))); + let _ = stream.set_write_timeout(Some(Duration::from_secs(60))); + let reader_stream = match stream.try_clone() { + Ok(s) => s, + Err(_) => return, + }; + let mut reader = BufReader::with_capacity(16 * 1024, reader_stream); + let mut writer = BufWriter::with_capacity(16 * 1024, stream); + + loop { + let head = match read_http_head(&mut reader) { + Ok(Some(h)) => h, + Ok(None) => return, // 客户端关闭 + Err(_) => { + let _ = respond_plain(&mut writer, "431 Request Header Fields Too Large", b"header too large"); + return; + } + }; + let req = match parse_http_head(&head) { + Ok(r) => r, + Err(_) => { + let _ = respond_plain(&mut writer, "400 Bad Request", b"bad request"); + return; + } + }; + if req.method == "OPTIONS" { + let _ = respond_options(&mut writer); + continue; + } + // 防 DNS rebinding:校验 Origin(规范 MUST) + if let Some(origin) = req.headers.get("origin") { + if !origin_allowed(&app.cfg, origin) { + let _ = respond_plain(&mut writer, "403 Forbidden", b"origin not allowed"); + continue; + } + } + if let Some(tok) = &app.cfg.token { + let ok = req + .headers + .get("authorization") + .map(|a| a == &format!("Bearer {tok}")) + .unwrap_or(false); + if !ok { + let _ = respond_401(&mut writer); + continue; + } + } + if req.path != "/mcp" { + let _ = respond_plain(&mut writer, "404 Not Found", b"not found"); + return; + } + match req.method.as_str() { + "GET" => { + sse_stream(&mut writer, &app); + return; + } + "POST" => { + let len: usize = req + .headers + .get("content-length") + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + if len > MAX_HTTP_BODY { + let _ = respond_plain(&mut writer, "413 Payload Too Large", b"body too large"); + return; + } + let mut body = vec![0u8; len]; + if reader.read_exact(&mut body).is_err() { + return; + } + let v: Value = match serde_json::from_slice(&body) { + Ok(v) => v, + Err(_) => { + let e = rpc_error(&Value::Null, -32700, "parse error"); + let _ = respond_json(&mut writer, "400 Bad Request", &e.to_string()); + return; + } + }; + if !v.is_object() { + let e = rpc_error(&Value::Null, -32600, "invalid request"); + let _ = respond_json(&mut writer, "400 Bad Request", &e.to_string()); + return; + } + let wants_sse = req + .headers + .get("accept") + .map(|a| a.contains("text/event-stream")) + .unwrap_or(false); + match process_message(&v, &app, false) { + Some(resp) => { + let s = resp.to_string(); + if wants_sse { + let _ = write_sse_event(&mut writer, &s); + return; // 规范:响应后关闭 SSE 流 + } + let _ = respond_json(&mut writer, "200 OK", &s); + } + None => { + // 通知 → 202 Accepted,无 body + let _ = respond_plain(&mut writer, "202 Accepted", b""); + } + } + if req.conn_close { + return; + } + } + _ => { + let _ = respond_plain(&mut writer, "405 Method Not Allowed", b"method not allowed"); + return; + } + } + } +} + +#[cfg(unix)] +fn run_http(addr: &str, app: Shared) -> io::Result<()> { + use std::os::unix::io::AsRawFd; + let listener = TcpListener::bind(addr)?; + eprintln!("[bolt] {VERSION} 监听 http://{addr}/mcp (Streamable HTTP)"); + log_policy(&app); + let fd = listener.as_raw_fd(); + let mut pollfd = libc::pollfd { + fd, + events: libc::POLLIN, + revents: 0, + }; + loop { + if app.exit_flag.load(Ordering::SeqCst) { + break; + } + let r = unsafe { libc::poll(&mut pollfd, 1, 200) }; + if r > 0 { + match listener.accept() { + Ok((s, _)) => { + let _ = s.set_nodelay(true); + let app2 = app.clone(); + thread::Builder::new() + .name("bolt-http".into()) + .spawn(move || conn_loop(s, app2)) + .unwrap(); + } + Err(e) => eprintln!("[bolt] accept 错误: {e}"), + } + } else if r < 0 { + let err = io::Error::last_os_error(); + if err.kind() != io::ErrorKind::Interrupted { + eprintln!("[bolt] poll 错误: {err}"); + break; + } + } + } + Ok(()) +} + +#[cfg(windows)] +fn run_http(addr: &str, app: Shared) -> io::Result<()> { + let listener = TcpListener::bind(addr)?; + eprintln!("[bolt] {VERSION} 监听 http://{addr}/mcp (Streamable HTTP)"); + log_policy(&app); + loop { + match listener.accept() { + Ok((s, _)) => { + let _ = s.set_nodelay(true); + let app2 = app.clone(); + thread::Builder::new() + .name("bolt-http".into()) + .spawn(move || conn_loop(s, app2)) + .unwrap(); + } + Err(e) => eprintln!("[bolt] accept 错误: {e}"), + } + } +} + +fn log_policy(app: &Shared) { + if app.cfg.token.is_some() { + eprintln!("[bolt] 认证: Bearer Token 已启用"); + } + if !app.cfg.allow.is_empty() { + eprintln!("[bolt] 白名单: {}", app.cfg.allow.join(", ")); + } + if !app.cfg.deny.is_empty() { + eprintln!("[bolt] 黑名单: {}", app.cfg.deny.join(", ")); + } +} + +// --------------------------------------------------------------------------- +// 信号与清理 +// --------------------------------------------------------------------------- + +#[cfg(unix)] +fn install_signal_flag(flag: Arc) { + unsafe { + let mut set: libc::sigset_t = std::mem::zeroed(); + libc::sigemptyset(&mut set); + libc::sigaddset(&mut set, libc::SIGINT); + libc::sigaddset(&mut set, libc::SIGTERM); + libc::pthread_sigmask(libc::SIG_BLOCK, &set, std::ptr::null_mut()); + thread::spawn(move || { + let mut sig: libc::c_int = 0; + libc::sigwait(&set, &mut sig); + flag.store(true, Ordering::SeqCst); + eprintln!("[bolt] 收到信号 {sig},正在清理并退出…"); + }); + } +} + +#[cfg(not(unix))] +fn install_signal_flag(_flag: Arc) {} + +/// 退出前杀掉所有进行中的命令与会话,避免孤儿进程 +fn shutdown_cleanup(app: &Shared) { + let inflight: Vec> = app.inflight.lock().unwrap().values().cloned().collect(); + for k in inflight { + k.kill(true); + } + let sessions: Vec> = app.sessions.lock().unwrap().drain().map(|(_, s)| s).collect(); + for s in sessions { + kill_tree(s.pid); + } +} + +// --------------------------------------------------------------------------- +// 入口 +// --------------------------------------------------------------------------- + +fn main() { + let Some(cfg) = parse_args() else { + print_help(); + return; + }; + let app = Arc::new(App::new(cfg)); + install_signal_flag(app.exit_flag.clone()); + let res = match &app.cfg.mode { + Mode::Stdio => run_stdio(app.clone()), + Mode::Http(addr) => run_http(addr, app.clone()), + }; + if let Err(e) = res { + eprintln!("[bolt] 致命错误: {e}"); + } + shutdown_cleanup(&app); +} + +// --------------------------------------------------------------------------- +// 单元测试 +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_size() { + assert_eq!(parse_size("1024"), Some(1024)); + assert_eq!(parse_size("64K"), Some(64 * 1024)); + assert_eq!(parse_size("1M"), Some(1024 * 1024)); + assert_eq!(parse_size("2G"), Some(2 * 1024 * 1024 * 1024)); + assert_eq!(parse_size("x"), None); + } + + #[test] + fn test_policy() { + let cfg = Config { + mode: Mode::Stdio, + token: None, + allow: vec!["git".into(), "docker".into()], + deny: vec!["rm".into()], + timeout_ms: 1000, + cwd: None, + max_output: 1024, + allow_origins: vec![], + }; + assert!(policy_check(&cfg, "git").is_ok()); + assert!(policy_check(&cfg, "git status").is_ok()); + assert!(policy_check(&cfg, "docker ps").is_ok()); + assert!(policy_check(&cfg, "curl").is_err()); + let cfg2 = Config { + allow: vec![], + deny: vec!["rm".into()], + ..cfg + }; + assert!(policy_check(&cfg2, "rm -rf /").is_err()); + assert!(policy_check(&cfg2, "ls").is_ok()); + } + + #[test] + fn test_read_line_capped() { + let data = b"hello\nworld\r\n"; + let mut r = BufReader::new(&data[..]); + let mut buf = Vec::new(); + assert!(read_line_capped(&mut r, 1024, &mut buf).unwrap()); + assert_eq!(buf, b"hello"); + buf.clear(); + assert!(read_line_capped(&mut r, 1024, &mut buf).unwrap()); + assert_eq!(buf, b"world\r"); + buf.clear(); + assert!(!read_line_capped(&mut r, 1024, &mut buf).unwrap()); + // 超限:返回 true 但截断 + let mut r2 = BufReader::new(&b"abcdef"[..]); + let mut buf2 = Vec::new(); + assert!(read_line_capped(&mut r2, 3, &mut buf2).unwrap()); + assert_eq!(buf2.len(), 3); + } + + #[test] + fn test_http_head_parse() { + let head = b"POST /mcp HTTP/1.1\r\nHost: x\r\nContent-Length: 10\r\nConnection: close\r\nAccept: application/json, text/event-stream\r\n\r\n"; + let req = parse_http_head(head).unwrap(); + assert_eq!(req.method, "POST"); + assert_eq!(req.path, "/mcp"); + assert!(req.conn_close); + assert_eq!(req.headers.get("content-length").unwrap(), "10"); + assert!(req.headers.get("accept").unwrap().contains("text/event-stream")); + } + + #[test] + fn test_origin() { + let cfg = Config { + mode: Mode::Stdio, + token: None, + allow: vec![], + deny: vec![], + timeout_ms: 1000, + cwd: None, + max_output: 1024, + allow_origins: vec!["https://example.com".into()], + }; + assert!(origin_allowed(&cfg, "http://localhost:3000")); + assert!(origin_allowed(&cfg, "https://127.0.0.1")); + assert!(origin_allowed(&cfg, "https://example.com")); + assert!(!origin_allowed(&cfg, "https://evil.com")); + } + + #[test] + fn test_dispatch_basic() { + let app = Arc::new(App::new(Config { + mode: Mode::Stdio, + token: None, + allow: vec![], + deny: vec![], + timeout_ms: 1000, + cwd: None, + max_output: 1024, + allow_origins: vec![], + })); + // ping + let v: Value = serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#).unwrap(); + let r = process_message(&v, &app, false).unwrap(); + assert_eq!(r["result"], json!({})); + assert_eq!(r["id"], json!(1)); + // initialize + let v: Value = serde_json::from_str( + r#"{"jsonrpc":"2.0","id":2,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}"#, + ) + .unwrap(); + let r = process_message(&v, &app, false).unwrap(); + assert_eq!(r["result"]["protocolVersion"], "2024-11-05"); + assert_eq!(r["result"]["serverInfo"]["name"], "bolt"); + // 未知方法 + let v: Value = serde_json::from_str(r#"{"jsonrpc":"2.0","id":3,"method":"nope"}"#).unwrap(); + let r = process_message(&v, &app, false).unwrap(); + assert_eq!(r["error"]["code"], -32601); + // 未知工具 + let v: Value = serde_json::from_str( + r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"nope","arguments":{}}}"#, + ) + .unwrap(); + let r = process_message(&v, &app, false).unwrap(); + assert_eq!(r["error"]["code"], -32602); + // 参数错误 + let v: Value = serde_json::from_str( + r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"run_command","arguments":{}}}"#, + ) + .unwrap(); + let r = process_message(&v, &app, false).unwrap(); + assert_eq!(r["error"]["code"], -32602); + // 解析失败走 PARSE_ERROR + assert!(PARSE_ERROR.contains("-32700")); + } + + #[test] + fn test_run_command_basic() { + let app = Arc::new(App::new(Config { + mode: Mode::Stdio, + token: None, + allow: vec![], + deny: vec![], + timeout_ms: 5000, + cwd: None, + max_output: 1024 * 1024, + allow_origins: vec![], + })); + let args = json!({"command": "echo hello && printf err >&2 && exit 3"}); + let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap(); + assert_eq!(r["isError"], Value::Null); + let sc = &r["structuredContent"]; + assert_eq!(sc["exit_code"], 3); + assert_eq!(sc["stdout"], "hello\n"); + assert_eq!(sc["stderr"], "err"); + assert!(!sc["timed_out"].as_bool().unwrap()); + // 超时 + let args = json!({"command": "sleep 10", "timeout_ms": 300}); + let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap(); + assert_eq!(r["structuredContent"]["timed_out"], true); + // shell=false 直接执行 + let args = json!({"command": "echo", "args": ["direct"], "shell": false}); + let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap(); + assert_eq!(r["structuredContent"]["stdout"], "direct\n"); + // 白名单拦截 + let app2 = Arc::new(App::new(Config { + allow: vec!["echo".into()], + ..Config { + mode: Mode::Stdio, + token: None, + allow: vec![], + deny: vec![], + timeout_ms: 5000, + cwd: None, + max_output: 1024, + allow_origins: vec![], + } + })); + let r = exec_tool_call(&app2, "run_command", &json!({"command": "ls"}), None, None).unwrap(); + assert_eq!(r["isError"], true); + // 经 shell 执行不存在的命令:sh 启动成功,内部报错,exit 127 + let r = exec_tool_call(&app, "run_command", &json!({"command": "definitely_not_a_cmd_xyz"}), None, None).unwrap(); + assert_eq!(r["structuredContent"]["exit_code"], 127); + assert!(!r["structuredContent"]["stderr"].as_str().unwrap().is_empty()); + // shell=false 直接执行不存在的程序:spawn 失败 → isError + let r = exec_tool_call(&app, "run_command", &json!({"command": "definitely_not_a_cmd_xyz", "shell": false}), None, None).unwrap(); + assert_eq!(r["isError"], true); + } + + #[test] + fn test_shell_session() { + let app = Arc::new(App::new(Config { + mode: Mode::Stdio, + token: None, + allow: vec![], + deny: vec![], + timeout_ms: 5000, + cwd: None, + max_output: 1024, + allow_origins: vec![], + })); + let r = shell_open(&app, Some("sh"), None).unwrap(); + let sid = r["session_id"].as_u64().unwrap(); + assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo FIRST", true).is_ok()); + let r0 = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 2000, 1 << 20); + assert!(r0["output"].as_str().unwrap().contains("FIRST"), "output: {}", r0); + // 状态保留:cd 后 pwd(非交互 sh 无回显,分两步写、一次读) + assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "cd /tmp", true).is_ok()); + assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo PROBE && pwd", true).is_ok()); + // 输出分块到达,循环读取直到拿到全部(工具语义:返回当前可用的新输出) + let mut probe_out = String::new(); + for _ in 0..10 { + let r = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 500, 1 << 20); + probe_out.push_str(r["output"].as_str().unwrap()); + if probe_out.contains("PROBE") && probe_out.contains("/tmp") { + break; + } + } + assert!(probe_out.contains("PROBE"), "output: {probe_out}"); + assert!(probe_out.contains("/tmp"), "output: {probe_out}"); + assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo STATE_OK && exit", true).is_ok()); + let mut out = String::new(); + for _ in 0..20 { + let r = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 1000, 1 << 20); + out.push_str(r["output"].as_str().unwrap()); + if r["exited"].as_bool().unwrap() { + break; + } + } + assert!(out.contains("STATE_OK"), "out: {out}"); + let r = shell_close(&app, sid).unwrap(); + assert_eq!(r["closed"], true); + } +}