diff --git a/inc/usersim/ex.h b/inc/usersim/ex.h
index c0fee2b..cb11664 100644
--- a/inc/usersim/ex.h
+++ b/inc/usersim/ex.h
@@ -198,6 +198,80 @@ USERSIM_API _Ret_maybenull_ void*
 ExAllocatePoolWithTagCPP(
     _In_ __drv_strictTypeMatch(__drv_typeExpr) POOL_TYPE pool_type, SIZE_T number_of_bytes, ULONG tag);
 
+/**
+ * @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);
+
+    /**
+ * @brief Allocate memory.
+ * @param[in] size Size of memory to allocate.
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size);
+
+/**
+ * @brief Reallocate memory.
+ * @param[in] memory Allocation to be reallocated.
+ * @param[in] old_size Old size of memory to reallocate.
+ * @param[in] new_size New size of memory to reallocate.
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+__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);
+
+/**
+ * @brief Reallocate memory with tag.
+ * @param[in] memory Allocation to be reallocated.
+ * @param[in] old_size Old size of memory to reallocate.
+ * @param[in] new_size New size of memory to reallocate.
+ * @param[in] tag Pool tag to use.
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+__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);
+
+/**
+ * @brief Free memory.
+ * @param[in] memory Allocation to be freed.
+ */
+void
+usersim_free(_Frees_ptr_opt_ void* memory);
+
+/**
+ * @brief Allocate memory that has a starting address that is cache aligned.
+ * @param[in] size Size of memory to allocate
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+USERSIM_API
+__drv_allocatesMem(Mem) _Must_inspect_result_
+    _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned(size_t size);
+
+/**
+ * @brief Allocate memory that has a starting address that is cache aligned with tag.
+ * @param[in] size Size of memory to allocate
+ * @param[in] tag Pool tag to use.
+ * @returns Pointer to memory block allocated, or null on failure.
+ */
+__drv_allocatesMem(Mem) _Must_inspect_result_
+    _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned_with_tag(size_t size, uint32_t tag);
+
+/**
+ * @brief Free memory that has a starting address that is cache aligned.
+ * @param[in] memory Allocation to be freed.
+ */
+void
+usersim_free_cache_aligned(_Frees_ptr_opt_ void* memory);
+
 USERSIM_API _Ret_maybenull_ void*
 ExAllocatePoolUninitializedCPP(_In_ POOL_TYPE pool_type, _In_ size_t number_of_bytes, _In_ unsigned long tag);
 
@@ -207,4 +281,26 @@ ExRaiseAccessViolationCPP();
 USERSIM_API void
 ExRaiseDatatypeMisalignmentCPP();
 
+void usersim_initialize_ex(bool leak_detector);
+void usersim_clean_up_ex();
+
+#ifdef __cplusplus
+#include <memory>
+namespace usersim_helper {
+
+struct _usersim_free_functor
+{
+    void
+    operator()(void* memory)
+    {
+        usersim_free(memory);
+    }
+};
+
+typedef std::unique_ptr<void, _usersim_free_functor> usersim_memory_ptr;
+
+} // namespace usersim_helper
+
+#endif
+
 #endif
diff --git a/inc/usersim/io.h b/inc/usersim/io.h
index 1c403f1..331ebb8 100644
--- a/inc/usersim/io.h
+++ b/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)
 }
diff --git a/inc/usersim/wdf.h b/inc/usersim/wdf.h
index 76f504a..1335744 100644
--- a/inc/usersim/wdf.h
+++ b/inc/usersim/wdf.h
@@ -13,7 +13,7 @@ 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;
@@ -21,8 +21,6 @@ extern "C"
 #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)
 
@@ -78,6 +76,22 @@ extern "C"
     void
     WDF_DRIVER_CONFIG_INIT(_Out_ PWDF_DRIVER_CONFIG config, _In_opt_ PFN_WDF_DRIVER_DEVICE_ADD evt_driver_device_add);
 
