Locate the kernel once and for all

README.md

@@ -25,7 +25,6 @@ Limitations
 -----------
 
 - 64-bit only.
-- Kernel stack traces will be wrong if ETW events are not generated by calling `EtwWrite`. This should never happen unless you are adapting the code for a different provider.
 
 Building requirements
 ---------------------

mitimon/src/data.cpp

@@ -9,6 +9,8 @@
 
 std::unordered_map<uint32_t, ProcessData> ProcessData::processMap;
 
+ImageData ProcessData::kernelImageData;
+
 bool ProcessData::add(uint32_t pid, const std::wstring& imageName)
 {
     auto [it, isNew] = processMap.emplace(std::make_pair(pid, ProcessData(pid, imageName)));
@@ -30,6 +32,10 @@ ProcessData& ProcessData::get(uint32_t pid)
     return processMap.at(pid);
 }
 
+bool ProcessData::addImage(const ImageData& imageData)
+{
+    return addImage(ImageData(imageData));
+}
 
 bool ProcessData::addImage(ImageData&& imageData)
 {

mitimon/src/data.h

@@ -7,7 +7,42 @@
 #include <string>
 #include <unordered_map>
 
-class ImageData;
+class ImageData {
+public:
+    static bool add(uint32_t pid, void* imageBase, size_t imageSize, const std::wstring& imageName);
+    static bool remove(uint32_t pid, void* imageBase);
+
+    static std::wstring nameFromEtwName(const std::wstring& imageName);
+    static std::wstring pathFromEtwName(const std::wstring& imageName);
+
+public:
+    ImageData() :
+        mBase{ nullptr },
+        mSize{ 0 },
+        mName{},
+        mPath{}
+    {
+    }
+
+    ImageData(void* base, size_t size, const std::wstring& name) :
+        mBase{ base },
+        mSize{ size },
+        mName{ nameFromEtwName(name) },
+        mPath{ pathFromEtwName(name) }
+    {
+    }
+
+    void* base() const { return mBase; }
+    size_t size() const { return mSize; }
+    const std::wstring& name() const { return mName; }
+    const std::wstring& path() const { return mPath; }
+
+private:
+    void* mBase;
+    size_t mSize;
+    std::wstring mName;
+    std::wstring mPath;
+};
 
 class ProcessData {
 public:
@@ -16,20 +51,35 @@ class ProcessData {
         mImageName{ imageName },
         mImageMap{}
     {
+        if (kernelImageData.base()) {
+            addImage(kernelImageData);
+        }
     };
 
     static bool add(uint32_t pid, const std::wstring& imageName);
     static bool remove(uint32_t pid);
     static bool exists(uint32_t pid);
     static ProcessData& get(uint32_t pid);
 
+    static void setKernelImage(ImageData && imageData)
+    {
+        kernelImageData = std::move(imageData);
+    }
+
+    static const ImageData& kernelImage()
+    {
+        return kernelImageData;
+    }
+
 private:
     static std::unordered_map<uint32_t, ProcessData> processMap;
+    static ImageData kernelImageData;
 
 public:
     uint32_t pid() const { return mPid; }
     const std::wstring& imageName() const { return mImageName; }
 
+    bool addImage(const ImageData& imageData);
     bool addImage(ImageData&& imageData);
     bool removeImage(void* imageBase);
     ImageData& getImage(void* imageBase);
@@ -42,35 +92,4 @@ class ProcessData {
     std::unordered_map<void*, ImageData> mImageMap;
 };
 
-
-
-class ImageData {
-public:
-    static bool add(uint32_t pid, void* imageBase, size_t imageSize, const std::wstring& imageName);
-    static bool remove(uint32_t pid, void* imageBase);
-
-    static std::wstring nameFromEtwName(const std::wstring& imageName);
-    static std::wstring pathFromEtwName(const std::wstring& imageName);
-
-public:
-    ImageData(void* base, size_t size, const std::wstring& name) :
-        mBase{ base },
-        mSize{ size },
-        mName{ nameFromEtwName(name) },
-        mPath{ pathFromEtwName(name) }
-    {
-    }
-
-    void* base() const { return mBase; }
-    size_t size() const { return mSize; }
-    const std::wstring& name() const { return mName; }
-    const std::wstring& path() const { return mPath; }
-
-private:
-    void* mBase;
-    size_t mSize;
-    std::wstring mName;
-    std::wstring mPath;
-};
-
 #endif // DATA_H

mitimon/src/main.cpp

@@ -71,13 +71,70 @@ inline std::wstring stringify(const EVENT_RECORD& record, krabs::parser& parser,
     return result;
 }
 
+void locateKernel()
+{
+    Tracer tracer(SESSION_NAME);
+    std::atomic<bool> canStop{ false };
+
+    // Use ACG failures originating from this process to guess the kernel address.
+    tracer.addCustomProvider(MITIGATIONS_PROVIDER, MITIGATIONS_ANY,
+        [&canStop](const EVENT_RECORD& record, const krabs::trace_context& traceContext) {
+
+        if (canStop.load()) {
+                return;
+        }
+
+        krabs::schema schema(record, traceContext.schema_locator);
+        auto pid = schema.process_id();
+        if (pid != GetProcessId(GetCurrentProcess())) {
+            return;
+        }
+
+        // Locate the kernel based on the assumption that the first return address points somewhere in EtwWrite.
+        auto stackTrace = schema.stack_trace();
+        Symbolicator symbolicator{ ProcessData{ pid, L"self" }, SYM_DIR, SYM_PATH };
+        ProcessData::setKernelImage(symbolicator.guessImageFromSymbol(
+            L"C:\\Windows\\System32\\ntoskrnl.exe", L"EtwWrite", reinterpret_cast<void*>(stackTrace[0])
+        ));
+
+        canStop.store(true);
+    });
+
+    // Provoke ACG failures originating from this process.
+    std::future<void> backgroundTask = std::async(std::launch::async, [&canStop, &tracer]() {
+        while (!canStop.load()) {
+            PROCESS_MITIGATION_DYNAMIC_CODE_POLICY policy{};
+            policy.AuditProhibitDynamicCode = 1;
+            ::SetProcessMitigationPolicy(ProcessDynamicCodePolicy, &policy, sizeof policy);
+            void* address = ::VirtualAlloc(nullptr, 0x1000, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
+            if (address) {
+                ::VirtualFree(address, 0, MEM_RELEASE);
+            }
+            ::Sleep(500);
+        }
+        tracer.stop();
+    });
+
+    try {
+        tracer.start();
+    }
+    catch (std::runtime_error e) {
+        std::cout << e.what() << std::endl;
+    }
+}
+
 int main()
 {
-    std::vector<std::future<void>> backgroundTasks;
+    std::wcout << L"Please wait while the kernel base address is being guessed..." << std::endl;
+
+    locateKernel();
+
+    std::wcout << L"Guessed kernel base address: " << ProcessData::kernelImage().base() << L"." << std::endl << std::endl;
 
     std::wofstream sout{OUTPUT_FILE};
     std::mutex soutMutex;
 
+    std::vector<std::future<void>> backgroundTasks;
     Tracer tracer(SESSION_NAME);
 
     // The process provider will track process creation and image loading,
@@ -122,10 +179,6 @@ int main()
                 sout << std::format(L"ThreadId 0x{:08x}", tid) << std::endl;
                 sout << std::endl;
 
-                // Locate the kernel based on the assumption that the first return address points somewhere in EtwWrite.
-                symbolicator.loadWithHint(L"ntoskrnl", L"C:\\Windows\\System32\\ntoskrnl.exe",
-                    L"EtwWrite", reinterpret_cast<void*>(stackTrace[0]));
-
                 sout << L"Call Stack:" << std::endl;
                 for (auto& return_address : stackTrace)
                 {
@@ -142,6 +195,9 @@ int main()
             }));
         }
     );
+
+    std::wcout << L"Ready to catch events! You may now start the processes you wish to monitor." << std::endl << std::endl;
+
     try {
         tracer.start();
     }

mitimon/src/symbols.cpp

@@ -119,12 +119,12 @@ bool Symbolicator::load(const ImageData& imageData)
     return true;
 }
 
