Merge pull request #9 from noita-player/feature/peb-teb-mcp

Add process.info MCP tool for PEB/TEB enumeration and peb/tebs API for providers to implement
This commit is contained in:
IChooseYou
2026-03-07 09:42:51 -07:00
committed by GitHub
5 changed files with 178 additions and 0 deletions

View File

@@ -19,6 +19,60 @@
#include <tlhelp32.h>
#include <psapi.h>
#include <shellapi.h>
typedef struct _UNICODE_STRING { USHORT Length, MaximumLength; PWSTR Buffer; } UNICODE_STRING;
typedef struct _CLIENT_ID { HANDLE UniqueProcess; HANDLE UniqueThread; } CLIENT_ID;
typedef struct _SYSTEM_THREAD_INFORMATION {
LARGE_INTEGER KernelTime, UserTime, CreateTime;
ULONG WaitTime; PVOID StartAddress; CLIENT_ID ClientId;
LONG Priority, BasePriority; ULONG ContextSwitches, ThreadState, WaitReason;
} SYSTEM_THREAD_INFORMATION;
typedef struct _SYSTEM_PROCESS_INFORMATION {
ULONG NextEntryOffset; // 0x000
ULONG NumberOfThreads; // 0x004
LARGE_INTEGER WorkingSetPrivateSize; // 0x008
ULONG HardFaultCount; // 0x010
ULONG NumberOfThreadsHighWatermark; // 0x014
ULONGLONG CycleTime; // 0x018
LARGE_INTEGER CreateTime; // 0x020
LARGE_INTEGER UserTime; // 0x028
LARGE_INTEGER KernelTime; // 0x030
UNICODE_STRING ImageName; // 0x038
LONG BasePriority; // 0x048
HANDLE UniqueProcessId; // 0x050
PVOID InheritedFromUniqueProcessId; // 0x058
ULONG HandleCount; // 0x060
ULONG SessionId; // 0x064
ULONG_PTR UniqueProcessKey; // 0x068
SIZE_T PeakVirtualSize; // 0x070
SIZE_T VirtualSize; // 0x078
ULONG PageFaultCount; // 0x080
ULONG _pad0; // 0x084
SIZE_T PeakWorkingSetSize; // 0x088
SIZE_T WorkingSetSize; // 0x090
SIZE_T QuotaPeakPagedPoolUsage; // 0x098
SIZE_T QuotaPagedPoolUsage; // 0x0A0
SIZE_T QuotaPeakNonPagedPoolUsage; // 0x0A8
SIZE_T QuotaNonPagedPoolUsage; // 0x0B0
SIZE_T PagefileUsage; // 0x0B8
SIZE_T PeakPagefileUsage; // 0x0C0
SIZE_T PrivatePageCount; // 0x0C8
LARGE_INTEGER ReadOperationCount; // 0x0D0
LARGE_INTEGER WriteOperationCount; // 0x0D8
LARGE_INTEGER OtherOperationCount; // 0x0E0
LARGE_INTEGER ReadTransferCount; // 0x0E8
LARGE_INTEGER WriteTransferCount; // 0x0F0
LARGE_INTEGER OtherTransferCount; // 0x0F8
} SYSTEM_PROCESS_INFORMATION; // sizeof = 0x100
typedef struct alignas(8) _THREAD_BASIC_INFORMATION {
NTSTATUS ExitStatus; // 0x00
ULONG _pad; // 0x04
PVOID TebBaseAddress; // 0x08
CLIENT_ID ClientId; // 0x10
ULONG_PTR AffinityMask; // 0x20
LONG Priority; // 0x28
LONG BasePriority; // 0x2C
} THREAD_BASIC_INFORMATION;
#elif defined(__linux__)
#include <climits>
#include <sys/types.h>
@@ -61,6 +115,17 @@ ProcessMemoryProvider::ProcessMemoryProvider(uint32_t pid, const QString& proces
BOOL isWow64 = FALSE;
if (IsWow64Process(m_handle, &isWow64) && isWow64)
m_pointerSize = 4;
// Query PEB address via NtQueryInformationProcess
{
typedef NTSTATUS(NTAPI* NtQIP_t)(HANDLE, ULONG, PVOID, ULONG, PULONG);
static NtQIP_t pNtQIP = (NtQIP_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess");
if (pNtQIP) {
struct { PVOID r1; PVOID PebBaseAddress; PVOID r2[2]; ULONG_PTR pid; PVOID r3; } pbi = {};
ULONG retLen = 0;
if (pNtQIP(m_handle, /*ProcessBasicInformation*/0, &pbi, sizeof(pbi), &retLen) >= 0 && pbi.PebBaseAddress)
m_peb = (uint64_t)(uintptr_t)pbi.PebBaseAddress;
}
}
cacheModules();
}
}
@@ -426,6 +491,58 @@ int ProcessMemoryProvider::size() const
#endif
}
QVector<rcx::Provider::ThreadInfo> ProcessMemoryProvider::tebs() const
{
#ifdef _WIN32
QVector<ThreadInfo> result;
if (!m_handle || !m_peb) return result;
typedef NTSTATUS(NTAPI* NtQSI_t)(ULONG, PVOID, ULONG, PULONG);
typedef NTSTATUS(NTAPI* NtQIT_t)(HANDLE, ULONG, PVOID, ULONG, PULONG);
static auto pNtQSI = (NtQSI_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQuerySystemInformation");
static auto pNtQIT = (NtQIT_t)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationThread");
if (!pNtQSI || !pNtQIT) return result;
// Enumerate threads via SystemProcessInformation (class 5)
ULONG retLen = 0;
ULONG bufSize = 1 << 20;
QByteArray buf(bufSize, 0);
NTSTATUS qsiSt;
for (int attempt = 0; attempt < 8; ++attempt) {
qsiSt = pNtQSI(5, buf.data(), bufSize, &retLen);
if ((uint32_t)qsiSt != 0xC0000004u) break;
bufSize *= 2;
buf.resize(bufSize);
}
if (qsiSt < 0) return result;
// Walk process entries to find ours
auto* proc = (SYSTEM_PROCESS_INFORMATION*)buf.data();
for (;;) {
if ((uintptr_t)proc->UniqueProcessId == m_pid) {
auto* threads = (SYSTEM_THREAD_INFORMATION*)((char*)proc + sizeof(*proc));
for (ULONG i = 0; i < proc->NumberOfThreads; ++i) {
DWORD tid = (DWORD)(uintptr_t)threads[i].ClientId.UniqueThread;
HANDLE hThread = OpenThread(THREAD_QUERY_LIMITED_INFORMATION, FALSE, tid);
if (!hThread) continue;
THREAD_BASIC_INFORMATION tbi = {};
ULONG tbiLen = 0;
NTSTATUS qitSt = pNtQIT(hThread, 0, &tbi, sizeof(tbi), &tbiLen);
if (qitSt >= 0 && tbi.TebBaseAddress)
result.append({(uint64_t)(uintptr_t)tbi.TebBaseAddress, tid});
CloseHandle(hThread);
}
break;
}
if (!proc->NextEntryOffset) break;
proc = (SYSTEM_PROCESS_INFORMATION*)((char*)proc + proc->NextEntryOffset);
}
return result;
#else
return {};
#endif
}
// ──────────────────────────────────────────────────────────────────────────
// ProcessMemoryPlugin implementation
// ──────────────────────────────────────────────────────────────────────────

View File

@@ -41,6 +41,8 @@ public:
// Process-specific helpers
uint32_t pid() const { return m_pid; }
void refreshModules() { m_modules.clear(); cacheModules(); }
uint64_t peb() const override { return m_peb; }
QVector<ThreadInfo> tebs() const override;
private:
void cacheModules();
@@ -56,6 +58,7 @@ private:
bool m_writable;
uint64_t m_base;
int m_pointerSize = 8;
uint64_t m_peb = 0;
struct ModuleInfo {
QString name;

View File

@@ -447,6 +447,22 @@ QJsonObject McpBridge::handleToolsList(const QJsonValue& id) {
}}
});
// process.info
tools.append(QJsonObject{
{"name", "process.info"},
{"description", "Returns PEB address and enumerates all Thread Environment Blocks (TEBs) for the attached process. "
"TEBs are discovered via NtQuerySystemInformation and NtQueryInformationThread. "
"Each TEB entry includes: address, threadId. "
"Requires a live process provider with PEB support."},
{"inputSchema", QJsonObject{
{"type", "object"},
{"properties", QJsonObject{
{"tabIndex", QJsonObject{{"type", "integer"},
{"description", "MDI tab index (0-based). Omit for active tab."}}}
}}
}}
});
return okReply(id, QJsonObject{{"tools", tools}});
}
@@ -472,6 +488,7 @@ QJsonObject McpBridge::handleToolsCall(const QJsonValue& id, const QJsonObject&
else if (toolName == "ui.action") result = toolUiAction(args);
else if (toolName == "tree.search") result = toolTreeSearch(args);
else if (toolName == "node.history") result = toolNodeHistory(args);
else if (toolName == "process.info") result = toolProcessInfo(args);
else return errReply(id, -32601, "Unknown tool: " + toolName);
m_mainWindow->clearMcpStatus();
@@ -1327,6 +1344,39 @@ QJsonObject McpBridge::toolNodeHistory(const QJsonObject& args) {
QJsonDocument(result).toJson(QJsonDocument::Compact)));
}
// ════════════════════════════════════════════════════════════════════
// TOOL: process.info — PEB address + TEB enumeration
// ════════════════════════════════════════════════════════════════════
QJsonObject McpBridge::toolProcessInfo(const QJsonObject& args) {
auto* tab = resolveTab(args);
if (!tab) return makeTextResult("No active tab", true);
auto* prov = tab->doc->provider.get();
if (!prov) return makeTextResult("No data source attached", true);
if (!prov->isLive()) return makeTextResult("Not a live provider", true);
uint64_t pebAddr = prov->peb();
if (!pebAddr) return makeTextResult("PEB not available for this provider", true);
QJsonObject out;
out["peb"] = "0x" + QString::number(pebAddr, 16).toUpper();
auto tebList = prov->tebs();
QJsonArray tebArr;
for (const auto& t : tebList) {
tebArr.append(QJsonObject{
{"address", "0x" + QString::number(t.tebAddress, 16).toUpper()},
{"threadId", (qint64)t.threadId}
});
}
out["tebs"] = tebArr;
out["tebCount"] = tebArr.size();
return makeTextResult(QString::fromUtf8(QJsonDocument(out).toJson(QJsonDocument::Indented)));
}
// ════════════════════════════════════════════════════════════════════
// Notifications (call from MainWindow/Controller hooks)
// ════════════════════════════════════════════════════════════════════

View File

@@ -60,6 +60,7 @@ private:
QJsonObject toolUiAction(const QJsonObject& args);
QJsonObject toolTreeSearch(const QJsonObject& args);
QJsonObject toolNodeHistory(const QJsonObject& args);
QJsonObject toolProcessInfo(const QJsonObject& args);
// Helpers
QJsonObject makeTextResult(const QString& text, bool isError = false);

View File

@@ -73,6 +73,13 @@ public:
// Default: returns empty (scan engine falls back to [0, size())).
virtual QVector<MemoryRegion> enumerateRegions() const { return {}; }
// Process Environment Block address (x64 PEB VA in target process).
// Only meaningful for live process providers. Returns 0 if unavailable.
virtual uint64_t peb() const { return 0; }
struct ThreadInfo { uint64_t tebAddress; uint32_t threadId; };
virtual QVector<ThreadInfo> tebs() const { return {}; }
// --- Derived convenience (non-virtual, never override) ---
bool isValid() const { return size() > 0; }