yeet

driver/callbacks.c

@@ -5,46 +5,55 @@
 
 #include "queue.h"
 
-PQUEUE_HEAD report_queue = NULL;
-
-QUEUE_HEAD test_queue = { 0 };
+QUEUE_HEAD head = { 0 };
 
+/*
+* This mutex is to prevent a new item being pushed to the queue
+* while the HandlePeriodicCallbackReportQueue is iterating through
+* the objects. This can be an issue because the spinlock is released
+* after each report is placed in the IRP buffer which means a new report
+* can be pushed into the queue before the next iteration can take ownership
+* of the spinlock. 
+*/
 KGUARDED_MUTEX mutex;
 
-VOID InitCallbackReportQueue( PBOOLEAN Status )
+VOID InitCallbackReportQueue(
+	_In_ PBOOLEAN Status 
+)
 {
-	//report_queue = QueueCreate();
+	head.start = NULL;
+	head.end = NULL;
+	head.entries = 0;
 
-	test_queue.start = NULL;
-	test_queue.end = NULL;
-	test_queue.entries = 0;
-	KeInitializeSpinLock( &test_queue.lock );
+	KeInitializeSpinLock( &head.lock );
 	KeInitializeGuardedMutex( &mutex );
 
-	//if ( report_queue == NULL )
-	//{
-	//	*Status = FALSE;
-	//	return;
-	//}
-
 	*Status = TRUE;
 }
 
-//VOID DeleteCallbackReportQueueHead()
-//{
-//	ExFreePoolWithTag( report_queue, QUEUE_POOL_TAG );
-//}
-
 VOID InsertReportToQueue(
 	_In_ POPEN_HANDLE_FAILURE_REPORT Report
 )
 {
-	QueuePush( &test_queue, Report );
+	KeAcquireGuardedMutex( &mutex );
+	QueuePush( &head, Report );
+	KeReleaseGuardedMutex( &mutex );
 }
 
