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

plugins/ProcessMemory/ProcessMemoryPlugin.cpp

@@ -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
 // ──────────────────────────────────────────────────────────────────────────

plugins/ProcessMemory/ProcessMemoryPlugin.h

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

src/mcp/mcp_bridge.cpp

@@ -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)
 // ════════════════════════════════════════════════════════════════════

src/mcp/mcp_bridge.h

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

src/providers/provider.h

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