Cleanup on process termination

src/driver/driver_main.cpp

@@ -19,7 +19,7 @@ public:
 			  this->sleep_notification(type);
 		  }),
 		  process_callback_(
-			  [this](const HANDLE parent_id, const HANDLE process_id, const process_callback::type type)
+			  [this](const process_id parent_id, const process_id process_id, const process_callback::type type)
 			  {
 				  this->process_notification(parent_id, process_id, type);
 			  }),
@@ -66,16 +66,11 @@ private:
 		}
 	}
 
-	void process_notification(HANDLE /*parent_id*/, const HANDLE process_id, const process_callback::type type)
+	void process_notification(process_id /*parent_id*/, const process_id process_id, const process_callback::type type)
 	{
-		if (type == process_callback::type::create)
-		{
-			debug_log("Created process: %X\n", process_id);
-		}
-
 		if (type == process_callback::type::destroy)
 		{
-			debug_log("Destroyed process: %X\n", process_id);
+			this->hypervisor_.handle_process_termination(process_id);
 		}
 	}
 };

src/driver/ept.cpp

@@ -97,7 +97,7 @@ namespace vmx
 
 		void reset_all_watch_point_pages(utils::list<ept_code_watch_point>& watch_points)
 		{
-			for(const auto& watch_point : watch_points)
+			for (const auto& watch_point : watch_points)
 			{
 				if (watch_point.target_page)
 				{
@@ -121,6 +121,8 @@ namespace vmx
 
 	ept_hook::~ept_hook()
 	{
+		this->target_page->flags = this->original_entry.flags;
+
 		if (mapped_virtual_address)
 		{
 			memory::unmap_physical_memory(mapped_virtual_address, PAGE_SIZE);
@@ -143,10 +145,13 @@ namespace vmx
 		this->disable_all_hooks();
 	}
 
-	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, const process_id source_pid,
+	                            const process_id target_pid,
 	                            const ept_translation_hint* translation_hint)
 	{
 		auto* hook = this->get_or_create_ept_hook(destination, translation_hint);
+		hook->source_pid = source_pid;
+		hook->target_pid = target_pid;
 
 		const auto page_offset = ADDRMASK_EPT_PML1_OFFSET(reinterpret_cast<uint64_t>(destination));
 		memcpy(hook->fake_page + page_offset, source, length);
@@ -170,6 +175,7 @@ namespace vmx
 	}
 
 	void ept::install_hook(const void* destination, const void* source, const size_t length,
+	                       const process_id source_pid, const process_id target_pid,
 	                       const utils::list<ept_translation_hint>& hints)
 	{
 		auto current_destination = reinterpret_cast<uint64_t>(destination);
@@ -184,9 +190,8 @@ namespace vmx
 			const auto data_to_write = min(page_remaining, current_length);
 
 
-
 			const ept_translation_hint* relevant_hint = nullptr;
-			for(const auto& hint : hints)
+			for (const auto& hint : hints)
 			{
 				if (hint.virtual_base_address == aligned_destination)
 				{
@@ -196,7 +201,8 @@ namespace vmx
 			}
 
 			this->install_page_hook(reinterpret_cast<void*>(current_destination),
-			                        reinterpret_cast<const void*>(current_source), data_to_write, relevant_hint);
+			                        reinterpret_cast<const void*>(current_source), data_to_write, source_pid,
+			                        target_pid, relevant_hint);
 
 			current_length -= data_to_write;
 			current_destination += data_to_write;
@@ -204,12 +210,9 @@ namespace vmx
 		}
 	}
 
-	void ept::disable_all_hooks() const
+	void ept::disable_all_hooks()
 	{
-		for(auto& hook : this->ept_hooks)
+		this->ept_hooks.clear();
-		{
-			hook.target_page->flags = hook.original_entry.flags;
-		}
 	}
 
 	void ept::handle_violation(guest_context& guest_context)
@@ -333,7 +336,8 @@ namespace vmx
 		}
 	}
 
-	void ept::install_code_watch_point(const uint64_t physical_page)
+	void ept::install_code_watch_point(const uint64_t physical_page, const process_id source_pid,
+	                                   const process_id target_pid)
 	{
 		const auto physical_base_address = reinterpret_cast<uint64_t>(PAGE_ALIGN(physical_page));
 
@@ -343,6 +347,8 @@ namespace vmx
 		}
 
 		auto& watch_point = this->allocate_ept_code_watch_point();
