Skip to content

Commit

Permalink
Support saving and restoring snapshots
Browse files Browse the repository at this point in the history
  • Loading branch information
@momo5502
momo5502 committed Sep 13, 2024
1 parent b1cbc5a commit b68892c
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 45 deletions.
54 changes: 48 additions & 6 deletions src/emulator/emulator.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)>;
Expand Down Expand Up @@ -62,6 +72,8 @@ class emulator : public memory_manager, public utils::serializable

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)
Expand All @@ -87,17 +99,35 @@ class emulator : public memory_manager, public utils::serializable

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)
{
Expand All @@ -110,6 +140,18 @@ class emulator : public memory_manager, public utils::serializable
});
}

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;
};
46 changes: 25 additions & 21 deletions src/emulator/memory_manager.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

std::vector<uint8_t> data{};

for (const auto& reserved_region : this->reserved_regions_)
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;
}

std::vector<uint8_t> data{};

for (const auto& reserved_region : this->reserved_regions_)
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);
}
}
}
Expand Down
16 changes: 2 additions & 14 deletions src/emulator/memory_manager.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,22 +75,10 @@ class memory_manager
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;

Expand All @@ -100,6 +88,6 @@ class memory_manager
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);
};
59 changes: 55 additions & 4 deletions src/unicorn_emulator/unicorn_x64_emulator.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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_{};
Expand Down

0 comments on commit b68892c

Please sign in to comment.