This commit is contained in:
donnaskiez 2024-05-06 00:07:33 +10:00
parent e1a5a78db0
commit 08beb2e036

View file

@ -547,7 +547,7 @@ IsInstructionPointerInInvalidRegion(_In_ UINT64 RIP,
BOOLEAN BOOLEAN
IsInstructionPointerInsideSpecifiedModule(_In_ UINT64 Rip, IsInstructionPointerInsideSpecifiedModule(_In_ UINT64 Rip,
_In_ PRTL_MODULE_EXTENDED_INFO Module) _In_ PRTL_MODULE_EXTENDED_INFO Module)
{ {
UINT64 base = (UINT64)Module->ImageBase; UINT64 base = (UINT64)Module->ImageBase;
UINT64 end = base + Module->ImageSize; UINT64 end = base + Module->ImageSize;
@ -722,8 +722,8 @@ NmiCallback(_Inout_opt_ PVOID Context, _In_ BOOLEAN Handled)
{ {
UNREFERENCED_PARAMETER(Handled); UNREFERENCED_PARAMETER(Handled);
PNMI_CONTEXT context = (PNMI_CONTEXT)Context;
ULONG core = KeGetCurrentProcessorNumber(); ULONG core = KeGetCurrentProcessorNumber();
PNMI_CONTEXT context = &((PNMI_CONTEXT)Context)[core];
UINT64 kpcr = 0; UINT64 kpcr = 0;
TASK_STATE_SEGMENT_64* tss = NULL; TASK_STATE_SEGMENT_64* tss = NULL;
PMACHINE_FRAME machine_frame = NULL; PMACHINE_FRAME machine_frame = NULL;
@ -745,12 +745,12 @@ NmiCallback(_Inout_opt_ PVOID Context, _In_ BOOLEAN Handled)
machine_frame = GetIsrMachineFrame(tss); machine_frame = GetIsrMachineFrame(tss);
if (IsUserModeAddress(machine_frame->rip)) if (IsUserModeAddress(machine_frame->rip))
context[core].user_thread = TRUE; context->user_thread = TRUE;
context[core].interrupted_rip = machine_frame->rip; context->interrupted_rip = machine_frame->rip;
context[core].interrupted_rsp = machine_frame->rsp; context->interrupted_rsp = machine_frame->rsp;
context[core].kthread = PsGetCurrentThread(); context->kthread = PsGetCurrentThread();
context[core].callback_count++; context->callback_count++;
DEBUG_VERBOSE( DEBUG_VERBOSE(
"[NMI CALLBACK]: Core Number: %lx, Interrupted RIP: %llx, Interrupted RSP: %llx", "[NMI CALLBACK]: Core Number: %lx, Interrupted RIP: %llx, Interrupted RSP: %llx",
@ -802,9 +802,9 @@ HandleNmiIOCTL()
PAGED_CODE(); PAGED_CODE();
NTSTATUS status = STATUS_UNSUCCESSFUL; NTSTATUS status = STATUS_UNSUCCESSFUL;
PVOID callback_handle = NULL; PVOID handle = NULL;
SYSTEM_MODULES system_modules = {0}; SYSTEM_MODULES modules = {0};
PNMI_CONTEXT nmi_context = NULL; PNMI_CONTEXT context = NULL;
UINT32 size = ImpKeQueryActiveProcessorCount(0) * sizeof(NMI_CONTEXT); UINT32 size = ImpKeQueryActiveProcessorCount(0) * sizeof(NMI_CONTEXT);
@ -817,10 +817,10 @@ HandleNmiIOCTL()
if (!NT_SUCCESS(status)) if (!NT_SUCCESS(status))
DEBUG_ERROR("ValidateHalDispatchTables failed with status %x", status); DEBUG_ERROR("ValidateHalDispatchTables failed with status %x", status);
nmi_context = context =
ImpExAllocatePool2(POOL_FLAG_NON_PAGED, size, NMI_CONTEXT_POOL); ImpExAllocatePool2(POOL_FLAG_NON_PAGED, size, NMI_CONTEXT_POOL);
if (!nmi_context) { if (!context) {
UnsetNmiInProgressFlag(); UnsetNmiInProgressFlag();
return STATUS_MEMORY_NOT_ALLOCATED; return STATUS_MEMORY_NOT_ALLOCATED;
} }
@ -830,48 +830,47 @@ HandleNmiIOCTL()
* becomes harder for people to hook our callback and get up to some * becomes harder for people to hook our callback and get up to some
* funny business * funny business
*/ */
callback_handle = ImpKeRegisterNmiCallback(NmiCallback, nmi_context); handle = ImpKeRegisterNmiCallback(NmiCallback, context);
if (!callback_handle) { if (!handle) {
DEBUG_ERROR("KeRegisterNmiCallback failed with no status."); DEBUG_ERROR("KeRegisterNmiCallback failed with no status.");
ImpExFreePoolWithTag(nmi_context, NMI_CONTEXT_POOL); goto end;
UnsetNmiInProgressFlag();
return STATUS_UNSUCCESSFUL;
} }
/* /*
* We query the system modules each time since they can potentially * We query the system modules each time since they can potentially
* change at any time * change at any time
*/ */
status = GetSystemModuleInformation(&system_modules); status = GetSystemModuleInformation(&modules);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
ImpKeDeregisterNmiCallback(callback_handle);
ImpExFreePoolWithTag(nmi_context, NMI_CONTEXT_POOL);
DEBUG_ERROR("Error retriving system module information"); DEBUG_ERROR("Error retriving system module information");
UnsetNmiInProgressFlag(); goto end;
return status;
} }
status = LaunchNonMaskableInterrupt(); status = LaunchNonMaskableInterrupt();
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
DEBUG_ERROR("Error running NMI callbacks"); DEBUG_ERROR("Error running NMI callbacks");
ImpKeDeregisterNmiCallback(callback_handle); goto end;
ImpExFreePoolWithTag(system_modules.address, SYSTEM_MODULES_POOL);
ImpExFreePoolWithTag(nmi_context, NMI_CONTEXT_POOL);
UnsetNmiInProgressFlag();
return status;
} }
status = AnalyseNmiData(nmi_context, &system_modules); status = AnalyseNmiData(context, &modules);
if (!NT_SUCCESS(status)) if (!NT_SUCCESS(status))
DEBUG_ERROR("Error analysing nmi data"); DEBUG_ERROR("Error analysing nmi data");
ImpExFreePoolWithTag(system_modules.address, SYSTEM_MODULES_POOL); end:
ImpExFreePoolWithTag(nmi_context, NMI_CONTEXT_POOL);
ImpKeDeregisterNmiCallback(callback_handle); if (modules.address)
ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL);
if (context)
ImpExFreePoolWithTag(context, NMI_CONTEXT_POOL);
if (handle)
ImpKeDeregisterNmiCallback(handle);
UnsetNmiInProgressFlag(); UnsetNmiInProgressFlag();
return status; return status;
} }