some small tweaks

driver/apc.c

@@ -2,24 +2,28 @@
 
 #include "driver.h"
 #include "imports.h"
-
 #include "lib/stdlib.h"
 
 VOID
-GetApcContextByIndex(_Out_ PVOID* Context, _In_ INT Index)
+GetApcContextByIndex(_Out_ PVOID* Context, _In_ UINT32 Index)
 {
+    NT_ASSERT(Index <= MAXIMUM_APC_CONTEXTS);
     AcquireDriverConfigLock();
     *Context = (PVOID)GetApcContextArray()[Index];
     ReleaseDriverConfigLock();
 }
 
 VOID
-GetApcContext(_Out_ PVOID* Context, _In_ LONG ContextIdentifier)
+GetApcContext(_Out_ PVOID* Context, _In_ UINT32 ContextIdentifier)
 {
+    NT_ASSERT(ContextIdentifier <= MAXIMUM_APC_CONTEXTS);
+
+    PAPC_CONTEXT_HEADER header = NULL;
+
     AcquireDriverConfigLock();
 
-    for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+    for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
-        PAPC_CONTEXT_HEADER header = GetApcContextArray()[index];
+        header = GetApcContextArray()[index];
 
         if (!header)
             continue;
@@ -43,10 +47,12 @@ unlock:
 BOOLEAN
 FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context)
 {
-    DEBUG_VERBOSE("All APCs executed, freeing context structure");
+    NT_ASSERT(Context <= MAXIMUM_APC_CONTEXTS);
 
-    for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+    PUINT64 entry = NULL;
-        PUINT64 entry = GetApcContextArray();
+
+    for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+        entry = GetApcContextArray();
 
         if (entry[index] != (UINT64)Context)
             continue;
@@ -63,23 +69,28 @@ FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context)
 }
 
 VOID
-IncrementApcCount(_In_ LONG ContextId)
+IncrementApcCount(_In_ UINT32 ContextId)
 {
+    NT_ASSERT(ContextId <= MAXIMUM_APC_CONTEXTS);
+
     PAPC_CONTEXT_HEADER header = NULL;
+
     GetApcContext(&header, ContextId);
 
     if (!header)
         return;
 
-    /* i actually dont think we need this lock here */
     AcquireDriverConfigLock();
     header->count += 1;
     ReleaseDriverConfigLock();
 }
 
 VOID
-FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ LONG ContextId)
+FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ UINT32 ContextId)
 {
+    NT_ASSERT(Apc != NULL);
+    NT_ASSERT(ContextId <= MAXIMUM_APC_CONTEXTS);
+
     PAPC_CONTEXT_HEADER context = NULL;
 
     ImpExFreePoolWithTag(Apc, POOL_TAG_APC);
@@ -123,10 +134,11 @@ FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ LONG ContextId)
 NTSTATUS
 QueryActiveApcContextsForCompletion()
 {
+    PAPC_CONTEXT_HEADER entry = NULL;
+
     AcquireDriverConfigLock();
 
-    for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+    for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
-        PAPC_CONTEXT_HEADER entry = NULL;
         GetApcContextByIndex(&entry, index);
 
         if (!entry)
@@ -137,8 +149,7 @@ QueryActiveApcContextsForCompletion()
 
         switch (entry->context_id) {
         case APC_CONTEXT_ID_STACKWALK:
-            FreeApcStackwalkApcContextInformation(
+            FreeApcStackwalkApcContextInformation(entry);
-                (PAPC_STACKWALK_CONTEXT)entry);
             FreeApcContextStructure(entry);
             break;
         }
@@ -151,13 +162,17 @@ QueryActiveApcContextsForCompletion()
 VOID
 InsertApcContext(_In_ PVOID Context)
 {
+    NT_ASSERT(Context != NULL);
+
+    PUINT64 entry = NULL;
+
     if (IsDriverUnloading())
         return;
 
     AcquireDriverConfigLock();
 
-    for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+    for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
-        PUINT64 entry = GetApcContextArray();
+        entry = GetApcContextArray();
 
         if (entry[index] == NULL) {
             entry[index] = (UINT64)Context;
@@ -198,17 +213,26 @@ end:
 BOOLEAN
 DrvUnloadFreeAllApcContextStructures()
 {
+    PUINT64 entry = NULL;
+    PAPC_CONTEXT_HEADER context = NULL;
+    LARGE_INTEGER delay = {.QuadPart = -ABSOLUTE(SECONDS(1))};
+
     AcquireDriverConfigLock();
 
-    for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
+    for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
-        PUINT64 entry = GetApcContextArray();
+        entry = GetApcContextArray();
 
         if (entry[index] == NULL)
             continue;
 
-        PAPC_CONTEXT_HEADER context = entry[index];
+        context = entry[index];
 
         if (context->count > 0) {
+            DEBUG_VERBOSE(
+                "Still active APCs: Index: %lx, Count: %lx",
+                index,
+                context->count);
+            KeDelayExecutionThread(KernelMode, FALSE, &delay);
             ReleaseDriverConfigLock();
             return FALSE;
         }

driver/apc.h

@@ -9,19 +9,19 @@
 #include "imports.h"
 
 VOID
-GetApcContextByIndex(_Out_ PVOID* Context, _In_ INT Index);
+GetApcContextByIndex(_Out_ PVOID* Context, _In_ UINT32 Index);
 
 VOID
-GetApcContext(_Out_ PVOID* Context, _In_ LONG ContextIdentifier);
+GetApcContext(_Out_ PVOID* Context, _In_ UINT32 ContextIdentifier);
 
 BOOLEAN
 FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context);
 
 VOID
-IncrementApcCount(_In_ LONG ContextId);
+IncrementApcCount(_In_ UINT32 ContextId);
 
 VOID
-FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ LONG ContextId);
+FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ UINT32 ContextId);
 
 NTSTATUS
 QueryActiveApcContextsForCompletion();

driver/callbacks.c

@@ -1,20 +1,17 @@
 #include "callbacks.h"
 
-#include "driver.h"
+#include "containers/map.h"
-
+#include "containers/tree.h"
 #include "crypt.h"
+#include "driver.h"
 #include "imports.h"
+#include "lib/stdlib.h"
 #include "modules.h"
 #include "pool.h"
 #include "session.h"
 #include "thread.h"
 #include "util.h"
 
-#include "lib/stdlib.h"
-
-#include "containers/map.h"
-#include "containers/tree.h"
-
 #define PROCESS_HASHMAP_BUCKET_COUNT 101
 
 STATIC
@@ -51,33 +48,29 @@ UnregisterProcessCreateNotifyRoutine()
 VOID
 UnregisterImageLoadNotifyRoutine()
 {
-    PDRIVER_LIST_HEAD list = GetDriverList();
+    InterlockedExchange(&GetDriverList()->active, FALSE);
-    InterlockedExchange(&list->active, FALSE);
     PsRemoveLoadImageNotifyRoutine(ImageLoadNotifyRoutineCallback);
 }
 
 VOID
 UnregisterThreadCreateNotifyRoutine()
 {
-    PRB_TREE tree = GetThreadTree();
+    InterlockedExchange(&GetThreadTree()->active, FALSE);
-    InterlockedExchange(&tree->active, FALSE);
     ImpPsRemoveCreateThreadNotifyRoutine(ThreadCreateNotifyRoutine);
 }
 
 VOID
 CleanupThreadListOnDriverUnload()
 {
-    PRB_TREE tree = GetThreadTree();
+    RtlRbTreeEnumerate(GetThreadTree(), CleanupThreadListFreeCallback, NULL);
-    DEBUG_VERBOSE("Freeing thread list!");
+    RtlRbTreeDeleteTree(GetThreadTree());
-    RtlRbTreeEnumerate(tree, CleanupThreadListFreeCallback, NULL);
-    RtlRbTreeDeleteTree(tree);
 }
 
 VOID
 CleanupDriverListOnDriverUnload()
 {
-    PDRIVER_LIST_HEAD head = GetDriverList();
     PLIST_ENTRY entry = NULL;
+    PDRIVER_LIST_HEAD head = GetDriverList();
     PDRIVER_LIST_ENTRY driver = NULL;
 
     ImpKeAcquireGuardedMutex(&head->lock);
@@ -95,19 +88,20 @@ VOID
 EnumerateDriverListWithCallbackRoutine(
     _In_ DRIVERLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context)
 {
+    NT_ASSERT(CallbackRoutine != NULL);
+
     PDRIVER_LIST_HEAD head = GetDriverList();
-    PLIST_ENTRY list_entry = NULL;
+    PLIST_ENTRY entry = NULL;
-    PDRIVER_LIST_ENTRY driver_entry = NULL;
+    PDRIVER_LIST_ENTRY driver = NULL;
 
     ImpKeAcquireGuardedMutex(&head->lock);
 
     if (CallbackRoutine) {
-        list_entry = head->list_entry.Flink;
+        entry = head->list_entry.Flink;
-        while (list_entry != &head->list_entry) {
+        while (entry != &head->list_entry) {
-            driver_entry =
+            driver = CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, list_entry);
-                CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry);
+            CallbackRoutine(driver, Context);
-            CallbackRoutine(driver_entry, Context);
+            entry = entry->Flink;
-            list_entry = list_entry->Flink;
         }
     }
 
@@ -197,7 +191,6 @@ InitialiseDriverList()
     }
 
     KeReleaseGuardedMutex(&head->lock);
-
     head->active = TRUE;
 
     if (modules.address)
@@ -215,25 +208,27 @@ VOID
 FindDriverEntryByBaseAddress(
     _In_ PVOID ImageBase, _Out_ PDRIVER_LIST_ENTRY* Entry)
 {
-    PDRIVER_LIST_HEAD head = GetDriverList();
+    NT_ASSERT(ImageBase != NULL);
-    PLIST_ENTRY list_entry = NULL;
+    NT_ASSERT(Entry != NULL);
-    PDRIVER_LIST_ENTRY driver_entry = NULL;
+
+    PDRIVER_LIST_HEAD head = GetDriverList();
+    PLIST_ENTRY entry = NULL;
+    PDRIVER_LIST_ENTRY driver = NULL;
 
-    ImpKeAcquireGuardedMutex(&head->lock);
     *Entry = NULL;
 
-    list_entry = head->list_entry.Flink;
+    ImpKeAcquireGuardedMutex(&head->lock);
+    entry = head->list_entry.Flink;
 
-    while (list_entry != &head->list_entry) {
+    while (entry != &head->list_entry) {
-        driver_entry =
+        driver = CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, list_entry);
-            CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry);
 
-        if (driver_entry->ImageBase == ImageBase) {
+        if (driver->ImageBase == ImageBase) {
-            *Entry = driver_entry;
+            *Entry = driver;
             goto unlock;
         }
 
-        list_entry = list_entry->Flink;
+        entry = entry->Flink;
     }
 
 unlock:
@@ -244,6 +239,9 @@ STATIC
 BOOLEAN
 ProcessHashmapCompareFunction(_In_ PVOID Struct1, _In_ PVOID Struct2)
 {
+    NT_ASSERT(Struct1 != NULL);
+    NT_ASSERT(Struct2 != NULL);
+
     HANDLE h1 = *((PHANDLE)Struct1);
     HANDLE h2 = *((PHANDLE)Struct2);
 
@@ -404,6 +402,7 @@ FreeProcessEntryModuleList(
     _In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context)
 {
     UNREFERENCED_PARAMETER(Context);
+    NT_ASSERT(Entry != NULL);
 
     PRTL_HASHMAP map = GetProcessHashmap();
     PLIST_ENTRY list = NULL;
@@ -460,6 +459,8 @@ VOID
 FindOurUserModeModuleEntry(
     _In_ PROCESS_MODULE_CALLBACK Callback, _In_opt_ PVOID Context)
 {
+    NT_ASSERT(Callback != NULL);
+
     INT32 index = 0;
     PRTL_HASHMAP map = GetProcessHashmap();
     PPROCESS_LIST_ENTRY entry = NULL;
@@ -524,7 +525,6 @@ CleanupProcessHashmap()
     }
 
     context = map->context;
-
     ExDeleteLookasideListEx(&context->pool);
     ExFreePoolWithTag(map->context, POOL_TAG_HASHMAP);
     RtlHashmapDelete(map);
@@ -583,6 +583,9 @@ STATIC
 UINT32
 ThreadListTreeCompare(_In_ PVOID Key, _In_ PVOID Object)
 {
+    NT_ASSERT(Key != NULL);
+    NT_ASSERT(Object != NULL);
+
     HANDLE tid_1 = *((PHANDLE)Object);
     HANDLE tid_2 = *((PHANDLE)Key);
 
@@ -603,8 +606,10 @@ InitialiseThreadList()
     status =
         RtlRbTreeCreate(ThreadListTreeCompare, sizeof(THREAD_LIST_ENTRY), tree);
 
-    if (!NT_SUCCESS(status))
+    if (!NT_SUCCESS(status)) {
         DEBUG_ERROR("RtlRbTreeCreate: %x", status);
+        return status;
+    }
 
     tree->active = TRUE;
     return status;
@@ -630,6 +635,7 @@ CanInitiateDeferredHashing(_In_ LPCSTR ProcessName, _In_ PDRIVER_LIST_HEAD Head)
                : FALSE;
 }
 
+#ifdef DEBUG
 STATIC
 VOID
 PrintHashmapCallback(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context)
@@ -656,6 +662,7 @@ EnumerateAndPrintProcessHashmap()
 {
     RtlHashmapEnumerate(GetProcessHashmap(), PrintHashmapCallback, NULL);
 }
+#endif
 
 VOID
 ProcessCreateNotifyRoutine(
@@ -852,16 +859,10 @@ ObPreOpCallbackRoutine(
     _In_ POB_PRE_OPERATION_INFORMATION OperationInformation)
 {
     PAGED_CODE();
-
     UNREFERENCED_PARAMETER(RegistrationContext);
 
     /* access mask to completely strip permissions */
     ACCESS_MASK deny_access = SYNCHRONIZE | PROCESS_TERMINATE;
-
-    /*
-     * This callback routine is executed in the context of the thread that
-     * is requesting to open said handle
-     */
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PEPROCESS process_creator = PsGetCurrentProcess();
     PEPROCESS protected_process = NULL;
@@ -1178,8 +1179,8 @@ EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context)
 {
     /* Handles are stored in pageable memory */
     PAGED_CODE();
-
     UNREFERENCED_PARAMETER(Context);
+    NT_ASSERT(Entry != NULL);
 
     if (!Entry)
         return STATUS_INVALID_PARAMETER;
@@ -1214,6 +1215,8 @@ VOID
 TimerObjectValidateProcessModuleCallback(
     _In_ PPROCESS_MAP_MODULE_ENTRY Entry, _In_opt_ PVOID Context)
 {
+    NT_ASSERT(Entry != NULL);
+
     CHAR hash[SHA_256_HASH_LENGTH] = {0};
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PACTIVE_SESSION session = (PACTIVE_SESSION)Context;
@@ -1297,6 +1300,7 @@ TimerObjectCallbackRoutine(
     UNREFERENCED_PARAMETER(Dpc);
     UNREFERENCED_PARAMETER(SystemArgument1);
     UNREFERENCED_PARAMETER(SystemArgument2);
+    NT_ASSERT(DeferredContext != NULL);
 
     if (!HasDriverLoaded() || !ARGUMENT_PRESENT(DeferredContext))
         return;

driver/crypt.c

@@ -2,13 +2,11 @@
 
 #include "driver.h"
 #include "imports.h"
+#include "lib/stdlib.h"
 #include "session.h"
-#include "util.h"
-
 #include "types/tpm20.h"
 #include "types/tpmptp.h"
-
+#include "util.h"
-#include "lib/stdlib.h"
 
 #include <bcrypt.h>
 #include <immintrin.h>

driver/driver.c

@@ -9,13 +9,12 @@
 #include "imports.h"
 #include "integrity.h"
 #include "io.h"
+#include "lib/stdlib.h"
 #include "modules.h"
 #include "pool.h"
 #include "session.h"
 #include "thread.h"
 
-#include "lib/stdlib.h"
-
 #include <immintrin.h>
 
 STATIC
@@ -466,7 +465,6 @@ DriverUnload(_In_ PDRIVER_OBJECT DriverObject)
     UnregisterProcessCreateNotifyRoutine();
     UnregisterImageLoadNotifyRoutine();
 
-
     DrvUnloadFreeThreadList();
     DrvUnloadFreeProcessList();
     DrvUnloadFreeDriverList();
@@ -1049,7 +1047,7 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath)
 
     SetDriverLoadedFlag();
     TpmExtractEndorsementKey();
-    //PoolScanForManualMappedDrivers();
+    // PoolScanForManualMappedDrivers();
 
     DEBUG_INFO("Driver Entry Complete.");
     return STATUS_SUCCESS;

driver/hv.c

@@ -3,10 +3,10 @@
 #include "common.h"
 #include "imports.h"
 #include "io.h"
-#include <intrin.h>
-
 #include "lib/stdlib.h"
 
+#include <intrin.h>
+
 #ifdef ALLOC_PRAGMA
 #    pragma alloc_text(PAGE, PerformVirtualizationDetection)
 #endif

driver/hw.c

@@ -2,9 +2,8 @@
 
 #include "crypt.h"
 #include "imports.h"
-#include "modules.h"
-
 #include "lib/stdlib.h"
+#include "modules.h"
 
 #define PCI_VENDOR_ID_OFFSET 0x00
 #define PCI_DEVICE_ID_OFFSET 0x02
@@ -283,11 +282,12 @@ ReportBlacklistedPcieDevice(
     _In_ PDEVICE_OBJECT DeviceObject, _In_ PPCI_COMMON_HEADER Header)
 {
     NTSTATUS status = STATUS_UNSUCCESSFUL;
-    UINT32 packet_size = CryptRequestRequiredBufferLength(
+    UINT32 len = 0;
-        sizeof(BLACKLISTED_PCIE_DEVICE_REPORT));
+    PBLACKLISTED_PCIE_DEVICE_REPORT report = NULL;
 
-    PBLACKLISTED_PCIE_DEVICE_REPORT report =
+    len = CryptRequestRequiredBufferLength(
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
+        sizeof(BLACKLISTED_PCIE_DEVICE_REPORT));
+    report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG);
 
     if (!report)
         return;
@@ -298,15 +298,15 @@ ReportBlacklistedPcieDevice(
     report->device_id = Header->DeviceID;
     report->vendor_id = Header->VendorID;
 
-    status = CryptEncryptBuffer(report, packet_size);
+    status = CryptEncryptBuffer(report, len);
 
     if (!NT_SUCCESS(status)) {
         DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
-        ImpExFreePoolWithTag(report, packet_size);
+        ImpExFreePoolWithTag(report, len);
         return;
     }
 
-    IrpQueueSchedulePacket(report, packet_size);
+    IrpQueueSchedulePacket(report, len);
 }
 
 STATIC
@@ -338,13 +338,6 @@ PciDeviceQueryCallback(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context)
             header.DeviceID);
         ReportBlacklistedPcieDevice(DeviceObject, &header);
     }
-    else {
-        DEBUG_VERBOSE(
-            "Device: %llx, DeviceID: %lx, VendorID: %lx",
-            DeviceObject,
-            header.DeviceID,
-            header.VendorID);
-    }
 
     return status;
 }

driver/integrity.c

@@ -6,48 +6,61 @@
 #include "driver.h"
 #include "imports.h"
 #include "io.h"
+#include "lib/stdlib.h"
 #include "modules.h"
 #include "pe.h"
 #include "session.h"
 #include "util.h"
 
-#include "lib/stdlib.h"
-
 #include <bcrypt.h>
 #include <devpkey.h>
 #include <initguid.h>
 
-/* Header for a buffer that contains an array of sections copied from a module
+// clang-format off
- */
+
 typedef struct _INTEGRITY_CHECK_HEADER {
+
+    /* Count of total sections contained within the buffer */
     UINT32 section_count;
+
+    /* Total size of the buffer */
     UINT32 total_size;
 
 } INTEGRITY_CHECK_HEADER, *PINTEGRITY_CHECK_HEADER;
 
+
 typedef struct _PROCESS_MODULE_INFORMATION {
+    /* Pointer to the base of the module*/
     PVOID module_base;
+
+    /* Total size of the module */
     SIZE_T module_size;
+
+    /* Path to the modules executable image*/
     WCHAR module_path[MAX_MODULE_PATH];
 
 } PROCESS_MODULE_INFORMATION, *PPROCESS_MODULE_INFORMATION;
 
-/* Structure representing the data passed back to user-mode after validating a
- * process module sections*/
 typedef struct _PROCESS_MODULE_VALIDATION_RESULT {
+
+    /* Boolean value of whether or not the module image is valid */
     UINT32 is_module_valid;
 
 } PROCESS_MODULE_VALIDATION_RESULT, *PPROCESS_MODULE_VALIDATION_RESULT;
 
 typedef struct _VAL_INTEGRITY_HEADER {
+
+    /* Header containing information pertaining to the buffer */
     INTEGRITY_CHECK_HEADER integrity_check_header;
+
+    /* Section header */
     IMAGE_SECTION_HEADER section_header;
+
+    /* Pointer to the start of the sections image */
     CHAR section_base[];
 
 } VAL_INTEGRITY_HEADER, *PVAL_INTEGRITY_HEADER;
 
-// clang-format off
-
 STATIC
 NTSTATUS
 InitiateEptFunctionAddressArrays();
@@ -136,10 +149,12 @@ GetDriverImageSize(_Inout_ PIRP Irp)
 {
     PAGED_CODE();
 
+    NT_ASSERT(Irp != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
-    LPCSTR driver_name = GetDriverName();
+    LPCSTR name = GetDriverName();
     SYSTEM_MODULES modules = {0};
-    PRTL_MODULE_EXTENDED_INFO driver_info = NULL;
+    PRTL_MODULE_EXTENDED_INFO driver = NULL;
 
     status = GetSystemModuleInformation(&modules);
 
@@ -148,27 +163,27 @@ GetDriverImageSize(_Inout_ PIRP Irp)
         return status;
     }
 
-    driver_info = FindSystemModuleByName(driver_name, &modules);
+    driver = FindSystemModuleByName(name, &modules);
 
-    if (!driver_info) {
+    if (!driver) {
         DEBUG_ERROR("FindSystemModuleByName failed with no status code");
         ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL);
         return STATUS_NOT_FOUND;
     }
 
-    status = ValidateIrpOutputBuffer(Irp, sizeof(ULONG));
+    status = ValidateIrpOutputBuffer(Irp, sizeof(UINT32));
 
     if (!NT_SUCCESS(status)) {
         DEBUG_ERROR("ValidateIrpOutputBuffer failed with status %x", status);
         goto end;
     }
 
-    Irp->IoStatus.Information = sizeof(ULONG);
+    Irp->IoStatus.Information = sizeof(UINT32);
 
     IntCopyMemory(
         Irp->AssociatedIrp.SystemBuffer,
-        &driver_info->ImageSize,
+        &driver->ImageSize,
-        sizeof(ULONG));
+        sizeof(UINT32));
 
 end:
 
@@ -185,6 +200,8 @@ GetModuleInformationByName(
 {
     PAGED_CODE();
 
+    NT_ASSERT(ModuleName != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     LPCSTR driver_name = GetDriverName();
     SYSTEM_MODULES modules = {0};
@@ -271,6 +288,10 @@ StoreModuleExecutableRegionsInBuffer(
 {
     PAGED_CODE();
 
+    NT_ASSERT(Buffer != NULL);
+    NT_ASSERT(ModuleBase != NULL);
+    NT_ASSERT(BytesWritten != NULL);
+
     UINT32 total_packet_size = 0;
     UINT32 num_sections = 0;
     UINT32 num_executable_sections = 0;
@@ -382,6 +403,11 @@ MapDiskImageIntoVirtualAddressSpace(
 {
     PAGED_CODE();
 
+    NT_ASSERT(SectionHandle != NULL);
+    NT_ASSERT(Section != NULL);
+    NT_ASSERT(Path != NULL);
+    NT_ASSERT(Size != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     HANDLE handle = NULL;
     OBJECT_ATTRIBUTES oa = {0};
@@ -464,6 +490,8 @@ RetrieveInMemoryModuleExecutableSections(_Inout_ PIRP Irp)
 {
     PAGED_CODE();
 
+    NT_ASSERT(Irp != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     SIZE_T bytes_written = NULL;
     PVOID buffer = NULL;
@@ -532,11 +560,17 @@ GetNextSMBIOSStructureInTable(_Inout_ PSMBIOS_TABLE_HEADER* CurrentStructure)
 {
     PAGED_CODE();
 
-    PCHAR string_section_start =
+    NT_ASSERT(CurrentStructure != NULL);
+
+    PCHAR string_section_start = NULL;
+    PCHAR current_char_in_strings = NULL;
+    PCHAR next_char_in_strings = NULL;
+
+    string_section_start =
         (PCHAR)((UINT64)*CurrentStructure + (*CurrentStructure)->Length);
 
-    PCHAR current_char_in_strings = string_section_start;
+    current_char_in_strings = string_section_start;
-    PCHAR next_char_in_strings = string_section_start + 1;
+    next_char_in_strings = string_section_start + 1;
 
     for (;;) {
         if (*current_char_in_strings == NULL_TERMINATOR &&
@@ -574,10 +608,14 @@ GetStringAtIndexFromSMBIOSTable(
 {
     PAGED_CODE();
 
+    NT_ASSERT(Table != NULL);
+    NT_ASSERT(Buffer != NULL);
+
     UINT32 current_string_char_index = 0;
     UINT32 string_count = 0;
     PCHAR current_string_char = (PCHAR)((UINT64)Table + Table->Length);
     PCHAR next_string_char = current_string_char + 1;
+    UINT64 dest = 0;
 
     for (;;) {
         if (*current_string_char == NULL_TERMINATOR &&
@@ -591,7 +629,7 @@ GetStringAtIndexFromSMBIOSTable(
             if (*current_string_char == NULL_TERMINATOR)
                 return STATUS_SUCCESS;
 
-            UINT64 dest = (UINT64)Buffer + current_string_char_index;
+            dest = (UINT64)Buffer + current_string_char_index;
 
             IntCopyMemory(dest, current_string_char, sizeof(CHAR));
             current_string_char_index++;
@@ -636,6 +674,8 @@ ParseSMBIOSTable(
 {
     PAGED_CODE();
 
+    NT_ASSERT(Buffer != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PVOID buffer = NULL;
     ULONG buffer_size = 0;
@@ -727,6 +767,13 @@ ComputeHashOfSections(
     _Out_ PVOID* MemoryHash,
     _Out_ PULONG MemoryHashSize)
 {
+    NT_ASSERT(DiskSection != NULL);
+    NT_ASSERT(MemorySection != NULL);
+    NT_ASSERT(DiskHash != NULL);
+    NT_ASSERT(DiskHashSize != NULL);
+    NT_ASSERT(MemoryHash != NULL);
+    NT_ASSERT(MemoryHashSize != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
 
     if (DiskSection->SizeOfRawData != MemorySection->SizeOfRawData) {
@@ -830,6 +877,8 @@ ValidateProcessLoadedModule(_Inout_ PIRP Irp)
 {
     PAGED_CODE();
 
+    NT_ASSERT(Irp != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PROCESS_MODULE_VALIDATION_RESULT validation_result = {0};
     PPROCESS_MODULE_INFORMATION module_info = NULL;
@@ -954,6 +1003,9 @@ HashUserModule(
 {
     PAGED_CODE();
 
+    NT_ASSERT(Entry != NULL);
+    NT_ASSERT(OutBuffer != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     KAPC_STATE apc_state = {0};
     PVAL_INTEGRITY_HEADER memory_buffer = NULL;
@@ -1050,6 +1102,8 @@ GetHardDiskDriveSerialNumber(
 {
     PAGED_CODE();
 
+    NT_ASSERT(ConfigDrive0Serial != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     HANDLE handle = NULL;
     OBJECT_ATTRIBUTES attributes = {0};
@@ -1171,11 +1225,14 @@ ScanForSignature(
 {
     PAGED_CODE();
 
+    NT_ASSERT(BaseAddress != NULL);
+    NT_ASSERT(Signature != NULL);
+
     CHAR current_char = 0;
     CHAR current_sig_char = 0;
 
-    for (INT index = 0; index < MaxLength; index++) {
+    for (UINT32 index = 0; index < MaxLength; index++) {
-        for (INT sig = 0; sig < SignatureLength + 1; sig++) {
+        for (UINT32 sig = 0; sig < SignatureLength + 1; sig++) {
             current_char = *(PCHAR)((UINT64)BaseAddress + index + sig);
             current_sig_char = Signature[sig];
 
@@ -1199,6 +1256,7 @@ STATIC
 UINT64
 MeasureInstructionRead(_In_ PVOID InstructionAddress)
 {
+    NT_ASSERT(InstructionAddress != NULL);
     CONST UINT64 start = __readmsr(IA32_APERF_MSR) << 32;
     CHAR value = *(PCHAR)InstructionAddress;
     return (__readmsr(IA32_APERF_MSR) << 32) - start;
@@ -1210,6 +1268,9 @@ STATIC
 UINT64
 MeasureReads(_In_ PVOID Address, _In_ ULONG Count)
 {
+    NT_ASSERT(Address != NULL);
+    NT_ASSERT(Count > 0);
+
     UINT64 read_average = 0;
     KIRQL irql = {0};
 
@@ -1253,6 +1314,9 @@ NTSTATUS
 GetAverageReadTimeAtRoutine(
     _In_ PVOID RoutineAddress, _Out_ PUINT64 AverageTime)
 {
+    NT_ASSERT(RoutineAddress != NULL);
+    NT_ASSERT(AverageTime != NULL);
+
     if (!RoutineAddress || !AverageTime)
         return STATUS_UNSUCCESSFUL;
 
@@ -1324,6 +1388,7 @@ InitiateEptFunctionAddressArrays()
         CONTROL_FUNCTION_ADDRESSES[index] =
             ImpMmGetSystemRoutineAddress(&current_function);
 
+        NT_ASSERT(CONTROL_FUNCTION_ADDRESSES[index] != NULL);
         if (!CONTROL_FUNCTION_ADDRESSES[index])
             return STATUS_UNSUCCESSFUL;
     }
@@ -1333,6 +1398,7 @@ InitiateEptFunctionAddressArrays()
         PROTECTED_FUNCTION_ADDRESSES[index] =
             ImpMmGetSystemRoutineAddress(&current_function);
 
+        NT_ASSERT(PROTECTED_FUNCTION_ADDRESSES[index] != NULL);
         if (!PROTECTED_FUNCTION_ADDRESSES[index])
             return STATUS_UNSUCCESSFUL;
     }
@@ -1464,6 +1530,9 @@ DetectEptHooksInKeyFunctions()
 VOID
 FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Node, _In_opt_ PVOID Context)
 {
+    NT_ASSERT(Node != NULL);
+    NT_ASSERT(Context != NULL);
+
     LPCSTR process_name = NULL;
     PEPROCESS* process = (PEPROCESS*)Context;
 
@@ -1483,6 +1552,10 @@ StoreModuleExecutableRegionsx86(
     _In_ PVOID* Buffer,
     _In_ PULONG BufferSize)
 {
+    NT_ASSERT(Module != NULL);
+    NT_ASSERT(Buffer != NULL);
+    NT_ASSERT(BufferSize != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PEPROCESS process = NULL;
     KAPC_STATE apc_state = {0};
@@ -1570,6 +1643,9 @@ end:
 NTSTATUS
 HashModule(_In_ PRTL_MODULE_EXTENDED_INFO Module, _Out_ PVOID Hash)
 {
+    NT_ASSERT(Module != NULL);
+    NT_ASSERT(Hash != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     ANSI_STRING ansi_string = {0};
     UNICODE_STRING path = {0};
@@ -1675,6 +1751,8 @@ STATIC
 VOID
 ReportModifiedSystemImage(_In_ PRTL_MODULE_EXTENDED_INFO Module)
 {
+    NT_ASSERT(Module != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     UINT32 len = 0;
     PSYSTEM_MODULE_INTEGRITY_CHECK_REPORT report = NULL;
@@ -1711,6 +1789,8 @@ ReportModifiedSystemImage(_In_ PRTL_MODULE_EXTENDED_INFO Module)
 VOID
 ValidateSystemModule(_In_ PRTL_MODULE_EXTENDED_INFO Module)
 {
+    NT_ASSERT(Module != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     PDRIVER_LIST_ENTRY entry = NULL;
     PVOID hash = NULL;
@@ -1774,6 +1854,8 @@ STATIC
 VOID
 ReportModifiedSelfDriverImage(_In_ PRTL_MODULE_EXTENDED_INFO Module)
 {
+    NT_ASSERT(Module != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     UINT32 len = 0;
     PDRIVER_SELF_INTEGRITY_CHECK_REPORT packet = NULL;
@@ -1942,11 +2024,12 @@ SystemModuleVerificationDispatchFunction(
 {
     UNREFERENCED_PARAMETER(DeviceObject);
 
+    NT_ASSERT(Context != NULL);
+
     UINT32 count = 0;
     UINT32 max = 0;
 
     IncrementActiveThreadCount(Context);
-
     count = GetCurrentVerificationIndex(Context);
 
     /*
@@ -1959,6 +2042,12 @@ SystemModuleVerificationDispatchFunction(
     max = GetCurrentVerificationMaxIndex(Context, count);
 
     for (; count < max && count < Context->total_count; count++) {
+        DEBUG_VERBOSE(
+            "ThrId: %lx, Count: %lx, Max: %lx, Total Count: %lx",
+            PsGetCurrentThreadId(),
+            count,
+            max,
+            Context->total_count);
         if (!InterlockedCompareExchange(
                 &Context->dispatcher_info[count].validated,
                 TRUE,
@@ -2011,6 +2100,8 @@ STATIC
 NTSTATUS
 InitialiseSystemModuleVerificationContext(PSYS_MODULE_VAL_CONTEXT Context)
 {
+    NT_ASSERT(Context != NULL);
+
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     SYSTEM_MODULES modules = {0};
     PMODULE_DISPATCHER_HEADER dispatcher = NULL;
@@ -2042,7 +2133,9 @@ InitialiseSystemModuleVerificationContext(PSYS_MODULE_VAL_CONTEXT Context)
 VOID
 FreeWorkItems(_In_ PSYS_MODULE_VAL_CONTEXT Context)
 {
-    for (INT index = 0; index < VERIFICATION_THREAD_COUNT; index++) {
+    NT_ASSERT(Context != NULL);
+
+    for (UINT32 index = 0; index < VERIFICATION_THREAD_COUNT; index++) {
         if (Context->work_items[index]) {
             ImpIoFreeWorkItem(Context->work_items[index]);
             Context->work_items[index] = NULL;
@@ -2054,6 +2147,8 @@ STATIC
 VOID
 FreeModuleVerificationItems(_In_ PSYS_MODULE_VAL_CONTEXT Context)
 {
+    NT_ASSERT(Context != NULL);
+
     /* if a thread hasnt completed by this point, something catastrophic has
      * gone wrong and maybe its better not to yield..*/
     while (Context->active_thread_count)
@@ -2171,40 +2266,6 @@ GetOsVersionInformation(_Out_ PRTL_OSVERSIONINFOW VersionInfo)
     return status;
 }
 
-#define KPCR_KPRCB_OFFSET        0x180
-#define KPCRB_IDLE_THREAD_OFFSET 0x018
-#define KTHREAD_IDLE_TIME_OFFSET 0x28c
-#define KPCRB_KERNEL_TIME_OFFSET 0x7e84
-#define KPCRB_USER_TIME_OFFSET   0x7e88
-
-UINT32
-CalculateCpuCoreUsage(_In_ UINT32 Core)
-{
-    PVOID kpcr = NULL;
-    PVOID kpcrb = NULL;
-    PVOID idle_thread = NULL;
-    UINT32 idle_time = 0;
-    UINT32 kernel_time = 0;
-    UINT32 user_time = 0;
-
-    KeSetSystemAffinityThread(1ull << Core);
-
-    while (Core != KeGetCurrentProcessorNumber())
-        YieldProcessor();
-
-    kpcr = __readmsr(IA32_GS_BASE);
-    kpcrb = (UINT64)kpcr + KPCR_KPRCB_OFFSET;
-    idle_thread = *(UINT64*)((UINT64)kpcrb + KPCRB_IDLE_THREAD_OFFSET);
-
-    idle_time = *(UINT32*)((UINT64)idle_thread + KTHREAD_IDLE_TIME_OFFSET);
-    kernel_time = *(UINT32*)((UINT64)kpcrb + KPCRB_KERNEL_TIME_OFFSET);
-    user_time = *(UINT32*)((UINT64)kpcrb + KPCRB_USER_TIME_OFFSET);
-
-    return (
-        100 - (UINT32)(UInt32x32To64(idle_time, 100) /
-                       (UINT64)(kernel_time + user_time)));
-}
-
 BOOLEAN
 ValidateOurDriversDispatchRoutines()
 {
@@ -2369,6 +2430,8 @@ HeartbeatWorkItem(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context)
 {
     UNREFERENCED_PARAMETER(DeviceObject);
 
+    NT_ASSERT(Context != NULL);
+
     if (!ARGUMENT_PRESENT(Context))
         return;
 

driver/io.c

@@ -1,20 +1,17 @@
 #include "io.h"
 
 #include "callbacks.h"
+#include "containers/map.h"
 #include "driver.h"
+#include "hv.h"
+#include "hw.h"
+#include "imports.h"
 #include "integrity.h"
+#include "lib/stdlib.h"
 #include "modules.h"
 #include "pool.h"
-#include "thread.h"
-
-#include "hv.h"
-#include "imports.h"
-
-#include "containers/map.h"
-#include "hw.h"
 #include "session.h"
-
+#include "thread.h"
-#include "lib/stdlib.h"
 
 STATIC
 NTSTATUS
@@ -106,6 +103,9 @@ PIRP
 IrpQueuePeekNextEntry(_In_ PIO_CSQ Csq, _In_ PIRP Irp, _In_ PVOID Context)
 {
     UNREFERENCED_PARAMETER(Context);
+
+    NT_ASSERT(Irp != NULL);
+
     PIRP_QUEUE_HEAD queue = GetIrpQueueHead();
 
     if (queue->irp_count == 0)
@@ -171,6 +171,8 @@ STATIC
 NTSTATUS
 IrpQueueCompleteDeferredPacket(_In_ PDEFERRED_REPORT Report, _In_ PIRP Irp)
 {
+    NT_ASSERT(Report != NULL);
+
     NTSTATUS status = ValidateIrpOutputBuffer(Irp, Report->buffer_size);
     PIRP_QUEUE_HEAD queue = GetIrpQueueHead();
     UINT16 type = GetPacketType(Report->buffer);
@@ -196,6 +198,8 @@ STATIC
 NTSTATUS
 IrpQueueQueryPendingPackets(_In_ PIRP Irp)
 {
+    NT_ASSERT(Irp != NULL);
+
     PIRP_QUEUE_HEAD queue = GetIrpQueueHead();
     PDEFERRED_REPORT report = NULL;
     NTSTATUS status = STATUS_UNSUCCESSFUL;
@@ -253,6 +257,9 @@ STATIC
 PDEFERRED_REPORT
 IrpQueueAllocateDeferredPacket(_In_ PVOID Buffer, _In_ UINT32 BufferSize)
 {
+    NT_ASSERT(Buffer != NULL);
+    NT_ASSERT(BufferSize != 0);
+
     PDEFERRED_REPORT report = ImpExAllocatePool2(
         POOL_FLAG_NON_PAGED,
         sizeof(DEFERRED_REPORT),
@@ -273,6 +280,9 @@ VOID
 IrpQueueDeferPacket(
     _In_ PIRP_QUEUE_HEAD Queue, _In_ PVOID Buffer, _In_ UINT32 BufferSize)
 {
+    NT_ASSERT(Queue != NULL);
+    NT_ASSERT(Buffer != NULL);
+
     PDEFERRED_REPORT report = NULL;
     /*
      * arbitrary number, if we ever do have 100 deferred reports, theres

driver/modules.c

@@ -53,7 +53,7 @@ CHAR WHITELISTED_MODULES[WHITELISTED_MODULE_COUNT][MODULE_MAX_STRING_SIZE] = {
 
 typedef struct _WHITELISTED_REGIONS {
     UINT64 base;
-    UINT64 end;
+    UINT64 size;
 
 } WHITELISTED_REGIONS, *PWHITELISTED_REGIONS;
 
@@ -201,7 +201,7 @@ PopulateWhitelistedModuleBuffer(
 
         region = &Whitelist[index];
         region->base = (UINT64)module->ImageBase;
-        region->end = (UINT64)module->ImageBase + module->ImageSize;
+        region->size = (UINT64)module->ImageBase + module->ImageSize;
     }
 }
 
@@ -267,7 +267,7 @@ DoesDriverHaveInvalidDispatchRoutine(
          */
         for (UINT32 index = 0; index < WHITELISTED_MODULE_COUNT; index++) {
             if (dispatch_function >= Regions[index].base &&
-                dispatch_function <= Regions[index].end)
+                dispatch_function <= Regions[index].size)
                 return FALSE;
         }