Support saving and restoring snapshots

This commit is contained in:
momo5502
2024-09-13 09:52:11 +02:00
parent b1cbc5a7ad
commit b68892ceca
4 changed files with 130 additions and 45 deletions

View File

@@ -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<void(const basic_block& current_block,
const basic_block& previous_block)>;
using instruction_hook_callback = std::function<instruction_hook_continuation()>;
using interrupt_hook_callback = std::function<void(int interrupt)>;
@@ -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<std::byte> 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;
};

View File

@@ -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<uint8_t> data{};
return;
}
for (const auto& reserved_region : this->reserved_regions_)
std::vector<uint8_t> 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<uint8_t> data{};
return;
}
for (const auto& reserved_region : this->reserved_regions_)
std::vector<uint8_t> 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);
}
}
}

View File

@@ -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<uint64_t, reserved_region>;
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);
};

View File

@@ -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<void, uc_engine*, uc_tb*, uc_tb*> 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<hook_container>();
uce(uc_hook_add(*this, hook.make_reference(), UC_HOOK_EDGE_GENERATED, wrapper.get_function(),
wrapper.get_user_data(), 0, std::numeric_limits<pointer_type>::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<void, uc_engine*, int> 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<std::unique_ptr<hook_object>> hooks_{};