+		watch_point.source_pid = source_pid;
+		watch_point.target_pid = target_pid;
 
 		this->split_large_page(physical_base_address);
 
@@ -420,7 +426,7 @@ namespace vmx
 
 	pml1* ept::find_pml1_table(const uint64_t physical_address)
 	{
-		for(auto& split : this->ept_splits)
+		for (auto& split : this->ept_splits)
 		{
 			if (memory::get_physical_address(&split.pml1[0]) == physical_address)
 			{
@@ -461,7 +467,7 @@ namespace vmx
 
 	ept_code_watch_point* ept::find_ept_code_watch_point(const uint64_t physical_address)
 	{
-		for(auto& watch_point : this->ept_code_watch_points)
+		for (auto& watch_point : this->ept_code_watch_points)
 		{
 			if (watch_point.physical_base_address == physical_address)
 			{
@@ -628,4 +634,37 @@ namespace vmx
 		*count = i;
 		return this->access_records;
 	}
+
+	bool ept::handle_process_termination(const process_id process)
+	{
+		bool changed = false;
+
+		for (auto i = this->ept_hooks.begin(); i != this->ept_hooks.end();)
+		{
+			if (i->source_pid == process || i->target_pid == process)
+			{
+				i = this->ept_hooks.erase(i);
+				changed = true;
+			}
+			else
+			{
+				++i;
+			}
+		}
+
+		for (auto i = this->ept_code_watch_points.begin(); i != this->ept_code_watch_points.end();)
+		{
+			if (i->source_pid == process || i->target_pid == process)
+			{
+				i = this->ept_code_watch_points.erase(i);
+				changed = true;
+			}
+			else
+			{
+				++i;
+			}
+		}
+
+		return changed;
+	}
 }

src/driver/ept.hpp

@@ -26,6 +26,8 @@ namespace vmx
 	{
 		uint64_t physical_base_address{};
 		pml1* target_page{};
+		process_id source_pid{0};
+		process_id target_pid{0};
 	};
 
 	struct ept_hook
@@ -43,6 +45,9 @@ namespace vmx
 		pml1 original_entry{};
 		pml1 execute_entry{};
 		pml1 readwrite_entry{};
+
+		process_id source_pid{0};
+		process_id target_pid{0};
 	};
 
 	struct ept_translation_hint
@@ -68,11 +73,12 @@ namespace vmx
 
 		void initialize();
 
-		void install_code_watch_point(uint64_t physical_page);
+		void install_code_watch_point(uint64_t physical_page, process_id source_pid, process_id target_pid);
 
-		void install_hook(const void* destination, const void* source, size_t length,
+		void install_hook(const void* destination, const void* source, size_t length, process_id source_pid,
+		                  process_id target_pid,
 		                  const utils::list<ept_translation_hint>& hints = {});
-		void disable_all_hooks() const;
+		void disable_all_hooks();
 
 		void handle_violation(guest_context& guest_context);
 		void handle_misconfiguration(guest_context& guest_context) const;
@@ -84,6 +90,8 @@ namespace vmx
 
 		uint64_t* get_access_records(size_t* count);
 
+		bool handle_process_termination(process_id process);
+
 	private:
 		DECLSPEC_PAGE_ALIGN pml4 epml4[EPT_PML4E_ENTRY_COUNT];
 		DECLSPEC_PAGE_ALIGN pml3 epdpt[EPT_PDPTE_ENTRY_COUNT];
@@ -110,8 +118,8 @@ namespace vmx
 
 		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, process_id source_pid,
-		                       const ept_translation_hint* translation_hint = nullptr);
+		                       process_id target_pid, const ept_translation_hint* translation_hint = nullptr);
 
 		void record_access(uint64_t rip);
 	};

src/driver/hypervisor.cpp

@@ -164,11 +164,12 @@ bool hypervisor::is_enabled() const
 }
 
 bool hypervisor::install_ept_hook(const void* destination, const void* source, const size_t length,
+                                  const process_id source_pid, const process_id target_pid,
                                   const utils::list<vmx::ept_translation_hint>& hints)
 {
 	try
 	{
-		this->ept_->install_hook(destination, source, length, hints);
+		this->ept_->install_hook(destination, source, length, source_pid, target_pid, hints);
 	}
 	catch (std::exception& e)
 	{
@@ -181,23 +182,16 @@ bool hypervisor::install_ept_hook(const void* destination, const void* source, c
 		return false;
 	}
 
-	volatile long failures = 0;
+	this->invalidate_cores();
-	thread::dispatch_on_all_cores([&]
+	return true;
-	{
-		if (!this->try_install_ept_hook_on_core(destination, source, length, hints))
-		{
-			InterlockedIncrement(&failures);
-		}
-	});
-
-	return failures == 0;
 }
 
-bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, const bool invalidate) const
+bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, const process_id source_pid,
+                                              const process_id target_pid, const bool invalidate) const
 {
 	try
 	{
-		this->ept_->install_code_watch_point(physical_page);
+		this->ept_->install_code_watch_point(physical_page, source_pid, target_pid);
 	}
 	catch (std::exception& e)
 	{
@@ -221,12 +215,13 @@ bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, cons
 	return true;
 }
 
-bool hypervisor::install_ept_code_watch_points(const uint64_t* physical_pages, const size_t count) const
+bool hypervisor::install_ept_code_watch_points(const uint64_t* physical_pages, const size_t count,
+                                               const process_id source_pid, const process_id target_pid) const
 {
 	bool success = true;
 	for (size_t i = 0; i < count; ++i)
 	{
-		success &= this->install_ept_code_watch_point(physical_pages[i], false);
+		success &= this->install_ept_code_watch_point(physical_pages[i], source_pid, target_pid, false);
 	}
 
 	thread::dispatch_on_all_cores([&]
@@ -243,7 +238,7 @@ void hypervisor::disable_all_ept_hooks() const
 
 	thread::dispatch_on_all_cores([&]
 	{
-		auto* vm_state = this->get_current_vm_state();
+		const auto* vm_state = this->get_current_vm_state();
 		if (!vm_state)
 		{
 			return;
@@ -266,6 +261,17 @@ hypervisor* hypervisor::get_instance()
 	return instance;
 }
 
+void hypervisor::handle_process_termination(const process_id process)
+{
+	if (!this->ept_->handle_process_termination(process))
+	{
+		return;
+	}
+
+	debug_log("Handled termination of %X\n", process);
+	this->invalidate_cores();
+}
+
 void hypervisor::enable()
 {
 	const auto cr3 = __readcr3();
@@ -430,7 +436,7 @@ vmx::gdt_entry convert_gdt_entry(const uint64_t gdt_base, const uint16_t selecto
 	result.access_rights.granularity = gdt_entry->granularity;
 
 	result.access_rights.reserved1 = 0;
-	result.access_rights.unusable = !gdt_entry->present;
+	result.access_rights.unusable = ~gdt_entry->present;
 
 	return result;
 }
@@ -826,45 +832,16 @@ void hypervisor::free_vm_states()
 	}
 }
 
-bool hypervisor::try_install_ept_hook_on_core(const void* destination, const void* source, const size_t length,
+void hypervisor::invalidate_cores() const
-                                              const utils::list<vmx::ept_translation_hint>& hints)
 {
-	try
+	thread::dispatch_on_all_cores([&]
 	{
-		this->install_ept_hook_on_core(destination, source, length, hints);
+		const auto* vm_state = this->get_current_vm_state();
-		return true;
+		if (vm_state && this->is_enabled())
-	}
-	catch (std::exception& e)
-	{
-		debug_log("Failed to install ept hook on core %d: %s\n", thread::get_processor_index(), e.what());
-		return false;
-	}
-	catch (...)
-	{
-		debug_log("Failed to install ept hook on core %d.\n", thread::get_processor_index());
-		return false;
-	}
-}
-
-void hypervisor::install_ept_hook_on_core(const void* destination, const void* source, const size_t length,
-                                          const utils::list<vmx::ept_translation_hint>& hints)
-{
-	auto* vm_state = this->get_current_vm_state();
-	if (!vm_state)
-	{
-		throw std::runtime_error("No vm state available");
-	}
-
-	(void)destination;
-	(void)source;
-	(void)length;
-	(void)hints;
-	//vm_state->ept->install_hook(destination, source, length, hints);
-
-	if (this->is_enabled())
 		{
 			vm_state->ept->invalidate();
 		}
+	});
 }
 
 vmx::state* hypervisor::get_current_vm_state() const

src/driver/hypervisor.hpp

@@ -19,11 +19,13 @@ 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, process_id source_pid,
-	                      const utils::list<vmx::ept_translation_hint>& hints = {});
+	                      process_id target_pid, const utils::list<vmx::ept_translation_hint>& hints = {});
 
-	bool install_ept_code_watch_point(uint64_t physical_page, bool invalidate = true) const;
+	bool install_ept_code_watch_point(uint64_t physical_page, process_id source_pid, process_id target_pid,
-	bool install_ept_code_watch_points(const uint64_t* physical_pages, size_t count) const;
+	                                  bool invalidate = true) const;
+	bool install_ept_code_watch_points(const uint64_t* physical_pages, size_t count, process_id source_pid,
+	                                   process_id target_pid) const;
 
 	void disable_all_ept_hooks() const;
 
@@ -31,6 +33,8 @@ public:
 
 	static hypervisor* get_instance();
 
+	void handle_process_termination(process_id process);
+
 private:
 	uint32_t vm_state_count_{0};
 	vmx::state** vm_states_{nullptr};
@@ -43,10 +47,7 @@ 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,
+	void invalidate_cores() const;
-	                                  const utils::list<vmx::ept_translation_hint>& hints = {});
-	void install_ept_hook_on_core(const void* destination, const void* source, size_t length,
-	                              const utils::list<vmx::ept_translation_hint>& hints = {});
 
 	vmx::state* get_current_vm_state() const;
 };

