Quick and dirty code watching bridge for usermode/kernelmode

src/driver/hypervisor.cpp

@@ -211,6 +211,11 @@ void hypervisor::disable_all_ept_hooks() const
 	});
 }
 
+vmx::ept& hypervisor::get_ept() const
+{
+	return *this->ept_;
+}
+
 hypervisor* hypervisor::get_instance()
 {
 	return instance;

src/driver/hypervisor.hpp

@@ -23,6 +23,8 @@ public:
 	                      vmx::ept_translation_hint* translation_hint = nullptr);
 	void disable_all_ept_hooks() const;
 
+	vmx::ept& get_ept() const;
+
 	static hypervisor* get_instance();
 
 private:

src/driver/irp.cpp

@@ -114,6 +114,9 @@ namespace
 
 	void try_apply_hook(const PIO_STACK_LOCATION irp_sp)
 	{
+		memory::assert_readability(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer,
+		                           irp_sp->Parameters.DeviceIoControl.InputBufferLength);
+
 		if (irp_sp->Parameters.DeviceIoControl.InputBufferLength < sizeof(hook_request))
 		{
 			throw std::runtime_error("Invalid hook request");
@@ -126,10 +129,103 @@ namespace
 		apply_hook(request);
 	}
 
+	void watch_regions(const watch_request& watch_request)
+	{
+		auto* hypervisor = hypervisor::get_instance();
+		if (!hypervisor)
+		{
+			throw std::runtime_error("Hypervisor not installed");
+		}
+
+		std::unique_ptr<watch_region[]> buffer(new watch_region[watch_request.watch_region_count]);
+		if (!buffer)
+		{
+			throw std::runtime_error("Failed to copy buffer");
+		}
+
+		memcpy(buffer.get(), watch_request.watch_regions, watch_request.watch_region_count * sizeof(watch_region));
+
+		auto watch_request_copy = watch_request;
+		watch_request_copy.watch_regions = buffer.get();
+
+		thread::kernel_thread t([watch_request_copy, hypervisor]
+		{
+			debug_log("Looking up process: %d\n", watch_request_copy.process_id);
+
+			const auto process_handle = process::find_process_by_id(watch_request_copy.process_id);
+			if (!process_handle || !process_handle.is_alive())
+			{
+				debug_log("Bad process\n");
+				return;
+			}
+
+			const auto name = process_handle.get_image_filename();
+			if (name)
+			{
+				debug_log("Attaching to %s\n", name);
+			}
+
+			process::scoped_process_attacher attacher{process_handle};
+
+			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)));
+
+				for (auto current = start; current < end; current += PAGE_SIZE)
+				{
+					const auto physical_address = memory::get_physical_address(const_cast<uint8_t*>(current));
+					if (physical_address)
+					{
+						hypervisor->get_ept().install_code_watch_point(physical_address);
+					}
+				}
+			}
+		});
+
+		t.join();
+	}
+
+	void try_watch_regions(const PIO_STACK_LOCATION irp_sp)
+	{
+		memory::assert_readability(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer,
+		                           irp_sp->Parameters.DeviceIoControl.InputBufferLength);
+
+		if (irp_sp->Parameters.DeviceIoControl.InputBufferLength < sizeof(watch_request))
+		{
+			throw std::runtime_error("Invalid watch request");
+		}
+
+		const auto& request = *static_cast<watch_request*>(irp_sp->Parameters.DeviceIoControl.Type3InputBuffer);
+		memory::assert_readability(request.watch_regions, request.watch_region_count * sizeof(watch_region));
+
+		watch_regions(request);
+	}
+
+	void get_records(const PIRP irp, const PIO_STACK_LOCATION irp_sp)
+	{
+		auto* hypervisor = hypervisor::get_instance();
+		if (!hypervisor)
+		{
+			throw std::runtime_error("Hypervisor not installed");
+		}
+
+		size_t count{};
+		const auto records = hypervisor->get_ept().get_access_records(&count);
+
+		memset(irp->UserBuffer, 0, irp_sp->Parameters.DeviceIoControl.OutputBufferLength);
+		memcpy(irp->UserBuffer, records, min(irp_sp->Parameters.DeviceIoControl.OutputBufferLength, count * 8));
+	}
+
 	void handle_irp(const PIRP irp)
 	{
 		irp->IoStatus.Information = 0;
-		irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
+		irp->IoStatus.Status = STATUS_SUCCESS;
 
 		const auto irp_sp = IoGetCurrentIrpStackLocation(irp);
 
@@ -145,6 +241,12 @@ namespace
 			case UNHOOK_DRV_IOCTL:
 				unhook();
 				break;
+			case WATCH_DRV_IOCTL:
+				try_watch_regions(irp_sp);
+				break;
+			case GET_RECORDS_DRV_IOCTL:
+				get_records(irp, irp_sp);
+				break;
 			default:
 				debug_log("Invalid IOCTL Code: 0x%X\n", ioctr_code);
 				irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;

src/runner/driver_device.cpp

@@ -24,21 +24,35 @@ bool driver_device::send(const DWORD ioctl_code, const data& input) const
 }
 
 bool driver_device::send(const DWORD ioctl_code, const data& input, data& output) const
