From 49d259f3ae53eb14ca59c0f0640fe8176233b145 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sat, 26 Mar 2022 20:01:13 +0100
Subject: [PATCH] More cleanup

---
 src/driver/driver_main.cpp | 25 +++++++++++++++----------
 src/driver/irp.cpp         |  4 ++--
 src/driver/new.cpp         | 24 ++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/src/driver/driver_main.cpp b/src/driver/driver_main.cpp
index d928811..4f95bce 100644
--- a/src/driver/driver_main.cpp
+++ b/src/driver/driver_main.cpp
@@ -55,26 +55,31 @@ private:
 
 global_driver* global_driver_instance{nullptr};
 
-extern "C" void __cdecl __std_terminate()
-{
-	KeBugCheckEx(DRIVER_VIOLATION, 14, 0, 0, 0);
-}
-
 _Function_class_(DRIVER_UNLOAD) void unload(const PDRIVER_OBJECT driver_object)
 {
-	if (global_driver_instance)
+	try
 	{
-		global_driver_instance->pre_destroy(driver_object);
-		delete global_driver_instance;
+		if (global_driver_instance)
+		{
+			global_driver_instance->pre_destroy(driver_object);
+			delete global_driver_instance;
+		}
+	}
+	catch (std::exception& e)
+	{
+		debug_log("Destruction error occured: %s\n", e.what());
+	}
+	catch (...)
+	{
+		debug_log("Unknown destruction error occured. This should not happen!");
 	}
 }
 
 extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT driver_object, PUNICODE_STRING /*registry_path*/)
 {
-	driver_object->DriverUnload = unload;
-
 	try
 	{
+		driver_object->DriverUnload = unload;
 		global_driver_instance = new global_driver(driver_object);
 	}
 	catch (std::exception& e)
diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp
index 82393bf..0bbec9b 100644
--- a/src/driver/irp.cpp
+++ b/src/driver/irp.cpp
@@ -55,10 +55,10 @@ namespace
 			switch (ioctr_code)
 			{
 			case HELLO_DRV_IOCTL:
-				debug_log("[< HelloDriver >] Hello from the Driver!\n");
+				debug_log("Hello from the Driver!\n");
 				break;
 			default:
-				debug_log("[-] Invalid IOCTL Code: 0x%X\n", ioctr_code);
+				debug_log("Invalid IOCTL Code: 0x%X\n", ioctr_code);
 				irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
 				break;
 			}
diff --git a/src/driver/new.cpp b/src/driver/new.cpp
index bb1fcb8..f5cefaa 100644
--- a/src/driver/new.cpp
+++ b/src/driver/new.cpp
@@ -1,14 +1,29 @@
 #include "std_include.hpp"
 #include "new.hpp"
+#include "exception.hpp"
+
+namespace
+{
+	void* perform_allocation(const size_t size, const POOL_TYPE pool, const unsigned long tag)
+	{
+		auto* memory = ExAllocatePoolWithTag(pool, size, tag);
+		if (!memory)
+		{
+			throw std::runtime_error("Memory allocation failed");
+		}
+
+		return memory;
+	}
+}
 
 void* __cdecl operator new(const size_t size, const POOL_TYPE pool, const unsigned long tag)
 {
-	return ExAllocatePoolWithTag(pool, size, tag);
+	return perform_allocation(size, pool, tag);
 }
 
 void* __cdecl operator new[](const size_t size, const POOL_TYPE pool, const unsigned long tag)
 {
-	return ExAllocatePoolWithTag(pool, size, tag);
+	return perform_allocation(size, pool, tag);
 }
 
 void* __cdecl operator new(const size_t size)
@@ -46,3 +61,8 @@ void __cdecl operator delete[](void* ptr)
 {
 	ExFreePool(ptr);
 }
+
+extern "C" void __cdecl __std_terminate()
+{
+	KeBugCheckEx(DRIVER_VIOLATION, 14, 0, 0, 0);
+}