Finish context saving/restoration

This commit is contained in:
momo5502
2024-09-12 17:53:41 +02:00
parent d94a92df45
commit b1cbc5a7ad
11 changed files with 458 additions and 79 deletions

View File

@@ -84,17 +84,52 @@ static void serialize(utils::buffer_serializer& buffer, const memory_manager::re
static void deserialize(utils::buffer_deserializer& buffer, memory_manager::reserved_region& region)
{
region.length = static_cast<size_t>(buffer.read<uint64_t>());
region.committed_regions = buffer.read_map<uint64_t, memory_manager::committed_region>();
buffer.read_map(region.committed_regions);
}
void memory_manager::serialize_memory_state(utils::buffer_serializer& buffer) const
{
buffer.write_map(this->reserved_regions_);
if (!this->use_in_place_serialization())
{
std::vector<uint8_t> data{};
for (const auto& reserved_region : this->reserved_regions_)
{
for (const auto& region : reserved_region.second.committed_regions)
{
data.resize(region.second.length);
this->read_memory(region.first, data.data(), region.second.length);
buffer.write(data.data(), region.second.length);
}
}
}
}
void memory_manager::deserialize_memory_state(utils::buffer_deserializer& buffer)
{
this->reserved_regions_ = buffer.read_map<uint64_t, reserved_region>();
buffer.read_map(this->reserved_regions_);
if (!this->use_in_place_serialization())
{
std::vector<uint8_t> data{};
for (const auto& reserved_region : this->reserved_regions_)
{
for (const auto& region : reserved_region.second.committed_regions)
{
data.resize(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);
}
}
}
}
bool memory_manager::protect_memory(const uint64_t address, const size_t size, const memory_permission permissions,

View File

@@ -45,8 +45,8 @@ public:
this->write_memory(address, &value, sizeof(value));
}
virtual void read_memory(uint64_t address, void* data, size_t size) = 0;
virtual bool try_read_memory(uint64_t address, void* data, size_t size) = 0;
virtual void read_memory(uint64_t address, void* data, size_t size) const = 0;
virtual bool try_read_memory(uint64_t address, void* data, size_t size) const = 0;
virtual void write_memory(uint64_t address, const void* data, size_t size) = 0;
bool protect_memory(uint64_t address, size_t size, memory_permission permissions,
@@ -75,10 +75,22 @@ 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;

View File

@@ -2,6 +2,7 @@
#include <span>
#include <vector>
#include <string_view>
#include <stdexcept>
#include <cstring>
@@ -20,32 +21,52 @@ namespace utils
namespace detail
{
template <typename, typename = void>
struct has_serialize_function : std::false_type {};
struct has_serialize_function : std::false_type
{
};
template <typename T>
struct has_serialize_function<T, std::void_t<decltype(serialize(std::declval<buffer_serializer&>(), std::declval<const T&>()))>>
: std::true_type {};
struct has_serialize_function<T, std::void_t<decltype(serialize(std::declval<buffer_serializer&>(),
std::declval<const T&>()))>>
: std::true_type
{
};
template <typename, typename = void>
struct has_deserialize_function : std::false_type {};
struct has_deserialize_function : std::false_type
{
};
template <typename T>
struct has_deserialize_function<T, std::void_t<decltype(deserialize(std::declval<buffer_deserializer&>(), std::declval<T&>()))>>
: std::true_type {};
struct has_deserialize_function<T, std::void_t<decltype(deserialize(
std::declval<buffer_deserializer&>(), std::declval<T&>()))>>
: std::true_type
{
};
}
class buffer_deserializer
{
public:
template <typename T>
buffer_deserializer(const std::span<T>& buffer)
buffer_deserializer(const std::span<T> buffer)
: buffer_(reinterpret_cast<const std::byte*>(buffer.data()), buffer.size() * sizeof(T))
{
static_assert(std::is_trivially_copyable_v<T>, "Type must be trivially copyable");
}
template <typename T>
buffer_deserializer(const std::vector<T>& buffer)
: buffer_deserializer(std::span(buffer))
{
}
std::span<const std::byte> read_data(const size_t length)
{
#ifndef NDEBUG
const uint64_t real_old_size = this->offset_;
#endif
if (this->offset_ + length > this->buffer_.size())
{
throw std::runtime_error("Out of bounds read from byte buffer");
@@ -54,6 +75,23 @@ namespace utils
const std::span result(this->buffer_.data() + this->offset_, length);
this->offset_ += length;
#ifndef NDEBUG
uint64_t old_size{};
if (this->offset_ + sizeof(old_size) > this->buffer_.size())
{
throw std::runtime_error("Out of bounds read from byte buffer");
}
memcpy(&old_size, this->buffer_.data() + this->offset_, sizeof(old_size));
if (old_size != real_old_size)
{
throw std::runtime_error("Reading from serialized buffer mismatches written data!");
}
this->offset_ += sizeof(old_size);
#endif
return result;
}
@@ -64,10 +102,8 @@ namespace utils
}
template <typename T>
T read()
void read(T& object)
{
T object{};
if constexpr (std::is_base_of_v<serializable, T>)
{
object.deserialize(*this);
@@ -78,63 +114,98 @@ namespace utils
}
else if constexpr (std::is_trivially_copyable_v<T>)
{
this->read(&object, sizeof(object));
union
{
T* type_{};
void* void_;
} pointers;
pointers.type_ = &object;
this->read(pointers.void_, sizeof(object));
}
else
{
static_assert(std::false_type::value, "Key must be trivially copyable or implement serializable!");
std::abort();
}
}
template <typename T>
T read()
{
T object{};
this->read(object);
return object;
}
template <typename T>
void read_vector(std::vector<T>& result)
{
const auto size = this->read<uint64_t>();
result.clear();
result.reserve(size);
for (uint64_t i = 0; i < size; ++i)
{
result.emplace_back(this->read<T>());
}
}
template <typename T>
std::vector<T> read_vector()
{
static_assert(std::is_trivially_copyable_v<T>, "Type must be trivially copyable");
std::vector<T> result{};
const auto size = this->read<uint64_t>();
const auto totalSize = size * sizeof(T);
if (this->offset_ + totalSize > this->buffer_.size())
{
throw std::runtime_error("Out of bounds read from byte buffer");
}
result.resize(size);
this->read(result.data(), totalSize);
this->read_vector(result);
return result;
}
template <typename Key, typename Value>
std::map<Key, Value> read_map()
template <typename Map>
void read_map(Map& map)
{
using key_type = typename Map::key_type;
using value_type = typename Map::mapped_type;
map.clear();
const auto size = this->read<uint64_t>();
std::map<Key, Value> map{};
for (uint64_t i = 0; i < size; ++i)
{
auto key = this->read<Key>();
auto value = this->read<Value>();
auto key = this->read<key_type>();
auto value = this->read<value_type>();
map[std::move(key)] = std::move(value);
}
}
template <typename Map>
Map read_map()
{
Map map{};
this->read_map(map);
return map;
}
std::string read_string()
template <typename T = char>
void read_string(std::basic_string<T>& result)
{
std::string result{};
const auto size = this->read<uint64_t>();
const auto span = this->read_data(size);
result.resize(size);
memcpy(result.data(), span.data(), size);
result.clear();
result.reserve(size);
for (uint64_t i = 0; i < size; ++i)
{
result.push_back(this->read<T>());
}
}
template <typename T= char>
std::basic_string<T> read_string()
{
std::basic_string<T> result{};
this->read_string(result);
return result;
}
@@ -165,7 +236,17 @@ namespace utils
void write(const void* buffer, const size_t length)
{
this->buffer_.append(static_cast<const char*>(buffer), length);
#ifndef NDEBUG
const uint64_t old_size = this->buffer_.size();
#endif
const auto* byte_buffer = static_cast<const std::byte*>(buffer);
this->buffer_.insert(this->buffer_.end(), byte_buffer, byte_buffer + length);
#ifndef NDEBUG
const auto* security_buffer = reinterpret_cast<const std::byte*>(&old_size);
this->buffer_.insert(this->buffer_.end(), security_buffer, security_buffer + sizeof(old_size));
#endif
}
void write(const buffer_serializer& object)
@@ -187,7 +268,15 @@ namespace utils
}
else if constexpr (std::is_trivially_copyable_v<T>)
{
this->write(&object, sizeof(object));
union
{
const T* type_{};
const void* void_;
} pointers;
pointers.type_ = &object;
this->write(pointers.void_, sizeof(object));
}
else
{
@@ -207,8 +296,26 @@ namespace utils
}
}
template <typename Key, typename Value>
void write_map(const std::map<Key, Value>& map)
template <typename T>
void write_vector(const std::vector<T> vec)
{
this->write_span(std::span(vec));
}
template <typename T>
void write_string(const std::basic_string_view<T> str)
{
this->write_span<const T>(str);
}
template <typename T>
void write_string(const std::basic_string<T>& str)
{
this->write_string(std::basic_string_view<T>(str));
}
template <typename Map>
void write_map(const Map& map)
{
this->write<uint64_t>(map.size());
@@ -219,17 +326,41 @@ namespace utils
}
}
const std::string& get_buffer() const
const std::vector<std::byte>& get_buffer() const
{
return this->buffer_;
}
std::string move_buffer()
std::vector<std::byte> move_buffer()
{
return std::move(this->buffer_);
}
private:
std::string buffer_{};
std::vector<std::byte> buffer_{};
};
template <>
inline void buffer_deserializer::read<std::string>(std::string& object)
{
object = this->read_string<char>();
}
template <>
inline void buffer_deserializer::read<std::wstring>(std::wstring& object)
{
object = this->read_string<wchar_t>();
}
template <>
inline void buffer_serializer::write<std::string>(const std::string& object)
{
this->write_string(object);
}
template <>
inline void buffer_serializer::write<std::wstring>(const std::wstring& object)
{
this->write_string(object);
}
}

