From 49a8435f6b50f1419bb6d9ff3ff09594b8f0de6b Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Mon, 18 Nov 2024 16:58:23 +0100
Subject: [PATCH] Prepare MMIO support

---
 src/emulator/memory_manager.hpp               |  5 ++
 src/unicorn-emulator/function_wrapper.hpp     |  2 +
 src/unicorn-emulator/unicorn_x64_emulator.cpp | 59 ++++++++++++-------
 3 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/emulator/memory_manager.hpp b/src/emulator/memory_manager.hpp
index d91c1297..c59ff3ce 100644
--- a/src/emulator/memory_manager.hpp
+++ b/src/emulator/memory_manager.hpp
@@ -13,6 +13,9 @@ struct region_info : basic_memory_region
 	bool is_committed{};
 };
 
+using mmio_read_callback = std::function<uint64_t(uint64_t addr, size_t size)>;
+using mmio_write_callback = std::function<void(uint64_t addr, size_t size, uint64_t data)>;
+
 class memory_manager
 {
 public:
@@ -109,11 +112,13 @@ private:
 	reserved_region_map::iterator find_reserved_region(uint64_t address);
 	bool overlaps_reserved_region(uint64_t address, size_t size) const;
 
+	virtual void map_mmio(uint64_t address, size_t size, mmio_read_callback read_cb, mmio_write_callback write_cb) = 0;
 	virtual void map_memory(uint64_t address, size_t size, memory_permission permissions) = 0;
 	virtual void unmap_memory(uint64_t address, size_t size) = 0;
 
 	virtual void apply_memory_protection(uint64_t address, size_t size, memory_permission permissions) = 0;
 
+
 protected:
 	void serialize_memory_state(utils::buffer_serializer& buffer, bool is_snapshot) const;
 	void deserialize_memory_state(utils::buffer_deserializer& buffer, bool is_snapshot);
diff --git a/src/unicorn-emulator/function_wrapper.hpp b/src/unicorn-emulator/function_wrapper.hpp
index ec543114..47bb6876 100644
--- a/src/unicorn-emulator/function_wrapper.hpp
+++ b/src/unicorn-emulator/function_wrapper.hpp
@@ -13,6 +13,8 @@ public:
 	using c_function_type = ReturnType(Args..., user_data_pointer);
 	using functor_type = std::function<ReturnType(Args...)>;
 
+	function_wrapper() = default;
+
 	function_wrapper(functor_type functor)
 		: functor_(std::make_unique<functor_type>(std::move(functor)))
 	{
diff --git a/src/unicorn-emulator/unicorn_x64_emulator.cpp b/src/unicorn-emulator/unicorn_x64_emulator.cpp
index 769c75a8..f6848195 100644
--- a/src/unicorn-emulator/unicorn_x64_emulator.cpp
+++ b/src/unicorn-emulator/unicorn_x64_emulator.cpp
@@ -106,6 +106,15 @@ namespace unicorn
 			std::vector<hook_entry> hooks_;
 		};
 
+		struct mmio_callbacks
+		{
+			using read_wrapper = function_wrapper<uint64_t, uc_engine*, uint64_t, unsigned>;
+			using write_wrapper = function_wrapper<void, uc_engine*, uint64_t, unsigned, uint64_t>;
+
+			read_wrapper read{};
+			write_wrapper write{};
+		};
+
 		class uc_context_serializer
 		{
 		public:
@@ -302,6 +311,27 @@ namespace unicorn
 				}
 			}
 
+			void map_mmio(const uint64_t address, const size_t size, mmio_read_callback read_cb,
+			              mmio_write_callback write_cb) override
+			{
+				mmio_callbacks cb{
+					.read = [c = std::move(read_cb)](uc_engine*, const uint64_t addr, const uint32_t s)
+					{
+						return c(addr, s);
+					},
+					.write = [c = std::move(write_cb)](uc_engine*, const uint64_t addr, const uint32_t s,
+					                                   const uint64_t value)
+					{
+						c(addr, s, value);
+					}
+				};
+
+				uce(uc_mmio_map(*this, address, size, cb.read.get_c_function(), cb.read.get_user_data(),
+				                cb.write.get_c_function(), cb.write.get_user_data()));
+
+				this->mmio_[address] = std::move(cb);
+			}
+
 			void map_memory(const uint64_t address, const size_t size, memory_permission permissions) override
 			{
 				uce(uc_mem_map(*this, address, size, static_cast<uint32_t>(permissions)));
@@ -310,6 +340,12 @@ namespace unicorn
 			void unmap_memory(const uint64_t address, const size_t size) override
 			{
 				uce(uc_mem_unmap(*this, address, size));
+
+				const auto mmio_entry = this->mmio_.find(address);
+				if (mmio_entry != this->mmio_.end())
+				{
+					this->mmio_.erase(mmio_entry);
+				}
 			}
 
 			bool try_read_memory(const uint64_t address, void* data, const size_t size) const override
@@ -333,28 +369,6 @@ namespace unicorn
 				uce(uc_mem_protect(*this, address, size, static_cast<uint32_t>(permissions)));
 			}
 
-			/*std::vector<memory_region> get_memory_regions() override
-			{
-				const unicorn_memory_regions regions{*this};
-				const auto region_span = regions.get_span();
-
-				std::vector<memory_region> result{};
-				result.reserve(region_span.size());
-
-				for (const auto region : region_span)
-				{
-					memory_region reg{};
-					reg.start = region.begin;
-					reg.length = region.end - region.begin;
-					reg.committed = true;
-					reg.pemissions = static_cast<memory_permission>(region.perms) & memory_permission::all;
-
-					result.push_back(reg);
-				}
-
-				return result;
-			}*/
-
 			emulator_hook* hook_instruction(int instruction_type,
 			                                instruction_hook_callback callback) override
 			{
@@ -614,6 +628,7 @@ namespace unicorn
 			uc_engine* uc_{};
 			bool has_violation_{false};
 			std::vector<std::unique_ptr<hook_object>> hooks_{};
+			std::unordered_map<uint64_t, mmio_callbacks> mmio_{};
 		};
 	}