More cleanup

src/driver/driver_main.cpp

@@ -3,9 +3,41 @@
 #include "sleep_callback.hpp"
 #include "irp.hpp"
 #include "exception.hpp"
-#include "finally.hpp"
 
-sleep_callback* sleep_cb{nullptr};
+#define DOS_DEV_NAME L"\\DosDevices\\HelloDev"
+#define DEV_NAME L"\\Device\\HelloDev"
+
+class global_driver
+{
+public:
+	global_driver(const PDRIVER_OBJECT driver_object)
+		: sleep_callback_([this](const sleep_callback::type type)
+		  {
+			  this->sleep_notification(type);
+		  })
+		  , irp_(driver_object, DEV_NAME, DOS_DEV_NAME)
+	{
+		debug_log("Driver started\n");
+	}
+
+	~global_driver()
+	{
+		debug_log("Unloading driver\n");
+	}
+
+	global_driver(global_driver&&) noexcept = delete;
+	global_driver& operator=(global_driver&&) noexcept = delete;
+
+	global_driver(const global_driver&) = delete;
+	global_driver& operator=(const global_driver&) = delete;
+
+	void pre_destroy(const PDRIVER_OBJECT /*driver_object*/)
+	{
+	}
+
+private:
+	sleep_callback sleep_callback_{};
+	irp irp_{};
 
 	void sleep_notification(const sleep_callback::type type)
 	{
@@ -19,34 +51,31 @@ void sleep_notification(const sleep_callback::type type)
 			debug_log("Waking up!");
 		}
 	}
+};
+
+global_driver* global_driver_instance{nullptr};
 
 extern "C" void __cdecl __std_terminate()
 {
 	KeBugCheckEx(DRIVER_VIOLATION, 14, 0, 0, 0);
 }
 