View File

@@ -109,10 +109,12 @@ namespace unicorn
class uc_context_serializer
{
public:
uc_context_serializer(uc_engine* uc)
uc_context_serializer(uc_engine* uc, const bool in_place)
: uc_(uc)
, size_(uc_context_size(uc))
{
uc_ctl_context_mode(uc, UC_CTL_CONTEXT_CPU | (in_place ? UC_CTL_CONTEXT_MEMORY : 0));
this->size_ = uc_context_size(uc);
uce(uc_context_alloc(uc, &this->context_));
}
@@ -296,12 +298,12 @@ namespace unicorn
uce(uc_mem_unmap(*this, address, size));
}
bool try_read_memory(const uint64_t address, void* data, const size_t size) override
bool try_read_memory(const uint64_t address, void* data, const size_t size) const override
{
return uc_mem_read(*this, address, data, size) == UC_ERR_OK;
}
void read_memory(const uint64_t address, void* data, const size_t size) override
void read_memory(const uint64_t address, void* data, const size_t size) const override
{
uce(uc_mem_read(*this, address, data, size));
}
@@ -498,13 +500,13 @@ namespace unicorn
void serialize_state(utils::buffer_serializer& buffer) const override
{
const uc_context_serializer serializer(this->uc_);
const uc_context_serializer serializer(this->uc_, this->use_in_place_serialization());
serializer.serialize(buffer);
}
void deserialize_state(utils::buffer_deserializer& buffer) override
{
const uc_context_serializer serializer(this->uc_);
const uc_context_serializer serializer(this->uc_, this->use_in_place_serialization());
serializer.deserialize(buffer);
}

