More cleanup

src/driver/driver_main.cpp

@@ -55,12 +55,9 @@ 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)
+{
+	try
 	{
 		if (global_driver_instance)
 		{
@@ -68,13 +65,21 @@ _Function_class_(DRIVER_UNLOAD) void unload(const PDRIVER_OBJECT 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)

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;
 			}

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);
+}