Prepare code watching

src/driver/ept.cpp

@@ -77,7 +77,7 @@ namespace vmx
 
 		void update_fake_page(ept_hook& hook)
 		{
-			if(!hook.mapped_virtual_address)
+			if (!hook.mapped_virtual_address)
 			{
 				return;
 			}
@@ -85,9 +85,9 @@ namespace vmx
 			uint8_t page_copy[PAGE_SIZE];
 			memcpy(page_copy, hook.mapped_virtual_address, PAGE_SIZE);
 
-			for(size_t i = 0; i < PAGE_SIZE; ++i)
+			for (size_t i = 0; i < PAGE_SIZE; ++i)
 			{
-				if(hook.diff_page[i] != page_copy[i])
+				if (hook.diff_page[i] != page_copy[i])
 				{
 					hook.diff_page[i] = page_copy[i];
 					hook.fake_page[i] = page_copy[i];
@@ -121,6 +121,8 @@ namespace vmx
 		memset(this->epml4, 0, sizeof(this->epml4));
 		memset(this->epdpt, 0, sizeof(this->epdpt));
 		memset(this->epde, 0, sizeof(this->epde));
+
+		memset(this->access_records, 0, sizeof(this->access_records));
 	}
 
 	ept::~ept()
@@ -140,6 +142,14 @@ namespace vmx
 			hook = hook->next_hook;
 			memory::free_aligned_object(current_hook);
 		}
+
+		auto* watch_point = this->ept_code_watch_points;
+		while (watch_point)
+		{
+			auto* current_watch_point = watch_point;
+			watch_point = watch_point->next_watch_point;
+			memory::free_aligned_object(current_watch_point);
+		}
 	}
 
 	void ept::install_page_hook(void* destination, const void* source, const size_t length,
@@ -151,6 +161,33 @@ namespace vmx
 		memcpy(hook->fake_page + page_offset, source, length);
 	}
 
+	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)
+			{
+				access_record = rip;
+				return;
+			}
+
+			if (access_record == rip)
+			{
+				return;
+			}
+		}
+	}
+
 	void ept::install_hook(const void* destination, const void* source, const size_t length,
 	                       ept_translation_hint* translation_hint)
 	{
@@ -197,7 +234,7 @@ namespace vmx
 		}
 	}
 
-	void ept::handle_violation(guest_context& guest_context) const
+	void ept::handle_violation(guest_context& guest_context)
 	{
 		vmx_exit_qualification_ept_violation violation_qualification{};
 		violation_qualification.flags = guest_context.exit_qualification;
@@ -208,6 +245,32 @@ namespace vmx
 		}
 
 		const auto physical_base_address = reinterpret_cast<uint64_t>(PAGE_ALIGN(guest_context.guest_physical_address));
+
+		// watch-point stuff
+
+		auto* watch_point = this->find_ept_code_watch_point(physical_base_address);
+		if (watch_point)
+		{
+			if (!violation_qualification.ept_executable && violation_qualification.execute_access)
+			{
+				watch_point->target_page->execute_access = 1;
+				watch_point->target_page->read_access = 0;
+				guest_context.increment_rip = false;
+			}
+
+			if (violation_qualification.ept_executable && violation_qualification.read_access)
+			{
+				watch_point->target_page->execute_access = 0;
+				watch_point->target_page->read_access = 1;
+				guest_context.increment_rip = false;
+				this->record_access(guest_context.guest_rip);
+			}
+
+			return;
+		}
+
+		// ept-hooking stuff
+
 		auto* hook = this->find_ept_hook(physical_base_address);
 		if (!hook)
 		{
@@ -283,6 +346,30 @@ namespace vmx
 		}
 	}
 
+	void ept::install_code_watch_point(const uint64_t physical_page)
+	{
+		const auto physical_base_address = reinterpret_cast<uint64_t>(PAGE_ALIGN(physical_page));
+
+		if (this->find_ept_code_watch_point(physical_base_address))
+		{
+			return;
+		}
+
+		auto* watch_point = this->allocate_ept_code_watch_point();
+		if (!watch_point)
+		{
+			throw std::runtime_error("Failed to allocate watch point");
+		}
+
+		this->split_large_page(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");
+		}
+	}
+
 	ept_pointer ept::get_ept_pointer() const
 	{
 		const auto ept_pml4_physical_address = memory::get_physical_address(const_cast<pml4*>(&this->epml4[0]));
@@ -404,6 +491,36 @@ namespace vmx
 		return nullptr;
 	}
 