+{
+	size_t out_len = output.size();
+	if (this->send(ioctl_code, input.data(), input.size(), output.data(), &out_len))
+	{
+		output.resize(out_len);
+		return true;
+	}
+
+	return false;
+}
+
+bool driver_device::send(const DWORD ioctl_code, const void* input, const size_t input_length, void* output,
+                         size_t* output_length) const
 {
 	DWORD size_returned = 0;
 	const auto success = DeviceIoControl(this->device_,
 	                                     ioctl_code,
-	                                     const_cast<uint8_t*>(input.data()),
-	                                     static_cast<DWORD>(input.size()),
-	                                     output.data(),
-	                                     static_cast<DWORD>(output.size()),
+	                                     const_cast<void*>(input),
+	                                     static_cast<DWORD>(input_length),
+	                                     output,
+	                                     static_cast<DWORD>(*output_length),
 	                                     &size_returned,
 	                                     nullptr
 	) != FALSE;
 
-	if (success && size_returned < output.size())
+	*output_length = 0;
+	if (success)
 	{
-		output.resize(size_returned);
+		*output_length = size_returned;
 	}
 
 	return success;

src/runner/driver_device.hpp

@@ -16,6 +16,7 @@ public:
 	using data = std::vector<uint8_t>;
 	bool send(DWORD ioctl_code, const data& input) const;
 	bool send(DWORD ioctl_code, const data& input, data& output) const;
+	bool send(DWORD ioctl_code, const void* input, size_t input_length, void* output, size_t* output_length) const;
 
 private:
 	native_handle device_{};

src/runner/main.cpp

@@ -4,7 +4,7 @@
 #include <filesystem>
 #include <conio.h>
 #include <fstream>
-
+#include <set>
 
 #include "driver.hpp"
 #include "driver_device.hpp"
@@ -107,16 +107,114 @@ std::vector<std::pair<size_t, size_t>> find_executable_regions(const std::string
 	return regions;
 }
 
-void unsafe_main(const int /*argc*/, char* /*argv*/[])
+uint32_t get_process_id()
 {
 	std::string pid_str{};
 	printf("Please enter the pid: ");
 	std::getline(std::cin, pid_str);
 
-	const auto pid = atoi(pid_str.data());
+	return atoi(pid_str.data());
+}
+
+void watch_regions(const driver_device& driver_device, const uint32_t pid, const HMODULE module,
+                   const std::vector<std::pair<size_t, size_t>>& regions)
+{
+	std::vector<watch_region> watch_regions{};
+	watch_regions.reserve(regions.size());
+
+	for (const auto& region : regions)
+	{
+		watch_region watch_region{};
+		watch_region.virtual_address = reinterpret_cast<uint8_t*>(module) + region.first;
+		watch_region.length = region.second;
+		watch_regions.push_back(watch_region);
+	}
+
+	watch_request request{};
+	request.process_id = pid;
+	request.watch_regions = watch_regions.data();
+	request.watch_region_count = watch_regions.size();
+
+	driver_device::data out{};
+	size_t out_len = 0;
+	driver_device.send(WATCH_DRV_IOCTL, &request, sizeof(request), out.data(), &out_len);
+}
+
+std::vector<uint64_t> query_records(const driver_device& driver_device, const size_t current_size = 0)
+{
+	std::vector<uint64_t> result{};
+	result.resize(std::max(size_t(1024), current_size * 2));
+
+	while (true)
+	{
+		char in[1];
+		constexpr auto element_len = sizeof(decltype(result)::value_type);
+		const size_t initial_len = result.size() * element_len;
+		size_t out_len = initial_len;
+		if (!driver_device.send(GET_RECORDS_DRV_IOCTL, in, 0, result.data(), &out_len))
+		{
+			return {};
+		}
+
+		//if (out_len <= initial_len)
+		if (result.back() == 0)
+		{
+			//result.resize(out_len / element_len);
+			break;
+		}
+
+		//result.resize((out_len / element_len) + 10);
+		const auto new_size = result.size() * 2;
+
+		result = {};
+		result.resize(result.size() * 2);
+	}
+
+	// Shrink
+	size_t i;
+	for (i = result.size(); i > 0; --i)
+	{
+		if (result[i - 1] != 0)
+		{
+			break;
+		}
+	}
+
+	result.resize(i);
+
+	return result;
+}
+
+void report_records(const std::atomic_bool& flag, const driver_device& driver_device)
+{
+	std::set<uint64_t> access_addresses{};
+
+	while (flag)
+	{
+		std::this_thread::sleep_for(std::chrono::seconds(1));
+		const auto new_records = query_records(driver_device, access_addresses.size());
+
+		for (const auto& new_record : new_records)
+		{
+			if (access_addresses.emplace(new_record).second)
+			{
+				printf("%p\n", reinterpret_cast<void*>(new_record));
+			}
+		}
+	}
+}
+
+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)"};
+
+	const auto pid = get_process_id();
 
 	printf("Opening process...\n");
