From c0a1bd4f7504e8c290bb7e73d83d26a80d4ae616 Mon Sep 17 00:00:00 2001
From: lhodges1 <lachiehodges7785@gmail.com>
Date: Mon, 21 Aug 2023 14:45:33 +1000
Subject: [PATCH] eee

---
 driver/callbacks.c   |  5 ++---
 driver/queue.c       |  2 +-
 user/km/driver.cpp   | 12 ++++++++++++
 user/km/driver.h     |  4 +++-
 user/km/kmanager.cpp |  2 +-
 user/main.cpp        |  7 +------
 6 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/driver/callbacks.c b/driver/callbacks.c
index 59e4f48..a6ec6b1 100644
--- a/driver/callbacks.c
+++ b/driver/callbacks.c
@@ -62,9 +62,8 @@ NTSTATUS HandlePeriodicCallbackReportQueue(
 	if ( report == NULL )
 	{
 		DEBUG_LOG( "callback report queue is empty, returning" );
-		KeReleaseGuardedMutex( &mutex );
-		Irp->IoStatus.Information = NULL;
-		return STATUS_SUCCESS;
+		Irp->IoStatus.Information = sizeof( OPEN_HANDLE_FAILURE_REPORT_HEADER );
+		goto end;
 	}
 
 	Irp->IoStatus.Information = sizeof( OPEN_HANDLE_FAILURE_REPORT ) * MAX_HANDLE_REPORTS_PER_IRP + 
diff --git a/driver/queue.c b/driver/queue.c
index 182c630..dfa208b 100644
--- a/driver/queue.c
+++ b/driver/queue.c
@@ -62,7 +62,7 @@ PVOID QueuePop(
 	if ( temp == NULL )
 		goto end;
 
-	Head->entries -= 1;
+	Head->entries = Head->entries - 1;
 
 	data = temp->data;
 	Head->start = temp->next;
diff --git a/user/km/driver.cpp b/user/km/driver.cpp
index ee861aa..d54efcf 100644
--- a/user/km/driver.cpp
+++ b/user/km/driver.cpp
@@ -182,6 +182,9 @@ void kernelmode::Driver::QueryReportQueue()
 	if ( !header )
 		goto end;
 
+	if ( header->count == 0 )
+		goto end;
+
 	for ( int i = 0; i < header->count; i++ )
 	{
 		global::report_structures::OPEN_HANDLE_FAILURE_REPORT* report =
@@ -198,6 +201,15 @@ end:
 	free( buffer );
 }
 
+void kernelmode::Driver::RunCallbackReportQueue()
+{
+	while ( true )
+	{
+		this->QueryReportQueue();
+		std::this_thread::sleep_for( std::chrono::seconds( 10 ) );
+	}
+}
+
 void kernelmode::Driver::NotifyDriverOnProcessLaunch()
 {
 	BOOLEAN status;
diff --git a/user/km/driver.h b/user/km/driver.h
index ab60d30..75b63d3 100644
--- a/user/km/driver.h
+++ b/user/km/driver.h
@@ -21,13 +21,15 @@ namespace kernelmode
 		HANDLE driver_handle;
 		LPCWSTR driver_name;
 		std::shared_ptr<global::Report> report_interface;
+
+		void QueryReportQueue();
 	public:
 
 		Driver(LPCWSTR DriverName, std::shared_ptr<global::Report> ReportInterface );
 
 		void RunNmiCallbacks();
 		void VerifySystemModules();
-		void QueryReportQueue();
+		void RunCallbackReportQueue();
 		void NotifyDriverOnProcessLaunch();
 		void CompleteQueuedCallbackReports();
 		void EnableProcessLoadNotifyCallbacks();
diff --git a/user/km/kmanager.cpp b/user/km/kmanager.cpp
index 9e2e964..5b8d908 100644
--- a/user/km/kmanager.cpp
+++ b/user/km/kmanager.cpp
@@ -18,5 +18,5 @@ void kernelmode::KManager::VerifySystemModules()
 
 void kernelmode::KManager::MonitorCallbackReports()
 {
-	this->thread_pool->QueueJob( [ this ]() { this->driver_interface->QueryReportQueue(); } );
+	this->thread_pool->QueueJob( [ this ]() { this->driver_interface->RunCallbackReportQueue(); } );
 }
diff --git a/user/main.cpp b/user/main.cpp
index 64233fb..d4d0983 100644
--- a/user/main.cpp
+++ b/user/main.cpp
@@ -28,15 +28,10 @@ DWORD WINAPI Init(HINSTANCE hinstDLL)
     usermode::UManager umanager( thread_pool, report_interface );
     kernelmode::KManager kmanager( driver_name, thread_pool, report_interface);
 
+    kmanager.MonitorCallbackReports();
     //kmanager.RunNmiCallbacks();
     //kmanager.VerifySystemModules();
 
-    while ( true )
-    {
-        kmanager.MonitorCallbackReports();
-        std::this_thread::sleep_for( std::chrono::seconds( 10 ) );
-    }
-
     //umanager.ValidateProcessModules();
     //umanager.ValidateProcessMemory();