Cleanup user-mode side

src/driver/CMakeLists.txt

@@ -1,5 +1,5 @@
 wdk_add_driver(driver
-    main.cpp
+    driver_main.cpp
     thread.cpp
     new.cpp
 )

src/runner/CMakeLists.txt

@@ -1,5 +1,13 @@
 add_executable(runner
 	main.cpp
+	driver.cpp
+	service_handle.cpp
+	native_handle.cpp
+	driver_device.cpp
+)
+
+target_precompile_headers(runner
+	PRIVATE std_include.hpp
 )
 
 set_property(TARGET runner APPEND_STRING PROPERTY LINK_FLAGS " /MANIFESTUAC:\"level='requireAdministrator'\"")

src/runner/driver.cpp

@@ -0,0 +1,46 @@
+#include "std_include.hpp"
+#include "driver.hpp"
+
+driver::driver(const std::filesystem::path& driver_file, const std::string& service_name)
+{
+	this->manager_ = OpenSCManagerA(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);
+	if (!this->manager_)
+	{
+		throw std::runtime_error("Unable to open SC manager");
+	}
+
+	this->service_ = OpenServiceA(this->manager_, service_name.data(), SERVICE_ALL_ACCESS);
+
+	if (!this->service_)
+	{
+		this->service_ = CreateServiceA(this->manager_, service_name.data(),
+		                                service_name.data(), SERVICE_ALL_ACCESS, SERVICE_KERNEL_DRIVER,
+		                                SERVICE_DEMAND_START, SERVICE_ERROR_NORMAL,
+		                                driver_file.generic_string().data(), nullptr, nullptr,
+		                                nullptr, nullptr, nullptr);
+	}
+
+	if (!this->service_)
+	{
+		this->service_ = OpenServiceA(this->manager_, service_name.data(), SERVICE_ALL_ACCESS);
+	}
+
+	if (!this->service_)
+	{
+		this->manager_ = {};
+		throw std::runtime_error("Unable to create service");
+	}
+
+	StartServiceA(this->service_, 0, nullptr);
+}
+
+driver::~driver()
+{
+	if (this->service_)
+	{
+		SERVICE_STATUS status{};
+		ControlService(this->service_, SERVICE_CONTROL_STOP, &status);
+
+		DeleteService(this->service_);
+	}
+}

src/runner/driver.hpp

@@ -0,0 +1,18 @@
+#include "service_handle.hpp"
+
+class driver
+{
+public:
+	driver(const std::filesystem::path& driver_file, const std::string& service_name);
+	~driver();
+	
+	driver(const driver&) = delete;
+	driver& operator=(const driver&) = delete;
+
+	driver(driver&& obj) noexcept = default;;
+	driver& operator=(driver&& obj) noexcept = default;
+
+private:
+	service_handle manager_{};
+	service_handle service_{};
+};

src/runner/driver_device.cpp

@@ -0,0 +1,46 @@
+#include "std_include.hpp"
+#include "driver_device.hpp"
+
+driver_device::driver_device(const std::string& driver_device)
+{
+	this->device_ = CreateFileA(driver_device.data(),
+	                            GENERIC_READ | GENERIC_WRITE,
+	                            NULL,
+	                            nullptr,
+	                            OPEN_EXISTING,
+	                            NULL,
+	                            nullptr
+	);
+
+	if (!this->device_)
+	{
+		throw std::runtime_error("Unable to access device");
+	}
+}
+
+bool driver_device::send(const DWORD ioctl_code, const data& input) const
+{
+	data output{};
+	return this->send(ioctl_code, input, output);
+}
+
+bool driver_device::send(const DWORD ioctl_code, const data& input, data& output) const
+{
+	DWORD size_returned = 0;
+	const auto success = DeviceIoControl(this->device_,
+	                                     ioctl_code,
+	                                     const_cast<uint8_t*>(input.data()),
+	                                     input.size(),
+	                                     output.data(),
+	                                     output.size(),
+	                                     &size_returned,
+	                                     nullptr
+	) != FALSE;
+
+	if (success && size_returned < output.size())
+	{
+		output.resize(size_returned);
+	}
+
+	return success;
+}