src/driver/irp.cpp

@@ -92,7 +92,7 @@ namespace
 		}
 
 		hypervisor->install_ept_hook(request.target_address, buffer.get(), request.source_data_size,
-		                             translation_hints);
+		                             process::get_current_process_id(), request.process_id, translation_hints);
 	}
 
 	void unhook()
@@ -160,7 +160,7 @@ namespace
 			throw std::runtime_error("Failed to copy buffer");
 		}
 
-		thread::kernel_thread t([watch_request_copy, hypervisor, &index, &page_buffer]
+		thread::kernel_thread t([watch_request_copy, &index, &page_buffer]
 		{
 			debug_log("Looking up process: %d\n", watch_request_copy.process_id);
 
@@ -210,7 +210,8 @@ namespace
 		t.join();
 
 		debug_log("Installing watch points...\n");
-		(void)hypervisor->install_ept_code_watch_points(page_buffer.get(), index);
+		(void)hypervisor->install_ept_code_watch_points(page_buffer.get(), index, process::get_current_process_id(),
+		                                                watch_request_copy.process_id);
 		debug_log("Watch points installed\n");
 	}
 
@@ -232,7 +233,7 @@ namespace
 
 	void get_records(const PIRP irp, const PIO_STACK_LOCATION irp_sp)
 	{
-		auto* hypervisor = hypervisor::get_instance();
+		const auto* hypervisor = hypervisor::get_instance();
 		if (!hypervisor)
 		{
 			throw std::runtime_error("Hypervisor not installed");
@@ -251,7 +252,6 @@ namespace
 		irp->IoStatus.Status = STATUS_SUCCESS;
 
 		const auto irp_sp = IoGetCurrentIrpStackLocation(irp);
-
 		if (irp_sp)
 		{
 			const auto ioctr_code = irp_sp->Parameters.DeviceIoControl.IoControlCode;

src/driver/list.hpp

@@ -36,9 +36,9 @@ namespace utils
 				return *this->entry_->entry;
 			}
 
-			T& operator->() const
+			T* operator->() const
 			{
-				return *this->entry_->entry;
+				return this->entry_->entry;
 			}
 
 			bool operator==(const iterator& i) const
@@ -94,9 +94,9 @@ namespace utils
 				return *this->entry_->entry;
 			}
 
-			const T& operator->() const
+			const T* operator->() const
 			{
-				return *this->entry_->entry;
+				return this->entry_->entry;
 			}
 
 			bool operator==(const const_iterator& i) const

src/driver/process.cpp

@@ -37,7 +37,7 @@ namespace process
 
 	process_handle::process_handle(const process_handle& obj)
 	{
-		this->operator=(std::move(obj));
+		this->operator=(obj);
 	}
 
 	process_handle& process_handle::operator=(const process_handle& obj)
@@ -80,14 +80,14 @@ namespace process
 		return KeWaitForSingleObject(this->handle_, Executive, KernelMode, FALSE, &zero_time) != STATUS_WAIT_0;
 	}
 
-	uint32_t process_handle::get_id() const
+	process_id process_handle::get_id() const
 	{
 		if (!this->handle_)
 		{
 			return 0;
 		}
 
-		return uint32_t(uint64_t(PsGetProcessId(this->handle_)));
+		return process_id_from_handle(PsGetProcessId(this->handle_));
 	}
 
 	const char* process_handle::get_image_filename() const
@@ -111,16 +111,25 @@ namespace process
 		this->own_ = false;
 	}
 
-	process_handle find_process_by_id(const uint32_t process_id)
+	process_id process_id_from_handle(HANDLE handle)
 	{
-		PEPROCESS process{};
+		return process_id(size_t(handle));
-		const uint64_t process_id_long = process_id;
+	}
-		if (PsLookupProcessByProcessId(HANDLE(process_id_long), &process) != STATUS_SUCCESS)
+
+	HANDLE handle_from_process_id(const process_id process)
+	{
+		return HANDLE(size_t(process));
+	}
+
+	process_handle find_process_by_id(const process_id process)
+	{
+		PEPROCESS process_obj{};
+		if (PsLookupProcessByProcessId(handle_from_process_id(process), &process_obj) != STATUS_SUCCESS)
 		{
 			return {};
 		}
 
-		return process_handle{process, true};
+		return process_handle{process_obj, true};
 	}
 
 	process_handle get_current_process()
@@ -128,6 +137,11 @@ namespace process
 		return process_handle{PsGetCurrentProcess(), false};
 	}
 
