From 76fcdd0d8e2d4c2a3c9006fa1b0cf1271538c875 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sun, 15 May 2022 13:07:17 +0200
Subject: [PATCH] Reflect memory writes to shadow page

---
 src/driver/ept.cpp    | 46 ++++++++++++++++++++++++++++++++++++++++---
 src/driver/ept.hpp    |  7 ++++++-
 src/driver/memory.cpp | 18 ++++++++++++++++-
 src/driver/memory.hpp |  7 +++++++
 4 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/src/driver/ept.cpp b/src/driver/ept.cpp
index 2d39b8a..8af2564 100644
--- a/src/driver/ept.cpp
+++ b/src/driver/ept.cpp
@@ -74,6 +74,44 @@ namespace vmx
 
 			return candidate_memory_type;
 		}
+
+		void update_fake_page(ept_hook& hook)
+		{
+			if(!hook.mapped_virtual_address)
+			{
+				return;
+			}
+
+			uint8_t page_copy[PAGE_SIZE];
+			memcpy(page_copy, hook.mapped_virtual_address, PAGE_SIZE);
+
+			for(size_t i = 0; i < PAGE_SIZE; ++i)
+			{
+				if(hook.diff_page[i] != page_copy[i])
+				{
+					hook.diff_page[i] = page_copy[i];
+					hook.fake_page[i] = page_copy[i];
+				}
+			}
+		}
+	}
+
+	ept_hook::ept_hook(const uint64_t physical_base)
+		: physical_base_address(physical_base)
+		  , mapped_virtual_address(memory::map_physical_memory(physical_base_address, PAGE_SIZE))
+	{
+		if (!mapped_virtual_address)
+		{
+			throw std::runtime_error("Failed to map physical memory");
+		}
+	}
+
+	ept_hook::~ept_hook()
+	{
+		if (mapped_virtual_address)
+		{
+			memory::unmap_physical_memory(mapped_virtual_address, PAGE_SIZE);
+		}
 	}
 
 	ept::ept()
@@ -178,6 +216,7 @@ namespace vmx
 
 		if (!violation_qualification.ept_executable && violation_qualification.execute_access)
 		{
+			update_fake_page(*hook);
 			hook->target_page->flags = hook->execute_entry.flags;
 			guest_context.increment_rip = false;
 		}
@@ -335,9 +374,9 @@ namespace vmx
 		return split;
 	}
 
-	ept_hook* ept::allocate_ept_hook()
+	ept_hook* ept::allocate_ept_hook(const uint64_t physical_address)
 	{
-		auto* hook = memory::allocate_aligned_object<ept_hook>();
+		auto* hook = memory::allocate_aligned_object<ept_hook>(physical_address);
 		if (!hook)
 		{
 			throw std::runtime_error("Failed to allocate ept hook object");
@@ -400,7 +439,7 @@ namespace vmx
 			return hook;
 		}
 
-		hook = this->allocate_ept_hook();
+		hook = this->allocate_ept_hook(physical_base_address);
 
 		if (!hook)
 		{
@@ -411,6 +450,7 @@ namespace vmx
 
 		const auto* data_source = translation_hint ? &translation_hint->page[0] : virtual_target;
 		memcpy(&hook->fake_page[0], data_source, PAGE_SIZE);
+		memcpy(&hook->diff_page[0], data_source, PAGE_SIZE);
 		hook->physical_base_address = physical_base_address;
 
 		hook->target_page = this->get_pml1_entry(physical_address);
diff --git a/src/driver/ept.hpp b/src/driver/ept.hpp
index 3c876cf..ee49c35 100644
--- a/src/driver/ept.hpp
+++ b/src/driver/ept.hpp
@@ -26,9 +26,14 @@ namespace vmx
 
 	struct ept_hook
 	{
+		ept_hook(const uint64_t physical_base);
+		~ept_hook();
+
 		DECLSPEC_PAGE_ALIGN uint8_t fake_page[PAGE_SIZE]{};
+		DECLSPEC_PAGE_ALIGN uint8_t diff_page[PAGE_SIZE]{};
 
 		uint64_t physical_base_address{};
+		void* mapped_virtual_address{};
 
 		pml1* target_page{};
 		pml1 original_entry{};
@@ -89,7 +94,7 @@ namespace vmx
 		pml1* find_pml1_table(uint64_t physical_address) const;
 
 		ept_split* allocate_ept_split();
-		ept_hook* allocate_ept_hook();
+		ept_hook* allocate_ept_hook(uint64_t physical_address);
 		ept_hook* find_ept_hook(uint64_t physical_address) const;
 
 		ept_hook* get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint = nullptr);
diff --git a/src/driver/memory.cpp b/src/driver/memory.cpp
index cb0fe7a..6e0b34a 100644
--- a/src/driver/memory.cpp
+++ b/src/driver/memory.cpp
@@ -1,7 +1,6 @@
 #include "std_include.hpp"
 #include "memory.hpp"
 #include "string.hpp"
-#include "process.hpp"
 
 namespace memory
 {
@@ -84,6 +83,23 @@ namespace memory
 	_Must_inspect_result_
 	_IRQL_requires_max_(DISPATCH_LEVEL)
 
+	void* map_physical_memory(const uint64_t address, const size_t size)
+	{
+		PHYSICAL_ADDRESS physical_address{};
+		physical_address.QuadPart = static_cast<LONGLONG>(address);
+		return MmMapIoSpace(physical_address, size, MmNonCached);
+	}
+
+	_IRQL_requires_max_(DISPATCH_LEVEL)
+
+	void unmap_physical_memory(void* address, const size_t size)
+	{
+		MmUnmapIoSpace(address, size);
+	}
+
+	_Must_inspect_result_
+	_IRQL_requires_max_(DISPATCH_LEVEL)
+
 	void* allocate_non_paged_memory(const size_t size)
 	{
 		void* memory = ExAllocatePoolWithTag(NonPagedPool, size, 'MOMO');
diff --git a/src/driver/memory.hpp b/src/driver/memory.hpp
index 7d39c37..3fae1c6 100644
--- a/src/driver/memory.hpp
+++ b/src/driver/memory.hpp
@@ -13,6 +13,13 @@ namespace memory
 	uint64_t get_physical_address(void* address);
 	void* get_virtual_address(uint64_t address);
 
+	_Must_inspect_result_
+	_IRQL_requires_max_(DISPATCH_LEVEL)
+	void* map_physical_memory(const uint64_t address, const size_t size);
+
+	_IRQL_requires_max_(DISPATCH_LEVEL)
+	void unmap_physical_memory(void* address, const size_t size);
+
 	_Must_inspect_result_
 	_IRQL_requires_max_(DISPATCH_LEVEL)
 	void* allocate_non_paged_memory(size_t size);