src/runner/driver_device.hpp

@@ -0,0 +1,21 @@
+#include "native_handle.hpp"
+
+class driver_device
+{
+public:
+	driver_device(const std::string& driver_device);
+	~driver_device() = default;
+	
+	driver_device(const driver_device&) = delete;
+	driver_device& operator=(const driver_device&) = delete;
+
+	driver_device(driver_device&& obj) noexcept = default;
+	driver_device& operator=(driver_device&& obj) noexcept = default;
+
+	using data = std::vector<uint8_t>;
+	bool send(DWORD ioctl_code, const data& input) const;
+	bool send(DWORD ioctl_code, const data& input, data& output) const;
+
+private:
+	native_handle device_{};
+};

src/runner/main.cpp

@@ -1,32 +1,11 @@
-#include <Windows.h>
-#include <Shlwapi.h>
+#include "std_include.hpp"
 #include "finally.hpp"
-#include <filesystem>
+#include "driver.hpp"
+#include "driver_device.hpp"
 
 #pragma comment(lib, "Shlwapi.lib")
 
-#define SERVICE_NAME "MomoLul"
-
 #define HELLO_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)
-const char kDevName[] = "\\\\.\\HelloDev";
-
-HANDLE open_device(const char* device_name)
-{
-	HANDLE device = CreateFileA(device_name,
-	                            GENERIC_READ | GENERIC_WRITE,
-	                            NULL,
-	                            NULL,
-	                            OPEN_EXISTING,
-	                            NULL,
-	                            NULL
-	);
-	return device;
-}
-
-void close_device(HANDLE device)
-{
-	CloseHandle(device);
-}
 
 BOOL send_ioctl(HANDLE device, DWORD ioctl_code)
 {
@@ -63,67 +42,36 @@ std::filesystem::path get_current_path()
 	return selfdir;
 }
 
-int main(const int /*argc*/, char* /*argv*/[])
+void unsafe_main(const int /*argc*/, char* /*argv*/[])
 {
-	const auto manager = OpenSCManagerA(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);
-	if (manager == nullptr)
-	{
-		return 1;
-	}
+	driver driver{get_current_path() / "driver.sys", "MomoLul"};
+	driver_device driver_device{"\\\\.\\HelloDev"};
 
-	const auto _1 = utils::finally([&manager]()
-	{
-		CloseServiceHandle(manager);
-	});
+	driver_device::data input{};
+	input.resize(4);
 
-	auto service = OpenServiceA(manager, SERVICE_NAME, SERVICE_ALL_ACCESS);
-	const auto _2 = utils::finally([&service]()
-	{
-		if (service)
-		{
-			SERVICE_STATUS status;
-			ControlService(service, SERVICE_CONTROL_STOP, &status);
-
-			DeleteService(service);
-			CloseServiceHandle(service);
-		}
-	});
-
-	if (service == nullptr)
-	{
-		const auto driver_path = get_current_path() / "driver.sys";
-
-		service = CreateServiceA(manager, SERVICE_NAME,
-		                         SERVICE_NAME, SERVICE_ALL_ACCESS, SERVICE_KERNEL_DRIVER,
-		                         SERVICE_DEMAND_START, SERVICE_ERROR_NORMAL,
-		                         driver_path.generic_string().data(), nullptr, nullptr,
-		                         nullptr, nullptr, nullptr);
-	}
-
-	if (service == nullptr)
-	{
-		service = OpenServiceA(manager, SERVICE_NAME,
-		                       SERVICE_ALL_ACCESS);
-	}
-
-	if (service)
-	{
-		StartServiceA(service, 0, nullptr);
-
-		HANDLE dev = open_device(kDevName);
-		if (dev == INVALID_HANDLE_VALUE)
-		{
-			printf("Failed!\n");
-			return -1;
-		}
-
-		send_ioctl(dev, HELLO_DRV_IOCTL);
-
-		close_device(dev);
+	(void)driver_device.send(HELLO_DRV_IOCTL, input);
 
 	MessageBoxA(0, "Service started!", 0, 0);
-	}
+}
+
+int main(const int argc, char* argv[])
+{
+	try
+	{
+		unsafe_main(argc, argv);
 		return 0;
+	}
+	catch (std::exception& e)
+	{
+		printf("Error: %s\n", e.what());
+		return 1;
+	}
+	catch (...)
+	{
+		printf("An unknown error occurd!\n");
+		return 1;
+	}
 }
 
 int __stdcall WinMain(HINSTANCE, HINSTANCE, char*, int)

