diff --git a/driver/callbacks.c b/driver/callbacks.c index 5ba7bc6..98c3817 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -198,10 +198,10 @@ InitialiseDriverList() ListInit(&list->start, &list->lock); InitializeListHead(&list->deferred_list); - list->can_hash_x86 = FALSE; - list->deferred_work_item = IoAllocateWorkItem(GetDriverDeviceObject()); + list->can_hash_x86 = FALSE; + list->work_item = IoAllocateWorkItem(GetDriverDeviceObject()); - if (!list->deferred_work_item) + if (!list->work_item) return STATUS_INSUFFICIENT_RESOURCES; status = GetSystemModuleInformation(&modules); @@ -454,6 +454,15 @@ unlock: ImpKeReleaseGuardedMutex(&list->lock); } +FORCEINLINE +STATIC +BOOLEAN +CanInitiateDeferredHashing(_In_ LPCSTR ProcessName, _In_ PDRIVER_LIST_HEAD Head) +{ + return !strcmp(ProcessName, "winlogon.exe") && Head->work_item ? TRUE + : FALSE; +} + VOID ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, _In_ HANDLE ProcessId, @@ -477,6 +486,8 @@ ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, process_name = ImpPsGetProcessImageFileName(process); + DEBUG_INFO("process create notify: %s", process_name); + if (Create) { entry = ExAllocateFromLookasideListEx(&list->lookaside_list); @@ -495,10 +506,8 @@ ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, * Notify to our driver that we can hash x86 modules, and hash * any x86 modules that werent hashed. */ - if (!strcmp(process_name, "winlogon.exe") && - !driver_list->deferred_complete) { - driver_list->can_hash_x86 = TRUE; - IoQueueWorkItem(driver_list->deferred_work_item, + if (CanInitiateDeferredHashing(process_name, driver_list)) { + IoQueueWorkItem(driver_list->work_item, DeferredModuleHashingCallback, NormalWorkQueue, NULL); diff --git a/driver/common.h b/driver/common.h index 330b95d..14ac802 100644 --- a/driver/common.h +++ b/driver/common.h @@ -80,7 +80,7 @@ typedef struct _DRIVER_LIST_HEAD { KGUARDED_MUTEX lock; /* modules that need to be hashed later. */ - PIO_WORKITEM deferred_work_item; + PIO_WORKITEM work_item; LIST_ENTRY deferred_list; volatile BOOLEAN deferred_complete; volatile LONG can_hash_x86; diff --git a/driver/integrity.c b/driver/integrity.c index d5d5027..f43204d 100644 --- a/driver/integrity.c +++ b/driver/integrity.c @@ -1441,6 +1441,14 @@ StoreModuleExecutableRegionsx86(_In_ PRTL_MODULE_EXTENDED_INFO Module, return status; } +FORCEINLINE +STATIC +VOID +Enablex86Hashing(_In_ PDRIVER_LIST_HEAD Head) +{ + Head->can_hash_x86 = TRUE; +} + VOID DeferredModuleHashingCallback() { @@ -1451,7 +1459,8 @@ DeferredModuleHashingCallback() PLIST_ENTRY list_entry = NULL; PDRIVER_LIST_ENTRY entry = NULL; - driver_list->deferred_complete = TRUE; + Enablex86Hashing(driver_list); + list_entry = RemoveHeadList(deferred_head); if (list_entry == deferred_head) @@ -1480,7 +1489,8 @@ DeferredModuleHashingCallback() end: DEBUG_VERBOSE("All deferred modules hashed."); - ImpIoFreeWorkItem(driver_list->deferred_work_item); + ImpIoFreeWorkItem(driver_list->work_item); + driver_list->work_item = NULL; } NTSTATUS diff --git a/driver/modules.c b/driver/modules.c index 673f7ca..e5af0ac 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -769,20 +769,20 @@ LaunchNonMaskableInterrupt() { PAGED_CODE(); - PKAFFINITY_EX ProcAffinityPool = ImpExAllocatePool2( + PKAFFINITY_EX affinity = ImpExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof(KAFFINITY_EX), PROC_AFFINITY_POOL); - if (!ProcAffinityPool) + if (!affinity) return STATUS_MEMORY_NOT_ALLOCATED; LARGE_INTEGER delay = {0}; delay.QuadPart -= NMI_DELAY_TIME; for (ULONG core = 0; core < ImpKeQueryActiveProcessorCount(0); core++) { - ImpKeInitializeAffinityEx(ProcAffinityPool); - ImpKeAddProcessorAffinityEx(ProcAffinityPool, core); + ImpKeInitializeAffinityEx(affinity); + ImpKeAddProcessorAffinityEx(affinity, core); - HalSendNMI(ProcAffinityPool); + HalSendNMI(affinity); /* * Only a single NMI can be active at any given time, so @@ -792,7 +792,7 @@ LaunchNonMaskableInterrupt() ImpKeDelayExecutionThread(KernelMode, FALSE, &delay); } - ImpExFreePoolWithTag(ProcAffinityPool, PROC_AFFINITY_POOL); + ImpExFreePoolWithTag(affinity, PROC_AFFINITY_POOL); return STATUS_SUCCESS; } @@ -801,10 +801,10 @@ HandleNmiIOCTL() { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - PVOID handle = NULL; - SYSTEM_MODULES modules = {0}; - PNMI_CONTEXT context = NULL; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PVOID handle = NULL; + SYSTEM_MODULES modules = {0}; + PNMI_CONTEXT context = NULL; UINT32 size = ImpKeQueryActiveProcessorCount(0) * sizeof(NMI_CONTEXT); @@ -817,8 +817,7 @@ HandleNmiIOCTL() if (!NT_SUCCESS(status)) DEBUG_ERROR("ValidateHalDispatchTables failed with status %x", status); - context = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, size, NMI_CONTEXT_POOL); + context = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, size, NMI_CONTEXT_POOL); if (!context) { UnsetNmiInProgressFlag();