From ced805cfbac525e3f0cb1025e620b805a2b5d216 Mon Sep 17 00:00:00 2001
From: lhodges1 <lachiehodges7785@gmail.com>
Date: Mon, 21 Aug 2023 14:08:57 +1000
Subject: [PATCH] GOT IT WORKING LOL

---
 driver/callbacks.c | 38 +++++++++++++++++++++++---------------
 driver/driver.c    |  6 ++++--
 driver/queue.c     |  1 -
 user/km/driver.cpp |  9 ++++++---
 4 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/driver/callbacks.c b/driver/callbacks.c
index ded2d41..59e4f48 100644
--- a/driver/callbacks.c
+++ b/driver/callbacks.c
@@ -7,38 +7,45 @@
 
 PQUEUE_HEAD report_queue = NULL;
 
+QUEUE_HEAD test_queue = { 0 };
+
 KGUARDED_MUTEX mutex;
 
 VOID InitCallbackReportQueue( PBOOLEAN Status )
 {
-	report_queue = QueueCreate();
+	//report_queue = QueueCreate();
 
-	if ( report_queue == NULL )
-	{
-		*Status = FALSE;
-		return;
-	}
+	test_queue.start = NULL;
+	test_queue.end = NULL;
+	test_queue.entries = 0;
+	KeInitializeSpinLock( &test_queue.lock );
+
+	//if ( report_queue == NULL )
+	//{
+	//	*Status = FALSE;
+	//	return;
+	//}
 
 	KeInitializeGuardedMutex( &mutex );
 
 	*Status = TRUE;
 }
 
-VOID DeleteCallbackReportQueueHead()
-{
-	ExFreePoolWithTag( report_queue, QUEUE_POOL_TAG );
-}
+//VOID DeleteCallbackReportQueueHead()
+//{
+//	ExFreePoolWithTag( report_queue, QUEUE_POOL_TAG );
+//}
 
 VOID InsertReportToQueue( 
 	_In_ POPEN_HANDLE_FAILURE_REPORT Report 
 )
 {
-	QueuePush( report_queue, Report );
+	QueuePush( &test_queue, Report );
 }
 
 POPEN_HANDLE_FAILURE_REPORT PopFirstReportFromQueue()
 {
-	return QueuePop( report_queue );
+	return QueuePop( &test_queue );
 }
 
 NTSTATUS HandlePeriodicCallbackReportQueue( 
@@ -78,13 +85,12 @@ NTSTATUS HandlePeriodicCallbackReportQueue(
 		count += 1;
 	}
 
+end:
 	header.count = count;
 	RtlCopyMemory( Irp->AssociatedIrp.SystemBuffer, &header, sizeof( OPEN_HANDLE_FAILURE_REPORT_HEADER ));
+	KeReleaseGuardedMutex( &mutex );
 
 	DEBUG_LOG( "Moved all reports into the IRP, sending !" );
-
-end:
-	KeReleaseGuardedMutex( &mutex );
 	return STATUS_SUCCESS;
 }
 
@@ -162,6 +168,7 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 			if ( !report )
 				goto end;
 
+			KeAcquireGuardedMutex( &mutex );
 			report->report_code = REPORT_ILLEGAL_HANDLE_OPERATION;
 			report->desired_access = OperationInformation->Parameters->CreateHandleInformation.DesiredAccess;
 			report->is_kernel_handle = OperationInformation->KernelHandle;
@@ -170,6 +177,7 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 			memcpy( report->process_name, process_creator_name, HANDLE_REPORT_PROCESS_NAME_MAX_LENGTH );
 
 			InsertReportToQueue( report );
+			KeReleaseGuardedMutex( &mutex );
 		}
 	}
 
diff --git a/driver/driver.c b/driver/driver.c
index c5c959a..7a7479d 100644
--- a/driver/driver.c
+++ b/driver/driver.c
@@ -53,7 +53,8 @@ VOID DriverUnload(
 	_In_ PDRIVER_OBJECT DriverObject
 )
 {
-	ExUnregisterCallback( callback_registration_handle );
+	PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE );
+	ObUnRegisterCallbacks( callback_registration_handle );
 	IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK );
 	IoDeleteDevice( DriverObject->DeviceObject );
 }
@@ -140,6 +141,7 @@ NTSTATUS DriverEntry(
 
 	KeInitializeGuardedMutex( &mutex );
 
+	__debugbreak();
 	InitCallbackReportQueue(&flag);
 
 	if ( !flag )
@@ -163,7 +165,7 @@ NTSTATUS DriverEntry(
 	if ( !NT_SUCCESS( status ) )
 	{
 		DEBUG_ERROR( "failed to launch thread to start tings" );
-		DeleteCallbackReportQueueHead();
+		//DeleteCallbackReportQueueHead();
 		IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK );
 		IoDeleteDevice( DriverObject->DeviceObject );
 		return STATUS_FAILED_DRIVER_ENTRY;
diff --git a/driver/queue.c b/driver/queue.c
index db9cb8d..182c630 100644
--- a/driver/queue.c
+++ b/driver/queue.c
@@ -15,7 +15,6 @@ PQUEUE_HEAD QueueCreate()
 	head->start = NULL;
 	head->entries = 0;
 
-	__debugbreak();
 	KeInitializeSpinLock( &head->lock );
 
 	return head;
diff --git a/user/km/driver.cpp b/user/km/driver.cpp
index 51d145c..ee861aa 100644
--- a/user/km/driver.cpp
+++ b/user/km/driver.cpp
@@ -153,7 +153,9 @@ void kernelmode::Driver::QueryReportQueue()
 	LONG buffer_size;
 	global::report_structures::OPEN_HANDLE_FAILURE_REPORT report;
 
-	buffer_size = sizeof( global::report_structures::OPEN_HANDLE_FAILURE_REPORT ) * MAX_HANDLE_REPORTS_PER_IRP ;
+	buffer_size = sizeof( global::report_structures::OPEN_HANDLE_FAILURE_REPORT ) * MAX_HANDLE_REPORTS_PER_IRP + 
+		sizeof( global::report_structures::OPEN_HANDLE_FAILURE_REPORT_HEADER );
+
 	buffer = malloc( buffer_size );
 
 	status = DeviceIoControl(
@@ -170,6 +172,7 @@ void kernelmode::Driver::QueryReportQueue()
 	if ( status == NULL )
 	{
 		LOG_ERROR( "DeviceIoControl failed with status code 0x%x", GetLastError() );
+		free( buffer );
 		return;
 	}
 
@@ -177,7 +180,7 @@ void kernelmode::Driver::QueryReportQueue()
 		( global::report_structures::OPEN_HANDLE_FAILURE_REPORT_HEADER* )buffer;
 
 	if ( !header )
-		return;
+		goto end;
 
 	for ( int i = 0; i < header->count; i++ )
 	{
@@ -191,8 +194,8 @@ void kernelmode::Driver::QueryReportQueue()
 		this->report_interface->ReportViolation( report );
 	}
 
+end:
 	free( buffer );
-
 }
 
 void kernelmode::Driver::NotifyDriverOnProcessLaunch()