From b68892cecafcf3331bf70400f47e77a2d6a4ca8a Mon Sep 17 00:00:00 2001 From: momo5502 Date: Fri, 13 Sep 2024 09:52:11 +0200 Subject: [PATCH] Support saving and restoring snapshots --- src/emulator/emulator.hpp | 54 +++++++++++++++-- src/emulator/memory_manager.cpp | 46 ++++++++------- src/emulator/memory_manager.hpp | 16 +---- src/unicorn_emulator/unicorn_x64_emulator.cpp | 59 +++++++++++++++++-- 4 files changed, 130 insertions(+), 45 deletions(-) diff --git a/src/emulator/emulator.hpp b/src/emulator/emulator.hpp index 1c9b6d06..b67f5d57 100644 --- a/src/emulator/emulator.hpp +++ b/src/emulator/emulator.hpp @@ -27,6 +27,16 @@ enum class memory_violation_type : uint8_t protection, }; +struct basic_block +{ + uint64_t address; + size_t instruction_count; + size_t size; +}; + +using edge_generation_hook_callback = std::function; + using instruction_hook_callback = std::function; using interrupt_hook_callback = std::function; @@ -62,6 +72,8 @@ public: virtual emulator_hook* hook_interrupt(interrupt_hook_callback callback) = 0; + virtual emulator_hook* hook_edge_generation(edge_generation_hook_callback callback) = 0; + virtual void delete_hook(emulator_hook* hook) = 0; emulator_hook* hook_memory_violation(memory_violation_hook_callback callback) @@ -87,17 +99,35 @@ public: void serialize(utils::buffer_serializer& buffer) const final { - this->serialize_memory_state(buffer); - this->serialize_state(buffer); + this->perform_serialization(buffer, false); } void deserialize(utils::buffer_deserializer& buffer) final { - this->deserialize_memory_state(buffer); - this->deserialize_state(buffer); + this->perform_deserialization(buffer, false); + } + + void save_snapshot() + { + utils::buffer_serializer serializer{}; + this->perform_serialization(serializer, true); + this->last_snapshot_data_ = serializer.move_buffer(); + } + + void restore_snapshot() + { + if (this->last_snapshot_data_.empty()) + { + return; + } + + utils::buffer_deserializer deserializer{this->last_snapshot_data_}; + this->perform_deserialization(deserializer, true); } private: + std::vector last_snapshot_data_{}; + emulator_hook* hook_simple_memory_access(const uint64_t address, const size_t size, simple_memory_hook_callback callback, const memory_operation operation) { @@ -110,6 +140,18 @@ private: }); } - virtual void serialize_state(utils::buffer_serializer& buffer) const = 0; - virtual void deserialize_state(utils::buffer_deserializer& buffer) = 0; + void perform_serialization(utils::buffer_serializer& buffer, const bool is_snapshot) const + { + this->serialize_memory_state(buffer, is_snapshot); + this->serialize_state(buffer, is_snapshot); + } + + void perform_deserialization(utils::buffer_deserializer& buffer, const bool is_snapshot) + { + this->deserialize_memory_state(buffer, is_snapshot); + this->deserialize_state(buffer, is_snapshot); + } + + virtual void serialize_state(utils::buffer_serializer& buffer, bool is_snapshot) const = 0; + virtual void deserialize_state(utils::buffer_deserializer& buffer, bool is_snapshot) = 0; }; diff --git a/src/emulator/memory_manager.cpp b/src/emulator/memory_manager.cpp index d741993e..539f4c93 100644 --- a/src/emulator/memory_manager.cpp +++ b/src/emulator/memory_manager.cpp @@ -87,47 +87,51 @@ static void deserialize(utils::buffer_deserializer& buffer, memory_manager::rese buffer.read_map(region.committed_regions); } -void memory_manager::serialize_memory_state(utils::buffer_serializer& buffer) const +void memory_manager::serialize_memory_state(utils::buffer_serializer& buffer, const bool is_snapshot) const { buffer.write_map(this->reserved_regions_); - if (!this->use_in_place_serialization()) + if (is_snapshot) { - std::vector data{}; + return; + } - for (const auto& reserved_region : this->reserved_regions_) + std::vector data{}; + + for (const auto& reserved_region : this->reserved_regions_) + { + for (const auto& region : reserved_region.second.committed_regions) { - for (const auto& region : reserved_region.second.committed_regions) - { - data.resize(region.second.length); + data.resize(region.second.length); - this->read_memory(region.first, data.data(), region.second.length); + this->read_memory(region.first, data.data(), region.second.length); - buffer.write(data.data(), region.second.length); - } + buffer.write(data.data(), region.second.length); } } } -void memory_manager::deserialize_memory_state(utils::buffer_deserializer& buffer) +void memory_manager::deserialize_memory_state(utils::buffer_deserializer& buffer, const bool is_snapshot) { buffer.read_map(this->reserved_regions_); - if (!this->use_in_place_serialization()) + if (is_snapshot) { - std::vector data{}; + return; + } - for (const auto& reserved_region : this->reserved_regions_) + std::vector data{}; + + for (const auto& reserved_region : this->reserved_regions_) + { + for (const auto& region : reserved_region.second.committed_regions) { - for (const auto& region : reserved_region.second.committed_regions) - { - data.resize(region.second.length); + data.resize(region.second.length); - buffer.read(data.data(), region.second.length); + buffer.read(data.data(), region.second.length); - this->map_memory(region.first, region.second.length, region.second.pemissions); - this->write_memory(region.first, data.data(), region.second.length); - } + this->map_memory(region.first, region.second.length, region.second.pemissions); + this->write_memory(region.first, data.data(), region.second.length); } } } diff --git a/src/emulator/memory_manager.hpp b/src/emulator/memory_manager.hpp index 116ea7af..f24253b4 100644 --- a/src/emulator/memory_manager.hpp +++ b/src/emulator/memory_manager.hpp @@ -75,22 +75,10 @@ public: return allocation_base; } - bool use_in_place_serialization() const - { - return this->in_place_serialization_; - } - - void set_in_place_serialization(const bool value) - { - this->in_place_serialization_ = value; - } - private: using reserved_region_map = std::map; reserved_region_map reserved_regions_{}; - bool in_place_serialization_{false}; - reserved_region_map::iterator find_reserved_region(uint64_t address); bool overlaps_reserved_region(uint64_t address, size_t size) const; @@ -100,6 +88,6 @@ private: virtual void apply_memory_protection(uint64_t address, size_t size, memory_permission permissions) = 0; protected: - void serialize_memory_state(utils::buffer_serializer& buffer) const; - void deserialize_memory_state(utils::buffer_deserializer& buffer); + 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/unicorn_x64_emulator.cpp b/src/unicorn_emulator/unicorn_x64_emulator.cpp index 8e2677c0..44faea20 100644 --- a/src/unicorn_emulator/unicorn_x64_emulator.cpp +++ b/src/unicorn_emulator/unicorn_x64_emulator.cpp @@ -210,6 +210,17 @@ namespace unicorn container.add(std::move(wrapper), std::move(hook)); } + basic_block map_block(const uc_tb& translation_block) + { + basic_block block{}; + + block.address = translation_block.pc; + block.instruction_count = translation_block.icount; + block.size = translation_block.size; + + return block; + } + class unicorn_x64_emulator : public x64_emulator { public: @@ -378,6 +389,31 @@ namespace unicorn return result; } + emulator_hook* hook_edge_generation(edge_generation_hook_callback callback) override + { + function_wrapper wrapper( + [c = std::move(callback)](uc_engine*, const uc_tb* cur_tb, const uc_tb* prev_tb) + { + const auto current_block = map_block(*cur_tb); + const auto previous_block = map_block(*prev_tb); + + c(current_block, previous_block); + }); + + unicorn_hook hook{*this}; + auto container = std::make_unique(); + + uce(uc_hook_add(*this, hook.make_reference(), UC_HOOK_EDGE_GENERATED, wrapper.get_function(), + wrapper.get_user_data(), 0, std::numeric_limits::max()) + ); + + container->add(std::move(wrapper), std::move(hook)); + + auto* result = container->as_opaque_hook(); + this->hooks_.push_back(std::move(container)); + return result; + } + emulator_hook* hook_interrupt(interrupt_hook_callback callback) override { function_wrapper wrapper( @@ -498,19 +534,34 @@ namespace unicorn return this->uc_; } - void serialize_state(utils::buffer_serializer& buffer) const override + void serialize_state(utils::buffer_serializer& buffer, const bool is_snapshot) const override { - const uc_context_serializer serializer(this->uc_, this->use_in_place_serialization()); + if (this->has_snapshots_ && !is_snapshot) + { + // TODO: Investigate if this is really necessary + throw std::runtime_error("Unable to serialize after snapshot was taken!"); + } + + this->has_snapshots_ |= is_snapshot; + + const uc_context_serializer serializer(this->uc_, is_snapshot); serializer.serialize(buffer); } - void deserialize_state(utils::buffer_deserializer& buffer) override + void deserialize_state(utils::buffer_deserializer& buffer, const bool is_snapshot) override { - const uc_context_serializer serializer(this->uc_, this->use_in_place_serialization()); + if (this->has_snapshots_ && !is_snapshot) + { + // TODO: Investigate if this is really necessary + throw std::runtime_error("Unable to deserialize after snapshot was taken!"); + } + + const uc_context_serializer serializer(this->uc_, is_snapshot); serializer.deserialize(buffer); } private: + mutable bool has_snapshots_{false}; uc_engine* uc_{}; bool retry_after_violation_{false}; std::vector> hooks_{};