Fix watch points

src/driver/ept.cpp

@@ -102,6 +102,7 @@ namespace vmx
 				if (watch_point->target_page)
 				{
 					watch_point->target_page->read_access = 0;
+					watch_point->target_page->write_access = 0;
 					watch_point->target_page->execute_access = 1;
 				}
 
@@ -162,7 +163,7 @@ namespace vmx
 		{
 			auto* current_watch_point = watch_point;
 			watch_point = watch_point->next_watch_point;
-			memory::free_aligned_object(current_watch_point);
+			memory::free_non_paged_object(current_watch_point);
 		}
 	}
 
@@ -177,16 +178,6 @@ namespace vmx
 
 	void ept::record_access(const uint64_t rip)
 	{
-		const auto _ = utils::finally([&]
-		{
-			InterlockedExchange(&this->access_records_barrier, 0);
-		});
-
-		// Aaaahhh, fuck that xD
-		while (InterlockedExchange(&this->access_records_barrier, 1))
-		{
-		}
-
 		for (unsigned long long& access_record : this->access_records)
 		{
 			if (access_record == 0)
@@ -271,15 +262,21 @@ namespace vmx
 			{
 				watch_point->target_page->execute_access = 1;
 				watch_point->target_page->read_access = 0;
+				watch_point->target_page->write_access = 0;
 				guest_context.increment_rip = false;
 			}
-
-			if (violation_qualification.ept_executable && violation_qualification.read_access)
+			else if (violation_qualification.ept_executable && (violation_qualification.read_access ||
+				violation_qualification.
+				write_access))
 			{
 				watch_point->target_page->execute_access = 0;
 				watch_point->target_page->read_access = 1;
+				watch_point->target_page->write_access = 1;
 				guest_context.increment_rip = false;
-				this->record_access(guest_context.guest_rip);
+				if (violation_qualification.read_access)
+				{
+					this->record_access(guest_context.guest_rip);
+				}
 			}
 
 			return;
@@ -310,6 +307,7 @@ namespace vmx
 
 	void ept::handle_misconfiguration(guest_context& guest_context) const
 	{
+		// We can actually not recover from this, but this should not occur anyways
 		guest_context.increment_rip = false;
 		guest_context.exit_vm = true;
 	}
@@ -379,12 +377,14 @@ namespace vmx
 
 		this->split_large_page(physical_base_address);
 
+		watch_point->physical_base_address = physical_base_address;
 		watch_point->target_page = this->get_pml1_entry(physical_base_address);
 		if (!watch_point->target_page)
 		{
 			throw std::runtime_error("Failed to get PML1 entry for target address");
 		}
 
+		watch_point->target_page->write_access = 0;
 		watch_point->target_page->read_access = 0;
 	}
 
@@ -511,7 +511,7 @@ namespace vmx
 
 	ept_code_watch_point* ept::allocate_ept_code_watch_point()
 	{
-		auto* watch_point = memory::allocate_aligned_object<ept_code_watch_point>();
+		auto* watch_point = memory::allocate_non_paged_object<ept_code_watch_point>();
 		if (!watch_point)
 		{
 			throw std::runtime_error("Failed to allocate ept watch point object");

src/driver/ept.hpp

@@ -97,7 +97,6 @@ namespace vmx
 		DECLSPEC_PAGE_ALIGN pml2 epde[EPT_PDPTE_ENTRY_COUNT][EPT_PDE_ENTRY_COUNT];
 
 		uint64_t access_records[1024];
-		volatile long access_records_barrier{0};
 
 		ept_split* ept_splits{nullptr};
 		ept_hook* ept_hooks{nullptr};

src/driver/hypervisor.cpp

@@ -192,7 +192,7 @@ bool hypervisor::install_ept_hook(const void* destination, const void* source, c
 	return failures == 0;
 }
 
-bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page) const
+bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, bool invalidate) const
 {
 	try
 	{
@@ -209,12 +209,31 @@ bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page) cons
 		return false;
 	}
 
+	if (invalidate)
+	{
+		thread::dispatch_on_all_cores([&]
+		{
+			this->ept_->invalidate();
+		});
+	}
+
+	return true;
+}
+
+bool hypervisor::install_ept_code_watch_points(const uint64_t* physical_pages, const size_t count) const
+{
+	bool success = true;
+	for (size_t i = 0; i < count; ++i)
+	{
+		success &= this->install_ept_code_watch_point(physical_pages[i], false);
+	}
+
 	thread::dispatch_on_all_cores([&]
 	{
 		this->ept_->invalidate();
 	});
 
-	return true;
+	return success;
 }
 
 void hypervisor::disable_all_ept_hooks() const
