From 8bd6f3b60ace5f5b9f043b6490b261d02cba81fa Mon Sep 17 00:00:00 2001 From: yuanyuanxiang <962914132@qq.com> Date: Sun, 23 Aug 2026 10:37:55 +0200 Subject: [PATCH] Feature: Add list_registry MCP tool Add a read-only list_registry MCP tool for querying a remote Windows host's registry. An empty path returns the five root keys; a path such as HKEY_LOCAL_MACHINE plus subkeys returns that key's immediate subkeys and values (name, type, and formatted data). Client side: enumerate with KEY_READ instead of KEY_ALL_ACCESS and map the full value-type set (REG_SZ, REG_DWORD, REG_BINARY, REG_EXPAND_SZ, REG_MULTI_SZ, REG_QWORD, REG_NONE) so unknown types are no longer mis-reported as REG_SZ. The registry manager now always sends both the TOKEN_REG_PATH and TOKEN_REG_KEY packets (an empty packet when a part is empty), making the two-packet reply deterministic for the server. Server side: add the three-phase flow (COMMAND_REGEDIT, COMMAND_REG_FIND, then PATH plus KEY) with per-host pending state, parse the fixed-width wire format with bounds checks, format value data for JSON, and guard registry paths against exceeding MAX_PATH to protect the client stack. Co-Authored-By: deepseek-v4-pro --- client/RegisterManager.cpp | 19 ++ client/RegisterOperation.cpp | 39 +-- server/2015Remote/2015RemoteDlg.cpp | 24 ++ server/2015Remote/McpServer.cpp | 424 ++++++++++++++++++++++++++++ server/2015Remote/McpServer.h | 21 +- 5 files changed, 508 insertions(+), 19 deletions(-) diff --git a/client/RegisterManager.cpp b/client/RegisterManager.cpp index cf39f2f..594f50e 100644 --- a/client/RegisterManager.cpp +++ b/client/RegisterManager.cpp @@ -9,6 +9,14 @@ #if ENABLE_REGISTRY +// 与 RegisterOperation.cpp 的 REGMSG 同构（int count + DWORD size + DWORD valsize = 12 字节）， +// 此处仅用于构造"空结果包"（count=0），无需跨文件共享完整定义。 +struct REGMSG { + int count; //名字个数 + DWORD size; //名字大小 + DWORD valsize; //值大小 +}; + ////////////////////////////////////////////////////////////////////// // Construction/Destruction ////////////////////////////////////////////////////////////////////// @@ -46,17 +54,28 @@ VOID CRegisterManager::Find(char bToken, char *szPath) if(szPath!=NULL) { Opt.SetPath(szPath); } + // 子键包（TOKEN_REG_PATH）：无子键时也回传空包，保证服务端（含 MCP 无头查询） + // 能确定性收齐 PATH+KEY 两包，无需靠超时猜测是否有第二包。 char *szBuffer= Opt.FindPath(); if(szBuffer!=NULL) { m_ClientObject->Send2Server((char*)szBuffer, LocalSize(szBuffer)); //目录下的目录 LocalFree(szBuffer); + } else { + char empty[1 + sizeof(REGMSG)] = { 0 }; + empty[0] = TOKEN_REG_PATH; + m_ClientObject->Send2Server(empty, sizeof(empty)); } + // 值包（TOKEN_REG_KEY）：同上，无值时也回传空包。 szBuffer = Opt.FindKey(); if(szBuffer!=NULL) { //目录下的文件 m_ClientObject->Send2Server((char*)szBuffer, LocalSize(szBuffer)); LocalFree(szBuffer); + } else { + char empty[1 + sizeof(REGMSG)] = { 0 }; + empty[0] = TOKEN_REG_KEY; + m_ClientObject->Send2Server(empty, sizeof(empty)); } } diff --git a/client/RegisterOperation.cpp b/client/RegisterOperation.cpp index 3bb4d56..fbb1a94 100644 --- a/client/RegisterOperation.cpp +++ b/client/RegisterOperation.cpp @@ -21,10 +21,13 @@ enum MYKEY { enum KEYVALUE { - MREG_SZ, - MREG_DWORD, - MREG_BINARY, - MREG_EXPAND_SZ + MREG_SZ, // REG_SZ + MREG_DWORD, // REG_DWORD + MREG_BINARY, // REG_BINARY + MREG_EXPAND_SZ, // REG_EXPAND_SZ + MREG_MULTI_SZ, // REG_MULTI_SZ + MREG_QWORD, // REG_QWORD + MREG_NONE // REG_NONE / 未知类型 }; struct REGMSG { @@ -70,7 +73,7 @@ char* RegisterOperation::FindPath() char *szBuffer=NULL; HKEY hKey; //注册表返回句柄 /*打开注册表 User kdjfjkf\kdjfkdjf\ */ - if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_ALL_ACCESS,&hKey)==ERROR_SUCCESS) { //打开 + if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_READ,&hKey)==ERROR_SUCCESS) { //打开 DWORD dwIndex=0,NameCount,NameMaxLen; DWORD KeySize,KeyCount,KeyMaxLen,MaxDataLen; //这就是枚举了 @@ -93,7 +96,7 @@ char* RegisterOperation::FindPath() REGMSG msg; //数据头 msg.size=KeySize; msg.count=KeyCount; - memcpy(szBuffer+1,(void*)&msg,Size); + memcpy(szBuffer+1,(void*)&msg,sizeof(REGMSG)); char * szTemp=new char[KeySize]; for(dwIndex=0; dwIndexGetClientID(); + if (McpServer().IsPending(devId) && McpServer().OnRegeditReady(devId, ContextObject)) + break; g_2015RemoteDlg->SendMessage(WM_OPENREGISTERDIALOG, 0, (LPARAM)ContextObject); break; } + case TOKEN_REG_PATH: { // 注册表子键【P3：仅 MCP 无头路径会到达】 + // 正常注册表管理由 RegisterDlg 的 hDlg 接管，TOKEN_REG_PATH 不进此分派； + // 到达此处只有 MCP 无头查询一种情形。命中则存子键包并保持子链接（等 TOKEN_REG_KEY）。 + uint64_t devId = ContextObject->GetClientID(); + if (McpServer().IsPending(devId)) { + McpServer().TakeRegPath(devId, szBuffer, len); + } else { + ContextObject->CancelIO(); // 迟到/游离包，关闭子链接避免泄漏 + } + break; + } + case TOKEN_REG_KEY: { // 注册表值【P3：仅 MCP 无头路径会到达】 + uint64_t devId = ContextObject->GetClientID(); + if (McpServer().IsPending(devId)) { + McpServer().TakeRegKey(devId, szBuffer, len); + } + ContextObject->CancelIO(); // 值包为最后包，无论命中与否用完即关 + break; + } case TOKEN_SERVERLIST: { // 服务管理【x】 // P3：MCP 挂起时接管。服务列表为一次性子链接回传（CServicesManager 构造即发）， // 取走即 CancelIO 关子链接（无 MFC 对话框续用该子链接）。 diff --git a/server/2015Remote/McpServer.cpp b/server/2015Remote/McpServer.cpp index 0144d5a..05fee98 100644 --- a/server/2015Remote/McpServer.cpp +++ b/server/2015Remote/McpServer.cpp @@ -401,6 +401,210 @@ Json::Value ParseServiceList(const std::vector& data, UINT cp) { return arr; } +// ========== list_registry：注册表解析 ========== + +// 注册表 REGMSG 头（client RegisterOperation.cpp / server RegisterDlg.cpp 各自本地定义，此处同构）。 +struct RegMsgHeader { + int count; // 名字个数 + DWORD size; // 名字大小（定宽） + DWORD valsize; // 值大小（定宽） +}; + +// client KEYVALUE 枚举（RegisterOperation.cpp）：MREG_SZ=0 … MREG_NONE=6。 +enum RegValueType { + REG_T_SZ = 0, + REG_T_DWORD = 1, + REG_T_BINARY = 2, + REG_T_EXPAND_SZ = 3, + REG_T_MULTI_SZ = 4, + REG_T_QWORD = 5, + REG_T_NONE = 6, +}; + +// client MYKEY 枚举：根键 token。 +enum RegRootToken { + REG_ROOT_CLASSES_ROOT = 0, + REG_ROOT_CURRENT_USER = 1, + REG_ROOT_LOCAL_MACHINE = 2, + REG_ROOT_USERS = 3, + REG_ROOT_CURRENT_CONFIG = 4, +}; + +// ASCII 不区分大小写前缀匹配（注册表根键名大小写不敏感）。 +bool StrPrefixCI(const std::string& s, const char* prefix, size_t len) { + if (s.size() < len) return false; + for (size_t i = 0; i < len; ++i) { + char a = s[i], b = prefix[i]; + if (a >= 'A' && a <= 'Z') a += 32; + if (b >= 'A' && b <= 'Z') b += 32; + if (a != b) return false; + } + return true; +} + +// 值类型名（可读字符串） +const char* RegTypeName(BYTE type) { + switch (type) { + case REG_T_SZ: return "REG_SZ"; + case REG_T_DWORD: return "REG_DWORD"; + case REG_T_BINARY: return "REG_BINARY"; + case REG_T_EXPAND_SZ: return "REG_EXPAND_SZ"; + case REG_T_MULTI_SZ: return "REG_MULTI_SZ"; + case REG_T_QWORD: return "REG_QWORD"; + case REG_T_NONE: return "REG_NONE"; + default: return "REG_UNKNOWN"; + } +} + +// 将注册表值数据格式化为字符串。valSize 为定宽（client 按 MaxDataLen+1 填充）， +// 字符串类取到 '\0' 为止、数值类取定长前缀；REG_BINARY 长度无法从协议还原，按 valSize 输出 hex。 +std::string FormatRegData(BYTE type, const BYTE* val, size_t valSize, UINT cp) { + switch (type) { + case REG_T_SZ: + case REG_T_EXPAND_SZ: + return ToUtf8((const char*)val, cp); + case REG_T_MULTI_SZ: { // 多个 '\0' 结尾串、双 '\0' 结束 + std::string out; + size_t i = 0; + while (i < valSize && val[i] != '\0') { + const char* s = (const char*)(val + i); + size_t n = BoundedStrlen(s, valSize - i); + if (!out.empty()) out += "\n"; + out += ToUtf8(s, cp); + i += n + 1; + } + return out; + } + case REG_T_DWORD: { // 4 字节小端 + if (valSize < 4) return ""; + DWORD v = 0; + memcpy(&v, val, 4); + char buf[64]; + sprintf(buf, "0x%08lX (%lu)", (unsigned long)v, (unsigned long)v); + return buf; + } + case REG_T_QWORD: { // 8 字节小端 + if (valSize < 8) return ""; + uint64_t v = 0; + memcpy(&v, val, 8); + char buf[96]; + sprintf(buf, "0x%016I64X (%I64u)", (unsigned __int64)v, (unsigned __int64)v); + return buf; + } + case REG_T_BINARY: { // hex，尾部可能含 0 填充 + std::string out; + char buf[4]; + for (size_t i = 0; i < valSize; ++i) { + sprintf(buf, "%02X", val[i]); + out += buf; + } + return out; + } + case REG_T_NONE: + default: + return ""; + } +} + +// 解析 TOKEN_REG_PATH 缓冲：[token:1][RegMsgHeader:12][count * size 定宽子键名]。cp 为客户端 ANSI。 +Json::Value ParseRegPath(const std::vector& data, UINT cp) { + Json::Value arr(Json::arrayValue); + if (data.size() < 1 + sizeof(RegMsgHeader)) return arr; + const char* p = (const char*)data.data(); + size_t off = 1; // 跳过 TOKEN 字节 + RegMsgHeader hdr; + memcpy(&hdr, p + off, sizeof(hdr)); + off += sizeof(hdr); + int count = hdr.count; + size_t nameSize = hdr.size; + for (int i = 0; i < count; ++i) { + if (off + nameSize > data.size()) break; + const char* name = p + off; + off += nameSize; + if (name[0] == '\0') break; // 空名 = 尾部零填充，停止 + arr.append(ToUtf8(name, cp)); + } + return arr; +} + +// 解析 TOKEN_REG_KEY 缓冲：[token:1][RegMsgHeader:12][count * {type:1,name:size,data:valsize}]。 +Json::Value ParseRegKey(const std::vector& data, UINT cp) { + Json::Value arr(Json::arrayValue); + if (data.size() < 1 + sizeof(RegMsgHeader)) return arr; + const char* p = (const char*)data.data(); + size_t off = 1; // 跳过 TOKEN 字节 + RegMsgHeader hdr; + memcpy(&hdr, p + off, sizeof(hdr)); + off += sizeof(hdr); + int count = hdr.count; + size_t nameSize = hdr.size; + size_t valSize = hdr.valsize; + for (int i = 0; i < count; ++i) { + if (off + 1 + nameSize + valSize > data.size()) break; + BYTE type = (BYTE)p[off]; + off += 1; + const char* name = p + off; // 定宽 nameSize 内以 '\0' 结尾 + off += nameSize; + const BYTE* val = (const BYTE*)(p + off); + off += valSize; + + Json::Value item(Json::objectValue); + item["name"] = ToUtf8(name, cp); + item["type"] = RegTypeName(type); + item["data"] = FormatRegData(type, val, valSize, cp); + arr.append(item); + } + return arr; +} + +// 解析 "HKEY_LOCAL_MACHINE\Software\..." → rootToken + 相对子键路径（去掉根键名与首个 '\'）。 +// 返回 false 表示 path 为空（调用方列根键）或未匹配任何根键（调用方报错）。 +bool ParseRegistryPath(const std::string& path, BYTE& rootToken, std::string& relPath) { + struct RootMap { const char* name; BYTE token; }; + static const RootMap roots[] = { + { "HKEY_CLASSES_ROOT", REG_ROOT_CLASSES_ROOT }, + { "HKEY_CURRENT_USER", REG_ROOT_CURRENT_USER }, + { "HKEY_LOCAL_MACHINE", REG_ROOT_LOCAL_MACHINE }, + { "HKEY_USERS", REG_ROOT_USERS }, + { "HKEY_CURRENT_CONFIG", REG_ROOT_CURRENT_CONFIG }, + }; + if (path.empty()) return false; + for (const RootMap& r : roots) { + size_t n = strlen(r.name); + if (StrPrefixCI(path, r.name, n)) { + rootToken = r.token; + relPath = path.substr(n); + if (!relPath.empty() && relPath[0] == '\\') relPath.erase(0, 1); + return true; + } + } + return false; +} + +// 列根键（path 为空）：返回 5 个固定根键，无需查询客户端。 +std::string BuildRegistryRoots(const Json::Value& id) { + Json::Value keys(Json::arrayValue); + const char* roots[] = { + "HKEY_CLASSES_ROOT", "HKEY_CURRENT_USER", "HKEY_LOCAL_MACHINE", + "HKEY_USERS", "HKEY_CURRENT_CONFIG" + }; + for (const char* r : roots) keys.append(r); + + Json::Value result(Json::objectValue); + Json::Value structuredContent(Json::objectValue); + structuredContent["keys"] = keys; + result["structuredContent"] = structuredContent; + + Json::Value content(Json::arrayValue); + Json::Value item(Json::objectValue); + item["type"] = "text"; + item["text"] = std::string(u8"共 ") + std::to_string((int)keys.size()) + std::string(u8" 个根键。"); + content.append(item); + result["content"] = content; + result["isError"] = false; + return BuildResult(id, result); +} + // 收集所有在线主机 JSON 数组（m_cs 锁内遍历，复用 BuildHostJson 序列化，方案 C） void CollectOnlineHosts(CMy2015RemoteDlg* parent, Json::Value& hosts) { if (!parent) return; @@ -622,6 +826,69 @@ Json::Value BuildServiceListOutputSchema() { return schema; } +// list_registry 的 inputSchema（id 必填，path 可选；path 缺省/空 = 列根键） +Json::Value BuildListRegistryInputSchema() { + Json::Value props(Json::objectValue); + + Json::Value idProp(Json::objectValue); + idProp["type"] = "string"; + idProp["description"] = u8"主机 id，取 list_online_hosts / search_hosts 返回的 id 字段"; + props["id"] = idProp; + + Json::Value pathProp(Json::objectValue); + pathProp["type"] = "string"; + pathProp["description"] = u8"注册表键路径（如 HKEY_LOCAL_MACHINE\\Software）；省略或传空则返回 5 个根键"; + props["path"] = pathProp; + + Json::Value schema(Json::objectValue); + schema["type"] = "object"; + schema["properties"] = props; + Json::Value required(Json::arrayValue); + required.append("id"); + schema["required"] = required; + return schema; +} + +// list_registry 的 outputSchema（keys 字符串数组 / values 对象数组） +Json::Value BuildListRegistryOutputSchema() { + Json::Value props(Json::objectValue); + + Json::Value pathProp(Json::objectValue); + pathProp["type"] = "string"; + props["path"] = pathProp; + + Json::Value keysProp(Json::objectValue); + keysProp["type"] = "array"; + Json::Value keyItems(Json::objectValue); + keyItems["type"] = "string"; + keysProp["items"] = keyItems; + props["keys"] = keysProp; + + Json::Value valsProp(Json::objectValue); + valsProp["type"] = "array"; + Json::Value valItems(Json::objectValue); + valItems["type"] = "object"; + Json::Value valProps(Json::objectValue); + const char* strFields[] = { "name", "type", "data" }; + for (const char* f : strFields) { + Json::Value s(Json::objectValue); + s["type"] = "string"; + valProps[f] = s; + } + valItems["properties"] = valProps; + valsProp["items"] = valItems; + props["values"] = valsProp; + + Json::Value schema(Json::objectValue); + schema["type"] = "object"; + schema["properties"] = props; + Json::Value required(Json::arrayValue); + required.append("keys"); + required.append("values"); + schema["required"] = required; + return schema; +} + // get_client_log 的 outputSchema（原始日志文本） Json::Value BuildClientLogOutputSchema() { Json::Value props(Json::objectValue); @@ -942,6 +1209,18 @@ std::string BuildToolsListResult(const Json::Value& id) { tools.append(tool); } + // 12) list_registry（P3：注册表查询，列根键/列子键+值，仅 Windows，一次性子链接） + { + Json::Value tool(Json::objectValue); + tool["name"] = "list_registry"; + tool["description"] = u8"查询指定在线 Windows 主机的注册表。省略或空 path 返回 5 个根键；给定 path（如 HKEY_LOCAL_MACHINE\\Software）返回该键一层内的子键（keys）与值（values，含 name/type/data）。只读。"; + + tool["inputSchema"] = BuildListRegistryInputSchema(); + tool["outputSchema"] = BuildListRegistryOutputSchema(); + + tools.append(tool); + } + result["tools"] = tools; return BuildResult(id, result); } @@ -1408,6 +1687,78 @@ std::string BuildListServices(const Json::Value& id, const Json::Value& args, CM return BuildResult(id, result); } +// tools/call：list_registry（主连接下发 COMMAND_REGEDIT → 子连接 TOKEN_REGEDIT → 下发 +// COMMAND_REG_FIND → TOKEN_REG_PATH + TOKEN_REG_KEY 两包；path 空 = 列根键，无需查询客户端） +std::string BuildListRegistry(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) { + uint64_t devId = 0; + std::string err; + if (!ParseHostIdArg(args, devId, err)) + return BuildError(id, -32602, err); + + context* ctx = FindMainContext(parent, devId); + if (!ctx) + return BuildError(id, -32002, "Host not found or offline: " + std::to_string(devId)); + + // 注册表仅 Windows 客户端实现（RegOpenKeyEx 等 A 接口）；LNX/MAC 无注册表。 + CString clientType = ctx->GetAdditionalData(RES_CLIENT_TYPE); + if (clientType == "LNX" || clientType == "MAC") + return BuildError(id, -32005, "list_registry is only supported on Windows hosts"); + + std::string path = GetStringArg(args, "path"); + BYTE rootToken = 0; + std::string relPath; + if (!ParseRegistryPath(path, rootToken, relPath)) { + if (path.empty()) + return BuildRegistryRoots(id); // 空 path → 列根键 + return BuildError(id, -32602, + "Invalid path: expected a root key like HKEY_LOCAL_MACHINE\\Software"); + } + + // 客户端 RegisterOperation 以 char KeyPath[MAX_PATH] 承载相对子键路径（SetPath→strcpy 无界）， + // 超长会栈溢出。按 UTF-8 字节数做保守上限（GBK/936 字节数 ≤ UTF-8，故只查 UTF-8 即可）。 + if (relPath.size() >= MAX_PATH) + return BuildError(id, -32602, "Registry path too long (exceeds MAX_PATH)"); + + // 注册表键路径为客户端 ANSI（GBK/936），下发前转好；rootToken 编码进 path[0] 供 OnRegeditReady 拆分。 + UINT cp = (clientType == "LNX" || clientType == "MAC") ? CP_UTF8 : 936; + std::string rootAndPath = std::string(1, (char)rootToken) + ToAnsi(relPath, cp); + + CMcpServer& mcp = CMcpServer::Instance(); + if (!mcp.BeginPending(devId, "list_registry", rootAndPath)) + return BuildError(id, -32003, "Device busy: another request is pending for this host"); + + BYTE cmd = COMMAND_REGEDIT; + if (!ctx->Send2Client(&cmd, 1)) { + mcp.ClearPending(devId); + return BuildError(id, -32004, "Failed to send command to host"); + } + + std::vector subkeys, values; + if (!mcp.WaitPendingRegistry(devId, subkeys, values, kMcpToolTimeoutMs)) + return BuildError(id, -32001, "Timeout waiting for registry data"); + + Json::Value keysArr = ParseRegPath(subkeys, cp); + Json::Value valsArr = ParseRegKey(values, cp); + + Json::Value result(Json::objectValue); + Json::Value structuredContent(Json::objectValue); + structuredContent["path"] = path; + structuredContent["keys"] = keysArr; + structuredContent["values"] = valsArr; + result["structuredContent"] = structuredContent; + + Json::Value content(Json::arrayValue); + Json::Value item(Json::objectValue); + item["type"] = "text"; + item["text"] = std::string(u8"共 ") + std::to_string((int)keysArr.size()) + std::string(u8" 个子键、") + + std::to_string((int)valsArr.size()) + std::string(u8" 个值。"); + content.append(item); + result["content"] = content; + result["isError"] = false; + + return BuildResult(id, result); +} + // tools/call：get_client_log（主连接下发 COMMAND_QUERY_LOG，子连接回传 TOKEN_REPORT_LOG） std::string BuildGetClientLog(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) { uint64_t devId = 0; @@ -1521,6 +1872,7 @@ std::string BuildToolsCall(const Json::Value& root, CMy2015RemoteDlg* parent) { if (toolName == "list_services") return BuildListServices(id, args, parent); if (toolName == "get_client_log") return BuildGetClientLog(id, args, parent); if (toolName == "get_audit_log") return BuildGetAuditLog(id, parent); + if (toolName == "list_registry") return BuildListRegistry(id, args, parent); return BuildError(id, -32602, "Unknown tool: " + (toolName.empty() ? std::string("(empty)") : toolName)); @@ -1735,6 +2087,78 @@ bool CMcpServer::OnDriveList(uint64_t device_id, context* subCtx, const BYTE* bu return false; // 继续等 TOKEN_FILE_LIST，调用方不 CancelIO } +// ===== P3：list_registry 扩展 ===== + +bool CMcpServer::OnRegeditReady(uint64_t device_id, context* subCtx) { + std::string rootAndPath; + { + std::lock_guard lk(m_PendingMutex); + auto it = m_Pending.find(device_id); + if (it == m_Pending.end()) return false; // 无挂起 → 回落 MFC + if (it->second.tool != "list_registry") return false; + rootAndPath = it->second.path; // [rootToken:1][相对子键路径...] + } + if (rootAndPath.empty()) return false; // 不应发生：list_registry 必有 rootToken + + BYTE rootToken = (BYTE)rootAndPath[0]; + std::string relPath = rootAndPath.substr(1); + + // 下发 COMMAND_REG_FIND：布局与 RegisterDlg::OnTvnSelchangedTree 一致 + // [COMMAND_REG_FIND][rootToken][relPath...]['\0']。子链接保持，等 TOKEN_REG_PATH + KEY。 + std::vector pkt; + pkt.reserve(2 + relPath.size() + 1); + pkt.push_back((BYTE)COMMAND_REG_FIND); + pkt.push_back(rootToken); + pkt.insert(pkt.end(), relPath.begin(), relPath.end()); + pkt.push_back(0); + subCtx->Send2Client(pkt.data(), (ULONG)pkt.size()); + return true; // 接管子链接，不打开 MFC 对话框 +} + +void CMcpServer::TakeRegPath(uint64_t device_id, const BYTE* data, ULONG len) { + std::lock_guard lk(m_PendingMutex); + auto it = m_Pending.find(device_id); + if (it == m_Pending.end()) return; // 已超时清理 → 迟到数据丢弃 + if (it->second.tool != "list_registry") return; + it->second.regPath.assign(data, data + len); + it->second.pathDone = true; + if (it->second.keyDone) { + it->second.done = true; + m_PendingCv.notify_one(); + } +} + +void CMcpServer::TakeRegKey(uint64_t device_id, const BYTE* data, ULONG len) { + std::lock_guard lk(m_PendingMutex); + auto it = m_Pending.find(device_id); + if (it == m_Pending.end()) return; + if (it->second.tool != "list_registry") return; + it->second.regKey.assign(data, data + len); + it->second.keyDone = true; + if (it->second.pathDone) { + it->second.done = true; + m_PendingCv.notify_one(); + } +} + +bool CMcpServer::WaitPendingRegistry(uint64_t device_id, std::vector& subkeys, + std::vector& values, int timeoutMs) { + std::unique_lock lk(m_PendingMutex); + auto it = m_Pending.find(device_id); + if (it == m_Pending.end()) return false; + + bool signaled = m_PendingCv.wait_for(lk, std::chrono::milliseconds(timeoutMs), + [&] { return it->second.done; }); + if (!signaled) { + m_Pending.erase(it); // 超时 → 清理 + return false; + } + subkeys = std::move(it->second.regPath); + values = std::move(it->second.regKey); + m_Pending.erase(it); + return true; +} + // rand_s：Windows CRT 加密安全随机源（基于系统 CSPRNG）。其声明需在 前 // 定义 _CRT_RAND_S；为避免依赖 PCH 的包含顺序，这里手动声明其导出原型（errno_t == int）。 extern "C" int __cdecl rand_s(unsigned int* randomValue); diff --git a/server/2015Remote/McpServer.h b/server/2015Remote/McpServer.h index a60b6b3..4d67734 100644 --- a/server/2015Remote/McpServer.h +++ b/server/2015Remote/McpServer.h @@ -83,6 +83,21 @@ public: // 等 TOKEN_FILE_LIST（调用方不关子链接）。 bool OnDriveList(uint64_t device_id, context* subCtx, const BYTE* buf, ULONG len); + // ===== P3：list_registry（注册表查询，三阶段：COMMAND_REGEDIT → TOKEN_REGEDIT → + // COMMAND_REG_FIND → TOKEN_REG_PATH + TOKEN_REG_KEY）===== + + // MessageHandle(TOKEN_REGEDIT) 调用：挂起命中且为 list_registry 则下发 COMMAND_REG_FIND + // 并返回 true（接管子链接，不打开 MFC 对话框）；否则 false（回落 MFC）。 + bool OnRegeditReady(uint64_t device_id, context* subCtx); + + // MessageHandle(TOKEN_REG_PATH / TOKEN_REG_KEY) 调用：两包都到齐后置 done 并唤醒等待者。 + void TakeRegPath(uint64_t device_id, const BYTE* data, ULONG len); + void TakeRegKey(uint64_t device_id, const BYTE* data, ULONG len); + + // 工具线程：等待注册表两包（子键 + 值）；成功返回 true 并把两包分别写入 out，超时返回 false。 + bool WaitPendingRegistry(uint64_t device_id, std::vector& subkeys, + std::vector& values, int timeoutMs); + private: CMcpServer(); ~CMcpServer(); @@ -101,9 +116,13 @@ private: // 挂起请求注册表（受 m_PendingMutex 保护，键 = device_id）。 struct PendingRequest { std::string tool; - std::string path; // list_files：目录路径；空 = 只列盘 + std::string path; // list_files：目录路径；list_registry：rootToken(1B)+相对子键路径 uint16_t expectedReqId = 0; // get_screenshot：期望的预览 reqId（0 = 未设置） std::vector data; + std::vector regPath; // list_registry：TOKEN_REG_PATH 子键包（含 token 字节） + std::vector regKey; // list_registry：TOKEN_REG_KEY 值包（含 token 字节） + bool pathDone = false; // list_registry：已收到 TOKEN_REG_PATH + bool keyDone = false; // list_registry：已收到 TOKEN_REG_KEY bool done = false; }; std::mutex m_PendingMutex;