Add ept translation hints

src/driver/ept.cpp

@@ -2,6 +2,8 @@
 #include "ept.hpp"
 
 #include "assembly.hpp"
+#include "finally.hpp"
+#include "logging.hpp"
 #include "memory.hpp"
 #include "vmx.hpp"
 
@@ -127,15 +129,17 @@ namespace vmx
 		}
 	}
 
-	void ept::install_page_hook(void* destination, const void* source, const size_t length)
+	void ept::install_page_hook(void* destination, const void* source, const size_t length,
+	                            ept_translation_hint* translation_hint)
 	{
-		auto* hook = this->get_or_create_ept_hook(destination);
+		auto* hook = this->get_or_create_ept_hook(destination, translation_hint);
 
 		const auto page_offset = ADDRMASK_EPT_PML1_OFFSET(reinterpret_cast<uint64_t>(destination));
 		memcpy(hook->fake_page + page_offset, source, length);
 	}
 
-	void ept::install_hook(const void* destination, const void* source, const size_t length)
+	void ept::install_hook(const void* destination, const void* source, const size_t length,
+	                       ept_translation_hint* translation_hint)
 	{
 		auto current_destination = reinterpret_cast<uint64_t>(destination);
 		auto current_source = reinterpret_cast<uint64_t>(source);
@@ -143,12 +147,26 @@ namespace vmx
 
 		while (current_length != 0)
 		{
+			const auto aligned_destination = PAGE_ALIGN(current_destination);
 			const auto page_offset = ADDRMASK_EPT_PML1_OFFSET(current_destination);
 			const auto page_remaining = PAGE_SIZE - page_offset;
 			const auto data_to_write = min(page_remaining, current_length);
 
+			ept_translation_hint* relevant_hint = nullptr;
+			ept_translation_hint* current_hint = translation_hint;
+			while (current_hint)
+			{
+				if (current_hint->virtual_base_address == aligned_destination)
+				{
+					relevant_hint = current_hint;
+					break;
+				}
+
+				current_hint = current_hint->next_hint;
+			}
+
 			this->install_page_hook(reinterpret_cast<void*>(current_destination),
-			                        reinterpret_cast<const void*>(current_source), data_to_write);
+			                        reinterpret_cast<const void*>(current_source), data_to_write, relevant_hint);
 
 			current_length -= data_to_write;
 			current_destination += data_to_write;
@@ -381,10 +399,20 @@ namespace vmx
 		return nullptr;
 	}
 
-	ept_hook* ept::get_or_create_ept_hook(void* destination)
+	ept_hook* ept::get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint)
 	{
 		const auto virtual_target = PAGE_ALIGN(destination);
-		const auto physical_address = memory::get_physical_address(virtual_target);
+
+		uint64_t physical_address = 0;
+		if (translation_hint)
+		{
+			physical_address = translation_hint->physical_base_address + ADDRMASK_EPT_PML1_OFFSET(
+				reinterpret_cast<uint64_t>(destination));
+		}
+		else
+		{
+			physical_address = memory::get_physical_address(virtual_target);
+		}
 
 		if (!physical_address)
 		{
@@ -412,7 +440,9 @@ namespace vmx
 
 		this->split_large_page(physical_address);
 
-		memcpy(&hook->fake_page[0], virtual_target, PAGE_SIZE);
+		const auto* data_source = translation_hint ? &translation_hint->page[0] : virtual_target;
+
+		memcpy(&hook->fake_page[0], data_source, PAGE_SIZE);
 		hook->physical_base_address = physical_base_address;
 
 		hook->target_page = this->get_pml1_entry(physical_address);
@@ -480,4 +510,61 @@ namespace vmx
 
 		target_entry->flags = new_pointer.flags;
 	}
+
+	ept_translation_hint* ept::generate_translation_hints(const void* destination, const size_t length)
+	{
+		auto current_destination = reinterpret_cast<uint64_t>(destination);
+		auto current_length = length;
+
+		 ept_translation_hint* current_hints = nullptr;
+
+		auto destructor = utils::finally([&current_hints]()
+		{
+			ept::free_translation_hints(current_hints);
+		});
+
+		while (current_length != 0)
+		{
+			const auto aligned_destination = PAGE_ALIGN(current_destination);
+			const auto page_offset = ADDRMASK_EPT_PML1_OFFSET(current_destination);
+			const auto page_remaining = PAGE_SIZE - page_offset;
+			const auto data_to_write = min(page_remaining, current_length);
+
+			auto* new_hint = memory::allocate_non_paged_object<ept_translation_hint>();
+			if(!new_hint)
+			{
+				throw std::runtime_error("Failed to allocate hint");
+			}
+
+			new_hint->next_hint = current_hints;
+			current_hints = new_hint;
+			current_hints->virtual_base_address = aligned_destination;
+			current_hints->physical_base_address = memory::get_physical_address(aligned_destination);
+
+			if(!current_hints->physical_base_address)
+			{
+				throw std::runtime_error("Failed to resolve physical address");
+			}
+
+			memcpy(&current_hints->page[0], aligned_destination, PAGE_SIZE);			
+
+			current_length -= data_to_write;
+			current_destination += data_to_write;
+		}
+
+		destructor.cancel();
+
+		return current_hints;
+	}
+
+	void ept::free_translation_hints(ept_translation_hint* hints)
+	{
+		auto* hint = hints;
+		while (hint)
+		{
+			auto* current_hint = hint;
+			hint = hint->next_hint;
+			memory::free_non_paged_object(current_hint);
+		}
+	}
 }