-void destroy_sleep_callback()
-{
-	delete sleep_cb;
-}
-
 _Function_class_(DRIVER_UNLOAD) void unload(const PDRIVER_OBJECT driver_object)
 {
-	irp::uninitialize(driver_object);
+	if (global_driver_instance)
-	destroy_sleep_callback();
+	{
+		global_driver_instance->pre_destroy(driver_object);
+		delete global_driver_instance;
+	}
 }
 
 extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT driver_object, PUNICODE_STRING /*registry_path*/)
 {
 	driver_object->DriverUnload = unload;
 
-	auto sleep_destructor = utils::finally(&destroy_sleep_callback);
-
 	try
 	{
-		sleep_cb = new sleep_callback(sleep_notification);
+		global_driver_instance = new global_driver(driver_object);
-		irp::initialize(driver_object);
-		sleep_destructor.cancel();
 	}
 	catch (std::exception& e)
 	{

src/driver/irp.cpp

@@ -3,13 +3,8 @@
 #include "logging.hpp"
 #include "exception.hpp"
 
-#define DOS_DEV_NAME L"\\DosDevices\\HelloDev"
-#define DEV_NAME L"\\Device\\HelloDev"
-
 #define HELLO_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)
 
-namespace irp
-{
 namespace
 {
 	UNICODE_STRING get_unicode_string(const wchar_t* string)
@@ -19,18 +14,7 @@ namespace irp
 		return unicode_string;
 	}
 
-		UNICODE_STRING get_device_name()
+	_Function_class_(DRIVER_DISPATCH) NTSTATUS not_supported_handler(PDEVICE_OBJECT /*device_object*/, const PIRP irp)
-		{
-			return get_unicode_string(DEV_NAME);
-		}
-
-		UNICODE_STRING get_dos_device_name()
-		{
-			return get_unicode_string(DOS_DEV_NAME);
-		}
-
-		_Function_class_(DRIVER_DISPATCH) NTSTATUS not_supported_handler(
-			PDEVICE_OBJECT /*device_object*/, const PIRP irp)
 	{
 		PAGED_CODE()
 
@@ -86,34 +70,23 @@ namespace irp
 	}
 }
 
-	_Function_class_(DRIVER_DISPATCH) void uninitialize(const PDRIVER_OBJECT driver_object)
+irp::irp(const PDRIVER_OBJECT driver_object, const wchar_t* device_name, const wchar_t* dos_device_name)
 {
 	PAGED_CODE()
 
-		auto dos_device_name = get_dos_device_name();
+	this->device_name_ = get_unicode_string(device_name);
+	this->dos_device_name_ = get_unicode_string(dos_device_name);
 
-		IoDeleteSymbolicLink(&dos_device_name);
+	auto destructor = utils::finally([this]()
-		IoDeleteDevice(driver_object->DeviceObject);
-	}
-
-	void initialize(const PDRIVER_OBJECT driver_object)
 	{
-		PAGED_CODE()
+		if (this->device_object_)
-
-		auto device_name = get_device_name();
-		auto dos_device_name = get_dos_device_name();
-
-		PDEVICE_OBJECT device_object{};
-		auto destructor = utils::finally([&device_object]()
 		{
-			if (device_object)
+			IoDeleteDevice(this->device_object_);
-			{
-				IoDeleteDevice(device_object);
 		}
 	});
 
-		auto status = IoCreateDevice(driver_object, 0, &device_name, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN,
+	auto status = IoCreateDevice(driver_object, 0, &this->device_name_, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN,
-		                             FALSE, &device_object);
+	                             FALSE, &this->device_object_);
 	if (!NT_SUCCESS(status))
 	{
 		throw std::runtime_error("Unable to create device");
@@ -128,10 +101,10 @@ namespace irp
 	driver_object->MajorFunction[IRP_MJ_CLOSE] = success_handler;
 	driver_object->MajorFunction[IRP_MJ_DEVICE_CONTROL] = io_ctl_handler;
 
-		device_object->Flags |= DO_DIRECT_IO;
+	this->device_object_->Flags |= DO_DIRECT_IO;
-		device_object->Flags &= ~DO_DEVICE_INITIALIZING;
+	this->device_object_->Flags &= ~DO_DEVICE_INITIALIZING;
 
-		status = IoCreateSymbolicLink(&dos_device_name, &device_name);
+	status = IoCreateSymbolicLink(&this->dos_device_name_, &this->device_name_);
 	if (!NT_SUCCESS(status))
 	{
 		throw std::runtime_error("Unable to create symbolic link");
@@ -139,4 +112,36 @@ namespace irp
 
 	destructor.cancel();
 }
+
+irp::~irp()
+{
+	PAGED_CODE()
+
+	if (this->device_object_)
+	{
+		IoDeleteSymbolicLink(&this->dos_device_name_);
+		IoDeleteDevice(this->device_object_);
+	}
+}
+
+irp::irp(irp&& obj) noexcept
+	: irp()
+{
+	this->operator=(std::move(obj));
+}
+
+irp& irp::operator=(irp&& obj) noexcept
+{
+	if (this != &obj)
+	{
+		this->~irp();
+
+		this->device_name_ = obj.device_name_;
+		this->dos_device_name_ = obj.dos_device_name_;
+		this->device_object_ = obj.device_object_;
+
+		obj.device_object_ = nullptr;
+	}
+
+	return *this;
 }

src/driver/irp.hpp

@@ -1,7 +1,20 @@
 #pragma once
 
-namespace irp
+class irp
 {
-	void initialize(PDRIVER_OBJECT driver_object);
+public:
-	void uninitialize(PDRIVER_OBJECT driver_object);
+	irp() = default;
-}
+	irp(PDRIVER_OBJECT driver_object, const wchar_t* device_name, const wchar_t* dos_device_name);
+	~irp();
+
+	irp(irp&& obj) noexcept;
+	irp& operator=(irp&& obj) noexcept;
+
+	irp(const irp&) = delete;
+	irp& operator=(const irp&) = delete;
+
+private:
+	UNICODE_STRING device_name_{};
+	UNICODE_STRING dos_device_name_{};
+	PDEVICE_OBJECT device_object_{};
+};

src/driver/sleep_callback.hpp

@@ -12,6 +12,7 @@ public:
 
 	using callback_function = std::function<void(type)>;
 
+	sleep_callback() = default;
 	sleep_callback(callback_function&& callback);
 	~sleep_callback();