Cleanup KUSD MMIO

momo5502 · Nov 24, 2024 · 5f49b30 · 5f49b30
1 parent d3cf88a
commit 5f49b30
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 50 deletions.
diff --git a/src/windows-emulator/emulator_utils.hpp b/src/windows-emulator/emulator_utils.hpp
@@ -36,8 +36,10 @@ class object_wrapper
 };
 
 class windows_emulator;
+struct process_context;
 
 using x64_emulator_wrapper = object_wrapper<x64_emulator>;
+using process_context_wrapper = object_wrapper<process_context>;
 using windows_emulator_wrapper = object_wrapper<windows_emulator>;
 
 template <typename T>

diff --git a/src/windows-emulator/kusd_mmio.cpp b/src/windows-emulator/kusd_mmio.cpp
@@ -99,41 +99,30 @@ inline void deserialize(utils::buffer_deserializer& buffer, KUSER_SHARED_DATA& k
 	buffer.read(&kusd, KUSD_SIZE);
 }
 
-kusd_mmio::kusd_mmio(windows_emulator& win_emu, const bool use_relative_time, const bool perform_registration)
-	: use_relative_time_(use_relative_time)
-	  , win_emu_(&win_emu)
+kusd_mmio::kusd_mmio(x64_emulator& emu, process_context& process)
+	: emu_(&emu)
+	  , process_(&process)
 {
-	setup_kusd(this->kusd_, use_relative_time);
-	this->start_time_ = convert_from_ksystem_time(this->kusd_.SystemTime);
-
-	if (perform_registration)
-	{
-		this->register_mmio();
-	}
 }
 
 kusd_mmio::~kusd_mmio()
 {
 	this->deregister_mmio();
 }
 
-kusd_mmio::kusd_mmio(kusd_mmio&& obj) // throws!
-	: use_relative_time_(obj.use_relative_time_)
-	  , win_emu_(obj.win_emu_)
-	  , start_time_(obj.start_time_)
+kusd_mmio::kusd_mmio(utils::buffer_deserializer& buffer)
+	: kusd_mmio(buffer.read<x64_emulator_wrapper>(), buffer.read<process_context_wrapper>())
 {
-	memcpy(&this->kusd_, &obj.kusd_, sizeof(this->kusd_));
-
-	if (obj.registered_)
-	{
-		obj.deregister_mmio();
-		this->register_mmio();
-	}
 }
 
-kusd_mmio::kusd_mmio(utils::buffer_deserializer& buffer)
-	: kusd_mmio(buffer.read<windows_emulator_wrapper>().get(), true, false)
+void kusd_mmio::setup(const bool use_relative_time)
 {
+	this->use_relative_time_ = use_relative_time;
+
+	setup_kusd(this->kusd_, use_relative_time);
+	this->start_time_ = convert_from_ksystem_time(this->kusd_.SystemTime);
+
+	this->register_mmio();
 }
 
 void kusd_mmio::serialize(utils::buffer_serializer& buffer) const
@@ -149,6 +138,7 @@ void kusd_mmio::deserialize(utils::buffer_deserializer& buffer)
 	buffer.read(this->kusd_);
 	buffer.read(this->start_time_);
 
+	this->deregister_mmio();
 	this->register_mmio();
 }
 
@@ -183,18 +173,13 @@ uint64_t kusd_mmio::address()
 	return KUSD_ADDRESS;
 }
 
