09018e2b4d
Add the busy/closed flags to ScreenCtrlSession plus BeginScreenCtrlAction / EndScreenCtrlAction, mirroring the terminal's busy discipline so an in-flight injection cannot race with session teardown (review finding #6). While an injection is in flight (busy=true), OnScreenControlClosed marks the session closed instead of erasing it (the injection thread, which still holds subCtx, cleans up in EndScreenCtrlAction), SweepIdleScreenCtrl skips it, and CloseScreenCtrlSession defers the erase. The subCtx is looked up under m_ScreenCtrlMutex and the injection is sent outside the lock, matching the established terminal pattern; the screen sub-connection's CONTEXT_OBJECT::Send2Client already serializes internally via SendLock. Co-Authored-By: deepseek-v4-pro
286 lines
16 KiB
C++
286 lines
16 KiB
C++
#pragma once
|
||
|
||
#include <Windows.h>
|
||
#include <string>
|
||
#include <thread>
|
||
#include <atomic>
|
||
#include <chrono>
|
||
#include <mutex>
|
||
#include <condition_variable>
|
||
#include <map>
|
||
#include <vector>
|
||
#include <cstdint>
|
||
#include <ctime>
|
||
|
||
// httplib 与 Windows 头部的 min/max 宏冲突，按 file_server.h 的既有约定处理。
|
||
#undef min
|
||
#undef max
|
||
#include "httplib.h"
|
||
#ifndef max
|
||
#define max(a,b) (((a) > (b)) ? (a) : (b))
|
||
#endif
|
||
#ifndef min
|
||
#define min(a,b) (((a) < (b)) ? (a) : (b))
|
||
#endif
|
||
|
||
#pragma comment(lib, "ws2_32.lib")
|
||
|
||
class CMy2015RemoteDlg;
|
||
class context;
|
||
|
||
// MCP (Model Context Protocol) 服务端：httplib 封装 + JSON-RPC 2.0 分发 + 静态 token 校验。
|
||
// 与 CWebService 平级、互不依赖；默认禁用，经「扩展 → MCP设置」开启后监听
|
||
// （默认 127.0.0.1:6544，仅本机回环）。见 docs/Mcp_Design.md。
|
||
class CMcpServer {
|
||
public:
|
||
static CMcpServer& Instance();
|
||
|
||
void SetParentDlg(CMy2015RemoteDlg* pDlg) { m_parent = pDlg; }
|
||
void SetToken(const std::string& token) { m_token = token; }
|
||
|
||
bool Start(const std::string& bind, int port);
|
||
void Stop();
|
||
bool IsRunning() const { return m_running.load(); }
|
||
|
||
// ===== P2b：无请求 id 的一次性响应 → 每 host 单飞行 =====
|
||
// 进程/窗口列表（TOKEN_PSLIST/TOKEN_WSLIST，一次性子链接）与历史活动
|
||
// （TOKEN_REPORT_ACTIVITY，主连接 RPC）都不带请求关联，同一 host 同一时刻
|
||
// 只允许一个挂起请求，避免同 host 并发请求响应归属歧义。
|
||
// 注意：单飞行只约束 MCP 侧同 host 并发，不约束与 MFC 对话框并存（各用独立子链接）。
|
||
|
||
// 该 host 是否已有挂起请求（MessageHandle 在 TOKEN_PSLIST/TOKEN_WSLIST 分支调用）。
|
||
bool IsPending(uint64_t device_id);
|
||
|
||
// 拷贝响应缓冲到挂起结果并唤醒等待者（在 MessageHandle 内同步调用，数据此时仍在
|
||
// context 缓冲中、IO 线程被阻塞，拷贝是安全的）。
|
||
void TakeMainResponse(uint64_t device_id, const BYTE* data, ULONG len);
|
||
|
||
// 工具线程：登记挂起（false = 该 host 已有挂起请求，设备忙）。
|
||
bool BeginPending(uint64_t device_id, const std::string& tool);
|
||
|
||
// 工具线程：等待响应；成功返回 true 并把缓冲写入 out，超时/失败返回 false（并清理）。
|
||
bool WaitPending(uint64_t device_id, std::vector<BYTE>& out, int timeoutMs);
|
||
|
||
// 工具线程：清理挂起状态（发送失败等提前退出路径）。
|
||
void ClearPending(uint64_t device_id);
|
||
|
||
// ===== P2c：list_files / get_screenshot 扩展 =====
|
||
|
||
// 登记挂起（带目录路径；list_files 专用，path 为空表示只列盘）。
|
||
bool BeginPending(uint64_t device_id, const std::string& tool, const std::string& path);
|
||
|
||
// get_screenshot：工具线程生成 16 位 reqId（跳过 0，0 表示未设置），用于丢弃过期响应。
|
||
uint16_t NextPreviewReqId();
|
||
|
||
// get_screenshot：登记期望的 reqId（必须在发送 COMMAND_SCREEN_PREVIEW_REQ 前调用）。
|
||
void SetPendingReqId(uint64_t device_id, uint16_t reqId);
|
||
|
||
// get_screenshot：MessageHandle(TOKEN_SCREEN_PREVIEW_RSP) 在 IO 线程同步调用；
|
||
// 挂起命中且 reqId 一致则拷贝响应并返回 true，否则 false（回落到 MFC 预览路径）。
|
||
bool TakePreviewResponse(uint64_t device_id, uint16_t reqId, const BYTE* data, ULONG len);
|
||
|
||
// list_files：MessageHandle(TOKEN_DRIVE_LIST) 调用。
|
||
// 返回 true = 已取走盘列表（调用方 CancelIO）；false = 已转向下发 COMMAND_LIST_FILES、
|
||
// 等 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<BYTE>& subkeys,
|
||
std::vector<BYTE>& values, int timeoutMs);
|
||
|
||
// ===== P3：exec_command（Windows 一次性远程命令，复刻 Web 终端链路）=====
|
||
// 流程：工具线程 BeginTermPending → 主连接下发 COMMAND_SHELL → 客户端 shell 子连接回
|
||
// TOKEN_TERMINAL_START / TOKEN_SHELL_START → RegisterTerminalContext 接管 → 发
|
||
// COMMAND_NEXT(+PTY resize) 启动输出回流 → 工具线程发哨兵命令行 → OnTerminalData 收集
|
||
// 并做哨兵检测 → WaitTerminalDone 取 stdout + exit_code → CancelIO 关子链接。
|
||
// 单设备单终端（MVP）：BeginTermPending 命中已有会话即拒绝，避免同设备并发归属歧义。
|
||
|
||
// 该 host 是否已有待接管的终端会话（MessageHandle TOKEN_*_START 分支调用）。
|
||
bool IsTermPending(uint64_t device_id);
|
||
|
||
// 终端子连接就绪：接管（登记 subCtx/isPty、清 pending、发 COMMAND_NEXT 唤醒回流）。
|
||
void RegisterTerminalContext(uint64_t device_id, context* subCtx, bool isPty);
|
||
|
||
// 终端子连接是否已被 MCP 接管（MessageHandle 顶部据此把 shell 输出路由到 OnTerminalData）。
|
||
bool IsTerminalContext(context* subCtx);
|
||
|
||
// 追加 shell 输出并做哨兵检测（命中置 done 并唤醒等待者）。
|
||
void OnTerminalData(context* subCtx, const BYTE* data, ULONG len);
|
||
|
||
// shell 进程退出（TOKEN_TERMINAL_CLOSE）。
|
||
void OnTerminalClosed(context* subCtx);
|
||
|
||
// 工具线程：登记终端挂起（false = 该 host 已有终端会话）。
|
||
bool BeginTermPending(uint64_t device_id, const std::string& command, const std::string& nonce);
|
||
|
||
// 工具线程：等待子连接回 START 并接管（true = 已接管，输出 subCtx/isPty）。
|
||
bool WaitTerminalReady(uint64_t device_id, context*& subCtx, bool& isPty, int timeoutMs);
|
||
|
||
// 工具线程：等待输出收集完成（哨兵命中 / 进程退出 / 超时）。
|
||
// 成功 true：out=截断到哨兵前的原始字节、exitCode=0/1（进程退出无哨兵时 -1）、closed=是否进程退出。
|
||
bool WaitTerminalDone(uint64_t device_id, std::vector<BYTE>& out, int& exitCode, bool& closed, int timeoutMs);
|
||
|
||
// 工具线程：清理终端挂起（发送失败等提前退出路径）。
|
||
void ClearTermPending(uint64_t device_id);
|
||
|
||
// ===== P4：持久远程终端（terminal_open / terminal_exec / terminal_close）=====
|
||
void SetTerminalEnabled(bool enabled) { m_terminalEnabled = enabled; }
|
||
bool IsTerminalEnabled() const { return m_terminalEnabled; }
|
||
|
||
// 登记持久终端挂起（复用单设备单终端约束；sessionId 由调用方生成）。
|
||
bool BeginTermOpen(uint64_t device_id, const std::string& sessionId);
|
||
|
||
// 复位单条命令字段并置 busy；输出 subCtx/isPty。false = 不存在/不匹配/未就绪/忙碌/已关。
|
||
bool BeginTermCommand(uint64_t device_id, const std::string& sessionId,
|
||
const std::string& command, const std::string& nonce,
|
||
context*& subCtx, bool& isPty);
|
||
|
||
// 等待本条命令完成（done || closed || 超时）。done=保持会话；closed/超时=清理会话+路由。
|
||
bool WaitTermCommand(uint64_t device_id, std::vector<BYTE>& out, int& exitCode,
|
||
bool& closed, int timeoutMs);
|
||
|
||
// 关闭持久终端（校验 sessionId；CancelIO + 清理路由/会话）。
|
||
// 返回：0=已关闭；1=会话不存在（幂等）；2=sessionId 不匹配（调用方报错）。
|
||
int CloseTermSession(uint64_t device_id, const std::string& sessionId);
|
||
|
||
// idle 回收：关闭 lastActiveAt 超时且非 busy 的持久会话。返回回收数量。
|
||
int SweepIdleTerminals(time_t idleTimeoutSec);
|
||
|
||
// ===== P5：MCP 远程控制（remote_open / remote_close / remote_mouse / remote_keyboard）=====
|
||
// 屏幕子连接由 WebService::StartRemoteDesktop 建立的隐藏 CScreenSpyDlg 持有（复用 Web
|
||
// 路径，见 docs/Mcp_RemoteControl_Design.md §6.3）；会话仅记录逻辑状态 + 子连接指针。
|
||
// 注入走该子连接（COMMAND_SCREEN_CONTROL），与终端一样单设备单会话。
|
||
void SetRemoteControlEnabled(bool enabled) { m_remoteControlEnabled = enabled; }
|
||
bool IsRemoteControlEnabled() const { return m_remoteControlEnabled; }
|
||
|
||
// 该 host 的屏幕子连接是否被 MCP 远程控制会话持有（OfflineProc 反查用）。
|
||
bool IsScreenCtrlContext(context* subCtx);
|
||
|
||
// 登记控制会话（单设备单会话；sessionId 由调用方生成）。false = 该设备已有会话。
|
||
bool BeginScreenCtrlOpen(uint64_t device_id, const std::string& sessionId);
|
||
|
||
// 子连接就绪（HasActiveSession && 分辨率到达）后调用：填 subCtx/分辨率并置 started。
|
||
// false = 会话已被并发关闭（调用方应放弃本次 open）。
|
||
bool MarkScreenCtrlReady(uint64_t device_id, const std::string& sessionId,
|
||
context* subCtx, int screenW, int screenH);
|
||
|
||
// 注入前登记（校验 sessionId/started/!busy/!closed）。输出 subCtx/screenW/H。
|
||
// 注入期间 busy=true，保证 close/sweep/断线不与注入并发擦会话（镜像终端 busy 模式，
|
||
// 评审发现 #6）。返回：0=ok；1=不存在/不匹配/未就绪/已断（-32002）；2=忙（-32003）。
|
||
int BeginScreenCtrlAction(uint64_t device_id, const std::string& sessionId,
|
||
context*& subCtx, int& screenW, int& screenH);
|
||
|
||
// 注入后复位 busy；若注入期间屏幕子连接已断（closed），擦会话+路由（防泄漏）。
|
||
void EndScreenCtrlAction(uint64_t device_id, const std::string& sessionId);
|
||
|
||
// 关闭控制会话（校验 sessionId；擦路由 + 会话）。屏幕子连接的关闭由调用方锁外执行
|
||
// （CloseWebRemoteDesktopByClientID）。返回 0=已关；1=不存在（幂等）；2=sessionId 不匹配。
|
||
int CloseScreenCtrlSession(uint64_t device_id, const std::string& sessionId);
|
||
|
||
// idle 回收：关闭 lastActiveAt 超时的控制会话，并锁外关闭其隐藏对话框。返回回收数量。
|
||
int SweepIdleScreenCtrl(time_t idleTimeoutSec);
|
||
|
||
// 屏幕子连接断开（OfflineProc 调用）：擦除该子连接对应的会话 + 路由（幂等）。
|
||
void OnScreenControlClosed(context* subCtx);
|
||
|
||
// 安全配置（启动时由 CMy2015RemoteDlg 读 THIS_CFG 后设置）。
|
||
void SetReadonly(bool readonly) { m_readonly = readonly; }
|
||
void SetCmdWhitelist(const std::string& whitelist) { m_cmdWhitelist = whitelist; }
|
||
bool IsReadonly() const { return m_readonly; }
|
||
const std::string& GetCmdWhitelist() const { return m_cmdWhitelist; }
|
||
|
||
private:
|
||
CMcpServer();
|
||
~CMcpServer();
|
||
CMcpServer(const CMcpServer&) = delete;
|
||
CMcpServer& operator=(const CMcpServer&) = delete;
|
||
|
||
// POST /mcp 处理器：token 校验 + JSON-RPC 分发。
|
||
void HandleMcp(const httplib::Request& req, httplib::Response& res);
|
||
|
||
httplib::Server m_server;
|
||
std::thread m_thread;
|
||
std::atomic<bool> m_running{false};
|
||
std::string m_token;
|
||
CMy2015RemoteDlg* m_parent = nullptr;
|
||
|
||
// 挂起请求注册表（受 m_PendingMutex 保护，键 = device_id）。
|
||
struct PendingRequest {
|
||
std::string tool;
|
||
std::string path; // list_files：目录路径；list_registry：rootToken(1B)+相对子键路径
|
||
uint16_t expectedReqId = 0; // get_screenshot：期望的预览 reqId（0 = 未设置）
|
||
std::vector<BYTE> data;
|
||
std::vector<BYTE> regPath; // list_registry：TOKEN_REG_PATH 子键包（含 token 字节）
|
||
std::vector<BYTE> 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;
|
||
std::condition_variable m_PendingCv;
|
||
std::map<uint64_t, PendingRequest> m_Pending;
|
||
|
||
// get_screenshot 的 reqId 发生器（16 位，跳过 0）。与 MFC 预览的 m_PreviewReqId 各自独立计数。
|
||
std::atomic<uint16_t> m_PreviewReqId{1};
|
||
|
||
// ===== P3：exec_command 终端会话（受 m_TermMutex 保护）=====
|
||
struct TermSession {
|
||
bool started = false; // false=已发 COMMAND_SHELL 待 START；true=已接管
|
||
context* subCtx = nullptr; // shell 子连接上下文
|
||
bool isPty = false; // true=ConPTY(UTF-8)；false=老 cmd 管道(GBK)
|
||
std::string command; // 原始 UTF-8 命令（审计用）
|
||
std::string nonce; // 哨兵随机串
|
||
std::vector<BYTE> data; // 收集的原始 shell 输出
|
||
size_t sentPos = 0; // 哨兵在 data 中的位置
|
||
int exitCode = -1;
|
||
bool done = false; // 哨兵命中
|
||
bool closed = false; // TOKEN_TERMINAL_CLOSE（进程退出）
|
||
// ===== 新增（持久终端）=====
|
||
std::string sessionId; // 持久会话 token（一次性 exec 为空）
|
||
bool persistent = false;
|
||
bool busy = false; // 一条命令在飞（terminal_exec 复位→WaitTermCommand 返回）
|
||
time_t lastActiveAt = 0; // idle 回收用（秒）
|
||
};
|
||
std::mutex m_TermMutex;
|
||
std::condition_variable m_TermCv;
|
||
std::map<uint64_t, TermSession> m_TermSessions; // device_id → 会话
|
||
std::map<context*, uint64_t> m_TermContextToDevice; // subCtx → device_id（顶部路由）
|
||
|
||
// ===== P5：远程控制会话（受 m_ScreenCtrlMutex 保护；单设备单会话）=====
|
||
struct ScreenCtrlSession {
|
||
std::string sessionId; // 会话 token（每次 open 独立随机）
|
||
context* subCtx = nullptr; // 屏幕子连接上下文（就绪后填；用于注入）
|
||
bool started = false; // false=子连接建立中；true=已就绪
|
||
bool busy = false; // 一条注入在飞（BeginScreenCtrlAction→End 之间）
|
||
bool closed = false; // 注入期间屏幕子连接已断（由注入线程收尾）
|
||
time_t lastActiveAt = 0; // idle 回收用（秒）
|
||
int screenW = 0; // 物理捕获分辨率（来自 TOKEN_BITMAPINFO）
|
||
int screenH = 0;
|
||
};
|
||
std::mutex m_ScreenCtrlMutex;
|
||
std::map<uint64_t, ScreenCtrlSession> m_ScreenCtrlSessions; // device_id → 会话
|
||
std::map<context*, uint64_t> m_ScreenCtrlContextToDevice; // subCtx → device_id（OfflineProc 反查）
|
||
|
||
bool m_readonly = true;
|
||
std::string m_cmdWhitelist;
|
||
bool m_terminalEnabled = false; // 持久终端开关（默认关；要求 m_readonly=false）
|
||
bool m_remoteControlEnabled = false; // 远程控制开关（默认关；要求 m_readonly=false）
|
||
};
|
||
|
||
// 全局访问器（仿 WebService()，见 WebService.h 末尾）
|
||
inline CMcpServer& McpServer() { return CMcpServer::Instance(); }
|
||
|
||
// 生成 32 hex（128-bit）随机 token；供运行时兜底与「MCP设置」对话框预填复用。
|
||
std::string GenerateRandomToken();
|