+	ept_code_watch_point* ept::allocate_ept_code_watch_point()
+	{
+		auto* watch_point = memory::allocate_aligned_object<ept_code_watch_point>();
+		if (!watch_point)
+		{
+			throw std::runtime_error("Failed to allocate ept watch point object");
+		}
+
+		watch_point->next_watch_point = this->ept_code_watch_points;
+		this->ept_code_watch_points = watch_point;
+
+		return watch_point;
+	}
+
+	ept_code_watch_point* ept::find_ept_code_watch_point(const uint64_t physical_address) const
+	{
+		auto* watch_point = this->ept_code_watch_points;
+		while (watch_point)
+		{
+			if (watch_point->physical_base_address == physical_address)
+			{
+				return watch_point;
+			}
+
+			watch_point = watch_point->next_watch_point;
+		}
+
+		return nullptr;
+	}
+
 	ept_hook* ept::get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint)
 	{
 		const auto virtual_target = PAGE_ALIGN(destination);
@@ -575,4 +692,17 @@ namespace vmx
 			memory::free_non_paged_object(current_hint);
 		}
 	}
+
+	uint64_t* ept::get_access_records(size_t* count)
+	{
+		size_t i = 0;
+		for (const auto& record : this->access_records)
+		{
+			if (record == 0) break;
+			++i;
+		}
+
+		*count = i;
+		return this->access_records;
+	}
 }

src/driver/ept.hpp

@@ -23,6 +23,12 @@ namespace vmx
 		ept_split* next_split{nullptr};
 	};
 
+	struct ept_code_watch_point
+	{
+		uint64_t physical_base_address{};
+		pml1* target_page{};
+		ept_code_watch_point* next_watch_point{nullptr};
+	};
 
 	struct ept_hook
 	{
@@ -68,11 +74,13 @@ namespace vmx
 
 		void initialize();
 
+		void install_code_watch_point(uint64_t physical_page);
+
 		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;
+		void handle_violation(guest_context& guest_context);
 		void handle_misconfiguration(guest_context& guest_context) const;
 
 		ept_pointer get_ept_pointer() const;
@@ -81,13 +89,19 @@ namespace vmx
 		static ept_translation_hint* generate_translation_hints(const void* destination, size_t length);
 		static void free_translation_hints(ept_translation_hint* hints);
 
+		uint64_t* get_access_records(size_t* count);
+
 	private:
 		DECLSPEC_PAGE_ALIGN pml4 epml4[EPT_PML4E_ENTRY_COUNT];
 		DECLSPEC_PAGE_ALIGN pml3 epdpt[EPT_PDPTE_ENTRY_COUNT];
 		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};
+		ept_code_watch_point* ept_code_watch_points{nullptr};
 
 		pml2* get_pml2_entry(uint64_t physical_address);
 		pml1* get_pml1_entry(uint64_t physical_address);
@@ -97,11 +111,16 @@ namespace vmx
 		ept_hook* allocate_ept_hook(uint64_t physical_address);
 		ept_hook* find_ept_hook(uint64_t physical_address) const;
 
+		ept_code_watch_point* allocate_ept_code_watch_point();
+		ept_code_watch_point* find_ept_code_watch_point(uint64_t physical_address) const;
+
 		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,
 		                       ept_translation_hint* translation_hint = nullptr);
+
+		void record_access(uint64_t rip);
 	};
 }

src/runner/CMakeLists.txt

