diff --git a/driver/apc.c b/driver/apc.c index 6c00065..397527b 100644 --- a/driver/apc.c +++ b/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++) { - PAPC_CONTEXT_HEADER header = GetApcContextArray()[index]; + for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; 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 = GetApcContextArray(); + PUINT64 entry = NULL; + + 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++) { - PAPC_CONTEXT_HEADER entry = NULL; + for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { GetApcContextByIndex(&entry, index); if (!entry) @@ -137,8 +149,7 @@ QueryActiveApcContextsForCompletion() switch (entry->context_id) { case APC_CONTEXT_ID_STACKWALK: - FreeApcStackwalkApcContextInformation( - (PAPC_STACKWALK_CONTEXT)entry); + FreeApcStackwalkApcContextInformation(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++) { - PUINT64 entry = GetApcContextArray(); + for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { + 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++) { - PUINT64 entry = GetApcContextArray(); + for (UINT32 index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { + 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; } diff --git a/driver/apc.h b/driver/apc.h index 680a39f..e3da44a 100644 --- a/driver/apc.h +++ b/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(); diff --git a/driver/callbacks.c b/driver/callbacks.c index 9e881b9..af36e00 100644 --- a/driver/callbacks.c +++ b/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(&list->active, FALSE); + InterlockedExchange(&GetDriverList()->active, FALSE); PsRemoveLoadImageNotifyRoutine(ImageLoadNotifyRoutineCallback); } VOID UnregisterThreadCreateNotifyRoutine() { - PRB_TREE tree = GetThreadTree(); - InterlockedExchange(&tree->active, FALSE); + InterlockedExchange(&GetThreadTree()->active, FALSE); ImpPsRemoveCreateThreadNotifyRoutine(ThreadCreateNotifyRoutine); } VOID CleanupThreadListOnDriverUnload() { - PRB_TREE tree = GetThreadTree(); - DEBUG_VERBOSE("Freeing thread list!"); - RtlRbTreeEnumerate(tree, CleanupThreadListFreeCallback, NULL); - RtlRbTreeDeleteTree(tree); + RtlRbTreeEnumerate(GetThreadTree(), CleanupThreadListFreeCallback, NULL); + RtlRbTreeDeleteTree(GetThreadTree()); } 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; - PDRIVER_LIST_ENTRY driver_entry = NULL; + PLIST_ENTRY entry = NULL; + PDRIVER_LIST_ENTRY driver = NULL; ImpKeAcquireGuardedMutex(&head->lock); if (CallbackRoutine) { - list_entry = head->list_entry.Flink; - while (list_entry != &head->list_entry) { - driver_entry = - CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry); - CallbackRoutine(driver_entry, Context); - list_entry = list_entry->Flink; + entry = head->list_entry.Flink; + while (entry != &head->list_entry) { + driver = CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, list_entry); + CallbackRoutine(driver, Context); + entry = 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(); - PLIST_ENTRY list_entry = NULL; - PDRIVER_LIST_ENTRY driver_entry = NULL; + NT_ASSERT(ImageBase != NULL); + NT_ASSERT(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) { - driver_entry = - CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry); + while (entry != &head->list_entry) { + driver = CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, list_entry); - if (driver_entry->ImageBase == ImageBase) { - *Entry = driver_entry; + if (driver->ImageBase == ImageBase) { + *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; diff --git a/driver/crypt.c b/driver/crypt.c index 2228497..d3e79fe 100644 --- a/driver/crypt.c +++ b/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 "lib/stdlib.h" +#include "util.h" #include #include diff --git a/driver/driver.c b/driver/driver.c index 452b57f..fd74ce2 100644 --- a/driver/driver.c +++ b/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 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; diff --git a/driver/hv.c b/driver/hv.c index ecafbcd..7963cf9 100644 --- a/driver/hv.c +++ b/driver/hv.c @@ -3,10 +3,10 @@ #include "common.h" #include "imports.h" #include "io.h" -#include - #include "lib/stdlib.h" +#include + #ifdef ALLOC_PRAGMA # pragma alloc_text(PAGE, PerformVirtualizationDetection) #endif diff --git a/driver/hw.c b/driver/hw.c index 3b143d2..558de82 100644 --- a/driver/hw.c +++ b/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( - sizeof(BLACKLISTED_PCIE_DEVICE_REPORT)); + UINT32 len = 0; + PBLACKLISTED_PCIE_DEVICE_REPORT report = NULL; - PBLACKLISTED_PCIE_DEVICE_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength( + 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; } diff --git a/driver/integrity.c b/driver/integrity.c index bf0bcc9..5454733 100644 --- a/driver/integrity.c +++ b/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 #include #include -/* 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, - sizeof(ULONG)); + &driver->ImageSize, + 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; - PCHAR next_char_in_strings = string_section_start + 1; + current_char_in_strings = string_section_start; + 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 (INT sig = 0; sig < SignatureLength + 1; sig++) { + for (UINT32 index = 0; index < MaxLength; index++) { + 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(¤t_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(¤t_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; diff --git a/driver/io.c b/driver/io.c index c1a9d2a..024550a 100644 --- a/driver/io.c +++ b/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 "lib/stdlib.h" +#include "thread.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 diff --git a/driver/modules.c b/driver/modules.c index 1311cb7..58f58c6 100644 --- a/driver/modules.c +++ b/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; }