View File

@@ -2,21 +2,21 @@
#include "memory_utils.hpp"
template <typename T>
class emulator_object
class emulator_object : utils::serializable
{
public:
using value_type = T;
emulator_object() = default;
emulator_object(emulator& emu, const void* address)
: emulator_object(emu, reinterpret_cast<uint64_t>(address))
emulator_object(emulator& emu, const uint64_t address = 0)
: emu_(&emu)
, address_(address)
{
}
emulator_object(emulator& emu, const uint64_t address)
: emu_(&emu)
, address_(address)
emulator_object(emulator& emu, const void* address)
: emulator_object(emu, reinterpret_cast<uint64_t>(address))
{
}
@@ -68,15 +68,28 @@ public:
this->write(obj, index);
}
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->address_);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->address_);
}
private:
emulator* emu_{};
uint64_t address_{};
};
class emulator_allocator
class emulator_allocator : utils::serializable
{
public:
emulator_allocator() = default;
emulator_allocator(emulator& emu)
: emu_(&emu)
{
}
emulator_allocator(emulator& emu, const uint64_t address, const uint64_t size)
: emu_(&emu)
@@ -146,6 +159,20 @@ public:
return this->size_;
}
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->address_);
buffer.write(this->size_);
buffer.write(this->active_address_);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->address_);
buffer.read(this->size_);
buffer.read(this->active_address_);
}
private:
emulator* emu_{};
uint64_t address_{};