+    typedef enum _WDFFUNCENUM
+    {
+        WdfControlDeviceInitAllocateTableIndex = 25,
+        WdfDeviceInitSetDeviceTypeTableIndex = 66,
+        WdfDeviceInitAssignNameTableIndex = 67,
+        WdfDeviceInitSetCharacteristicsTableIndex = 70,
+        WdfDeviceInitSetFileObjectConfigTableIndex = 71,
+        WdfDeviceInitAssignWdmIrpPreprocessCallbackTableIndex = 73,
+        WdfDeviceCreateTableIndex = 75,
+        WdfDeviceCreateSymbolicLinkTableIndex = 80,
+        WdfDriverCreateTableIndex = 116,
+        WdfIoQueueCreateTableIndex = 152,
+        WdfObjectDeleteTableIndex = 208,
+        WdfFunctionTableNumEntries = 444,
+    } WDFFUNCENUM;
+
     void
     usersim_initialize_wdf();
 
diff --git a/src/etw.cpp b/src/etw.cpp
index 3a7ad1b..134e2b0 100644
--- a/src/etw.cpp
+++ b/src/etw.cpp
@@ -3,6 +3,7 @@
 
 #include "platform.h"
 #include "usersim/etw.h"
+#include "usersim/ex.h"
 
 typedef struct
 {
diff --git a/src/ex.cpp b/src/ex.cpp
index b0446be..797f07b 100644
--- a/src/ex.cpp
+++ b/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();
+    }
 }
\ No newline at end of file
diff --git a/src/io.cpp b/src/io.cpp
index 914f5df..9d9a8ee 100644
--- a/src/io.cpp
+++ b/src/io.cpp
@@ -3,6 +3,7 @@
 
 #include "platform.h"
 #include "kernel_um.h"
+#include "usersim/ex.h"
 #include "usersim/io.h"
 
 // Io* functions.
@@ -117,3 +118,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);
+}
\ No newline at end of file
diff --git a/src/ke.cpp b/src/ke.cpp
index 0767864..51d456d 100644
--- a/src/ke.cpp
+++ b/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;
     }
 }
diff --git a/src/kernel_um.h b/src/kernel_um.h
index d2130c7..b87e2e8 100644
--- a/src/kernel_um.h
+++ b/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)
diff --git a/src/leak_detector.cpp b/src/leak_detector.cpp
index 8416088..62b827b 100644
--- a/src/leak_detector.cpp
+++ b/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);
 }
 
diff --git a/src/platform.h b/src/platform.h
index 8cf1d10..e178fff 100644
--- a/src/platform.h
+++ b/src/platform.h
@@ -99,76 +99,6 @@ extern "C"
     void
     usersim_platform_terminate();
 
-    /**
-     * @brief Allocate memory.
-     * @param[in] size Size of memory to allocate.
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    __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.
-     * @param[in] old_size Old size of memory to reallocate.
-     * @param[in] new_size New size of memory to reallocate.
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    __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);
-
-    /**
-     * @brief Reallocate memory with tag.
-     * @param[in] memory Allocation to be reallocated.
-     * @param[in] old_size Old size of memory to reallocate.
-     * @param[in] new_size New size of memory to reallocate.
-     * @param[in] tag Pool tag to use.
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    __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);
-
-    /**
-     * @brief Free memory.
-     * @param[in] memory Allocation to be freed.
-     */
-    void
-    usersim_free(_Frees_ptr_opt_ void* memory);
-
-    /**
-     * @brief Allocate memory that has a starting address that is cache aligned.
-     * @param[in] size Size of memory to allocate
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    USERSIM_API 
-    __drv_allocatesMem(Mem) _Must_inspect_result_
-        _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned(size_t size);
-
-    /**
-     * @brief Allocate memory that has a starting address that is cache aligned with tag.
-     * @param[in] size Size of memory to allocate
-     * @param[in] tag Pool tag to use.
-     * @returns Pointer to memory block allocated, or null on failure.
-     */
-    __drv_allocatesMem(Mem) _Must_inspect_result_
-        _Ret_writes_maybenull_(size) void* usersim_allocate_cache_aligned_with_tag(size_t size, uint32_t tag);
-
-    /**
-     * @brief Free memory that has a starting address that is cache aligned.
-     * @param[in] memory Allocation to be freed.
-     */
-    void
-    usersim_free_cache_aligned(_Frees_ptr_opt_ void* memory);
-
     typedef enum _usersim_page_protection
     {
         USERSIM_PAGE_PROTECT_READ_ONLY,
@@ -775,22 +705,3 @@ extern "C"
 #ifdef __cplusplus
 }
 #endif
