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
This commit is contained in:
yuanyuanxiang
2026-08-23 10:37:55 +02:00
parent 8815c4f896
commit 8bd6f3b60a
5 changed files with 508 additions and 19 deletions

View File

@@ -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));
}
}

View File

@@ -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; dwIndex<KeyCount; dwIndex++) { //枚举项
@@ -124,7 +127,7 @@ char* RegisterOperation::FindKey()
char *szBuffer=NULL;
HKEY hKey; //注册表返回句柄
if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_ALL_ACCESS,&hKey)==ERROR_SUCCESS) { //打开
if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_READ,&hKey)==ERROR_SUCCESS) { //打开
DWORD dwIndex=0,NameSize,NameCount,NameMaxLen,Type;
DWORD KeyCount,KeyMaxLen,DataSize,MaxDataLen;
//这就是枚举了
@@ -167,17 +170,17 @@ char* RegisterOperation::FindKey()
RegEnumValue(hKey,dwIndex,szValueName,&NameSize,
NULL,&Type,szValueData,&DataSize);//读取键值
if(Type==REG_SZ) {
szTemp[0]=MREG_SZ;
}
if(Type==REG_DWORD) {
szTemp[0]=MREG_DWORD;
}
if(Type==REG_BINARY) {
szTemp[0]=MREG_BINARY;
}
if(Type==REG_EXPAND_SZ) {
szTemp[0]=MREG_EXPAND_SZ;
// 值类型映射REG_SZ/DWORD/BINARY/EXPAND_SZ 之外补 MULTI_SZ/QWORD/NONE
// 未识别类型一律 MREG_NONE避免误报为 REG_SZ)。
switch (Type) {
case REG_SZ: szTemp[0] = MREG_SZ; break;
case REG_EXPAND_SZ: szTemp[0] = MREG_EXPAND_SZ; break;
case REG_MULTI_SZ: szTemp[0] = MREG_MULTI_SZ; break;
case REG_DWORD: szTemp[0] = MREG_DWORD; break;
case REG_QWORD: szTemp[0] = MREG_QWORD; break;
case REG_BINARY: szTemp[0] = MREG_BINARY; break;
case REG_NONE: szTemp[0] = MREG_NONE; break;
default: szTemp[0] = MREG_NONE; break;
}
szTemp+=sizeof(BYTE);
strcpy(szTemp,szValueName);

View File

@@ -6458,9 +6458,33 @@ VOID CMy2015RemoteDlg::MessageHandle(CONTEXT_OBJECT* ContextObject)
break;
}
case TOKEN_REGEDIT: { // 注册表管理【x】
// P3MCP 挂起时接管。TOKEN_REGEDIT 是客户端子链接构造即发CRegisterManager 构造),
// 命中则下发 COMMAND_REG_FIND 并保持子链接(等 TOKEN_REG_PATH/KEY不打开 MFC 对话框。
uint64_t devId = ContextObject->GetClientID();
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】
// P3MCP 挂起时接管。服务列表为一次性子链接回传CServicesManager 构造即发),
// 取走即 CancelIO 关子链接(无 MFC 对话框续用该子链接)。

View File

@@ -401,6 +401,210 @@ Json::Value ParseServiceList(const std::vector<BYTE>& 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.cppMREG_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<BYTE>& 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<BYTE>& 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 的 inputSchemaid 必填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 的 outputSchemakeys 字符串数组 / 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_registryP3注册表查询列根键/列子键+值,仅 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/calllist_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)");
// 注册表键路径为客户端 ANSIGBK/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<BYTE> 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/callget_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
}
// ===== P3list_registry 扩展 =====
bool CMcpServer::OnRegeditReady(uint64_t device_id, context* subCtx) {
std::string rootAndPath;
{
std::lock_guard<std::mutex> 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<BYTE> 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<std::mutex> 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<std::mutex> 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<BYTE>& subkeys,
std::vector<BYTE>& values, int timeoutMs) {
std::unique_lock<std::mutex> 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_sWindows CRT 加密安全随机源(基于系统 CSPRNG。其声明需在 <stdlib.h> 前
// 定义 _CRT_RAND_S为避免依赖 PCH 的包含顺序这里手动声明其导出原型errno_t == int
extern "C" int __cdecl rand_s(unsigned int* randomValue);

View File

@@ -83,6 +83,21 @@ public:
// 等 TOKEN_FILE_LIST调用方不关子链接
bool OnDriveList(uint64_t device_id, context* subCtx, const BYTE* buf, ULONG len);
// ===== P3list_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<BYTE>& subkeys,
std::vector<BYTE>& 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_registryrootToken(1B)+相对子键路径
uint16_t expectedReqId = 0; // get_screenshot期望的预览 reqId0 = 未设置)
std::vector<BYTE> data;
std::vector<BYTE> regPath; // list_registryTOKEN_REG_PATH 子键包(含 token 字节)
std::vector<BYTE> regKey; // list_registryTOKEN_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;