View File

@@ -68,7 +68,8 @@ constexpr handle make_pseudo_handle(const uint32_t id, const handle_types::type
}
template <handle_types::type Type, typename T>
class handle_store
requires(std::is_base_of_v<utils::serializable, T>)
class handle_store : utils::serializable
{
public:
handle store(T value)
@@ -134,6 +135,16 @@ public:
return this->erase(hh);
}
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write_map(this->store_);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read_map(this->store_);
}
private:
using value_map = std::map<uint32_t, T>;

View File

@@ -289,7 +289,7 @@ namespace
process_context setup_context(x64_emulator& emu, const std::filesystem::path& file)
{
process_context context{};
process_context context{emu};
setup_stack(emu, STACK_ADDRESS, STACK_SIZE);
setup_gdt(emu);

View File

@@ -2,6 +2,48 @@
#include "module_manager.hpp"
#include "module_mapping.hpp"
static void serialize(utils::buffer_serializer& buffer, const exported_symbol& sym)
{
buffer.write(sym.name);
buffer.write(sym.ordinal);
buffer.write(sym.rva);
buffer.write(sym.address);
}
static void deserialize(utils::buffer_deserializer& buffer, exported_symbol& sym)
{
buffer.read(sym.name);
buffer.read(sym.ordinal);
buffer.read(sym.rva);
buffer.read(sym.address);
}
static void serialize(utils::buffer_serializer& buffer, const mapped_module& mod)
{
buffer.write_string(mod.name);
buffer.write_string(mod.path.wstring());
buffer.write(mod.image_base);
buffer.write(mod.size_of_image);
buffer.write(mod.entry_point);
buffer.write_vector(mod.exports);
buffer.write_map(mod.address_names);
}
static void deserialize(utils::buffer_deserializer& buffer, mapped_module& mod)
{
mod.name = buffer.read_string();
mod.path = buffer.read_string<wchar_t>();
buffer.read(mod.image_base);
buffer.read(mod.size_of_image);
buffer.read(mod.entry_point);
buffer.read_vector(mod.exports);
buffer.read_map(mod.address_names);
}
module_manager::module_manager(emulator& emu)
: emu_(&emu)
{
@@ -30,3 +72,13 @@ mapped_module* module_manager::map_module(const std::filesystem::path& file)
const auto entry = this->modules_.try_emplace(image_base, std::move(*mod));
return &entry.first->second;
}
void module_manager::serialize(utils::buffer_serializer& buffer) const
{
buffer.write_map(this->modules_);
}
void module_manager::deserialize(utils::buffer_deserializer& buffer)
{
buffer.read_map(this->modules_);
}

View File

@@ -2,10 +2,9 @@
#include "mapped_module.hpp"
#include <emulator.hpp>
class module_manager
class module_manager : public utils::serializable
{
public:
module_manager() = default; // TODO: Get rid of that
module_manager(emulator& emu);
mapped_module* map_module(const std::filesystem::path& file);
@@ -21,6 +20,9 @@ public:
return nullptr;
}
void serialize(utils::buffer_serializer& buffer) const override;
void deserialize(utils::buffer_deserializer& buffer) override;
private:
emulator* emu_{};

View File

@@ -5,7 +5,9 @@
#include "module/module_manager.hpp"
#include <utils/nt_handle.hpp>
struct event
#include <x64_emulator.hpp>
struct event : utils::serializable
{
bool signaled{};
EVENT_TYPE type{};
@@ -21,30 +23,80 @@ struct event
return res;
}
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->signaled);
buffer.write(this->type);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->signaled);
buffer.read(this->type);
}
};
struct file
struct file : utils::serializable
{
utils::nt::handle<INVALID_HANDLE_VALUE> handle{};
std::wstring name{};
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->name);
// TODO: Serialize handle
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->name);
this->handle = INVALID_HANDLE_VALUE;
}
};
struct semaphore
struct semaphore : utils::serializable
{
std::wstring name{};
volatile uint32_t current_count{};
uint32_t max_count{};
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->name);
buffer.write(this->current_count);
buffer.write(this->max_count);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->name);
buffer.read(this->current_count);
buffer.read(this->max_count);
}
};
struct process_context
struct process_context : utils::serializable
{
uint64_t executed_instructions{0};
emulator_object<TEB> teb{};
emulator_object<PEB> peb{};
emulator_object<RTL_USER_PROCESS_PARAMETERS> process_params{};
emulator_object<KUSER_SHARED_DATA> kusd{};
process_context(x64_emulator& emu)
: emu(&emu)
, teb(emu)
, peb(emu)
, process_params(emu)
, kusd(emu)
, module_manager(emu)
, gs_segment(emu)
{
}
module_manager module_manager{};
x64_emulator* emu{};
uint64_t executed_instructions{0};
emulator_object<TEB> teb;
emulator_object<PEB> peb;
emulator_object<RTL_USER_PROCESS_PARAMETERS> process_params;
emulator_object<KUSER_SHARED_DATA> kusd;
module_manager module_manager;
mapped_module* executable{};
mapped_module* ntdll{};
@@ -56,7 +108,53 @@ struct process_context
handle_store<handle_types::file, file> files{};
handle_store<handle_types::semaphore, semaphore> semaphores{};
std::map<uint16_t, std::wstring> atoms{};
emulator_allocator gs_segment{};
emulator_allocator gs_segment;
bool verbose{false};
void serialize(utils::buffer_serializer& buffer) const override
{
buffer.write(this->executed_instructions);
buffer.write(this->teb);
buffer.write(this->peb);
buffer.write(this->process_params);
buffer.write(this->kusd);
buffer.write(this->module_manager);
buffer.write(this->executable->image_base);
buffer.write(this->ntdll->image_base);
buffer.write(this->win32u->image_base);
buffer.write(this->shared_section_size);
buffer.write(this->events);
buffer.write(this->files);
buffer.write(this->semaphores);
buffer.write_map(this->atoms);
buffer.write(this->gs_segment);
}
void deserialize(utils::buffer_deserializer& buffer) override
{
buffer.read(this->executed_instructions);
buffer.read(this->teb);
buffer.read(this->peb);
buffer.read(this->process_params);
buffer.read(this->kusd);
buffer.read(this->module_manager);
const auto executable_base = buffer.read<uint64_t>();
const auto ntdll_base = buffer.read<uint64_t>();
const auto win32u_base = buffer.read<uint64_t>();
this->executable = this->module_manager.find_by_address(executable_base);
this->ntdll = this->module_manager.find_by_address(ntdll_base);
this->win32u = this->module_manager.find_by_address(win32u_base);
buffer.read(this->shared_section_size);
buffer.read(this->events);
buffer.read(this->files);
buffer.read(this->semaphores);
buffer.read_map(this->atoms);
buffer.read(this->gs_segment);
}
};