-
-#ifdef __cplusplus
-#include <memory>
-namespace usersim_helper {
-
-struct _usersim_free_functor
-{
-    void
-    operator()(void* memory)
-    {
-        usersim_free(memory);
-    }
-};
-
-typedef std::unique_ptr<void, _usersim_free_functor> usersim_memory_ptr;
-
-} // namespace usersim_helper
-
-#endif
diff --git a/src/platform_user.cpp b/src/platform_user.cpp
index 304ff2c..d3128e6 100644
--- a/src/platform_user.cpp
+++ b/src/platform_user.cpp
@@ -37,12 +37,9 @@ bool _usersim_platform_is_preemptible = true;
 int32_t _usersim_platform_initiate_count = 0;
 
 extern "C" bool usersim_fuzzing_enabled = false;
-extern "C" size_t usersim_fuzzing_memory_limit = MAXSIZE_T;
 
 static EX_RUNDOWN_REF _usersim_platform_preemptible_work_items_rundown;
 
-usersim_leak_detector_ptr _usersim_leak_detector_ptr;
-
 /**
  * @brief Environment variable to enable fault injection testing.
  *
@@ -204,9 +201,9 @@ usersim_platform_initiate()
     try {
         _usersim_platform_maximum_group_count = GetMaximumProcessorGroupCount();
         _usersim_platform_maximum_processor_count = GetMaximumProcessorCount(ALL_PROCESSOR_GROUPS);
-        auto fault_injection_stack_depth =
+        bool fault_injection_stack_depth =
             _get_environment_variable_as_size_t(USERSIM_FAULT_INJECTION_SIMULATION_ENVIRONMENT_VARIABLE_NAME);
-        auto leak_detector = _get_environment_variable_as_bool(USERSIM_MEMORY_LEAK_DETECTION_ENVIRONMENT_VARIABLE_NAME);
+        bool leak_detector = _get_environment_variable_as_bool(USERSIM_MEMORY_LEAK_DETECTION_ENVIRONMENT_VARIABLE_NAME);
         if (fault_injection_stack_depth || leak_detector) {
             _usersim_symbol_decoder_initialize();
         }
@@ -218,9 +215,7 @@ usersim_platform_initiate()
             usersim_fuzzing_enabled = true;
         }
 
-        if (leak_detector) {
-            _usersim_leak_detector_ptr = std::make_unique<usersim_leak_detector_t>();
-        }
+        usersim_initialize_ex(leak_detector);
 
         result = usersim_initialize_irql();
         if (result != STATUS_SUCCESS) {
@@ -269,10 +264,7 @@ usersim_platform_terminate()
     usersim_clean_up_dpcs();
     usersim_clean_up_irql();
     _clean_up_thread_pool();
-    if (_usersim_leak_detector_ptr) {
-        _usersim_leak_detector_ptr->dump_leaks();
-        _usersim_leak_detector_ptr.reset();
-    }
+    usersim_clean_up_ex();
 
     int32_t count = InterlockedDecrement((volatile long*)&_usersim_platform_initiate_count);
     if (count < 0) {
@@ -294,114 +286,6 @@ usersim_get_code_integrity_state(_Out_ usersim_code_integrity_state_t* state)
     USERSIM_RETURN_RESULT(STATUS_SUCCESS);
 }
 
-__drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_writes_maybenull_(size) void* usersim_allocate(size_t size)
-{
-    return usersim_allocate_with_tag(size, 'tset', true);
-}
-
-__drv_allocatesMem(Mem) _Must_inspect_result_
-    _Ret_writes_maybenull_(size) void* usersim_allocate_with_tag(size_t size, uint32_t tag, bool initialize)
-{
-    UNREFERENCED_PARAMETER(tag);
-    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;
-    }
-
-    void* memory = calloc(size, 1);
-    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)
-{
-    UNREFERENCED_PARAMETER(old_size);
-    if (new_size > usersim_fuzzing_memory_limit) {
-        return nullptr;
-    }
-
-    if (usersim_fault_injection_inject_fault()) {
-        return nullptr;
-    }
-
-    void* p = realloc(memory, 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_(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);
-
-    return usersim_reallocate(memory, old_size, new_size);
-}
-
-void
-usersim_free(_Frees_ptr_opt_ void* memory)
-{
-    if (_usersim_leak_detector_ptr) {
-        _usersim_leak_detector_ptr->unregister_allocation(reinterpret_cast<uintptr_t>(memory));
-    }
-    free(memory);
-}
-
-__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);
-}
-
 struct _usersim_ring_descriptor
 {
     void* primary_view;
diff --git a/src/se.cpp b/src/se.cpp
index 59b7ec2..c4ada13 100644
--- a/src/se.cpp
+++ b/src/se.cpp
@@ -4,6 +4,7 @@
 #include "fault_injection.h"
 #include "kernel_um.h"
 #include "platform.h"
+#include "usersim/ex.h"
 #include "usersim/se.h"
 #include "utilities.h"
 #include <processthreadsapi.h>
diff --git a/src/wdf.cpp b/src/wdf.cpp
index 8195a50..1eb609b 100644
--- a/src/wdf.cpp
+++ b/src/wdf.cpp
@@ -1,10 +1,13 @@
 // Copyright (c) Microsoft Corporation
 // SPDX-License-Identifier: MIT
 
+#include "fault_injection.h"
 #include "framework.h"
+#include "platform.h"
+#include "usersim/ex.h"
 #include "usersim/wdf.h"
 
-static WDFDRIVER g_CurrentDriver = {};
+WDF_DRIVER_GLOBALS g_UsersimWdfDriverGlobals = {0};
 
 static NTSTATUS
 _WdfDriverCreate(
@@ -15,14 +18,19 @@ _WdfDriverCreate(
     _In_ PWDF_DRIVER_CONFIG driver_config,
     _Out_opt_ WDFDRIVER* driver)
 {
-    UNREFERENCED_PARAMETER(driver_globals);
-    UNREFERENCED_PARAMETER(driver_object);
     UNREFERENCED_PARAMETER(registry_path);
     UNREFERENCED_PARAMETER(driver_attributes);
 
-    g_CurrentDriver.config = *driver_config;
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (driver_globals->Driver != nullptr) {
+        return STATUS_DRIVER_INTERNAL_ERROR;
+    }
+    g_UsersimWdfDriverGlobals.Driver = driver_object;
+    driver_object->config = *driver_config;
     if (driver != nullptr) {
-        *driver = g_CurrentDriver;
+        *driver = driver_object;
     }
     return STATUS_SUCCESS;
 }
@@ -30,34 +38,226 @@ _WdfDriverCreate(
 static NTSTATUS
 _WdfDeviceCreate(
     _In_ WDF_DRIVER_GLOBALS* driver_globals,
-    _Inout_ PWDFDEVICE_INIT* device_init, _In_opt_ PWDF_OBJECT_ATTRIBUTES device_attributes, _Out_ WDFDEVICE* device)
+    _Inout_ PWDFDEVICE_INIT* device_init,
+    _In_opt_ PWDF_OBJECT_ATTRIBUTES device_attributes,
+    _Out_ WDFDEVICE* device)
 {
-    UNREFERENCED_PARAMETER(driver_globals);
     UNREFERENCED_PARAMETER(device_init);
     UNREFERENCED_PARAMETER(device_attributes);
 
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+
     *device = nullptr;
     return STATUS_SUCCESS;
 }
 
-WDF_DRIVER_GLOBALS g_UsersimWdfDriverGlobals = {
-    .Driver = g_CurrentDriver,
-};
+typedef struct _WDF_FILEOBJECT_CONFIG WDF_FILEOBJECT_CONFIG, *PWDF_FILEOBJECT_CONFIG;
+typedef struct _FN_WDFDEVICE_WDM_IRP_PREPROCESS FN_WDFDEVICE_WDM_IRP_PREPROCESS, *PFN_WDFDEVICE_WDM_IRP_PREPROCESS;
+#define IRP_MJ_MAXIMUM_FUNCTION 0x1b
+
+typedef struct _WDFDEVICE_INIT
+{
+    DEVICE_TYPE device_type;
+    ULONG device_characteristics;
+    UNICODE_STRING device_name;
+    PWDF_FILEOBJECT_CONFIG file_object_config;
+    PWDF_OBJECT_ATTRIBUTES file_object_attributes;
+    PFN_WDFDEVICE_WDM_IRP_PREPROCESS evt_device_wdm_irp_preprocess;
+    PUCHAR minor_functions[IRP_MJ_MAXIMUM_FUNCTION];
+    ULONG num_minor_functions[IRP_MJ_MAXIMUM_FUNCTION];
+} WDFDEVICE_INIT;
+
+static
+_Must_inspect_result_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+PWDFDEVICE_INIT
+_WdfControlDeviceInitAllocate(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ WDFDRIVER driver,
+    _In_ CONST UNICODE_STRING* sddl_string)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    UNREFERENCED_PARAMETER(driver);
+    UNREFERENCED_PARAMETER(sddl_string);
+
+    PWDFDEVICE_INIT device_init = (PWDFDEVICE_INIT)usersim_allocate(sizeof(*device_init));
+    return device_init;
+}
+
+static
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+_WdfDeviceInitSetDeviceType(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals, 
+    _In_ PWDFDEVICE_INIT device_init,
+    _In_ DEVICE_TYPE device_type)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    device_init->device_type = device_type;
+}
+
+static
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+_WdfDeviceInitSetCharacteristics(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals, 
+    _In_ PWDFDEVICE_INIT device_init,
+    _In_ ULONG device_characteristics,
+    _In_ BOOLEAN or_in_values)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    if (!or_in_values) {
+        device_init->device_characteristics = 0;
+    }
+    device_init->device_characteristics |= device_characteristics;
+}
+
+static _Must_inspect_result_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+NTSTATUS
+_WdfDeviceInitAssignName(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ PWDFDEVICE_INIT device_init,
+    _In_opt_ PCUNICODE_STRING device_name)
+{
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    if (device_name != nullptr) {
+        device_init->device_name = *device_name;
+    }
+    return STATUS_SUCCESS;
+}
 
-typedef enum _WDFFUNCENUM
+static
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+_WdfDeviceInitSetFileObjectConfig(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ PWDFDEVICE_INIT device_init,
+    _In_ PWDF_FILEOBJECT_CONFIG file_object_config,
+    _In_opt_ PWDF_OBJECT_ATTRIBUTES file_object_attributes)
 {
-    WdfDeviceCreateTableIndex = 75,
-    WdfDriverCreateTableIndex = 116,
-    WdfFunctionTableNumEntries = 444,
-} WDFFUNCENUM;
+    UNREFERENCED_PARAMETER(driver_globals);
+    device_init->file_object_config = file_object_config;
+    device_init->file_object_attributes = file_object_attributes;
+}
+
+static _Must_inspect_result_
+_IRQL_requires_max_(DISPATCH_LEVEL)
+NTSTATUS
+_WdfDeviceInitAssignWdmIrpPreprocessCallback(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ PWDFDEVICE_INIT device_init,
+    _In_ PFN_WDFDEVICE_WDM_IRP_PREPROCESS evt_device_wdm_irp_preprocess,
+    _In_ UCHAR major_function,
+    _When_(num_minor_functions > 0, _In_reads_bytes_(num_minor_functions)) _When_(num_minor_functions == 0, _In_opt_)
+        PUCHAR minor_functions,
+    _In_ ULONG num_minor_functions)
+{
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (device_init->minor_functions[major_function] != nullptr) {
+        return STATUS_INVALID_DEVICE_REQUEST;
+    }
+    if (major_function >= IRP_MJ_MAXIMUM_FUNCTION) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+
+    device_init->evt_device_wdm_irp_preprocess = evt_device_wdm_irp_preprocess;
+    device_init->minor_functions[major_function] = minor_functions;
+    device_init->num_minor_functions[major_function] = num_minor_functions;
+    return STATUS_SUCCESS;
+}
+
+static _Must_inspect_result_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+NTSTATUS
+_WdfDeviceCreateSymbolicLink(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ WDFDEVICE device,
+    _In_ PCUNICODE_STRING symbolic_link_name)
+{
+    UNREFERENCED_PARAMETER(device);
+    UNREFERENCED_PARAMETER(symbolic_link_name);
+
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    return STATUS_SUCCESS;
+}
+
+typedef struct _WDF_IO_QUEUE_CONFIG WDF_IO_QUEUE_CONFIG, *PWDF_IO_QUEUE_CONFIG;
+typedef struct _WDFQUEUE
+{
+    int tbd;
+} WDFQUEUE;
+
+static
+_Must_inspect_result_
+_IRQL_requires_max_(DISPATCH_LEVEL)
+NTSTATUS
+_WdfIoQueueCreate(
+    _In_ PWDF_DRIVER_GLOBALS driver_globals,
+    _In_ WDFDEVICE device,
+    _In_ PWDF_IO_QUEUE_CONFIG config,
+    _In_opt_ PWDF_OBJECT_ATTRIBUTES queue_attributes,
+    _Out_opt_ WDFQUEUE* queue)
+{
+    UNREFERENCED_PARAMETER(device);
+    UNREFERENCED_PARAMETER(config);
+    UNREFERENCED_PARAMETER(queue_attributes);
+
+    if (driver_globals != &g_UsersimWdfDriverGlobals) {
+        return STATUS_INVALID_PARAMETER;
+    }
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    if (queue != nullptr) {
+        memset(queue, 0, sizeof(*queue));
+    }
+    return STATUS_SUCCESS;
+}
+
+typedef HANDLE WDFOBJECT;
+
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+_WdfObjectDelete(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFOBJECT object)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    UNREFERENCED_PARAMETER(object);
+}
 
 WDFFUNC g_UsersimWdfFunctions[WdfFunctionTableNumEntries];
 
 void
 usersim_initialize_wdf()
 {
+    g_UsersimWdfFunctions[WdfControlDeviceInitAllocateTableIndex] = (WDFFUNC)_WdfControlDeviceInitAllocate;
+    g_UsersimWdfFunctions[WdfDeviceInitSetDeviceTypeTableIndex] = (WDFFUNC)_WdfDeviceInitSetDeviceType;
+    g_UsersimWdfFunctions[WdfDeviceInitAssignNameTableIndex] = (WDFFUNC)_WdfDeviceInitAssignName;
+    g_UsersimWdfFunctions[WdfDeviceInitSetCharacteristicsTableIndex] = (WDFFUNC)_WdfDeviceInitSetCharacteristics;
+    g_UsersimWdfFunctions[WdfDeviceInitSetFileObjectConfigTableIndex] = (WDFFUNC)_WdfDeviceInitSetFileObjectConfig;
+    g_UsersimWdfFunctions[WdfDeviceInitAssignWdmIrpPreprocessCallbackTableIndex] =
+        (WDFFUNC)_WdfDeviceInitAssignWdmIrpPreprocessCallback;
     g_UsersimWdfFunctions[WdfDeviceCreateTableIndex] = (WDFFUNC)_WdfDeviceCreate;
+    g_UsersimWdfFunctions[WdfDeviceCreateSymbolicLinkTableIndex] = (WDFFUNC)_WdfDeviceCreateSymbolicLink;
     g_UsersimWdfFunctions[WdfDriverCreateTableIndex] = (WDFFUNC)_WdfDriverCreate;
+    g_UsersimWdfFunctions[WdfIoQueueCreateTableIndex] = (WDFFUNC)_WdfIoQueueCreate;
+    g_UsersimWdfFunctions[WdfObjectDeleteTableIndex] = (WDFFUNC)_WdfObjectDelete;
 }
 
 extern "C"
@@ -69,7 +269,7 @@ extern "C"
 }
 
 WDFDRIVER
-WdfGetDriver() { return g_CurrentDriver; }
+WdfGetDriver() { return g_UsersimWdfDriverGlobals.Driver; }
 
 void
 WDF_DRIVER_CONFIG_INIT(_Out_ PWDF_DRIVER_CONFIG config, _In_opt_ PFN_WDF_DRIVER_DEVICE_ADD evt_driver_device_add)
diff --git a/tests/ex_test.cpp b/tests/ex_test.cpp
index 633f7a2..4de3bd9 100644
--- a/tests/ex_test.cpp
+++ b/tests/ex_test.cpp
@@ -10,12 +10,21 @@
 
 TEST_CASE("ExAllocatePool", "[ex]")
 {
+    // Try an allocation that need not be cache aligned.
     uint64_t* buffer = (uint64_t*)ExAllocatePoolUninitialized(NonPagedPool, 8, 'tset');
     REQUIRE(buffer != nullptr);
     REQUIRE(*buffer != 0);
     *buffer = 0;
     ExFreePool(buffer);
 
+    // Try an allocation that must be cache aligned.
+    buffer = (uint64_t*)ExAllocatePoolUninitialized(NonPagedPoolNxCacheAligned, 8, 'tset');
+    REQUIRE(buffer != nullptr);
+    REQUIRE(*buffer != 0);
+    REQUIRE((((uintptr_t)buffer) % 64) == 0);
+    *buffer = 0;
+    ExFreePool(buffer);
+
     buffer = (uint64_t*)ExAllocatePoolWithTag(NonPagedPool, 8, 'tset');
     REQUIRE(buffer != nullptr);
     REQUIRE(*buffer == 0);
diff --git a/tests/wdf_test.cpp b/tests/wdf_test.cpp
index f93869e..093a62c 100644
--- a/tests/wdf_test.cpp
+++ b/tests/wdf_test.cpp
@@ -7,6 +7,7 @@
 #else
 #include <catch2/catch.hpp>
 #endif
+#include "usersim/wdf.h"
 
 TEST_CASE("DriverEntry", "[wdf]")
 {
@@ -15,3 +16,30 @@ TEST_CASE("DriverEntry", "[wdf]")
 
     FreeLibrary(module);
 }
+
+extern "C"
+{
+    __declspec(dllimport) const WDFFUNC* UsersimWdfFunctions;
+    __declspec(dllimport) PWDF_DRIVER_GLOBALS UsersimWdfDriverGlobals;
+}
+
+TEST_CASE("WdfDriverCreate", "[wdf]")
+{
+    DRIVER_OBJECT driver_object = {0};
+    WDF_DRIVER_CONFIG config;
+    WDF_DRIVER_CONFIG_INIT(&config, nullptr);
+
+    typedef NTSTATUS(FN_WdfDriverCreate)(
+        _In_ WDF_DRIVER_GLOBALS* driver_globals,
+        _In_ PDRIVER_OBJECT driver_object,
+        _In_ PCUNICODE_STRING registry_path,
+        _In_opt_ PWDF_OBJECT_ATTRIBUTES driver_attributes,
+        _In_ PWDF_DRIVER_CONFIG driver_config,
+        _Out_opt_ WDFDRIVER* driver);
+    FN_WdfDriverCreate* WdfDriverCreate = (FN_WdfDriverCreate*)UsersimWdfFunctions[WdfDriverCreateTableIndex];
+    WDFDRIVER driver;
+    NTSTATUS status = WdfDriverCreate(UsersimWdfDriverGlobals, &driver_object, nullptr, nullptr, &config, &driver);
+    REQUIRE(status == STATUS_SUCCESS);
+    REQUIRE(driver == &driver_object);
+    UsersimWdfDriverGlobals->Driver = nullptr;
+}
diff --git a/usersim_dll_skeleton/dll_skeleton.c b/usersim_dll_skeleton/dll_skeleton.c
index 40b94ae..0711440 100644
--- a/usersim_dll_skeleton/dll_skeleton.c
+++ b/usersim_dll_skeleton/dll_skeleton.c
@@ -22,11 +22,6 @@ __declspec(dllimport) const WDFFUNC* UsersimWdfFunctions;
 PWDF_DRIVER_GLOBALS WdfDriverGlobals = NULL;
 const WDFFUNC* WdfFunctions_01015 = NULL;
 
-    struct _driver_object
-    {
-        int dummy;
-    };
-
     NTSTATUS
     UsersimStartDriver()
     {
@@ -41,10 +36,11 @@ const WDFFUNC* WdfFunctions_01015 = NULL;
     void
     UsersimStopDriver()
     {
-        WDFDRIVER driver = WdfDriverGlobals->Driver;
-        if (driver.config.EvtDriverUnload != NULL) {
-            driver.config.EvtDriverUnload(driver);
+        DRIVER_OBJECT* driver = (DRIVER_OBJECT*)WdfDriverGlobals->Driver;
+        if (driver->config.EvtDriverUnload != NULL) {
+            driver->config.EvtDriverUnload(driver);
         }
+        WdfDriverGlobals->Driver = NULL;
     }
 
     BOOL APIENTRY