@@ -534,14 +553,13 @@ extern "C" [[ noreturn ]] void vm_exit_handler(CONTEXT* context)
 	if (guest_context.exit_vm)
 	{
 		context->Rcx = 0x43434343;
-		restore_descriptor_tables(vm_state->launch_context);
-
-		__writecr3(read_vmx(VMCS_GUEST_CR3));
-
 		context->Rsp = guest_context.guest_rsp;
 		context->Rip = guest_context.guest_rip;
 		context->EFlags = static_cast<uint32_t>(guest_context.guest_e_flags);
 
+		restore_descriptor_tables(vm_state->launch_context);
+
+		__writecr3(read_vmx(VMCS_GUEST_CR3));
 		__vmx_off();
 	}
 	else

src/driver/hypervisor.hpp

@@ -22,7 +22,8 @@ public:
 	bool install_ept_hook(const void* destination, const void* source, size_t length,
 	                      vmx::ept_translation_hint* translation_hint = nullptr);
 
-	bool install_ept_code_watch_point(uint64_t physical_page) const;
+	bool install_ept_code_watch_point(uint64_t physical_page, bool invalidate = true) const;
+	bool install_ept_code_watch_points(const uint64_t* physical_pages, size_t count) const;
 
 	void disable_all_ept_hooks() const;
 

src/driver/irp.cpp

@@ -148,7 +148,27 @@ namespace
 		auto watch_request_copy = watch_request;
 		watch_request_copy.watch_regions = buffer.get();
 
-		thread::kernel_thread t([watch_request_copy, hypervisor]
+		size_t page_count = 0;
+		for (size_t i = 0; i < watch_request_copy.watch_region_count; ++i)
+		{
+			const auto& watch_region = watch_request_copy.watch_regions[i];
+
+			auto start = static_cast<const uint8_t*>(watch_region.virtual_address);
+			auto end = start + watch_region.length;
+
+			start = static_cast<const uint8_t*>(PAGE_ALIGN(start));
+			end = static_cast<const uint8_t*>(PAGE_ALIGN(reinterpret_cast<uint64_t>(end) + (PAGE_SIZE - 1)));
+			page_count += (end - start) / PAGE_SIZE;
+		}
+
+		volatile long index = 0;
+		std::unique_ptr<uint64_t[]> page_buffer(new uint64_t[page_count]);
+		if (!page_buffer)
+		{
+			throw std::runtime_error("Failed to copy buffer");
+		}
+
+		thread::kernel_thread t([watch_request_copy, hypervisor, &index, &page_buffer]
 		{
 			debug_log("Looking up process: %d\n", watch_request_copy.process_id);
 
@@ -182,8 +202,10 @@ namespace
 					const auto physical_address = memory::get_physical_address(const_cast<uint8_t*>(current));
 					if (physical_address)
 					{
-						debug_log("Watching %p -> %llX\n", current, physical_address);
-						(void)hypervisor->install_ept_code_watch_point(physical_address);
+						debug_log("Resolved %p -> %llX\n", current, physical_address);
+						page_buffer.get()[index] = physical_address;
+						InterlockedIncrement(&index);
+						//(void)hypervisor->install_ept_code_watch_point(physical_address);
 					}
 					else
 					{
@@ -194,6 +216,10 @@ namespace
 		});
 
 		t.join();
+
+		debug_log("Installing watch points...\n");
+		(void)hypervisor->install_ept_code_watch_points(page_buffer.get(), index);
+		debug_log("Watch points installed\n");
 	}
 
 	void try_watch_regions(const PIO_STACK_LOCATION irp_sp)