View File

@@ -308,7 +308,10 @@ namespace
//return STATUS_NOT_SUPPORTED;
}
event e{initial_state != FALSE, event_type};
event e{};
e.type = event_type;
e.signaled = initial_state != FALSE;
const auto handle = c.proc.events.store(std::move(e));
event_handle.write(handle.bits);
@@ -414,7 +417,10 @@ namespace
return STATUS_FILE_INVALID;
}
const auto handle = c.proc.files.store(file{{}, std::move(filename)});
file f{};
f.name = std::move(filename);
const auto handle = c.proc.files.store(std::move(f));
section_handle.write(handle.bits);
return STATUS_SUCCESS;
@@ -1316,12 +1322,15 @@ namespace
return STATUS_NOT_SUPPORTED;
}
NTSTATUS handle_NtAlpcSendWaitReceivePort(const syscall_context& /*c*/, const uint64_t /*port_handle*/, const ULONG /*flags*/,
NTSTATUS handle_NtAlpcSendWaitReceivePort(const syscall_context& /*c*/, const uint64_t /*port_handle*/,
const ULONG /*flags*/,
const emulator_object<PORT_MESSAGE> /*send_message*/,
const emulator_object<ALPC_MESSAGE_ATTRIBUTES> /*send_message_attributes*/,
const emulator_object<ALPC_MESSAGE_ATTRIBUTES> /*send_message_attributes*/
,
const emulator_object<PORT_MESSAGE> receive_message,
const emulator_object<SIZE_T> /*buffer_length*/,
const emulator_object<ALPC_MESSAGE_ATTRIBUTES> /*receive_message_attributes*/,
const emulator_object<ALPC_MESSAGE_ATTRIBUTES>
/*receive_message_attributes*/,
const emulator_object<LARGE_INTEGER> /*timeout*/)
{
receive_message.access([](PORT_MESSAGE& msg)