src/driver/ept.hpp

@@ -38,6 +38,16 @@ namespace vmx
 		ept_hook* next_hook{nullptr};
 	};
 
+	struct ept_translation_hint
+	{
+		DECLSPEC_PAGE_ALIGN uint8_t page[PAGE_SIZE]{};
+
+		uint64_t physical_base_address{};
+		const void* virtual_base_address{};
+
+		ept_translation_hint* next_hint{nullptr};
+	};
+
 	struct guest_context;
 
 	class ept
@@ -53,7 +63,7 @@ namespace vmx
 
 		void initialize();
 
-		void install_hook(const void* destination, const void* source, size_t length);
+		void install_hook(const void* destination, const void* source, size_t length, ept_translation_hint* translation_hint = nullptr);
 		void disable_all_hooks() const;
 
 		void handle_violation(guest_context& guest_context) const;
@@ -62,6 +72,9 @@ namespace vmx
 		ept_pointer get_ept_pointer() const;
 		void invalidate() const;
 
+		static ept_translation_hint* generate_translation_hints(const void* destination, size_t length);
+		static void free_translation_hints(ept_translation_hint* hints);
+
 	private:
 		DECLSPEC_PAGE_ALIGN pml4 epml4[EPT_PML4E_ENTRY_COUNT]{};
 		DECLSPEC_PAGE_ALIGN pml3 epdpt[EPT_PDPTE_ENTRY_COUNT]{};
@@ -78,10 +91,10 @@ namespace vmx
 		ept_hook* allocate_ept_hook();
 		ept_hook* find_ept_hook(uint64_t physical_address) const;
 
-		ept_hook* get_or_create_ept_hook(void* destination);
+		ept_hook* get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint = nullptr);
 
 		void split_large_page(uint64_t physical_address);
 
-		void install_page_hook(void* destination, const void* source, size_t length);
+		void install_page_hook(void* destination, const void* source, size_t length, ept_translation_hint* translation_hint = nullptr);
 	};
 }

src/driver/hypervisor.cpp

@@ -159,12 +159,12 @@ bool hypervisor::is_enabled() const
 	return is_hypervisor_present();
 }
 
