refactor driver interface

src/driver/memory.c

@@ -2,10 +2,7 @@
 
 uint64_t GetProcessDirectoryTableBase(PEPROCESS pProcess) {
     PUCHAR process = (PUCHAR)pProcess;
-    uint64_t process_dirbase = *(uint64_t*)(process + 0x28); // DirectoryTableBase
-    if (process_dirbase == 0)
-        process_dirbase = *(uint64_t*)(process + 0x388); // UserDirectoryTableBase
-    return process_dirbase;
+    return *(uint64_t*)(process + 0x28); // DirectoryTableBase;
 }
 
 PVOID GetVirtualForPhysical(uint64_t PhysicalAddress) {

src/driver/thread.c

@@ -48,7 +48,6 @@ NTSTATUS FindProcessThread(PEPROCESS pProcess, PETHREAD* ppThread) {
     NTSTATUS status = EnumProcesses(&Processes, SystemProcessInformation);
     if (!NT_SUCCESS(status)) {
         Log("Failed to enumerate processes (0x%08X)", status);
-        ExFreePoolWithTag(Processes, POOL_TAG);
         return status;
     }
 

src/driver_interface/driver_interface.cpp

@@ -1,29 +1,31 @@
 #include "driver_interface.h"
 
-std::optional<std::reference_wrapper<fumo::DriverInterface>> fumo::DriverInterface::Open(LPCWSTR lpFileName) {
+std::shared_ptr<fumo::DriverInterface> fumo::DriverInterface::Open(LPCWSTR lpFileName) {
     HANDLE hDevice = CreateFileW(lpFileName, GENERIC_WRITE, FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, 0, nullptr);
     if (hDevice == INVALID_HANDLE_VALUE) {
-        return std::nullopt;
+        return nullptr;
     }
-    return std::optional<std::reference_wrapper<fumo::DriverInterface>>(*new DriverInterface(hDevice));
+    // not using make_shared because the constructor is private
+    return std::shared_ptr<fumo::DriverInterface>(new fumo::DriverInterface(hDevice));
 }
 
-std::optional<ULONG> fumo::DriverInterface::GetVersion() {
+VOID fumo::DriverInterface::Unload() {
+    DeviceIoControl(hDevice, IO_UNLOAD_REQUEST, 
+        nullptr, 0, 
+        nullptr, 0, 
+        nullptr, nullptr);
+}
+
+BOOL fumo::DriverInterface::GetVersion(PULONG pVersion) {
     IO_VERSION_RESPONSE_DATA version_response = {0};
     if (!DeviceIoControl(hDevice, IO_VERSION_REQUEST, 
         nullptr, 0, 
         &version_response, sizeof(version_response), 
         nullptr, nullptr)) {
-        return std::nullopt;
+        return FALSE;
     }
-    return version_response.Version;
-}
-
-VOID fumo::DriverInterface::Unload() {
-    DeviceIoControl(hDevice, IO_UNLOAD_REQUEST, 
-        nullptr, 0, 
-        nullptr, 0, 
-        nullptr, nullptr);
+    *pVersion = version_response.Version;
+    return TRUE;
 }
 
 PVOID fumo::DriverInterface::AllocateKernelMemory(ULONG size) {

src/driver_interface/driver_interface.h

@@ -14,9 +14,9 @@ class DriverInterface {
     DriverInterface(const DriverInterface&) = delete;
     DriverInterface& operator=(const DriverInterface&) = delete;
 public:
-    static std::optional<std::reference_wrapper<DriverInterface>> Open(LPCWSTR lpFileName);
-    std::optional<ULONG> GetVersion();
+    static std::shared_ptr<fumo::DriverInterface> Open(LPCWSTR lpFileName);
     VOID Unload();
+    BOOL GetVersion(PULONG pVersion);
     PVOID AllocateKernelMemory(ULONG size);
     BOOL ExposeKernelMemory(ULONG pid, PVOID address, ULONG size);
     BOOL ExecuteCode(ULONG pid, PVOID address, PVOID argument);

src/include/fomo_common.h

@@ -6,7 +6,7 @@
 
 #define FUMO_MAGIC 'OMUF'
 #define FUMO_DATA_VERSION 0x00000002
-#define FUMO_DRIVER_VERSION 0x00000002
+#define FUMO_DRIVER_VERSION 0x00000003
 #define FUMO_HOOKED_DRIVER_NAME L"\\Driver\\Null"
 #define FUMO_HOOKED_DRIVER_NAME_USER L"\\\\.\\NUL"
 #define FUMO_SECOND_STAGE_PROCESS_NAME L"explorer.exe"

src/stage1/fumo_preloader.cpp

@@ -1,7 +1,6 @@
 #include "fumo_preloader.h"
 #include <libKDU.h>
 #include <lazy_importer.hpp>
-#include <driver_interface.h>
 #include <fumo_drv_data.h>
 #include <stage2_data.h>
 
@@ -189,28 +188,27 @@ DWORD stage2_loader_shellcode(PSTAGE2_LOADER_DATA loader_data) {
 
 void stage2_loader_shellcode_end() {}
 
-int init_driver(DWORD osBuildNumber) {
-    auto driver = fumo::DriverInterface::Open(FUMO_HOOKED_DRIVER_NAME_USER);
-    if (!driver.has_value())
-        return fumo::error(ERR_STAGE1_FAILED_TO_OPEN_DRIVER, L"Failed to open driver");
-
-    auto& driver_ref = driver.value().get();
-
-    auto version = driver_ref.GetVersion();
-    if (!version.has_value()) {
+int init_driver(fumo::DriverInterface* pDriver, DWORD osBuildNumber, bool forceReload) {
+    ULONG version;
+    BOOL just_loaded = FALSE;
+    if (!pDriver->GetVersion(&version)) {
         auto error = load_driver(osBuildNumber);
         if (error != ERR_STAGE1_SUCCESS)
             return error;
 
-        version = driver_ref.GetVersion();
-        if (!version.has_value())
+        just_loaded = TRUE;
+        if (!pDriver->GetVersion(&version))
             return fumo::error(ERR_STAGE1_FAILED_TO_GET_DRIVER_VERSION, L"Failed to get driver version");
     }
 
-    if (version.value() != FUMO_DRIVER_VERSION) {
+    if (version != FUMO_DRIVER_VERSION || forceReload) {
+        // if the driver we just loaded reports a wrong version something has gone terribly wrong
+        if (just_loaded && !forceReload)
+            return fumo::error(ERR_STAGE1_LOADED_DERIVER_VERSION_MISMATCH, L"Driver version mismatch (expected: {}, found: {})", FUMO_DRIVER_VERSION, version);
+
         // unload the old driver and try again
-        driver_ref.Unload();
-        return init_driver(osBuildNumber);
+        pDriver->Unload();
+        return init_driver(pDriver, osBuildNumber, false);
     }
 
     return ERR_STAGE1_SUCCESS;

src/stage1/fumo_preloader.h

@@ -1,6 +1,7 @@
 #pragma once
 #include <util.h>
 #include <fomo_common.h>
+#include <driver_interface.h>
 #include <vector>
 
 #define ERR_STAGE1_SUCCESS 0
@@ -13,14 +14,15 @@
 #define ERR_STAGE1_FAILED_TO_MAP_DRIVER 100
 #define ERR_STAGE1_FAILED_TO_OPEN_DRIVER 101
 #define ERR_STAGE1_FAILED_TO_GET_DRIVER_VERSION 102
+#define ERR_STAGE1_LOADED_DERIVER_VERSION_MISMATCH 103
 #define ERR_STAGE1_FAILED_TO_FIND_PROCESS 150
 #define ERR_STAGE1_FAILED_TO_OPEN_PROCESS 151
 #define ERR_STAGE1_FAILED_TO_ALLOCATE_MEMORY 200
 #define ERR_STAGE1_FAILED_TO_RELOCATE_MODULE 201
 #define ERR_STAGE1_FAILED_TO_WRITE_MEMORY 202
 #define ERR_STAGE1_FAILED_TO_EXECUTE_SHELLCODE 203
 
-int init_driver(DWORD osBuildNumber);
+int init_driver(fumo::DriverInterface* pDriver, DWORD osBuildNumber, bool forceReload = false);
 int load_driver(DWORD osBuildNumber);
 std::wstring get_proces_name(HANDLE process);
 DWORD find_process_by_name(LPCWSTR lpProcessName);

src/stage1/stage1.cpp

@@ -87,7 +87,11 @@ int main(PFUMO_EMBEDDED_DATA embedded_data) {
     if(!get_debug_privileges())
         return fumo::error(ERR_STAGE1_FAILED_TO_GET_DEBUG_PRIVILEGES, L"Failed to get debug privileges");
 
-    auto error = init_driver(osv.dwBuildNumber);
+    auto driver = fumo::DriverInterface::Open(FUMO_HOOKED_DRIVER_NAME_USER);
+    if (!driver)
+        return fumo::error(ERR_STAGE1_FAILED_TO_OPEN_DRIVER, L"Failed to open driver");
+
+    auto error = init_driver(driver.get(), osv.dwBuildNumber, true);
     if (error != ERR_STAGE1_SUCCESS)
         return error;
 

src/stage2/stage2.cpp

@@ -23,17 +23,15 @@ int main(HANDLE loader_process) {
         return fumo::error(ERR_STAGE2_FAILED_TO_DELETE_LOADER, L"Failed to delete loader executable: {}", loader_process_name);
 
     auto driver = fumo::DriverInterface::Open(FUMO_HOOKED_DRIVER_NAME_USER);
-    if (!driver.has_value())
+    if (!driver)
         return fumo::error(ERR_STAGE2_FAILED_TO_OPEN_DRIVER, L"Failed to open driver");
 
-    auto& driver_ref = driver.value().get();
-
-    auto driver_version = driver_ref.GetVersion();
-    if (!driver_version.has_value())
+    ULONG driver_version;
+    if (!driver->GetVersion(&driver_version))
         return fumo::error(ERR_STAGE2_FAILED_TO_OPEN_DRIVER, L"Failed to get driver version");
 
-    if (driver_version.value() != FUMO_DRIVER_VERSION)
-        return fumo::error(ERR_STAGE2_INVALID_DRIVER_VERSION, L"Invalid driver version (expected: {}, found: {})", FUMO_DRIVER_VERSION, driver_version.value());
+    if (driver_version != FUMO_DRIVER_VERSION)
+        return fumo::error(ERR_STAGE2_INVALID_DRIVER_VERSION, L"Invalid driver version (expected: {}, found: {})", FUMO_DRIVER_VERSION, driver_version);
 
     PFUMO_DATA_HEADER header = (PFUMO_DATA_HEADER)loader_data.fumo_data_base;
 
@@ -107,7 +105,7 @@ int main(HANDLE loader_process) {
     for (auto& module : wait_for_modules) {
         std::wstring module_name = convert_to_wstring(module);
         WAIT_FOR_MODULE_DATA wait_for_module_data;
-        wait_for_module_data.driver_interface = &driver_ref;
+        wait_for_module_data.driver_interface = driver;
         wait_for_module_data.process_id = process_id;
         wait_for_module_data.module_base = 0;
         wait_for_module_data.module_name = module_name.c_str();
@@ -162,7 +160,7 @@ int main(HANDLE loader_process) {
         return fumo::error(ERR_STAGE2_FAILED_TO_DECOMPRESS_DATA, L"Failed to decompress data");
 
     // let the magic happen
-    auto error = MapImage(&driver_ref, process_id, decompressed_data.get());
+    auto error = MapImage(driver.get(), process_id, decompressed_data.get());
     if (error != ERROR_SUCCESS)
         return error;
 

src/stage2/stage2.h

@@ -27,7 +27,7 @@ typedef struct _WAIT_FOR_PROCESS_DATA {
 } WAIT_FOR_PROCESS_DATA, *PWAIT_FOR_PROCESS_DATA;
 
 typedef struct _WAIT_FOR_MODULE_DATA {
-    fumo::DriverInterface* driver_interface;
+    std::shared_ptr<fumo::DriverInterface> driver_interface;
     LPCWSTR module_name;
     DWORD process_id;
     PVOID module_base;