mirror of
https://github.com/donnaskiez/ac.git
synced 2024-11-21 22:24:08 +01:00
299 lines
No EOL
8.5 KiB
C
299 lines
No EOL
8.5 KiB
C
#include "nmi.h"
|
|
|
|
#include "modules.h"
|
|
#include "common.h"
|
|
|
|
#define NMI_DELAY 100 * 10000
|
|
|
|
typedef struct _NMI_POOLS
|
|
{
|
|
PVOID thread_data_pool;
|
|
PVOID stack_frames;
|
|
PVOID nmi_context;
|
|
|
|
}NMI_POOLS, * PNMI_POOLS;
|
|
|
|
PVOID nmi_callback_handle = NULL;
|
|
|
|
/* Global structure to hold pointers to required memory for the NMI's */
|
|
NMI_POOLS nmi_pools = { 0 };
|
|
|
|
NTSTATUS IsInstructionPointerInInvalidRegion(
|
|
_In_ UINT64 RIP,
|
|
_In_ PSYSTEM_MODULES SystemModules,
|
|
_Out_ PBOOLEAN Result
|
|
)
|
|
{
|
|
if ( !RIP || !SystemModules || !Result )
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
/* Note that this does not check for HAL or PatchGuard Execution */
|
|
for ( INT i = 0; i < SystemModules->module_count; i++ )
|
|
{
|
|
PRTL_MODULE_EXTENDED_INFO system_module = ( PRTL_MODULE_EXTENDED_INFO )(
|
|
( uintptr_t )SystemModules->address + i * sizeof( RTL_MODULE_EXTENDED_INFO ) );
|
|
|
|
UINT64 base = ( UINT64 )system_module->ImageBase;
|
|
UINT64 end = base + system_module->ImageSize;
|
|
|
|
if ( RIP >= base && RIP <= end )
|
|
{
|
|
*Result = TRUE;
|
|
return STATUS_SUCCESS;;
|
|
}
|
|
}
|
|
|
|
*Result = FALSE;
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
NTSTATUS AnalyseNmiData(
|
|
_In_ INT NumCores,
|
|
_In_ PSYSTEM_MODULES SystemModules,
|
|
_In_ PIRP Irp
|
|
)
|
|
{
|
|
if ( !NumCores || !SystemModules )
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
for ( INT core = 0; core < NumCores; core++ )
|
|
{
|
|
PNMI_CONTEXT context = ( PNMI_CONTEXT )( ( uintptr_t )nmi_pools.nmi_context + core * sizeof( NMI_CONTEXT ) );
|
|
|
|
/* Make sure our NMIs were run */
|
|
if ( !context->nmi_callbacks_run )
|
|
{
|
|
NMI_CALLBACK_FAILURE report;
|
|
report.report_code = REPORT_NMI_CALLBACK_FAILURE;
|
|
report.kthread_address = NULL;
|
|
report.invalid_rip = NULL;
|
|
report.were_nmis_disabled = TRUE;
|
|
|
|
Irp->IoStatus.Information = sizeof( NMI_CALLBACK_FAILURE );
|
|
|
|
RtlCopyMemory(
|
|
Irp->AssociatedIrp.SystemBuffer,
|
|
&report,
|
|
sizeof( NMI_CALLBACK_FAILURE )
|
|
);
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
PNMI_CALLBACK_DATA thread_data = ( PNMI_CALLBACK_DATA )(
|
|
( uintptr_t )nmi_pools.thread_data_pool + core * sizeof( NMI_CALLBACK_DATA ) );
|
|
|
|
DEBUG_LOG( "cpu number: %i callback count: %i", core, context->nmi_callbacks_run );
|
|
|
|
/* Walk the stack */
|
|
for ( INT frame = 0; frame < thread_data->num_frames_captured; frame++ )
|
|
{
|
|
BOOLEAN flag;
|
|
DWORD64 stack_frame = *( DWORD64* )(
|
|
( ( uintptr_t )nmi_pools.stack_frames + thread_data->stack_frames_offset + frame * sizeof( PVOID ) ) );
|
|
|
|
if ( !NT_SUCCESS( IsInstructionPointerInInvalidRegion( stack_frame, SystemModules, &flag ) ) )
|
|
{
|
|
DEBUG_ERROR( "errro checking RIP for current stack address" );
|
|
continue;
|
|
}
|
|
|
|
if ( flag == FALSE )
|
|
{
|
|
/*
|
|
* Note: for now, we only handle 1 report at a time so we stop the
|
|
* analysis once we receive a report since we only send a buffer
|
|
* large enough for 1 report. In the future this should be changed
|
|
* to a buffer that can hold atleast 4 reports (since the chance we
|
|
* get 4 reports with a single NMI would be impossible) so we can
|
|
* continue parsing the rest of the stack frames after receiving a
|
|
* single report.
|
|
*/
|
|
|
|
NMI_CALLBACK_FAILURE report;
|
|
report.report_code = REPORT_NMI_CALLBACK_FAILURE;
|
|
report.kthread_address = thread_data->kthread_address;
|
|
report.invalid_rip = stack_frame;
|
|
report.were_nmis_disabled = FALSE;
|
|
|
|
Irp->IoStatus.Information = sizeof( NMI_CALLBACK_FAILURE );
|
|
|
|
RtlCopyMemory(
|
|
Irp->AssociatedIrp.SystemBuffer,
|
|
&report,
|
|
sizeof( NMI_CALLBACK_FAILURE )
|
|
);
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
}
|
|
}
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
BOOLEAN NmiCallback(
|
|
_In_ PVOID Context,
|
|
_In_ BOOLEAN Handled
|
|
)
|
|
{
|
|
UNREFERENCED_PARAMETER( Handled );
|
|
|
|
ULONG proc_num = KeGetCurrentProcessorNumber();
|
|
PVOID current_thread = KeGetCurrentThread();
|
|
NMI_CALLBACK_DATA thread_data = { 0 };
|
|
|
|
/*
|
|
* Cannot allocate pool in this function as it runs at IRQL >= dispatch level
|
|
* so ive just allocated a global pool with size equal to 0x200 * num_procs
|
|
*/
|
|
INT num_frames_captured = RtlCaptureStackBackTrace(
|
|
NULL,
|
|
STACK_FRAME_POOL_SIZE,
|
|
( uintptr_t )nmi_pools.stack_frames + proc_num * STACK_FRAME_POOL_SIZE,
|
|
NULL
|
|
);
|
|
|
|
/*
|
|
* This function is run in the context of the interrupted thread hence we can
|
|
* gather any and all information regarding the thread that may be useful for analysis
|
|
*/
|
|
thread_data.kthread_address = ( UINT64 )current_thread;
|
|
thread_data.kprocess_address = ( UINT64 )PsGetCurrentProcess();
|
|
thread_data.stack_base = *( ( UINT64* )( ( uintptr_t )current_thread + KTHREAD_STACK_BASE_OFFSET ) );
|
|
thread_data.stack_limit = *( ( UINT64* )( ( uintptr_t )current_thread + KTHREAD_STACK_LIMIT_OFFSET ) );
|
|
thread_data.start_address = *( ( UINT64* )( ( uintptr_t )current_thread + KTHREAD_START_ADDRESS_OFFSET ) );
|
|
thread_data.cr3 = __readcr3();
|
|
thread_data.stack_frames_offset = proc_num * STACK_FRAME_POOL_SIZE;
|
|
thread_data.num_frames_captured = num_frames_captured;
|
|
|
|
RtlCopyMemory(
|
|
( ( uintptr_t )nmi_pools.thread_data_pool ) + proc_num * sizeof( thread_data ),
|
|
&thread_data,
|
|
sizeof( thread_data )
|
|
);
|
|
|
|
PNMI_CONTEXT context = ( PNMI_CONTEXT )( ( uintptr_t )Context + proc_num * sizeof( NMI_CONTEXT ) );
|
|
context->nmi_callbacks_run += 1;
|
|
DEBUG_LOG( "num nmis called: %i from addr: %llx", context->nmi_callbacks_run, ( uintptr_t )context );
|
|
|
|
return TRUE;
|
|
}
|
|
|
|
NTSTATUS LaunchNonMaskableInterrupt(
|
|
_In_ ULONG NumCores
|
|
)
|
|
{
|
|
if ( !NumCores )
|
|
return STATUS_INVALID_PARAMETER;
|
|
|
|
PKAFFINITY_EX ProcAffinityPool = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( KAFFINITY_EX ), PROC_AFFINITY_POOL );
|
|
|
|
if ( !ProcAffinityPool )
|
|
return STATUS_ABANDONED;
|
|
|
|
nmi_pools.stack_frames = ExAllocatePool2( POOL_FLAG_NON_PAGED, NumCores * STACK_FRAME_POOL_SIZE, STACK_FRAMES_POOL );
|
|
|
|
if ( !nmi_pools.stack_frames )
|
|
{
|
|
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
|
|
return STATUS_ABANDONED;
|
|
}
|
|
|
|
nmi_pools.thread_data_pool = ExAllocatePool2( POOL_FLAG_NON_PAGED, NumCores * sizeof( NMI_CALLBACK_DATA ), THREAD_DATA_POOL );
|
|
|
|
if ( !nmi_pools.thread_data_pool )
|
|
{
|
|
ExFreePoolWithTag( nmi_pools.stack_frames, STACK_FRAMES_POOL );
|
|
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
|
|
return STATUS_ABANDONED;
|
|
}
|
|
|
|
LARGE_INTEGER delay = { 0 };
|
|
delay.QuadPart -= NMI_DELAY;
|
|
|
|
for ( ULONG core = 0; core < NumCores; core++ )
|
|
{
|
|
KeInitializeAffinityEx( ProcAffinityPool );
|
|
KeAddProcessorAffinityEx( ProcAffinityPool, core );
|
|
|
|
DEBUG_LOG( "Sending NMI" );
|
|
HalSendNMI( ProcAffinityPool );
|
|
|
|
/*
|
|
* Only a single NMI can be active at any given time, so arbitrarily
|
|
* delay execution to allow time for the NMI to be processed
|
|
*/
|
|
KeDelayExecutionThread( KernelMode, FALSE, &delay );
|
|
}
|
|
|
|
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
NTSTATUS HandleNmiIOCTL(
|
|
_In_ PIRP Irp
|
|
)
|
|
{
|
|
NTSTATUS status = STATUS_SUCCESS;
|
|
SYSTEM_MODULES system_modules = { 0 };
|
|
ULONG num_cores = KeQueryActiveProcessorCountEx( 0 );
|
|
|
|
/* Fix annoying visual studio linting error */
|
|
RtlZeroMemory( &system_modules, sizeof( SYSTEM_MODULES ) );
|
|
RtlZeroMemory( &nmi_pools, sizeof( NMI_POOLS ) );
|
|
|
|
nmi_pools.nmi_context = ExAllocatePool2( POOL_FLAG_NON_PAGED, num_cores * sizeof( NMI_CONTEXT ), NMI_CONTEXT_POOL );
|
|
|
|
if ( !nmi_pools.nmi_context )
|
|
{
|
|
DEBUG_ERROR( "nmi_context ExAllocatePool2 failed" );
|
|
return STATUS_ABANDONED;
|
|
}
|
|
|
|
/*
|
|
* We want to register and unregister our callback each time so it becomes harder
|
|
* for people to hook our callback and get up to some funny business
|
|
*/
|
|
nmi_callback_handle = KeRegisterNmiCallback( NmiCallback, nmi_pools.nmi_context );
|
|
|
|
if ( !nmi_callback_handle )
|
|
{
|
|
DEBUG_ERROR( "KeRegisterNmiCallback failed" );
|
|
ExFreePoolWithTag( nmi_pools.nmi_context, NMI_CONTEXT_POOL );
|
|
return STATUS_ABANDONED;
|
|
}
|
|
|
|
/*
|
|
* We query the system modules each time since they can potentially
|
|
* change at any time
|
|
*/
|
|
status = GetSystemModuleInformation( &system_modules );
|
|
|
|
if ( !NT_SUCCESS( status ) )
|
|
{
|
|
DEBUG_ERROR( "Error retriving system module information" );
|
|
return status;
|
|
}
|
|
status = LaunchNonMaskableInterrupt( num_cores );
|
|
|
|
if ( !NT_SUCCESS( status ) )
|
|
{
|
|
DEBUG_ERROR( "Error running NMI callbacks" );
|
|
ExFreePoolWithTag( system_modules.address, SYSTEM_MODULES_POOL );
|
|
return status;
|
|
}
|
|
status = AnalyseNmiData( num_cores, &system_modules, Irp );
|
|
|
|
if ( !NT_SUCCESS( status ) )
|
|
DEBUG_ERROR( "Error analysing nmi data" );
|
|
|
|
ExFreePoolWithTag( system_modules.address, SYSTEM_MODULES_POOL );
|
|
ExFreePoolWithTag( nmi_pools.stack_frames, STACK_FRAMES_POOL );
|
|
ExFreePoolWithTag( nmi_pools.thread_data_pool, THREAD_DATA_POOL );
|
|
ExFreePoolWithTag( nmi_pools.nmi_context, NMI_CONTEXT_POOL );
|
|
KeDeregisterNmiCallback( nmi_callback_handle );
|
|
|
|
return status;
|
|
} |