Add Wdf* APIs required by ebpf-for-windows

Add cache-aligned allocations (fixes #31)

Signed-off-by: Dave Thaler <[email protected]>

inc/usersim/ex.h

@@ -198,6 +198,22 @@ USERSIM_API _Ret_maybenull_ void*
 ExAllocatePoolWithTagCPP(
     _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, SIZE_T number_of_bytes, ULONG tag);
 
+// TODO: clean up usersim_* apis that deal with allocation so other
+// file just use Ex*() APIs instead.
+/**
+ * @brief Allocate memory.
+ * @param[in] pool_type Pool type to use.
+ * @param[in] size Size of memory to allocate.
+ * @param[in] tag Pool tag to use.
+ * @param[in] initialize False to return "uninitialized" memory.
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate_with_tag(
+    _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type,
+    size_t size,
+    uint32_t tag,
+    bool initialize);
+
 USERSIM_API _Ret_maybenull_ void*
 ExAllocatePoolUninitializedCPP(_In_ POOL_TYPE pool_type, _In_ size_t number_of_bytes, _In_ unsigned long tag);
 
@@ -207,4 +223,7 @@ ExRaiseAccessViolationCPP();
 USERSIM_API void
 ExRaiseDatatypeMisalignmentCPP();
 
+void usersim_initialize_ex(bool leak_detector);
+void usersim_clean_up_ex();
+
 #endif

inc/usersim/io.h

@@ -68,7 +68,13 @@ extern "C"
     IoGetFileObjectGenericMapping();
 
     USERSIM_API
-    _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI PEPROCESS IoGetCurrentProcess(VOID);
+    _IRQL_requires_max_(DISPATCH_LEVEL) PEPROCESS IoGetCurrentProcess(VOID);
+
+    typedef struct _IRP* PIRP;
+
+    USERSIM_API
+    _IRQL_requires_max_(DISPATCH_LEVEL) VOID
+    IofCompleteRequest(_In_ PIRP irp, _In_ CCHAR priority_boost);
 
 #if defined(__cplusplus)
 }

inc/usersim/wdf.h

@@ -13,16 +13,14 @@ extern "C"
 #endif
 
     typedef HANDLE WDFDEVICE;
-    typedef struct _wdfdriver WDFDRIVER;
+    typedef HANDLE WDFDRIVER;
 
     typedef struct _WDF_OBJECT_ATTRIBUTES WDF_OBJECT_ATTRIBUTES, *PWDF_OBJECT_ATTRIBUTES;
     typedef struct _WDFDEVICE_INIT WDFDEVICE_INIT, *PWDFDEVICE_INIT;
 
 #define WDF_NO_OBJECT_ATTRIBUTES 0
 #define WDF_NO_HANDLE 0
 
-    typedef struct _wdfdriver WDFDRIVER;
-
     typedef struct _driver_object DRIVER_OBJECT, *PDRIVER_OBJECT;
 
     typedef NTSTATUS(DRIVER_INITIALIZE)(_In_ PDRIVER_OBJECT driver_object, _In_ PUNICODE_STRING registry_path);
@@ -44,10 +42,10 @@ extern "C"
         ULONG DriverPoolTag;
     } WDF_DRIVER_CONFIG, *PWDF_DRIVER_CONFIG;
 
-    typedef struct _wdfdriver
+    struct _driver_object
     {
         WDF_DRIVER_CONFIG config;
-    } WDFDRIVER;
+    };
 
 #define WDF_DRIVER_GLOBALS_NAME_LEN (32)
 

src/ex.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: MIT
 
 #include "fault_injection.h"
+#include "leak_detector.h"
 #include "platform.h"
 #include "kernel_um.h"
 #include "usersim/ex.h"
@@ -15,6 +16,9 @@
 
 // Ex* functions.
 
+extern "C" size_t usersim_fuzzing_memory_limit = MAXSIZE_T;
+usersim_leak_detector_ptr _usersim_leak_detector_ptr;
+
 /***
  * @brief This following class implements a mock of the Windows Kernel's rundown reference implementation.
  * 1) It uses a map to track the number of references to a given EX_RUNDOWN_REF structure.
@@ -299,7 +303,7 @@ ExAllocatePoolUninitializedCPP(
     if (tag == 0) {
         KeBugCheckExCPP(BAD_POOL_CALLER, 0x9B, pool_type, number_of_bytes, 0);
     }
-    return usersim_allocate_with_tag(number_of_bytes, tag, false);
+    return usersim_allocate_with_tag(pool_type, number_of_bytes, tag, false);
 }
 
 _Ret_maybenull_ void*
@@ -318,7 +322,170 @@ ExAllocatePoolWithTagCPP(
     if (tag == 0) {
         KeBugCheckExCPP(BAD_POOL_CALLER, 0x9B, pool_type, number_of_bytes, 0);
     }
-    return usersim_allocate_with_tag(number_of_bytes, tag, true);
+    return usersim_allocate_with_tag(pool_type, number_of_bytes, tag, true);
+}
+
+#define USERSIM_CACHE_LINE_SIZE 64
+
+typedef struct
+{
+    POOL_TYPE pool_type;
+    uint32_t tag;
+} usersim_allocation_header_t;
+
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void*
+usersim_allocate_with_tag(
+    _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, size_t size, uint32_t tag, bool initialize)
+{
+    if (size == 0) {
+        KeBugCheckEx(BAD_POOL_CALLER, 0x00, 0, 0, 0);
+    }
+    if (size > usersim_fuzzing_memory_limit) {
+        return nullptr;
+    }
+
+    if (usersim_fault_injection_inject_fault()) {
+        return nullptr;
+    }
+
+    // Allocate space with a usersim_allocation_header_t prepended.
+    void* memory;
+    if (pool_type == NonPagedPoolNxCacheAligned) {
+        // The pointer we return has to be cache aligned so we allocate
+        // enough extra space to fill a cache line, and put the
+        // usersim_allocation_header_t at the end of that space.
+        size_t full_size = USERSIM_CACHE_LINE_SIZE + size;
+        uint8_t* pointer = (uint8_t*)_aligned_malloc(full_size, USERSIM_CACHE_LINE_SIZE);
+        if (pointer == nullptr) {
+            return nullptr;
+        }
+        memory = pointer + USERSIM_CACHE_LINE_SIZE;
+    } else {
+        size_t full_size = sizeof(usersim_allocation_header_t) + size;
+        uint8_t* pointer = (uint8_t*)calloc(full_size, 1);
+        if (pointer == nullptr) {
+            return nullptr;
+        }
+        memory = pointer + sizeof(usersim_allocation_header_t);
+    }
+
+    // Do any initialization.
+    auto header = (usersim_allocation_header_t*)((uint8_t*)memory - sizeof(usersim_allocation_header_t));
+    header->pool_type = pool_type;
+    header->tag = tag;
+    if (!initialize) {
+        // The calloc call always zero-initializes memory.  To test
+        // returning uninitialized memory, we explicitly fill it with 0xcc.
+        memset(memory, 0xcc, size);
+    }
+
+    if (memory && _usersim_leak_detector_ptr) {
+        _usersim_leak_detector_ptr->register_allocation(reinterpret_cast<uintptr_t>(memory), size);
+    }
+
+    return memory;
+}
+
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void*
+usersim_reallocate(
+    _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size)
+{
+    return usersim_reallocate_with_tag(memory, old_size, new_size, 0);
+}
+
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(new_size) void*
+usersim_reallocate_with_tag(
+    _In_ _Post_invalid_ void* memory, size_t old_size, size_t new_size, uint32_t tag)
+{
+    UNREFERENCED_PARAMETER(tag);
+    UNREFERENCED_PARAMETER(old_size);
+
+    if (new_size > usersim_fuzzing_memory_limit) {
+        return nullptr;
+    }
+
+    if (usersim_fault_injection_inject_fault()) {
+        return nullptr;
+    }
+
+    void* p;
+    auto header = (usersim_allocation_header_t*)((uint8_t*)memory - sizeof(usersim_allocation_header_t));
+    if (header->pool_type == NonPagedPoolNxCacheAligned) {
+        uint8_t* pointer = ((uint8_t*)memory) - USERSIM_CACHE_LINE_SIZE;
+        p = _aligned_realloc(pointer, USERSIM_CACHE_LINE_SIZE + new_size, USERSIM_CACHE_LINE_SIZE);
+    } else {
+        uint8_t* pointer = ((uint8_t*)memory) - sizeof(usersim_allocation_header_t);
+        p = realloc(pointer, sizeof(usersim_allocation_header_t)  + new_size);
+    }
+
+    if (p && (new_size > old_size)) {
+        memset(((char*)p) + old_size, 0, new_size - old_size);
+    }
+
+    if (_usersim_leak_detector_ptr) {
+        _usersim_leak_detector_ptr->unregister_allocation(reinterpret_cast<uintptr_t>(memory));
+        _usersim_leak_detector_ptr->register_allocation(reinterpret_cast<uintptr_t>(p), new_size);
+    }
+
+    return p;
+}
+
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size)
+{
+    return usersim_allocate_with_tag(NonPagedPool, size, 'tset', true);
+}
+
+void
+usersim_free(_Frees_ptr_opt_ void* memory)
+{
+    if (memory == nullptr) {
+        return;
+    }
+    auto header = (usersim_allocation_header_t*)((uint8_t*)memory - sizeof(usersim_allocation_header_t));
+    if (_usersim_leak_detector_ptr) {
+        _usersim_leak_detector_ptr->unregister_allocation(reinterpret_cast<uintptr_t>(memory));
+    }
+    if (header->pool_type == NonPagedPoolNxCacheAligned) {
+        uint8_t* pointer = ((uint8_t*)memory) - USERSIM_CACHE_LINE_SIZE;
+        _aligned_free(pointer);
+    } else {
+        uint8_t* pointer = ((uint8_t*)memory) - sizeof(usersim_allocation_header_t);
+        free(pointer);
+    }
+}
+
+__drv_allocatesMem(Mem) _Must_inspect_result_
+_Ret_writes_maybenull_(size) void*
+usersim_allocate_cache_aligned(size_t size)
+{
+    if (size > usersim_fuzzing_memory_limit) {
+        return nullptr;
+    }
+
+    if (usersim_fault_injection_inject_fault()) {
+        return nullptr;
+    }
+
+    void* memory = _aligned_malloc(size, USERSIM_CACHE_LINE_SIZE);
+    if (memory) {
+        memset(memory, 0, size);
+    }
+    return memory;
+}
+
+__drv_allocatesMem(Mem) _Must_inspect_result_
+_Ret_writes_maybenull_(size) void*
+usersim_allocate_cache_aligned_with_tag(size_t size, uint32_t tag)
+{
+    UNREFERENCED_PARAMETER(tag);
+
+    return usersim_allocate_cache_aligned(size);
+}
+
+void
+usersim_free_cache_aligned(_Frees_ptr_opt_ void* memory)
+{
+    _aligned_free(memory);
 }
 
 _Ret_maybenull_ void*
@@ -409,4 +576,21 @@ void
 ExRaiseDatatypeMisalignment()
 {
     ExRaiseDatatypeMisalignmentCPP();
+}
+
+void
+usersim_initialize_ex(bool leak_detector)
+{
+    if (leak_detector) {
+        _usersim_leak_detector_ptr = std::make_unique<usersim_leak_detector_t>();
+    }
+}
+
+void
+usersim_clean_up_ex()
+{
+    if (_usersim_leak_detector_ptr) {
+        _usersim_leak_detector_ptr->dump_leaks();
+        _usersim_leak_detector_ptr.reset();
+    }
 }

src/io.cpp

@@ -117,3 +117,10 @@ _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI PEPROCESS IoGetCurrentProcess(VO
 {
     return (PEPROCESS)GetCurrentProcess();
 }
+
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+IofCompleteRequest(_In_ PIRP irp, _In_ CCHAR priority_boost)
+{
+    UNREFERENCED_PARAMETER(irp);
+    UNREFERENCED_PARAMETER(priority_boost);
+}

src/ke.cpp

@@ -266,7 +266,7 @@ usersim_free_semaphores()
         for (auto handle : *g_usersim_semaphore_handles) {
             ::CloseHandle(handle);
         }
-        usersim_free(g_usersim_semaphore_handles);
+        delete g_usersim_semaphore_handles;
         g_usersim_semaphore_handles = nullptr;
     }
 }
@@ -731,7 +731,7 @@ usersim_free_threadpool_timers()
         for (TP_TIMER* threadpool_timer : *g_usersim_threadpool_timers) {
             CloseThreadpoolTimer(threadpool_timer);
         }
-        usersim_free(g_usersim_threadpool_timers);
+        delete g_usersim_threadpool_timers;
         g_usersim_threadpool_timers = nullptr;
     }
 }

src/kernel_um.h

@@ -42,6 +42,7 @@ extern "C"
 #define STATUS_OBJECT_NAME_EXISTS ((NTSTATUS)0x40000000L)
 #define RPC_NT_CALL_FAILED ((NTSTATUS)0xC002001BL)
 #define STATUS_INVALID_IMAGE_FORMAT ((NTSTATUS)0xC000007BL)
+#define STATUS_DRIVER_INTERNAL_ERROR ((NTSTATUS)0xC0000183L)
 #define STATUS_TOO_MANY_NODES ((NTSTATUS)0xC000020EL)
 #define STATUS_GENERIC_COMMAND_FAILED ((NTSTATUS)0xC0150026L)
 #define STATUS_CONTENT_BLOCKED ((NTSTATUS)0xC0000804L)

src/leak_detector.cpp

@@ -3,6 +3,7 @@
 
 #include "leak_detector.h"
 #include "symbol_decoder.h"
+#include "usersim/ke.h"
 
 #include <iostream>
 #include <sstream>
@@ -30,6 +31,9 @@ void
 _usersim_leak_detector::unregister_allocation(uintptr_t address)
 {
     std::unique_lock<std::mutex> lock(_mutex);
+    if (!_allocations.contains(address)) {
+        KeBugCheck(0);
+    }
     _allocations.erase(address);
 }
 

src/platform.h

@@ -106,16 +106,6 @@ extern "C"
      */
     __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size);
 
-    /**
-     * @brief Allocate memory.
-     * @param[in] size Size of memory to allocate.
-     * @param[in] tag Pool tag to use.
-     * @param[in] initialize False to return "uninitialized" memory.
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    __drv_allocatesMem(Mem) _Must_inspect_result_
-        _Ret_writes_maybenull_(size) void* usersim_allocate_with_tag(size_t size, uint32_t tag, bool initialize);
-
     /**
      * @brief Reallocate memory.
      * @param[in] memory Allocation to be reallocated.