From 4c3e5d78ac356cac0a5bd6d7c81842fbe4f0e21f Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sat, 16 Apr 2022 10:15:59 +0200
Subject: [PATCH] Implement threading

---
 src/driver/driver_main.cpp |  6 +--
 src/driver/thread.cpp      | 88 ++++++++++++++++++++++++++++++++++++++
 src/driver/thread.hpp      | 22 ++++++++++
 3 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/src/driver/driver_main.cpp b/src/driver/driver_main.cpp
index be9ba45..b83f023 100644
--- a/src/driver/driver_main.cpp
+++ b/src/driver/driver_main.cpp
@@ -168,14 +168,14 @@ public:
 		  , irp_(driver_object, DEV_NAME, DOS_DEV_NAME)
 	{
 		debug_log("Driver started\n");
-		this->trampoline = HookCreateFile(this->hypervisor_);
+		this->trampoline_ = HookCreateFile(this->hypervisor_);
 	}
 
 	~global_driver()
 	{
 		debug_log("Unloading driver\n");
 		this->hypervisor_.disable_all_ept_hooks();
-		memory::free_non_paged_memory(this->trampoline);
+		memory::free_non_paged_memory(this->trampoline_);
 	}
 
 	global_driver(global_driver&&) noexcept = delete;
@@ -193,7 +193,7 @@ private:
 	hypervisor hypervisor_{};
 	sleep_callback sleep_callback_{};
 	irp irp_{};
-	void* trampoline{nullptr};
+	void* trampoline_{nullptr};
 
 	void sleep_notification(const sleep_callback::type type)
 	{
diff --git a/src/driver/thread.cpp b/src/driver/thread.cpp
index d77e3b1..261baaa 100644
--- a/src/driver/thread.cpp
+++ b/src/driver/thread.cpp
@@ -2,6 +2,7 @@
 #include "std_include.hpp"
 #include "logging.hpp"
 #include "exception.hpp"
+#include "finally.hpp"
 
 namespace thread
 {
@@ -66,6 +67,93 @@ namespace thread
 
 			KeSignalCallDpcDone(arg1);
 		}
+
+		void thread_starter(void* context)
+		{
+			auto* function_ptr = static_cast<std::function<void()>*>(context);
+			const auto function = std::move(*function_ptr);
+			delete function_ptr;
+
+			try
+			{
+				function();
+			}
+			catch (...)
+			{
+			}
+		}
+	}
+
+	kernel_thread::kernel_thread(std::function<void()>&& callback)
+	{
+		auto* function_object = new std::function(std::move(callback));
+
+		auto destructor = utils::finally([&function_object]()
+		{
+			delete function_object;
+		});
+
+
+		HANDLE handle{};
+		const auto status = PsCreateSystemThread(&handle, 0, nullptr, nullptr, nullptr, thread_starter,
+		                                         function_object);
+
+		if (status != STATUS_SUCCESS)
+		{
+			throw std::runtime_error("Failed to create thread!");
+		}
+
+		ObReferenceObjectByHandle(handle, THREAD_ALL_ACCESS, nullptr, KernelMode,
+		                          reinterpret_cast<void**>(&this->handle_), nullptr);
+
+		ZwClose(handle);
+
+		destructor.cancel();
+	}
+
+	kernel_thread::~kernel_thread()
+	{
+		this->join();
+	}
+
+	kernel_thread::kernel_thread(kernel_thread&& obj) noexcept
+	{
+		this->operator=(std::move(obj));
+	}
+
+	kernel_thread& kernel_thread::operator=(kernel_thread&& obj) noexcept
+	{
+		if (this != &obj)
+		{
+			this->join();
+			this->handle_ = obj.handle_;
+			obj.handle_ = nullptr;
+		}
+
+		return *this;
+	}
+
+	bool kernel_thread::joinable() const
+	{
+		return this->handle_ != nullptr;
+	}
+
+	void kernel_thread::join()
+	{
+		if (this->joinable())
+		{
+			KeWaitForSingleObject(this->handle_, Executive, KernelMode, FALSE, nullptr);
+			this->detach();
+		}
+	}
+
+	void kernel_thread::detach()
+	{
+		if (this->joinable())
+		{
+			ObDereferenceObject(this->handle_);
+			this->handle_ = nullptr;
+		}
 	}
 
 	uint32_t get_processor_count()
diff --git a/src/driver/thread.hpp b/src/driver/thread.hpp
index 3693bfb..88700d9 100644
--- a/src/driver/thread.hpp
+++ b/src/driver/thread.hpp
@@ -1,4 +1,5 @@
 #pragma once
+#include "functional.hpp"
 
 namespace thread
 {
@@ -26,4 +27,25 @@ namespace thread
 			(*static_cast<F*>(data))();
 		}, &callback, sequential);
 	}
+
+	class kernel_thread
+	{
+	public:
+		kernel_thread() = default;
+		kernel_thread(std::function<void()>&& callback);
+		~kernel_thread();
+
+		kernel_thread(kernel_thread&& obj) noexcept;
+		kernel_thread& operator=(kernel_thread&& obj) noexcept;
+
+		kernel_thread(const kernel_thread& obj) = delete;
+		kernel_thread& operator=(const kernel_thread& obj) = delete;
+
+		bool joinable() const;
+		void join();
+		void detach();
+
+	private:
+		PETHREAD handle_{nullptr};
+	};
 }