noury/mirror-ac
1
0
Fork 0
mirror of https://github.com/donnaskiez/ac.git synced 2024-11-21 22:24:08 +01:00
Code Issues Projects Releases Packages Wiki Activity

uopdate nmi callback

Browse source
This commit is contained in:
lhodges1 2023-09-13 19:46:28 +10:00
parent abfce7ddf7
commit 1a61c9f6f1
2 changed files with 68 additions and 77 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

139
driver/modules.c
View file

@ -2,13 +2,6 @@
#define WHITELISTED_MODULE_TAG 'whte'
PVOID nmi_callback_handle = NULL;
/* Global structure to hold pointers to required memory for the NMI's */
NMI_POOLS nmi_pools = { 0 };
volatile LONG lock;
#define NMI_DELAY 200 * 10000
#define WHITELISTED_MODULE_COUNT 3
@ -25,6 +18,21 @@ CHAR WHITELISTED_MODULES[ WHITELISTED_MODULE_COUNT ][ MODULE_MAX_STRING_SIZE ] =
"Wdf01000.sys",
};
typedef struct _NMI_CORE_CONTEXT
{
INT nmi_callbacks_run;
}NMI_CORE_CONTEXT, * PNMI_CORE_CONTEXT;
typedef struct _NMI_CONTEXT
{
PVOID thread_data_pool;
PVOID stack_frames;
PVOID nmi_core_context;
INT core_count;
}NMI_CONTEXT, * PNMI_CONTEXT;
/*
* TODO: this needs to be refactored to just return the entry not the whole fukin thing
*/
@ -582,17 +590,17 @@ NTSTATUS IsInstructionPointerInInvalidRegion(
}
NTSTATUS AnalyseNmiData(
_In_ INT NumCores,
_In_ PNMI_CONTEXT NmiContext,
_In_ PSYSTEM_MODULES SystemModules,
_In_ PIRP Irp
)
{
if ( !NumCores || !SystemModules )
if ( !NmiContext || !SystemModules )
return STATUS_INVALID_PARAMETER;
for ( INT core = 0; core < NumCores; core++ )
for ( INT core = 0; core < NmiContext->core_count; core++ )
{
PNMI_CONTEXT context = ( PNMI_CONTEXT )( ( uintptr_t )nmi_pools.nmi_context + core * sizeof( NMI_CONTEXT ) );
PNMI_CORE_CONTEXT context = ( PNMI_CORE_CONTEXT )( ( uintptr_t )NmiContext->nmi_core_context + core * sizeof( NMI_CORE_CONTEXT ) );
/* Make sure our NMIs were run */
if ( !context->nmi_callbacks_run )
@ -615,7 +623,7 @@ NTSTATUS AnalyseNmiData(
}
PNMI_CALLBACK_DATA thread_data = ( PNMI_CALLBACK_DATA )(
( uintptr_t )nmi_pools.thread_data_pool + core * sizeof( NMI_CALLBACK_DATA ) );
( uintptr_t )NmiContext->thread_data_pool + core * sizeof( NMI_CALLBACK_DATA ) );
DEBUG_LOG( "cpu number: %i callback count: %i", core, context->nmi_callbacks_run );
@ -624,7 +632,7 @@ NTSTATUS AnalyseNmiData(
{
BOOLEAN flag;
DWORD64 stack_frame = *( DWORD64* )(
( ( uintptr_t )nmi_pools.stack_frames + thread_data->stack_frames_offset + frame * sizeof( PVOID ) ) );
( ( uintptr_t )NmiContext->stack_frames + thread_data->stack_frames_offset + frame * sizeof( PVOID ) ) );
if ( !NT_SUCCESS( IsInstructionPointerInInvalidRegion( stack_frame, SystemModules, &flag ) ) )
{
@ -673,9 +681,10 @@ BOOLEAN NmiCallback(
{
UNREFERENCED_PARAMETER( Handled );
ULONG proc_num = KeGetCurrentProcessorNumber();
PVOID current_thread = KeGetCurrentThread();
NMI_CALLBACK_DATA thread_data = { 0 };
PNMI_CONTEXT nmi_context = ( PNMI_CONTEXT )Context;
ULONG proc_num = KeGetCurrentProcessorNumber();
/*
* Cannot allocate pool in this function as it runs at IRQL >= dispatch level
@ -684,7 +693,7 @@ BOOLEAN NmiCallback(
INT num_frames_captured = RtlCaptureStackBackTrace(
NULL,
STACK_FRAME_POOL_SIZE,
( uintptr_t )nmi_pools.stack_frames + proc_num * STACK_FRAME_POOL_SIZE,
( uintptr_t )nmi_context->stack_frames + proc_num * STACK_FRAME_POOL_SIZE,
NULL
);
@ -702,68 +711,62 @@ BOOLEAN NmiCallback(
thread_data.num_frames_captured = num_frames_captured;
RtlCopyMemory(
( ( uintptr_t )nmi_pools.thread_data_pool ) + proc_num * sizeof( thread_data ),
( ( uintptr_t )nmi_context->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 );
PNMI_CORE_CONTEXT core_context =
( PNMI_CORE_CONTEXT )( ( uintptr_t )nmi_context->nmi_core_context + proc_num * sizeof( NMI_CORE_CONTEXT ) );
InterlockedDecrement( &lock );
core_context->nmi_callbacks_run += 1;
DEBUG_LOG( "core number: %lx, num nmis run: %i", proc_num, core_context->nmi_callbacks_run );
return TRUE;
}
NTSTATUS LaunchNonMaskableInterrupt(
_In_ ULONG NumCores
_In_ PNMI_CONTEXT NmiContext
)
{
if ( !NumCores )
if ( !NmiContext )
return STATUS_INVALID_PARAMETER;
PKAFFINITY_EX proc_affinity = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( KAFFINITY_EX ), PROC_AFFINITY_POOL );
PKAFFINITY_EX ProcAffinityPool =
ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( KAFFINITY_EX ), PROC_AFFINITY_POOL );
if ( !proc_affinity )
return STATUS_ABANDONED;
if ( !ProcAffinityPool )
return STATUS_MEMORY_NOT_ALLOCATED;
nmi_pools.stack_frames = ExAllocatePool2( POOL_FLAG_NON_PAGED, NumCores * STACK_FRAME_POOL_SIZE, STACK_FRAMES_POOL );
NmiContext->stack_frames =
ExAllocatePool2( POOL_FLAG_NON_PAGED, NmiContext->core_count * STACK_FRAME_POOL_SIZE, STACK_FRAMES_POOL );
if ( !nmi_pools.stack_frames )
if ( !NmiContext->stack_frames )
{
ExFreePoolWithTag( proc_affinity, PROC_AFFINITY_POOL );
return STATUS_ABANDONED;
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
return STATUS_MEMORY_NOT_ALLOCATED;
}
nmi_pools.thread_data_pool = ExAllocatePool2( POOL_FLAG_NON_PAGED, NumCores * sizeof( NMI_CALLBACK_DATA ), THREAD_DATA_POOL );
NmiContext->thread_data_pool =
ExAllocatePool2( POOL_FLAG_NON_PAGED, NmiContext->core_count * sizeof( NMI_CALLBACK_DATA ), THREAD_DATA_POOL );
if ( !nmi_pools.thread_data_pool )
if ( !NmiContext->thread_data_pool )
{
ExFreePoolWithTag( nmi_pools.stack_frames, STACK_FRAMES_POOL );
ExFreePoolWithTag( proc_affinity, PROC_AFFINITY_POOL );
return STATUS_ABANDONED;
ExFreePoolWithTag( NmiContext->stack_frames, STACK_FRAMES_POOL );
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
return STATUS_MEMORY_NOT_ALLOCATED;
}
LARGE_INTEGER delay = { 0 };
delay.QuadPart -= NMI_DELAY;
delay.QuadPart -= 100 * 10000;
LONG ready = 0;
for ( ULONG core = 0; core < NumCores; core++ )
for ( ULONG core = 0; core < NmiContext->core_count; core++ )
{
KeInitializeAffinityEx( proc_affinity );
KeAddProcessorAffinityEx( proc_affinity, core );
InterlockedExchange( &ready, lock );
while ( ready > 0 )
InterlockedExchange( &ready, lock );
InterlockedIncrement( &lock );
KeInitializeAffinityEx( ProcAffinityPool );
KeAddProcessorAffinityEx( ProcAffinityPool, core );
DEBUG_LOG( "Sending NMI" );
HalSendNMI( proc_affinity );
HalSendNMI( ProcAffinityPool );
/*
* Only a single NMI can be active at any given time, so arbitrarily
@ -772,7 +775,7 @@ NTSTATUS LaunchNonMaskableInterrupt(
KeDelayExecutionThread( KernelMode, FALSE, &delay );
}
ExFreePoolWithTag( proc_affinity, PROC_AFFINITY_POOL );
ExFreePoolWithTag( ProcAffinityPool, PROC_AFFINITY_POOL );
return STATUS_SUCCESS;
}
@ -783,32 +786,26 @@ NTSTATUS HandleNmiIOCTL(
{
NTSTATUS status = STATUS_SUCCESS;
SYSTEM_MODULES system_modules = { 0 };
ULONG num_cores = KeQueryActiveProcessorCountEx( 0 );
NMI_CONTEXT nmi_context = { 0 };
PVOID callback_handle;
/* Fix annoying visual studio linting error */
RtlZeroMemory( &system_modules, sizeof( SYSTEM_MODULES ) );
RtlZeroMemory( &nmi_pools, sizeof( NMI_POOLS ) );
nmi_context.core_count = KeQueryActiveProcessorCountEx( 0 );
nmi_context.nmi_core_context =
ExAllocatePool2( POOL_FLAG_NON_PAGED, nmi_context.core_count * sizeof( NMI_CORE_CONTEXT ), NMI_CONTEXT_POOL );
KeInitializeSpinLock( &lock );
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;
}
if ( !nmi_context.nmi_core_context )
return STATUS_MEMORY_NOT_ALLOCATED;
/*
* 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 );
callback_handle = KeRegisterNmiCallback( NmiCallback, nmi_context.nmi_core_context );
if ( !nmi_callback_handle )
if ( !callback_handle )
{
DEBUG_ERROR( "KeRegisterNmiCallback failed" );
ExFreePoolWithTag( nmi_pools.nmi_context, NMI_CONTEXT_POOL );
ExFreePoolWithTag( nmi_context.nmi_core_context, NMI_CONTEXT_POOL );
return STATUS_ABANDONED;
}
@ -823,7 +820,7 @@ NTSTATUS HandleNmiIOCTL(
DEBUG_ERROR( "Error retriving system module information" );
return status;
}
status = LaunchNonMaskableInterrupt( num_cores );
status = LaunchNonMaskableInterrupt( &nmi_context );
if ( !NT_SUCCESS( status ) )
{
@ -831,16 +828,16 @@ NTSTATUS HandleNmiIOCTL(
ExFreePoolWithTag( system_modules.address, SYSTEM_MODULES_POOL );
return status;
}
status = AnalyseNmiData( num_cores, &system_modules, Irp );
status = AnalyseNmiData( &nmi_context, &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 );
ExFreePoolWithTag( nmi_context.stack_frames, STACK_FRAMES_POOL );
ExFreePoolWithTag( nmi_context.thread_data_pool, THREAD_DATA_POOL );
ExFreePoolWithTag( nmi_context.nmi_core_context, NMI_CONTEXT_POOL );
KeDeregisterNmiCallback( callback_handle );
return status;
}

6
driver/modules.h
View file

@ -34,12 +34,6 @@ typedef struct NMI_CALLBACK_FAILURE
}NMI_CALLBACK_FAILURE, * PNMI_CALLBACK_FAILURE;
typedef struct _NMI_CONTEXT
{
INT nmi_callbacks_run;
}NMI_CONTEXT, * PNMI_CONTEXT;
typedef struct _NMI_CALLBACK_DATA
{
UINT64 kthread_address;