-POPEN_HANDLE_FAILURE_REPORT PopFirstReportFromQueue()
+VOID FreeQueueObjectsAndCleanup()
 {
-	return QueuePop( &test_queue );
+	KeAcquireGuardedMutex( &mutex );
+
+	PVOID report = QueuePop(&head );
+
+	if ( report == NULL )
+		goto end;
+
+	while ( report != NULL )
+		report = QueuePop( &head );
+
+end:
+	KeReleaseGuardedMutex( &mutex );
 }
 
 NTSTATUS HandlePeriodicCallbackReportQueue( 
@@ -56,7 +65,7 @@ NTSTATUS HandlePeriodicCallbackReportQueue(
 	OPEN_HANDLE_FAILURE_REPORT_HEADER header;
 
 	KeAcquireGuardedMutex( &mutex );
-	report = PopFirstReportFromQueue();
+	report = QueuePop( &head );
 
 	if ( report == NULL )
 	{
@@ -79,11 +88,12 @@ NTSTATUS HandlePeriodicCallbackReportQueue(
 			sizeof( OPEN_HANDLE_FAILURE_REPORT )
 		);
 
-		report = PopFirstReportFromQueue();
+		report = QueuePop( &head );
 		count += 1;
 	}
 
 end:
+
 	header.count = count;
 	RtlCopyMemory( Irp->AssociatedIrp.SystemBuffer, &header, sizeof( OPEN_HANDLE_FAILURE_REPORT_HEADER ));
 	KeReleaseGuardedMutex( &mutex );

driver/callbacks.h

@@ -55,8 +55,9 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 	_In_ POB_PRE_OPERATION_INFORMATION OperationInformation
 );
 
-VOID InitCallbackReportQueue(PBOOLEAN Status);
-VOID DeleteCallbackReportQueueHead();
+VOID InitCallbackReportQueue(
+	_In_ PBOOLEAN Status
+);
 
 NTSTATUS HandlePeriodicCallbackReportQueue(
 	_In_ PIRP Irp
@@ -68,4 +69,6 @@ VOID ProcessCreateNotifyRoutine(
 	_In_ BOOLEAN Create
 );
 
+VOID FreeQueueObjectsAndCleanup();
+
 #endif

driver/driver.c

@@ -55,6 +55,7 @@ VOID DriverUnload(
 {
 	PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE );
 	ObUnRegisterCallbacks( callback_registration_handle );
+	FreeQueueObjectsAndCleanup();
 	IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK );
 	IoDeleteDevice( DriverObject->DeviceObject );
 }
@@ -105,7 +106,7 @@ NTSTATUS DriverEntry(
 {
 	UNREFERENCED_PARAMETER( RegistryPath );
 
-	BOOLEAN flag;
+	BOOLEAN flag = FALSE;
 	NTSTATUS status;
 	HANDLE handle;
 
@@ -164,7 +165,6 @@ NTSTATUS DriverEntry(
 	if ( !NT_SUCCESS( status ) )
 	{
 		DEBUG_ERROR( "failed to launch thread to start tings" );
-		//DeleteCallbackReportQueueHead();
 		IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK );
 		IoDeleteDevice( DriverObject->DeviceObject );
 		return STATUS_FAILED_DRIVER_ENTRY;

driver/ioctl.c

@@ -82,9 +82,9 @@ NTSTATUS DeviceControl(
 
 		/* KeWaitForSingleObject with infinite time must be called from IRQL <= APC_LEVEL */
 		PAGED_CODE();
-		DEBUG_LOG( "waiting for thread to finish" );
+
 		KeWaitForSingleObject( thread, Executive, KernelMode, FALSE, NULL );
-		DEBUG_LOG( "THREAD FINISHED" );
+
 		ZwClose( handle );
 		ObDereferenceObject( thread );
 
@@ -131,6 +131,7 @@ NTSTATUS DeviceCreate(
 	_In_ PIRP Irp
 )
 {
+	DEBUG_LOG( "Handle opened to DonnaAC" );
 	IoCompleteRequest( Irp, IO_NO_INCREMENT );
 	return Irp->IoStatus.Status;
 }

driver/nmi.c

@@ -3,6 +3,8 @@
 #include "modules.h"
 #include "common.h"
 
+#define NMI_DELAY 100 * 10000
+
 typedef struct _NMI_POOLS
 {
 	PVOID thread_data_pool;
@@ -208,7 +210,7 @@ NTSTATUS LaunchNonMaskableInterrupt(
 	}
 
 	LARGE_INTEGER delay = { 0 };
-	delay.QuadPart -= 100 * 10000;
+	delay.QuadPart -= NMI_DELAY;
 
 	for ( ULONG core = 0; core < NumCores; core++ )
 	{

driver/queue.c

@@ -1,24 +1,20 @@
 #include "queue.h"
 
-/*
-* Basic thread-safe queue implementation that can be used to create queues for any data.
-*/
-
-PQUEUE_HEAD QueueCreate()
-{
-	PQUEUE_HEAD head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( PQUEUE_HEAD ), QUEUE_POOL_TAG );
-
-	if ( !head )
-		return NULL;
-
-	head->end = NULL;
-	head->start = NULL;
-	head->entries = 0;
-
-	KeInitializeSpinLock( &head->lock );
-
-	return head;
-}
+//PQUEUE_HEAD QueueCreate()
+//{
+//	PQUEUE_HEAD head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( PQUEUE_HEAD ), QUEUE_POOL_TAG );
+//
+//	if ( !head )
+//		return NULL;
+//
+//	head->end = NULL;
+//	head->start = NULL;
+//	head->entries = 0;
+//
+//	KeInitializeSpinLock( &head->lock );
+//
+//	return head;
+//}
 
 VOID QueuePush( 
 	_In_ PQUEUE_HEAD Head,

driver/queue.h

@@ -21,8 +21,6 @@ typedef struct QUEUE_HEAD
 
 }QUEUE_HEAD, *PQUEUE_HEAD;
 
-PQUEUE_HEAD QueueCreate();
-
 VOID QueuePush(
 	_In_ PQUEUE_HEAD Head,
 	_In_ PVOID Data