+	process_id get_current_process_id()
+	{
+		return get_current_process().get_id();
+	}
+
 	scoped_process_attacher::scoped_process_attacher(const process_handle& process)
 	{
 		if (!process || !process.is_alive())

src/driver/process.hpp

@@ -19,7 +19,7 @@ namespace process
 		operator PEPROCESS() const;
 
 		bool is_alive() const;
-		uint32_t get_id() const;
+		process_id get_id() const;
 
 		const char* get_image_filename() const;
 
@@ -30,9 +30,14 @@ namespace process
 		void release();
 	};
 
-	process_handle find_process_by_id(uint32_t process_id);
+	process_id process_id_from_handle(HANDLE handle);
+	HANDLE handle_from_process_id(process_id process);
+
+	process_handle find_process_by_id(process_id process);
 	process_handle get_current_process();
 
+	process_id get_current_process_id();
+
 	class scoped_process_attacher
 	{
 	public:

src/driver/process_callback.cpp

@@ -1,5 +1,6 @@
 #include "std_include.hpp"
 #include "process_callback.hpp"
+#include "process.hpp"
 #include "list.hpp"
 #include "logging.hpp"
 
@@ -13,12 +14,13 @@ namespace process_callback
 			return list;
 		}
 
-		void process_notification_callback(const HANDLE parent_id, const HANDLE process_id, const BOOLEAN create)
+		void process_notification_callback(const HANDLE parent, const HANDLE process, const BOOLEAN create)
 		{
 			const auto& list = get_callback_list();
 			for (const auto& callback : list)
 			{
-				callback(parent_id, process_id, create == FALSE ? type::destroy : type::create);
+				callback(process::process_id_from_handle(parent), process::process_id_from_handle(process),
+				         create == FALSE ? type::destroy : type::create);
 			}
 		}
 

src/driver/process_callback.hpp

@@ -9,7 +9,7 @@ namespace process_callback
 		destroy,
 	};
 
-	using callback = void(HANDLE parent_id, HANDLE process_id, type type);
+	using callback = void(process_id parent_id, process_id process_id, type type);
 	using callback_function = std::function<callback>;
 
 	void* add(callback_function callback);

src/driver/std_include.hpp

@@ -16,3 +16,6 @@
 #include "nt_ext.hpp"
 #include "new.hpp"
 #include "exception.hpp"
+
+// Not sure if this is good, but fuck it.
+using process_id = uint32_t;