From 85904bfd13d77cb567e7379de354e2182010bd06 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sun, 24 Apr 2022 10:41:27 +0200
Subject: [PATCH] Small cleanup

---
 src/driver/irp.cpp         | 63 ++++++++++++++++++++++----------------
 src/driver/type_traits.hpp | 10 ++++++
 src/driver/unique_ptr.hpp  | 19 ++++++++----
 3 files changed, 59 insertions(+), 33 deletions(-)

diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp
index 6279af8..aaf02cf 100644
--- a/src/driver/irp.cpp
+++ b/src/driver/irp.cpp
@@ -37,34 +37,15 @@ namespace
 		return STATUS_SUCCESS;
 	}
 
-	void apply_hook(const hook_request& request)
+	vmx::ept_translation_hint* generate_translation_hints(uint32_t process_id, const void* target_address, size_t size)
 	{
-		auto* buffer = new uint8_t[request.source_data_size];
-		if (!buffer)
+		vmx::ept_translation_hint* translation_hints{nullptr};
+
+		thread::kernel_thread t([&translation_hints, process_id, target_address, size]
 		{
-			throw std::runtime_error("Failed to copy buffer");
-		}
+			debug_log("Looking up process: %d\n", process_id);
 
-		vmx::ept_translation_hint* translation_hints = nullptr;
-		auto destructor = utils::finally([&translation_hints, &buffer]()
-		{
-			delete[] buffer;
-			vmx::ept::free_translation_hints(translation_hints);
-		});
-
-		memcpy(buffer, request.source_data, request.source_data_size);
-
-		auto* hypervisor = hypervisor::get_instance();
-		if (!hypervisor)
-		{
-			throw std::runtime_error("Hypervisor not installed");
-		}
-
-		thread::kernel_thread t([&translation_hints, r = request]
-		{
-			debug_log("Pid: %d | Address: %p\n", r.process_id, r.target_address);
-
-			const auto process_handle = process::find_process_by_id(r.process_id);
+			const auto process_handle = process::find_process_by_id(process_id);
 			if (!process_handle || !process_handle.is_alive())
 			{
 				debug_log("Bad process\n");
@@ -78,18 +59,46 @@ namespace
 			}
 
 			process::scoped_process_attacher attacher{process_handle};
-			translation_hints = vmx::ept::generate_translation_hints(r.target_address, r.source_data_size);
+
+			debug_log("Generating translation hints for address: %p\n", target_address);
+			translation_hints = vmx::ept::generate_translation_hints(target_address, size);
 		});
 
 		t.join();
 
+		return translation_hints;
+	}
+
+	void apply_hook(const hook_request& request)
+	{
+		auto* hypervisor = hypervisor::get_instance();
+		if (!hypervisor)
+		{
+			throw std::runtime_error("Hypervisor not installed");
+		}
+
+		std::unique_ptr<uint8_t[]> buffer(new uint8_t[request.source_data_size]);
+		if (!buffer)
+		{
+			throw std::runtime_error("Failed to copy buffer");
+		}
+
+		vmx::ept_translation_hint* translation_hints = nullptr;
+		auto destructor = utils::finally([&translation_hints]()
+		{
+			vmx::ept::free_translation_hints(translation_hints);
+		});
+
+		memcpy(buffer.get(), request.source_data, request.source_data_size);
+		translation_hints = generate_translation_hints(request.process_id, request.target_address, request.source_data_size);
+
 		if (!translation_hints)
 		{
 			debug_log("Failed to generate tranlsation hints\n");
 			return;
 		}
 
-		hypervisor->install_ept_hook(request.target_address, buffer, request.source_data_size,
+		hypervisor->install_ept_hook(request.target_address, buffer.get(), request.source_data_size,
 		                             translation_hints);
 	}
 
diff --git a/src/driver/type_traits.hpp b/src/driver/type_traits.hpp
index 85b0710..ea7cc21 100644
--- a/src/driver/type_traits.hpp
+++ b/src/driver/type_traits.hpp
@@ -132,4 +132,14 @@ namespace std
 	struct is_array<T[N]> : std::true_type
 	{
 	};
+
+
+	template<class T>
+	struct remove_extent { typedef T type; };
+
+	template<class T>
+	struct remove_extent<T[]> { typedef T type; };
+
+	template<class T, std::size_t N>
+	struct remove_extent<T[N]> { typedef T type; };
 }
diff --git a/src/driver/unique_ptr.hpp b/src/driver/unique_ptr.hpp
index 3faafd6..09c04cc 100644
--- a/src/driver/unique_ptr.hpp
+++ b/src/driver/unique_ptr.hpp
@@ -7,9 +7,11 @@ namespace std
 	class unique_ptr
 	{
 	public:
+		using value_type = typename remove_extent<T>::type;
+
 		unique_ptr() = default;
 
-		unique_ptr(T* pointer)
+		unique_ptr(value_type* pointer)
 			: pointer_(pointer)
 		{
 		}
@@ -44,22 +46,27 @@ namespace std
 		unique_ptr(const unique_ptr<T>& obj) = delete;
 		unique_ptr& operator=(const unique_ptr<T>& obj) = delete;
 
-		T* operator->()
+		value_type* get()
 		{
 			return this->pointer_;
 		}
 
-		const T* operator->() const
+		value_type* operator->()
 		{
 			return this->pointer_;
 		}
 
-		T& operator*()
+		const value_type* operator->() const
+		{
+			return this->pointer_;
+		}
+
+		value_type& operator*()
 		{
 			return *this->pointer_;
 		}
 
-		const T& operator*() const
+		const value_type& operator*() const
 		{
 			return *this->pointer_;
 		}
@@ -71,7 +78,7 @@ namespace std
 
 	private:
 		static constexpr auto is_array_type = is_array<T>::value;
-		T* pointer_{nullptr};
+		value_type* pointer_{nullptr};
 
 		void delete_pointer() const
 		{