fix leak

driver/callbacks.c

@@ -45,13 +45,13 @@ VOID FreeQueueObjectsAndCleanup()
 {
 	KeAcquireGuardedMutex( &mutex );
 
-	PVOID report = QueuePop(&head );
-
-	if ( report == NULL )
-		goto end;
+	PVOID report = QueuePop( &head );
 
 	while ( report != NULL )
+	{
+		ExFreePoolWithTag( report, REPORT_POOL_TAG );
 		report = QueuePop( &head );
+	}
 
 end:
 	KeReleaseGuardedMutex( &mutex );
@@ -89,6 +89,8 @@ NTSTATUS HandlePeriodicCallbackReportQueue(
 			sizeof( OPEN_HANDLE_FAILURE_REPORT )
 		);
 
+		ExFreePoolWithTag( report, REPORT_POOL_TAG );
+
 		report = QueuePop( &head );
 		count += 1;
 	}
@@ -129,21 +131,25 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 	* is requesting to open said handle
 	*/
 	PEPROCESS process_creator = PsGetCurrentProcess();
+	PEPROCESS protected_process;
 	PEPROCESS target_process = ( PEPROCESS )OperationInformation->Object;
 	LONG target_process_id = PsGetProcessId( target_process );
 	LONG process_creator_id = PsGetProcessId( process_creator );
 	LONG protected_process_id = NULL;
 	LPCSTR process_creator_name;
 	LPCSTR target_process_name;
+	LPCSTR protected_process_name;
 
 	KeAcquireGuardedMutex( &configuration.mutex );
 
 	GetProtectedProcessId( &protected_process_id );
+	GetProtectedProcessEProcess( &protected_process );
 
 	process_creator_name = PsGetProcessImageFileName( process_creator );
 	target_process_name = PsGetProcessImageFileName( target_process );
+	protected_process_name = PsGetProcessImageFileName( protected_process );
 
-	if ( !strcmp( "notepad.exe", target_process_name) )
+	if ( !strcmp( protected_process_name, target_process_name) )
 	{
 		if ( !strcmp( process_creator_name, "lsass.exe" ) || !strcmp( process_creator_name, "csrss.exe" ) )
 		{

driver/callbacks.h

@@ -14,6 +14,8 @@
 
 #define MAX_HANDLE_REPORTS_PER_IRP 10
 
+VOID UnregisterCallbacksOnProcessTermination();
+
 typedef struct _OPEN_HANDLE_FAILURE_REPORT_HEADER
 {
 	INT count;

driver/ioctl.c

@@ -245,6 +245,7 @@ NTSTATUS DeviceClose(
 
 	FreeQueueObjectsAndCleanup();
 	ClearDriverConfigOnProcessTermination();
+	UnregisterCallbacksOnProcessTermination();
 
 	IoCompleteRequest( Irp, IO_NO_INCREMENT );
 	return Irp->IoStatus.Status;

driver/modules.c

@@ -7,7 +7,9 @@ PVOID nmi_callback_handle = NULL;
 /* Global structure to hold pointers to required memory for the NMI's */
 NMI_POOLS nmi_pools = { 0 };
 
-#define NMI_DELAY 500 * 10000
+volatile LONG lock;
+
+#define NMI_DELAY 200 * 10000
 
 #define WHITELISTED_MODULE_COUNT 3
 #define MODULE_MAX_STRING_SIZE 256
@@ -706,6 +708,8 @@ BOOLEAN NmiCallback(
 	context->nmi_callbacks_run += 1;
 	DEBUG_LOG( "num nmis called: %i from addr: %llx", context->nmi_callbacks_run, ( uintptr_t )context );
 
+	InterlockedDecrement( &lock );
+
 	return TRUE;
 }
 
@@ -741,11 +745,23 @@ NTSTATUS LaunchNonMaskableInterrupt(
 	LARGE_INTEGER delay = { 0 };
 	delay.QuadPart -= NMI_DELAY;
 
+	LONG ready = 0;
+
 	for ( ULONG core = 0; core < NumCores; core++ )
 	{
 		KeInitializeAffinityEx( proc_affinity );
 		KeAddProcessorAffinityEx( proc_affinity, core );
 
+		InterlockedExchange( &ready, lock );
+
+		if ( ready > 0 )
+		{
+			while ( ready > 0 )
+				InterlockedExchange( &ready, lock );
+		}
+
+		InterlockedIncrement( &lock );
+
 		DEBUG_LOG( "Sending NMI" );
 		HalSendNMI( proc_affinity );
 
@@ -773,6 +789,8 @@ NTSTATUS HandleNmiIOCTL(
 	RtlZeroMemory( &system_modules, sizeof( SYSTEM_MODULES ) );
 	RtlZeroMemory( &nmi_pools, sizeof( NMI_POOLS ) );
 
+	KeInitializeSpinLock( &lock );
+
 	nmi_pools.nmi_context = ExAllocatePool2( POOL_FLAG_NON_PAGED, num_cores * sizeof( NMI_CONTEXT ), NMI_CONTEXT_POOL );
 
 	if ( !nmi_pools.nmi_context )

driver/pool.c

@@ -352,6 +352,9 @@ VOID WalkKernelPageTables(PVOID AddressBuffer)
 
 				physical.QuadPart = pd_entry.Bits.PhysicalAddress << PAGE_4KB_SHIFT;
 
+				if ( !MmIsAddressValid( pd_entry.BitAddress ) )
+					continue;
+
 				pt_base = MmGetVirtualForPhysical( physical );
 
 				if ( !pt_base || !MmIsAddressValid( pt_base ) )

driver/queue.c

@@ -1,5 +1,7 @@
 #include "queue.h"
 
+#include "common.h"
+
 //PQUEUE_HEAD QueueCreate()
 //{
 //	PQUEUE_HEAD head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( QUEUE_HEAD ), QUEUE_POOL_TAG );