-void kusd_mmio::write(const uint64_t /*addr*/, const size_t /*size*/, const uint64_t /*data*/)
-{
-	// Unsupported!
-}
-
 void kusd_mmio::update()
 {
 	auto time = this->start_time_;
 
 	if (this->use_relative_time_)
 	{
-		const auto passed_time = this->win_emu_->process().executed_instructions;
+		const auto passed_time = this->process_->executed_instructions;
 		const auto clock_frequency = this->kusd_.QpcFrequency;
 
 		using duration = std::chrono::system_clock::duration;
@@ -217,22 +202,21 @@ void kusd_mmio::register_mmio()
 
 	this->registered_ = true;
 
-	this->win_emu_->emu().allocate_mmio( //
-		KUSD_ADDRESS, KUSD_BUFFER_SIZE,
-		[this](const uint64_t addr, const size_t size)
-		{
-			return this->read(addr, size);
-		}, [this](const uint64_t addr, const size_t size, const uint64_t data)
-		{
-			this->write(addr, size, data);
-		});
+	this->emu_->allocate_mmio(KUSD_ADDRESS, KUSD_BUFFER_SIZE,
+	                          [this](const uint64_t addr, const size_t size)
+	                          {
+		                          return this->read(addr, size);
+	                          }, [this](const uint64_t, const size_t, const uint64_t)
+	                          {
+		                          // Writing not supported!
+	                          });
 }
 
 void kusd_mmio::deregister_mmio()
 {
 	if (this->registered_)
 	{
 		this->registered_ = false;
-		this->win_emu_->emu().release_memory(KUSD_ADDRESS, KUSD_BUFFER_SIZE);
+		this->emu_->release_memory(KUSD_ADDRESS, KUSD_BUFFER_SIZE);
 	}
 }
diff --git a/src/windows-emulator/kusd_mmio.hpp b/src/windows-emulator/kusd_mmio.hpp
@@ -3,18 +3,20 @@
 #include "std_include.hpp"
 #include <serialization.hpp>
 
+#include "x64_emulator.hpp"
+
+struct process_context;
 class windows_emulator;
 
 class kusd_mmio
 {
 public:
-	kusd_mmio(windows_emulator& win_emu,  bool use_relative_time, bool perform_registration = true);
+	kusd_mmio(x64_emulator& emu, process_context& process);
 	~kusd_mmio();
 
-	kusd_mmio(kusd_mmio&& obj);
-
 	kusd_mmio(utils::buffer_deserializer& buffer);
 
+	kusd_mmio(kusd_mmio&&) = delete;
 	kusd_mmio(const kusd_mmio&) = delete;
 	kusd_mmio& operator=(kusd_mmio&& obj) = delete;
 	kusd_mmio& operator=(const kusd_mmio&) = delete;
@@ -34,15 +36,19 @@ class kusd_mmio
 
 	static uint64_t address();
 
+	void setup(bool use_relative_time);
+
 private:
+	x64_emulator* emu_{};
+	process_context* process_{};
+
 	bool registered_{};
 	bool use_relative_time_{};
-	windows_emulator* win_emu_{};
+
 	KUSER_SHARED_DATA kusd_{};
 	std::chrono::system_clock::time_point start_time_{};
 
 	uint64_t read(uint64_t addr, size_t size);
-	void write(uint64_t addr, size_t size, uint64_t data);
 
 	void update();
 

diff --git a/src/windows-emulator/process_context.hpp b/src/windows-emulator/process_context.hpp
@@ -365,6 +365,7 @@ struct process_context
 		: base_allocator(emu)
 		  , peb(emu)
 		  , process_params(emu)
+		  , kusd(emu, *this)
 		  , module_manager(emu)
 	{
 	}
@@ -382,7 +383,7 @@ struct process_context
 
 	emulator_object<PEB> peb;
 	emulator_object<RTL_USER_PROCESS_PARAMETERS> process_params;
-	std::optional<kusd_mmio> kusd{};
+	kusd_mmio kusd;
 
 	module_manager module_manager;
 
@@ -421,7 +422,7 @@ struct process_context
 		buffer.write(this->base_allocator);
 		buffer.write(this->peb);
 		buffer.write(this->process_params);
-		buffer.write_optional(this->kusd);
+		buffer.write(this->kusd);
 		buffer.write(this->module_manager);
 
 		buffer.write(this->executable->image_base);
@@ -459,7 +460,7 @@ struct process_context
 		buffer.read(this->base_allocator);
 		buffer.read(this->peb);
 		buffer.read(this->process_params);
-		buffer.read_optional(this->kusd);
+		buffer.read(this->kusd);
 		buffer.read(this->module_manager);
 
 		const auto executable_base = buffer.read<uint64_t>();

diff --git a/src/windows-emulator/syscalls.cpp b/src/windows-emulator/syscalls.cpp
@@ -37,7 +37,7 @@ namespace
 			{
 				performance_frequency.access([&](LARGE_INTEGER& value)
 				{
-					value.QuadPart = c.proc.kusd->get().QpcFrequency;
+					value.QuadPart = c.proc.kusd.get().QpcFrequency;
 				});
 			}
 
@@ -534,7 +534,7 @@ namespace
 			c.emu.allocate_memory(address,
 			                      c.proc.shared_section_size, memory_permission::read_write);
 
-			const std::wstring_view windows_dir = c.proc.kusd->get().NtSystemRoot.arr;
+			const std::wstring_view windows_dir = c.proc.kusd.get().NtSystemRoot.arr;
 			const auto windows_dir_size = windows_dir.size() * 2;
 
 			constexpr auto windows_dir_offset = 0x10;

diff --git a/src/windows-emulator/windows_emulator.cpp b/src/windows-emulator/windows_emulator.cpp
@@ -182,7 +182,7 @@ namespace
 
 		context.registry = registry_manager(settings.registry_directory);
 
-		context.kusd.emplace(win_emu, settings.use_relative_time);
+		context.kusd.setup(settings.use_relative_time);
 
 		context.base_allocator = create_allocator(emu, PEB_SEGMENT_SIZE);
 		auto& allocator = context.base_allocator;