src/runner/native_handle.cpp

@@ -0,0 +1,49 @@
+#include "std_include.hpp"
+#include "native_handle.hpp"
+
+native_handle::native_handle()
+	: native_handle(INVALID_HANDLE_VALUE)
+{
+}
+
+native_handle::native_handle(const HANDLE handle)
+	: handle_{handle}
+{
+}
+
+native_handle::~native_handle()
+{
+	if (this->operator bool())
+	{
+		CloseHandle(this->handle_);
+		this->handle_ = INVALID_HANDLE_VALUE;
+	}
+}
+
+native_handle::native_handle(native_handle&& obj) noexcept
+	: native_handle()
+{
+	this->operator=(std::move(obj));
+}
+
+native_handle& native_handle::operator=(native_handle&& obj) noexcept
+{
+	if (this != &obj)
+	{
+		this->~native_handle();
+		this->handle_ = obj.handle_;
+		obj.handle_ = INVALID_HANDLE_VALUE;
+	}
+
+	return *this;
+}
+
+native_handle::operator HANDLE() const
+{
+	return this->handle_;
+}
+
+native_handle::operator bool() const
+{
+	return this->handle_ != INVALID_HANDLE_VALUE;
+}

src/runner/native_handle.hpp

@@ -0,0 +1,19 @@
+class native_handle
+{
+public:
+	native_handle();
+	native_handle(HANDLE handle);
+	~native_handle();
+	
+	native_handle(const native_handle&) = delete;
+	native_handle& operator=(const native_handle&) = delete;
+
+	native_handle(native_handle&& obj) noexcept;
+	native_handle& operator=(native_handle&& obj) noexcept;
+
+	operator HANDLE() const;
+	operator bool() const;
+
+private:
+	HANDLE handle_{INVALID_HANDLE_VALUE};
+};

src/runner/service_handle.cpp

@@ -0,0 +1,44 @@
+#include "std_include.hpp"
+#include "service_handle.hpp"
+
+service_handle::service_handle()
+	: service_handle(nullptr)
+{
+}
+
+service_handle::service_handle(const SC_HANDLE handle)
+	: handle_{handle}
+{
+}
+
+service_handle::~service_handle()
+{
+	if (this->handle_)
+	{
+		CloseServiceHandle(this->handle_);
+		this->handle_ = nullptr;
+	}
+}
+
+service_handle::service_handle(service_handle&& obj) noexcept
+	: service_handle()
+{
+	this->operator=(std::move(obj));
+}
+
+service_handle& service_handle::operator=(service_handle&& obj) noexcept
+{
+	if (this != &obj)
+	{
+		this->~service_handle();
+		this->handle_ = obj.handle_;
+		obj.handle_ = nullptr;
+	}
+
+	return *this;
+}
+
+service_handle::operator SC_HANDLE() const
+{
+	return this->handle_;
+}

src/runner/service_handle.hpp

@@ -0,0 +1,18 @@
+class service_handle
+{
+public:
+	service_handle();
+	service_handle(SC_HANDLE handle);
+	~service_handle();
+	
+	service_handle(const service_handle&) = delete;
+	service_handle& operator=(const service_handle&) = delete;
+
+	service_handle(service_handle&& obj) noexcept;
+	service_handle& operator=(service_handle&& obj) noexcept;
+
+	operator SC_HANDLE() const;
+
+private:
+	SC_HANDLE handle_{nullptr};
+};

src/runner/std_include.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <string>
+#include <vector>
+#include <filesystem>
+
+#include <Windows.h>
+#include <Shlwapi.h>
+
+#pragma comment(lib, "Shlwapi.lib")