-bool Symbolicator::loadWithHint(const std::wstring& imageName, const std::wstring& imagePath, const std::wstring& symbolName, void* symbolAddress)
+ImageData Symbolicator::guessImageFromSymbol(const std::wstring& imagePath, const std::wstring& symbolName, void* symbolAddress)
 {
     SYMSRV_INDEX_INFOW indexInfo{};
     indexInfo.sizeofstruct = sizeof(indexInfo);
     if (!::SymSrvGetFileIndexInfoW(imagePath.c_str(), &indexInfo, 0)) {
-        return false;
+        return ImageData{};
     }
 
     std::wcout << L"Downloading symbols file " << indexInfo.pdbfile << L"..." << std::endl;
@@ -133,29 +133,29 @@ bool Symbolicator::loadWithHint(const std::wstring& imageName, const std::wstrin
     if (!::SymFindFileInPathW(mProcess, nullptr, indexInfo.pdbfile,
         &indexInfo.guid, indexInfo.age, 0, SSRVOPT_GUIDPTR, foundFile,
         nullptr, nullptr)) {
-        return false;
+        return ImageData{};
     }
 
     auto module_ = ::SymLoadModuleExW(
-        mProcess, nullptr, imagePath.c_str(), imageName.c_str(),
+        mProcess, nullptr, imagePath.c_str(), L"_temporary_guess_",
         0, 0, nullptr, 0);
     if (!module_) {
-        return false;
+        return ImageData{};
     }
 
     IMAGEHLP_MODULEW64 moduleInfo{};
     moduleInfo.SizeOfStruct = sizeof moduleInfo;
     if (!::SymGetModuleInfoW64(mProcess, module_, &moduleInfo)) {
         ::SymUnloadModule64(mProcess, module_);
-        return false;
+        return ImageData{};
     }
 
     SYMBOL_INFOW symbol{};
     symbol.SizeOfStruct = sizeof (symbol);
     symbol.MaxNameLen = 0;
-    if (!::SymFromNameW(mProcess, std::format(L"{}!{}", imageName, symbolName).c_str(), &symbol)) {
+    if (!::SymFromNameW(mProcess, std::format(L"_temporary_guess_!{}", symbolName).c_str(), &symbol)) {
         ::SymUnloadModule64(mProcess, module_);
-        return false;
+        return ImageData{};
     }
 
     bool isOnNextPage = (reinterpret_cast<DWORD64>(symbolAddress) & 0xFFFUi64) < (symbol.Address & 0xFFFUi64);
@@ -168,22 +168,5 @@ bool Symbolicator::loadWithHint(const std::wstring& imageName, const std::wstrin
 
     ::SymUnloadModule64(mProcess, module_);
 
-    ImageData imageData{ reinterpret_cast<void*>(guessedImageBase), moduleInfo.ImageSize, imagePath };
-    mProcessData.addImage(std::move(imageData));
-
-    module_ = ::SymLoadModuleExW(
-        mProcess, nullptr, imagePath.c_str(), imageName.c_str(),
-        guessedImageBase, 0, nullptr, 0);
-    if (!module_) {
-        return false;
-    }
-
-    moduleInfo.SizeOfStruct = sizeof moduleInfo;
-    if (!::SymGetModuleInfoW64(mProcess, module_, &moduleInfo)) {
-        ::SymUnloadModule64(mProcess, module_);
-        return false;
-    }
-
-    mModuleMap.emplace(std::make_pair(reinterpret_cast<void*>(guessedImageBase), module_));
-    return true;
+    return ImageData{ reinterpret_cast<void*>(guessedImageBase), moduleInfo.ImageSize, imagePath };
 }

mitimon/src/symbols.h

@@ -22,7 +22,7 @@ class Symbolicator {
 
     std::wstring symbolicate(void* address);
 
-    bool loadWithHint(const std::wstring& imageName, const std::wstring& imagePath, const std::wstring& symbolName, void* symbolAddress);
+    ImageData guessImageFromSymbol(const std::wstring& imagePath, const std::wstring& symbolName, void* symbolAddress);
 
 private:
     ProcessData mProcessData;

mitimon/src/trace.cpp

@@ -85,3 +85,8 @@ void Tracer::start()
 
     mTrace.start();
 }
+
+void Tracer::stop()
+{
+    mTrace.stop();
+}

mitimon/src/trace.h

@@ -30,6 +30,8 @@ class Tracer {
 
     void start();
 
+    void stop();
+
 private:
     krabs::user_trace mTrace;
     std::vector<krabs::provider<>> mProviders;