-	const auto proc = process::open(pid, PROCESS_QUERY_INFORMATION | PROCESS_VM_READ);
+	auto proc = process::open(pid, PROCESS_QUERY_INFORMATION | PROCESS_VM_READ);
 	if (!proc)
 	{
 		printf("Failed to open process...\n");
@@ -138,6 +236,9 @@ void unsafe_main(const int /*argc*/, char* /*argv*/[])
 		module_files.emplace_back(std::move(name));
 	}
 
+	// We don't need this anymore 
+	proc = {};
+
 	std::string module_str{};
 	printf("\nPlease enter the module number: ");
 	std::getline(std::cin, module_str);
@@ -151,7 +252,8 @@ void unsafe_main(const int /*argc*/, char* /*argv*/[])
 		return;
 	}
 
-	const auto module_base = reinterpret_cast<uint8_t*>(modules[module_num]);
+	const auto target_module = modules[module_num];
+	const auto module_base = reinterpret_cast<uint8_t*>(target_module);
 	const auto& file = module_files[module_num];
 	printf("Analyzing %s...\n", file.data());
 	const auto regions = find_executable_regions(file);
@@ -161,14 +263,21 @@ void unsafe_main(const int /*argc*/, char* /*argv*/[])
 		printf("%p - %zu\n", module_base + region.first, region.second);
 	}
 
+	watch_regions(driver_device, pid, target_module, regions);
+
+	std::atomic_bool terminate{false};
+	std::thread t([&]()
+	{
+		report_records(terminate, driver_device);
+	});
+
 	_getch();
+
+	terminate = true;
+	t.join();
+
 	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/shared/irp_data.hpp

@@ -24,6 +24,6 @@ struct watch_region
 struct watch_request
 {
 	uint32_t process_id{};
-	const void* watch_region{};
+	const watch_region* watch_regions{};
 	uint64_t watch_region_count{};
 };