From d191f8e66701e79ffc1e2b9d58cf76b6834eec7e Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sat, 5 Apr 2025 08:55:08 +0200 Subject: [PATCH] Optimize MMIO handling --- src/emulator/memory_interface.hpp | 4 +-- src/icicle-emulator/icicle_x64_emulator.cpp | 29 ++++------------- src/icicle/src/lib.rs | 8 ++--- src/unicorn-emulator/unicorn_x64_emulator.cpp | 32 +++++++++++++++---- src/windows-emulator/kusd_mmio.cpp | 25 ++++++--------- src/windows-emulator/kusd_mmio.hpp | 2 +- src/windows-emulator/memory_manager.hpp | 4 +-- 7 files changed, 49 insertions(+), 55 deletions(-) diff --git a/src/emulator/memory_interface.hpp b/src/emulator/memory_interface.hpp index 3cdcc895..2f735e9a 100644 --- a/src/emulator/memory_interface.hpp +++ b/src/emulator/memory_interface.hpp @@ -4,8 +4,8 @@ #include "memory_permission.hpp" -using mmio_read_callback = std::function; -using mmio_write_callback = std::function; +using mmio_read_callback = std::function; +using mmio_write_callback = std::function; class memory_manager; diff --git a/src/icicle-emulator/icicle_x64_emulator.cpp b/src/icicle-emulator/icicle_x64_emulator.cpp index 31b62ce6..d6fddef1 100644 --- a/src/icicle-emulator/icicle_x64_emulator.cpp +++ b/src/icicle-emulator/icicle_x64_emulator.cpp @@ -12,8 +12,8 @@ extern "C" using violation_func = int32_t(void*, uint64_t address, uint8_t operation, int32_t unmapped); using data_accessor_func = void(void* user, const void* data, size_t length); - using icicle_mmio_read_func = void(void* user, uint64_t address, size_t length, void* data); - using icicle_mmio_write_func = void(void* user, uint64_t address, size_t length, const void* data); + using icicle_mmio_read_func = void(void* user, uint64_t address, void* data, size_t length); + using icicle_mmio_write_func = void(void* user, uint64_t address, const void* data, size_t length); icicle_emulator* icicle_create_emulator(); int32_t icicle_protect_memory(icicle_emulator*, uint64_t address, uint64_t length, uint8_t permissions); @@ -185,31 +185,14 @@ namespace icicle auto* ptr = wrapper.get(); this->storage_.push_back(std::move(wrapper)); - auto* read_wrapper = +[](void* user, const uint64_t addr, const size_t length, void* data) { - constexpr auto limit = sizeof(uint64_t); + auto* read_wrapper = +[](void* user, const uint64_t addr, void* data, const size_t length) { const auto* w = static_cast(user); - - // TODO: Change interface to get rid of loop - for (size_t offset = 0; offset < length; offset += limit) - { - const auto max_read = std::min(limit, length - offset); - const auto value = w->read_cb(addr + offset - w->base, max_read); - memcpy(static_cast(data) + offset, &value, max_read); - } + w->read_cb(addr - w->base, data, length); }; - auto* write_wrapper = +[](void* user, const uint64_t addr, const size_t length, const void* data) { - constexpr auto limit = sizeof(uint64_t); + auto* write_wrapper = +[](void* user, const uint64_t addr, const void* data, const size_t length) { const auto* w = static_cast(user); - - // TODO: Change interface to get rid of loop - for (size_t offset = 0; offset < length; offset += limit) - { - uint64_t value{}; - const auto max_read = std::min(limit, length - offset); - memcpy(&value, static_cast(data) + offset, max_read); - w->write_cb(addr + offset - w->base, max_read, value); - } + w->write_cb(addr + w->base, data, length); }; icicle_map_mmio(this->emu_, address, size, read_wrapper, ptr, write_wrapper, ptr); diff --git a/src/icicle/src/lib.rs b/src/icicle/src/lib.rs index 1df8d39b..dab13cfb 100644 --- a/src/icicle/src/lib.rs +++ b/src/icicle/src/lib.rs @@ -38,8 +38,8 @@ pub fn icicle_stop(ptr: *mut c_void) { type RawFunction = extern "C" fn(*mut c_void); type PtrFunction = extern "C" fn(*mut c_void, u64); type DataFunction = extern "C" fn(*mut c_void, *const c_void, usize); -type MmioReadFunction = extern "C" fn(*mut c_void, u64, usize, *mut c_void); -type MmioWriteFunction = extern "C" fn(*mut c_void, u64, usize, *const c_void); +type MmioReadFunction = extern "C" fn(*mut c_void, u64, *mut c_void, usize); +type MmioWriteFunction = extern "C" fn(*mut c_void, u64, *const c_void, usize); type ViolationFunction = extern "C" fn(*mut c_void, u64, u8, i32) -> i32; #[unsafe(no_mangle)] @@ -57,12 +57,12 @@ pub fn icicle_map_mmio( let read_wrapper = Box::new(move |addr: u64, data: &mut [u8]| { let raw_pointer: *mut u8 = data.as_mut_ptr(); - read_cb(read_data, addr, data.len(), raw_pointer as *mut c_void); + read_cb(read_data, addr, raw_pointer as *mut c_void, data.len()); }); let write_wrapper = Box::new(move |addr: u64, data: &[u8]| { let raw_pointer: *const u8 = data.as_ptr(); - write_cb(write_data, addr, data.len(), raw_pointer as *const c_void); + write_cb(write_data, addr, raw_pointer as *const c_void, data.len()); }); let res = emulator.map_mmio(address, length, read_wrapper, write_wrapper); diff --git a/src/unicorn-emulator/unicorn_x64_emulator.cpp b/src/unicorn-emulator/unicorn_x64_emulator.cpp index f26dddda..de01afcf 100644 --- a/src/unicorn-emulator/unicorn_x64_emulator.cpp +++ b/src/unicorn-emulator/unicorn_x64_emulator.cpp @@ -243,6 +243,14 @@ namespace unicorn return block; } + void assert_64bit_limit(const size_t size) + { + if (size > sizeof(uint64_t)) + { + throw std::runtime_error("Exceeded uint64_t size limit"); + } + } + class unicorn_x64_emulator : public x64_emulator { public: @@ -370,13 +378,23 @@ 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 = mmio_callbacks::read_wrapper( - [c = std::move(read_cb)](uc_engine*, const uint64_t addr, const uint32_t s) { - return c(addr, s); - }), - .write = mmio_callbacks::write_wrapper( - [c = std::move(write_cb)](uc_engine*, const uint64_t addr, const uint32_t s, - const uint64_t value) { c(addr, s, value); })}; + auto read_wrapper = [c = std::move(read_cb)](uc_engine*, const uint64_t addr, const uint32_t s) { + assert_64bit_limit(s); + uint64_t value{}; + c(addr, &value, s); + return value; + }; + + auto write_wrapper = [c = std::move(write_cb)](uc_engine*, const uint64_t addr, const uint32_t s, + const uint64_t value) { + assert_64bit_limit(s); + c(addr, &value, s); + }; + + mmio_callbacks cb{ + .read = mmio_callbacks::read_wrapper(std::move(read_wrapper)), + .write = mmio_callbacks::write_wrapper(std::move(write_wrapper)), + }; 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())); diff --git a/src/windows-emulator/kusd_mmio.cpp b/src/windows-emulator/kusd_mmio.cpp index 136ecbc7..856d7187 100644 --- a/src/windows-emulator/kusd_mmio.cpp +++ b/src/windows-emulator/kusd_mmio.cpp @@ -83,13 +83,13 @@ namespace namespace utils { - inline void serialize(buffer_serializer& buffer, const KUSER_SHARED_DATA64& kusd) + static void serialize(buffer_serializer& buffer, const KUSER_SHARED_DATA64& kusd) { static_assert(KUSD_SIZE == sizeof(kusd)); buffer.write(&kusd, KUSD_SIZE); } - inline void deserialize(buffer_deserializer& buffer, KUSER_SHARED_DATA64& kusd) + static void deserialize(buffer_deserializer& buffer, KUSER_SHARED_DATA64& kusd) { buffer.read(&kusd, KUSD_SIZE); } @@ -130,30 +130,21 @@ void kusd_mmio::deserialize(utils::buffer_deserializer& buffer) this->register_mmio(); } -uint64_t kusd_mmio::read(const uint64_t addr, const size_t size) +void kusd_mmio::read(const uint64_t addr, void* data, const size_t size) { - uint64_t result{}; - this->update(); if (addr >= KUSD_SIZE) { - return result; + return; } const auto end = addr + size; const auto valid_end = std::min(end, static_cast(KUSD_SIZE)); const auto real_size = valid_end - addr; - if (real_size > sizeof(result)) - { - return result; - } - const auto* kusd_buffer = reinterpret_cast(&this->kusd_); - memcpy(&result, kusd_buffer + addr, real_size); - - return result; + memcpy(data, kusd_buffer + addr, real_size); } uint64_t kusd_mmio::address() @@ -178,8 +169,10 @@ void kusd_mmio::register_mmio() this->memory_->allocate_mmio( KUSD_ADDRESS, KUSD_BUFFER_SIZE, - [this](const uint64_t addr, const size_t size) { return this->read(addr, size); }, - [](const uint64_t, const size_t, const uint64_t) { + [this](const uint64_t addr, void* data, const size_t size) { + this->read(addr, data, size); // + }, + [](const uint64_t, const void*, const size_t) { // Writing not supported! }); } diff --git a/src/windows-emulator/kusd_mmio.hpp b/src/windows-emulator/kusd_mmio.hpp index 24c4df53..bba54287 100644 --- a/src/windows-emulator/kusd_mmio.hpp +++ b/src/windows-emulator/kusd_mmio.hpp @@ -48,7 +48,7 @@ class kusd_mmio KUSER_SHARED_DATA64 kusd_{}; - uint64_t read(uint64_t addr, size_t size); + void read(uint64_t addr, void* data, size_t size); void update(); diff --git a/src/windows-emulator/memory_manager.hpp b/src/windows-emulator/memory_manager.hpp index 02e0aa07..92daab60 100644 --- a/src/windows-emulator/memory_manager.hpp +++ b/src/windows-emulator/memory_manager.hpp @@ -20,8 +20,8 @@ struct region_info : basic_memory_region bool is_committed{}; }; -using mmio_read_callback = std::function; -using mmio_write_callback = std::function; +using mmio_read_callback = std::function; +using mmio_write_callback = std::function; class memory_manager : public memory_interface {