@@ -1,7 +1,7 @@
 file(GLOB_RECURSE runner_sources ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
 file(GLOB_RECURSE  runner_headers ${CMAKE_CURRENT_SOURCE_DIR}/*.hpp)
 
-add_executable(runner WIN32
+add_executable(runner #WIN32
 	${runner_sources}
 	${runner_headers}
 )

src/runner/main.cpp

@@ -1,15 +1,20 @@
+#include "std_include.hpp"
+
 #include <iostream>
 #include <filesystem>
 #include <conio.h>
 #include <fstream>
 
-#include "std_include.hpp"
+
 #include "driver.hpp"
 #include "driver_device.hpp"
+#include "process.hpp"
 
 #include <irp_data.hpp>
 
 #include "resource.hpp"
+#include "utils/io.hpp"
+#include "utils/nt.hpp"
 
 #pragma comment(lib, "Shlwapi.lib")
 
@@ -74,21 +79,96 @@ std::filesystem::path extract_driver()
 	return driver_file;
 }
 
+std::vector<std::pair<size_t, size_t>> find_executable_regions(const std::string& pe_file)
+{
+	std::string data{};
+	if (!utils::io::read_file(pe_file, &data))
+	{
+		return {};
+	}
+
+	const utils::nt::library library(reinterpret_cast<HMODULE>(data.data()));
+	if (!library.is_valid())
+	{
+		return {};
+	}
+
+	const auto section_headers = library.get_section_headers();
+
+	std::vector<std::pair<size_t, size_t>> regions{};
+	for (const auto& section_header : section_headers)
+	{
+		if (section_header->Characteristics & IMAGE_SCN_MEM_EXECUTE)
+		{
+			regions.emplace_back(section_header->VirtualAddress, section_header->Misc.VirtualSize);
+		}
+	}
+
+	return regions;
+}
+
 void unsafe_main(const int /*argc*/, char* /*argv*/[])
 {
-	const auto driver_file = extract_driver();
-
-	driver driver{driver_file, "MomoLul"};
-	const driver_device driver_device{R"(\\.\HelloDev)"};
-
-	//launcher().run();
-
 	std::string pid_str{};
 	printf("Please enter the pid: ");
 	std::getline(std::cin, pid_str);
 
 	const auto pid = atoi(pid_str.data());
 
+	printf("Opening process...\n");
+	const auto proc = process::open(pid, PROCESS_QUERY_INFORMATION | PROCESS_VM_READ);
+	if (!proc)
+	{
+		printf("Failed to open process...\n");
+		return;
+	}
+
+	printf("Reading modules...\n");
+	const auto modules = process::get_modules(proc);
+	printf("Found %zu modules\n", modules.size());
+
+	std::vector<std::string> module_files{};
+	module_files.reserve(modules.size());
+
+	int i = 0;
+	for (const auto& module : modules)
+	{
+		auto name = process::get_module_filename(proc, module);
+		printf("(%i)\t%p: %s\n", i++, static_cast<void*>(module), name.data());
+		module_files.emplace_back(std::move(name));
+	}
+
+	std::string module_str{};
+	printf("\nPlease enter the module number: ");
+	std::getline(std::cin, module_str);
+
+	const auto module_num = atoi(module_str.data());
+
+	if (module_num < 0 || static_cast<size_t>(module_num) >= modules.size())
+	{
+		printf("Invalid module num\n");
+		_getch();
+		return;
+	}
+
+	const auto module_base = reinterpret_cast<uint8_t*>(modules[module_num]);
+	const auto& file = module_files[module_num];
+	printf("Analyzing %s...\n", file.data());
+	const auto regions = find_executable_regions(file);
+
+	for (const auto& region : regions)
+	{
+		printf("%p - %zu\n", module_base + region.first, region.second);
+	}
+
+	_getch();
+	return;
+
+	const auto driver_file = extract_driver();
+
+	driver driver{driver_file, "MomoLul"};
+	const driver_device driver_device{R"(\\.\HelloDev)"};
+
 	/*
 	// IW5
 	insert_nop(driver_device, pid, 0x4488A8, 2); // Force calling CG_DrawFriendOrFoeTargetBoxes

src/runner/process.cpp

@@ -0,0 +1,81 @@
+#include "std_include.hpp"
+#include "process.hpp"
+
+namespace process
+{
+	native_handle open(const uint32_t process_id, const DWORD access)
+	{
+		const auto handle = ::OpenProcess(access, FALSE, process_id);
+		if (handle)
+		{
+			return {handle};
+		}
+
+		return {};
+	}
+
+	std::vector<HMODULE> get_modules(const native_handle& process)
+	{
+		if (!process)
+		{
+			return {};
+		}
+
+		DWORD needed = 1024;
+		std::vector<HMODULE> result{};
+
+		do
+		{
+			result.resize(needed);
+			if (!EnumProcessModulesEx(process, result.data(), static_cast<DWORD>(result.size()), &needed,
+			                          LIST_MODULES_ALL))
+			{
+				return {};
+			}
+		}
+		while (needed > result.size());
+
+		result.resize(needed);
+
+		// Remove duplicates
+		std::ranges::sort(result);
+		const auto last = std::ranges::unique(result).begin();
+		result.erase(last, result.end());
+
+		// Remove nullptr
+		for (auto i = result.begin(); i != result.end();)
+		{
+			if (*i == nullptr)
+			{
+				i = result.erase(i);
+				break;
+			}
+			else
+			{
+				++i;
+			}
+		}
+
+		return result;
+	}
+
+	std::string get_module_filename(const native_handle& process, const HMODULE module)
+	{
+		if (!process)
+		{
+			return {};
+		}
+
+		std::string buffer{};
+		buffer.resize(1024);
+
+		const auto length = GetModuleFileNameExA(process, module, &buffer[0], static_cast<DWORD>(buffer.size()));
+		if (length > 0)
+		{
+			buffer.resize(length);
+			return buffer;
+		}
+
+		return {};
+	}
+}

src/runner/process.hpp

@@ -0,0 +1,9 @@
+#pragma once
+#include "native_handle.hpp"
+
+namespace process
+{
+	native_handle open(uint32_t process_id, DWORD access);
+	std::vector<HMODULE> get_modules(const native_handle& process);
+	std::string get_module_filename(const native_handle& process, HMODULE module);
+}

src/runner/std_include.hpp

@@ -9,5 +9,6 @@
 #include <Windows.h>
 #include <Shlwapi.h>
 #include <ShlObj.h>
+#include <Psapi.h>
 
 #pragma comment(lib, "Shlwapi.lib")

src/runner/utils/io.cpp

@@ -0,0 +1,125 @@
+#include "io.hpp"
+#include "nt.hpp"
+#include <fstream>
+
+namespace utils::io
+{
+	bool remove_file(const std::string& file)
+	{
+		return DeleteFileA(file.data()) == TRUE;
+	}
+
+	bool move_file(const std::string& src, const std::string& target)
+	{
+		return MoveFileA(src.data(), target.data()) == TRUE;
+	}
+
+	bool file_exists(const std::string& file)
+	{
+		return std::ifstream(file).good();
+	}
+
+	bool write_file(const std::string& file, const std::string& data, const bool append)
+	{
+		const auto pos = file.find_last_of("/\\");
+		if (pos != std::string::npos)
+		{
+			create_directory(file.substr(0, pos));
+		}
+
+		std::ofstream stream(
+			file, std::ios::binary | std::ofstream::out | (append ? std::ofstream::app : 0));
+
+		if (stream.is_open())
+		{
+			stream.write(data.data(), data.size());
+			stream.close();
+			return true;
+		}
+
+		return false;
+	}
+
+	std::string read_file(const std::string& file)
+	{
+		std::string data;
+		read_file(file, &data);
+		return data;
+	}
+
+	bool read_file(const std::string& file, std::string* data)
+	{
+		if (!data) return false;
+		data->clear();
+
+		if (file_exists(file))
+		{
+			std::ifstream stream(file, std::ios::binary);
+			if (!stream.is_open()) return false;
+
+			stream.seekg(0, std::ios::end);
+			const std::streamsize size = stream.tellg();
+			stream.seekg(0, std::ios::beg);
+
+			if (size > -1)
+			{
+				data->resize(static_cast<uint32_t>(size));
+				stream.read(const_cast<char*>(data->data()), size);
+				stream.close();
+				return true;
+			}
+		}
+
+		return false;
+	}
+
+	size_t file_size(const std::string& file)
+	{
+		if (file_exists(file))
+		{
+			std::ifstream stream(file, std::ios::binary);
+
+			if (stream.good())
+			{
+				stream.seekg(0, std::ios::end);
+				return static_cast<size_t>(stream.tellg());
+			}
+		}
+
+		return 0;
+	}
+
+	bool create_directory(const std::string& directory)
+	{
+		return std::filesystem::create_directories(directory);
+	}
+
+	bool directory_exists(const std::string& directory)
+	{
+		return std::filesystem::is_directory(directory);
+	}
+
+	bool directory_is_empty(const std::string& directory)
+	{
+		return std::filesystem::is_empty(directory);
+	}
+
+	std::vector<std::string> list_files(const std::string& directory)
+	{
+		std::vector<std::string> files;
+
+		for (auto& file : std::filesystem::directory_iterator(directory))
+		{
+			files.push_back(file.path().generic_string());
+		}
+
+		return files;
+	}
+
+	void copy_folder(const std::filesystem::path& src, const std::filesystem::path& target)
+	{
+		std::filesystem::copy(src, target,
+		                      std::filesystem::copy_options::overwrite_existing |
+		                      std::filesystem::copy_options::recursive);
+	}
+}

src/runner/utils/io.hpp

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <string>
+#include <vector>
+#include <filesystem>
+
+namespace utils::io
+{
+	bool remove_file(const std::string& file);
+	bool move_file(const std::string& src, const std::string& target);
+	bool file_exists(const std::string& file);
+	bool write_file(const std::string& file, const std::string& data, bool append = false);
+	bool read_file(const std::string& file, std::string* data);
+	std::string read_file(const std::string& file);
+	size_t file_size(const std::string& file);
+	bool create_directory(const std::string& directory);
+	bool directory_exists(const std::string& directory);
+	bool directory_is_empty(const std::string& directory);
+	std::vector<std::string> list_files(const std::string& directory);
+	void copy_folder(const std::filesystem::path& src, const std::filesystem::path& target);
+}

src/shared/irp_data.hpp

@@ -2,6 +2,8 @@
 
 #define HOOK_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)
 #define UNHOOK_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_NEITHER, FILE_ANY_ACCESS)
+#define WATCH_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_NEITHER, FILE_ANY_ACCESS)
+#define GET_RECORDS_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x803, METHOD_NEITHER, FILE_ANY_ACCESS)
 
 static_assert(sizeof(void*) == 8);
 
@@ -12,3 +14,16 @@ struct hook_request
 	const void* source_data{};
 	uint64_t source_data_size{};
 };
+
+struct watch_region
+{
+	const void* virtual_address{};
+	size_t length{};
+};
+
+struct watch_request
+{
+	uint32_t process_id{};
+	const void* watch_region{};
+	uint64_t watch_region_count{};
+};