Merge pull request 'fix(frame): 加固 TCP 收发路径（大响应、分帧判定、超限回包、PostMessage 泄漏）' (#11) from home-lenovo-legion-r9000/home-ai-dev into main
Reviewed-on: #11
This commit was merged in pull request #11.
This commit is contained in:
gjm committed
commit
8c07c7cad6
3 files changed
+276
-30
No files matched your search
@@ -85,6 +85,13 @@ type: project
|
||||
DLL 代理写锁改为 `ScopedDocumentWriteLock` RAII）。**插件作者无需自行 try/catch**，
|
||||
但仍不要在异常路径上留下副作用。新增/改动工具后跑 `verify_mcp_fault_injection.py` 回归。
|
||||
|
||||
**同一根因的另一面：阻塞。** 工具响应由 **AutoCAD 主线程**发送
|
||||
（`WndProc` → `ExecuteTaskInMainThread` → `SendAll`），所以发送路径也必须有时限——
|
||||
`accept` 时同时设了 `SO_RCVTIMEO`/`SO_SNDTIMEO`（各 5s），否则慢客户端能把 CAD 的 UI 线程卡死。
|
||||
分帧规则：`\n` 优先，其次「暂时无更多数据、且累积内容已是完整 JSON」（兼容不发 `\n` 的旧客户端）；
|
||||
**不要再引入基于等待时长的分帧判定**（旧的 150ms select 探测会把分片间隔大的大报文误截断）。
|
||||
回归脚本：`verify_mcp_tcp_hardening.py`。
|
||||
|
||||
## 验证结论（2026-10-13 实测通过）
|
||||
|
||||
- `verify_mcp_frame.py` 4/4 PASS（换行分帧 / >2KB 报文不截断 / string id 原样回显 / 通知不回包 / isError）。
|
||||
|
||||
@@ -43,6 +43,27 @@ namespace
|
||||
std::string rawJson;
|
||||
};
|
||||
|
||||
// 完整发送一帧：send 可能只发出部分字节（截图这类大响应尤甚），必须循环发完。
|
||||
// 失败（对端已关闭 / 发送超时）即放弃——调用方已无法补救，但至少不能卡住调用线程。
|
||||
// 注意：主线程（WndProc -> ExecuteTaskInMainThread）也会调用本函数发送工具响应，
|
||||
// 因此 clientSocket 必须设置 SO_SNDTIMEO，否则慢客户端会把 AutoCAD 的 UI 线程拖死。
|
||||
void SendAll(SOCKET sock, const std::string &data)
|
||||
{
|
||||
size_t sent = 0;
|
||||
const size_t total = data.size();
|
||||
while (sent < total)
|
||||
{
|
||||
int n = send(sock, data.data() + sent, (int)(total - sent), 0);
|
||||
if (n <= 0)
|
||||
return; // 对端已关闭或超时：放弃
|
||||
sent += static_cast<size_t>(n);
|
||||
}
|
||||
}
|
||||
|
||||
// 帧超过上限时回给客户端的协议错误（不触碰任何 ARX API，可在 TCP 线程直接发送）
|
||||
const char *kFrameTooLargeResponse =
|
||||
R"({"jsonrpc":"2.0","error":{"code":-32700,"message":"Parse error: request frame exceeds size limit"},"id":null})";
|
||||
|
||||
// 处理工具查询请求（msgId 允许 string 或 number）
|
||||
std::string HandleToolsList(const nlohmann::json &msgId)
|
||||
{
|
||||
@@ -199,7 +220,7 @@ namespace
|
||||
if (pTask->clientSocket != INVALID_SOCKET && !responseStr.empty())
|
||||
{
|
||||
responseStr.push_back('\n');
|
||||
send(pTask->clientSocket, responseStr.c_str(), (int)responseStr.length(), 0);
|
||||
SendAll(pTask->clientSocket, responseStr); // 循环发完，避免大响应被截断
|
||||
}
|
||||
}
|
||||
|
||||
@@ -219,13 +240,23 @@ namespace
|
||||
return DefWindowProc(hWnd, message, wParam, lParam);
|
||||
}
|
||||
|
||||
// RecvOneFrame 的返回状态
|
||||
enum class FrameStatus
|
||||
{
|
||||
Ok, // 成功取到一帧，内容在 out
|
||||
TooLarge, // 超过帧上限，未取到完整帧
|
||||
None // 对端未发数据即关闭，或读取失败：无内容可回
|
||||
};
|
||||
|
||||
// 从 socket 读取一整帧 JSON-RPC 报文。
|
||||
// 消息边界规则：优先以 '\n' 作为帧结束标记（新客户端）；
|
||||
// 若客户端不发 '\n'（旧脚本），则在收到数据后用 select 短暂探测是否还有后续数据，
|
||||
// 无后续数据即视为整帧，从而保持向后兼容。
|
||||
// 返回 true 表示成功取到一帧（内容写入 out）；false 表示无数据或出错。
|
||||
// 注意：本函数运行在后台 TCP 线程，只做字节累积，不解析 JSON、不触碰任何 ARX API。
|
||||
bool RecvOneFrame(SOCKET sock, std::string &out)
|
||||
// 消息边界规则（按优先级）：
|
||||
// 1) 出现 '\n' —— 新客户端的标准分帧；
|
||||
// 2) 暂时无更多数据、且累积内容已是完整合法 JSON —— 兼容不发 '\n' 的旧客户端；
|
||||
// 3) 超过上限 / 对端关闭 / 读超时 —— 结束（可能是不完整报文，交由上层报错）。
|
||||
// 旧实现用"收完 select 探测 150ms 无后续数据即判定整帧"，真实网络上分片间隔不可控，
|
||||
// 大报文会被误截断；改为 JSON 完整性判定后与分片时机无关。
|
||||
// 注意：本函数运行在后台 TCP 线程，只做字节累积与（无异常的）JSON 完整性判定，不触碰 ARX API。
|
||||
FrameStatus RecvOneFrame(SOCKET sock, std::string &out)
|
||||
{
|
||||
const size_t kMaxFrameBytes = 8 * 1024 * 1024; // 8MB 上限，防御异常/恶意大报文
|
||||
char buf[4096];
|
||||
@@ -238,40 +269,48 @@ namespace
|
||||
{
|
||||
out.append(buf, n);
|
||||
|
||||
// 命中换行符即成帧（去掉 '\n' 及之后内容，并兼容 CRLF）
|
||||
// 规则 1：命中换行符即成帧（去掉 '\n' 及之后内容，并兼容 CRLF）
|
||||
size_t pos = out.find('\n');
|
||||
if (pos != std::string::npos)
|
||||
{
|
||||
out.resize(pos);
|
||||
if (!out.empty() && out.back() == '\r')
|
||||
out.pop_back();
|
||||
return true;
|
||||
return FrameStatus::Ok;
|
||||
}
|
||||
|
||||
if (out.size() > kMaxFrameBytes)
|
||||
return false; // 超过上限，判为非法报文
|
||||
return FrameStatus::TooLarge;
|
||||
|
||||
// 尚未见到 '\n'：短暂探测是否还有后续分片。
|
||||
// 无后续数据 => 判定为"整帧已到"（兼容不发 '\n' 的旧客户端）。
|
||||
// 规则 2：仅当内核里暂时没有更多数据、且累积内容已是完整 JSON 时才判定为整帧。
|
||||
// 用零超时 select 而非固定等待：数据仍在流式到达时不解析，既不会把分片间隔较大
|
||||
// 的大报文误截断，也避免大报文被反复解析（退化成 O(n^2)）。
|
||||
fd_set probe;
|
||||
FD_ZERO(&probe);
|
||||
FD_SET(sock, &probe);
|
||||
timeval probeTimeout = { 0, 150000 }; // 150ms
|
||||
int pr = select(0, &probe, nullptr, nullptr, &probeTimeout);
|
||||
if (pr <= 0)
|
||||
return true;
|
||||
timeval zero = { 0, 0 };
|
||||
if (select(0, &probe, nullptr, nullptr, &zero) <= 0)
|
||||
{
|
||||
// 仅在"看起来可能收尾"时才尝试解析（JSON-RPC 请求是对象，以 '}' 收尾）
|
||||
size_t last = out.find_last_not_of(" \t\r\n");
|
||||
if (last != std::string::npos && (out[last] == '}' || out[last] == ']') &&
|
||||
!nlohmann::json::parse(out, nullptr, false).is_discarded())
|
||||
return FrameStatus::Ok;
|
||||
}
|
||||
}
|
||||
else if (n == 0)
|
||||
{
|
||||
// 对端正常关闭：若已收到数据则视为整帧
|
||||
return !out.empty();
|
||||
return out.empty() ? FrameStatus::None : FrameStatus::Ok;
|
||||
}
|
||||
else
|
||||
{
|
||||
int err = WSAGetLastError();
|
||||
// 读超时/暂时无数据但已有内容：视为整帧（不完整报文由上层报 -32700）
|
||||
if ((err == WSAETIMEDOUT || err == WSAEWOULDBLOCK) && !out.empty())
|
||||
return true; // 超时但已有数据：视为整帧
|
||||
return false;
|
||||
return FrameStatus::Ok;
|
||||
// 其它错误（对端重置等）：有内容就当整帧，否则无内容可回
|
||||
return out.empty() ? FrameStatus::None : FrameStatus::Ok;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -306,33 +345,56 @@ namespace
|
||||
SOCKET clientSocket = accept(g_listenSocket, nullptr, nullptr);
|
||||
if (clientSocket != INVALID_SOCKET)
|
||||
{
|
||||
// 设置接收超时：保证插件卸载时后台线程能及时退出，不会永久阻塞在 recv 上
|
||||
DWORD recvTimeoutMs = 5000;
|
||||
// 收发超时：保证插件卸载时后台线程能及时退出，不会永久阻塞在 recv 上；
|
||||
// 发送侧超时尤其重要——主线程发送工具响应时若被慢客户端卡住，会拖死 CAD 的 UI。
|
||||
DWORD ioTimeoutMs = 5000;
|
||||
setsockopt(clientSocket, SOL_SOCKET, SO_RCVTIMEO,
|
||||
(const char *)&recvTimeoutMs, sizeof(recvTimeoutMs));
|
||||
(const char *)&ioTimeoutMs, sizeof(ioTimeoutMs));
|
||||
setsockopt(clientSocket, SOL_SOCKET, SO_SNDTIMEO,
|
||||
(const char *)&ioTimeoutMs, sizeof(ioTimeoutMs));
|
||||
|
||||
// 按消息边界读取一整帧（不再受 2048 字节单次 recv 限制）
|
||||
std::string frame;
|
||||
if (RecvOneFrame(clientSocket, frame) && !frame.empty())
|
||||
FrameStatus status = RecvOneFrame(clientSocket, frame);
|
||||
if (status == FrameStatus::Ok)
|
||||
{
|
||||
// 将数据拷贝到堆内存，准备派发给主线程
|
||||
TaskData *pTask = new TaskData;
|
||||
pTask->clientSocket = clientSocket; // 记录 Socket
|
||||
pTask->rawJson = frame;
|
||||
|
||||
// 关键转变：通过 PostMessage 发送给主线程拥有的隐式窗口
|
||||
if (g_hMsgWnd != NULL)
|
||||
{
|
||||
PostMessage(g_hMsgWnd, WM_MCP_EXECUTE_TASK, 0, reinterpret_cast<LPARAM>(pTask));
|
||||
}
|
||||
else
|
||||
// 关键转变：通过 PostMessage 发送给主线程拥有的隐式窗口。
|
||||
// 投递成功则所有权移交主线程（McpWndProc 负责 delete）；
|
||||
// 失败（窗口已销毁 / 消息队列满）则消息永不会被处理，必须在此释放
|
||||
//（TaskData 析构会关闭 socket），否则内存泄漏。
|
||||
if (g_hMsgWnd == NULL ||
|
||||
!PostMessage(g_hMsgWnd, WM_MCP_EXECUTE_TASK, 0, reinterpret_cast<LPARAM>(pTask)))
|
||||
{
|
||||
delete pTask;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// 空帧或读取失败：直接关闭连接
|
||||
// 帧超限：明确回一条 -32700，而不是静默断开，便于客户端定位问题。
|
||||
// 注意：接收缓冲里仍有未读数据时直接 closesocket 会触发 RST，可能把刚发出的
|
||||
// 响应丢掉；故先 shutdown 发送方向（发 FIN），再把剩余输入有界地读掉，最后关闭。
|
||||
if (status == FrameStatus::TooLarge)
|
||||
{
|
||||
SendAll(clientSocket, std::string(kFrameTooLargeResponse) + "\n");
|
||||
shutdown(clientSocket, SD_SEND);
|
||||
|
||||
const size_t kMaxDrainBytes = 64 * 1024 * 1024; // 有界排空，防恶意客户端无限灌数据
|
||||
char drainBuf[4096];
|
||||
size_t drained = 0;
|
||||
while (drained < kMaxDrainBytes)
|
||||
{
|
||||
int r = recv(clientSocket, drainBuf, sizeof(drainBuf), 0);
|
||||
if (r <= 0)
|
||||
break;
|
||||
drained += static_cast<size_t>(r);
|
||||
}
|
||||
}
|
||||
// 其余情况（对端未发数据即关闭 / 读取失败）无内容可回，直接关闭
|
||||
closesocket(clientSocket);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
"""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())
|
||||
Reference in new issue
Block a user