-bool hypervisor::install_ept_hook(const void* destination, const void* source, const size_t length)
+bool hypervisor::install_ept_hook(const void* destination, const void* source, const size_t length, vmx::ept_translation_hint* translation_hint)
 {
 	volatile long failures = 0;
 	thread::dispatch_on_all_cores([&]()
 	{
-		if (!this->try_install_ept_hook_on_core(destination, source, length))
+		if (!this->try_install_ept_hook_on_core(destination, source, length, translation_hint))
 		{
 			InterlockedIncrement(&failures);
 		}
@@ -1005,11 +1005,11 @@ void hypervisor::free_vm_states()
 	this->vm_state_count_ = 0;
 }
 
-bool hypervisor::try_install_ept_hook_on_core(const void* destination, const void* source, const size_t length)
+bool hypervisor::try_install_ept_hook_on_core(const void* destination, const void* source, const size_t length, vmx::ept_translation_hint* translation_hint)
 {
 	try
 	{
-		this->install_ept_hook_on_core(destination, source, length);
+		this->install_ept_hook_on_core(destination, source, length, translation_hint);
 		return true;
 	}
 	catch (std::exception& e)
@@ -1024,7 +1024,7 @@ bool hypervisor::try_install_ept_hook_on_core(const void* destination, const voi
 	}
 }
 
-void hypervisor::install_ept_hook_on_core(const void* destination, const void* source, const size_t length)
+void hypervisor::install_ept_hook_on_core(const void* destination, const void* source, const size_t length, vmx::ept_translation_hint* translation_hint)
 {
 	auto* vm_state = this->get_current_vm_state();
 	if (!vm_state)
@@ -1032,7 +1032,7 @@ void hypervisor::install_ept_hook_on_core(const void* destination, const void* s
 		throw std::runtime_error("No vm state available");
 	}
 
-	vm_state->ept.install_hook(destination, source, length);
+	vm_state->ept.install_hook(destination, source, length, translation_hint);
 
 	if (this->is_enabled())
 	{

src/driver/hypervisor.hpp

@@ -19,7 +19,7 @@ public:
 
 	bool is_enabled() const;
 
-	bool install_ept_hook(const void* destination, const void* source, size_t length);
+	bool install_ept_hook(const void* destination, const void* source, size_t length, vmx::ept_translation_hint* translation_hint = nullptr);
 	void disable_all_ept_hooks() const;
 
 	static hypervisor* get_instance();
@@ -35,8 +35,8 @@ private:
 	void allocate_vm_states();
 	void free_vm_states();
 
-	bool try_install_ept_hook_on_core(const void* destination, const void* source, size_t length);
+	bool try_install_ept_hook_on_core(const void* destination, const void* source, size_t length, vmx::ept_translation_hint* translation_hint = nullptr);
-	void install_ept_hook_on_core(const void* destination, const void* source, size_t length);
+	void install_ept_hook_on_core(const void* destination, const void* source, size_t length, vmx::ept_translation_hint* translation_hint = nullptr);
 
 	vmx::state* get_current_vm_state() const;
 };

src/driver/irp.cpp

@@ -38,9 +38,24 @@ namespace
 	}
 
 	// TODO: This is vulnerable as fuck. Optimize!
-	void apply_hook(hook_request* request)
+	void apply_hook(const hook_request* request)
 	{
-		thread::kernel_thread t([r = *request]()
+		auto* 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, &buffer]()
+		{
+			delete[] buffer;
+			vmx::ept::free_translation_hints(translation_hints);
+		});
+
+		memcpy(buffer, request->source_data, request->source_data_size);
+
+		thread::kernel_thread t([&translation_hints, r = *request]()
 		{
 			debug_log("Pid: %d | Address: %p\n", r.process_id, r.target_address);
 
@@ -59,36 +74,22 @@ namespace
 
 			debug_log("Level: %d\n", static_cast<int>(KeGetCurrentIrql()));
 
-			/*
-			auto buffer = new uint8_t[r.source_data_size];
-			if (!buffer)
-			{
-				debug_log("Failed to allocate buffer\n");
-				return;
-			}
-
-			auto destructor = utils::finally([buffer]()
-			{
-				delete[] buffer;
-			});
-
-			memcpy(buffer, r.source_data, r.source_data_size);
-			*/
 
 			process::scoped_process_attacher attacher{process_handle};
-
+			translation_hints = vmx::ept::generate_translation_hints(r.target_address, r.source_data_size);
-			debug_log("Original: %p\n", r.target_address);
-
-			uint8_t buffer = 0xEB;
-
-			//hypervisor::get_instance()->install_ept_hook(r.target_address, buffer, r.source_data_size);
-			hypervisor::get_instance()->install_ept_hook(r.target_address, &buffer, 1);
-
-			debug_log("Done1\n");
 		});
 
 		t.join();
-		debug_log("Done\n");
+
+		if(!translation_hints)
+		{
+			debug_log("Failed to generate tranlsation hints");
+			return;
+		}
+
+		hypervisor::get_instance()->install_ept_hook(request->target_address, buffer, request->source_data_size, translation_hints);
+
+		debug_log("Done1\n");
 	}
 
 	_Function_class_(DRIVER_DISPATCH) NTSTATUS io_ctl_handler(

src/driver/memory.hpp

@@ -25,6 +25,28 @@ namespace memory
 
 	void copy_physical_data(uint64_t address, void* destination, size_t length);
 
+	template <typename T, typename... Args>
+	T* allocate_non_paged_object(Args ... args)
+	{
+		auto* object = static_cast<T*>(allocate_non_paged_memory(sizeof(T)));
+		if (object)
+		{
+			new(object) T(std::forward<Args>(args)...);
+		}
+
+		return object;
+	}
+
+	template <typename T>
+	void free_non_paged_object(T* object)
+	{
+		if (object)
+		{
+			object->~T();
+			free_non_paged_memory(object);
+		}
+	}
+
 	template <typename T, typename... Args>
 	T* allocate_aligned_object(Args ... args)
 	{

src/driver/thread.cpp

@@ -78,8 +78,13 @@ namespace thread
 			{
 				function();
 			}
+			catch (std::exception& e)
+			{
+				debug_log("Kernel thread threw an exception: %s\n", e.what());
+			}
 			catch (...)
 			{
+				debug_log("Kernel thread threw an unknown exception\n");
 			}
 		}
 	}