cad-agent/verify_mcp_tcp_hardening.py

"""TCP 收发路径加固验证。

对应 mcp_server.cpp 的四项加固：
  1) SendAll 循环发送、检查返回值（大响应不被截断）；
  2) RecvOneFrame 用「零超时探测 + JSON 完整性」替代旧的 150ms select 探测；
  3) 帧超限回 -32700 而非静默断开；
  4) PostMessage 失败释放 TaskData（无法从外部观测，靠代码审查）。

前置条件：CAD 已启动并加载 cad_mcp_frame.arx，监听 127.0.0.1:8080。
运行：python verify_mcp_tcp_hardening.py
"""

import base64
import json
import socket
import sys
import time

from dwg_review import config

PASS = 0
FAIL = 0

PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
PNG_IEND = b"IEND\xaeB`\x82"


def check(name, cond, detail=""):
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"[PASS] {name}")
    else:
        FAIL += 1
        print(f"[FAIL] {name}  {detail!r}")


def connect():
    return socket.create_connection((config.HOST, config.PORT), timeout=20)


def read_frame(sock):
    """读到一帧（以换行结尾）；返回原始 bytes（空表示对端已关闭/无数据）。"""
    buf = b""
    while b"\n" not in buf:
        try:
            chunk = sock.recv(65536)
        except (socket.timeout, OSError):
            break
        if not chunk:
            break
        buf += chunk
    return buf


def parse(raw):
    text = raw.decode("utf-8", errors="replace").strip()
    try:
        return json.loads(text) if text else None
    except ValueError:
        return None


def rpc_once(payload):
    """发一帧（自带换行），读回一帧，返回解析后的 dict（无响应返回 None）。"""
    sock = connect()
    try:
        sock.sendall(payload + b"\n")
        return parse(read_frame(sock))
    finally:
        sock.close()


def rpc(method, params=None, rid=1):
    msg = {"jsonrpc": "2.0", "id": rid, "method": method}
    if params is not None:
        msg["params"] = params
    return rpc_once(json.dumps(msg, ensure_ascii=False).encode("utf-8"))


def main():
    print("=" * 68)
    print(" TCP 收发路径加固验证")
    print("=" * 68)

    # ---- T1 大响应完整性：截图 base64 必须解码为完整 PNG ----------------
    r = rpc("tools/call", {"name": "get_viewport_screenshot", "arguments": {}}, rid=1)
    content = ((r or {}).get("result") or {}).get("content") or []
    img = next((c for c in content if c.get("type") == "image"), None)
    check("T1 截图返回 image 内容块", img is not None, r)
    if img:
        raw = base64.b64decode(img.get("data", ""))
        check("T1 截图解码为完整 PNG（含文件头与 IEND）",
              raw.startswith(PNG_MAGIC) and raw.rstrip().endswith(PNG_IEND),
              f"len={len(raw)} head={raw[:8]!r} tail={raw[-12:]!r}")
        print(f"       截图 {len(raw)} 字节")

    # ---- T2 超限帧 -> -32700（而非静默断开）------------------------------
    padding = "A" * (8 * 1024 * 1024 + 4096)  # 超过 8MB 上限
    oversized = json.dumps({
        "jsonrpc": "2.0", "id": 2, "method": "tools/call",
        "params": {"name": "__pad__", "arguments": {"pad": padding}},
    }, ensure_ascii=False).encode("utf-8")
    check("T2 构造的报文确实超过 8MB", len(oversized) > 8 * 1024 * 1024, len(oversized))

    sock = connect()
    try:
        # 服务端可能在客户端发完前就回包并关闭，sendall 会抛异常——忽略，随后仍尝试读响应
        try:
            sock.sendall(oversized + b"\n")
        except OSError:
            pass
        resp = parse(read_frame(sock))
        code = ((resp or {}).get("error") or {}).get("code")
        check("T2 超限帧 -> 顶层 error -32700", code == -32700, resp)
    finally:
        sock.close()

    # ---- T3 不发换行的旧式请求仍可用（规则 2）----------------------------
    legacy = json.dumps({"jsonrpc": "2.0", "id": 3, "method": "tools/list",
                         "params": {}}, ensure_ascii=False).encode("utf-8")
    sock = connect()
    try:
        sock.sendall(legacy)  # 不带换行，且保持连接
        resp = parse(read_frame(sock))
        tools = ((resp or {}).get("result") or {}).get("tools") or []
        check("T3 无换行的旧式请求仍能取到 tools/list", len(tools) == 26, f"got {len(tools)}")
    finally:
        sock.close()

    # ---- T4 分片间隔 >150ms 不被截断（旧实现会在此误判为整帧）------------
    part1 = (b'{"jsonrpc":"2.0","id":4,"method":"tools/call",'
             b'"params":{"name":"get_drawing_info","arguments":{}}')
    part2 = b"}\n"
    sock = connect()
    try:
        sock.sendall(part1)
        time.sleep(0.45)  # 远超旧实现的 150ms 探测窗口
        sock.sendall(part2)
        resp = parse(read_frame(sock))
        ok = (resp or {}).get("result") is not None
        check("T4 分片间隔 450ms 的请求未被截断", ok, resp)
    finally:
        sock.close()

    # ---- T5 发送后立即断开：基座不崩 -------------------------------------
    for i in range(3):
        s = connect()
        try:
            s.sendall(json.dumps({
                "jsonrpc": "2.0", "id": 50 + i, "method": "tools/call",
                "params": {"name": "get_drawing_info", "arguments": {}},
            }).encode("utf-8") + b"\n")
        except OSError:
            pass
        s.close()  # 不读响应直接断开

    time.sleep(0.5)
    r = rpc("tools/list", {}, rid=60)
    tools = ((r or {}).get("result") or {}).get("tools") or []
    check("T5 客户端中途断开后基座仍存活（tools/list 正常）", len(tools) == 26, f"got {len(tools)}")

    # ---- T6 常规路径无回归 -----------------------------------------------
    r = rpc("tools/call", {"name": "get_drawing_info", "arguments": {}}, rid=70)
    check("T6 常规请求正常返回", (r or {}).get("result") is not None, r)

    print("-" * 68)
    print(f"通过 {PASS} / 失败 {FAIL}")
    if FAIL:
        print("结果：FAIL")
        return 1
    print("结果：全部通过")
    return 0


if __name__ == "__main__":
    sys.exit(main())