Skip to content

Commit

Permalink
Safely access buffer when mapping modules
Browse files Browse the repository at this point in the history
  • Loading branch information
@momo5502
momo5502 committed Sep 9, 2024
1 parent 73b1555 commit 9d21fd8
Show file tree
Hide file tree
Showing 3 changed files with 205 additions and 57 deletions.
118 changes: 118 additions & 0 deletions src/common/utils/buffer_accessor.hpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
#pragma once
#include <span>
#include <cstdint>
#include <stdexcept>

namespace utils
{
template <typename T, typename S = const uint8_t>
requires(std::is_trivially_copyable_v<T> && std::is_same_v<uint8_t, std::remove_cv_t<S>>)
class safe_object_accessor
{
public:
safe_object_accessor(const std::span<S> buffer, const size_t offset)
: buffer_(buffer)
, offset_(offset)
{
}

/*****************************************************************************
* Object is copied to make sure platform-dependent alignment requirements
* are respected
****************************************************************************/

T get(const size_t element_index = 0) const
{
T value{};
memcpy(&value, get_valid_pointer(element_index), size);
return value;
}

void set(const T value, const size_t element_index = 0) const
{
memcpy(get_valid_pointer(element_index), &value, size);
}

private:
static constexpr auto size = sizeof(T);

std::span<S> buffer_{};
size_t offset_{};

S* get_valid_pointer(const size_t element_index) const
{
const auto start_offset = offset_ + (size * element_index);
const auto end_offset = start_offset + size;
if (end_offset > buffer_.size())
{
throw std::runtime_error("Buffer accessor overflow");
}

return buffer_.data() + start_offset;
}
};

template <typename T>
requires(std::is_same_v<uint8_t, std::remove_cv_t<T>>)
class safe_buffer_accessor
{
public:
safe_buffer_accessor(const std::span<T> buffer)
: buffer_(buffer)
{
}

template <typename S>
safe_buffer_accessor(const safe_buffer_accessor<S>& obj)
: buffer_(obj.get_buffer())
{
}

template <typename S>
safe_object_accessor<S, T> as(const size_t offset) const
{
return {this->buffer_, offset};
}

T* get_pointer_for_range(const size_t offset, const size_t size) const
{
this->validate(offset, size);
return this->buffer_.data() + offset;
}

void validate(const size_t offset, const size_t size) const
{
const auto end = offset + size;
if (end > buffer_.size())
{
throw std::runtime_error("Buffer accessor overflow");
}
}

template <typename S = char>
std::basic_string<S> as_string(const size_t offset) const
{
safe_object_accessor<S> string_accessor{this->buffer_, offset};
std::basic_string<S> result{};

while (true)
{
auto value = string_accessor.get(result.size());
if (!value)
{
return result;
}

result.push_back(std::move(value));
}
}

std::span<T> get_buffer() const
{
return this->buffer_;
}

private:
const std::span<T> buffer_{};
};
}
142 changes: 86 additions & 56 deletions src/windows_emulator/module/module_mapping.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,38 +2,51 @@
#include "module_mapping.hpp"
#include <address_utils.hpp>

#include <utils/buffer_accessor.hpp>

namespace
{
void collect_exports(emulator& emu, mapped_module& binary, const IMAGE_OPTIONAL_HEADER& optional_header)
uint64_t get_first_section_offset(const IMAGE_NT_HEADERS& nt_headers, const uint64_t nt_headers_offset)
{
auto& export_directory_entry = optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
if (export_directory_entry.VirtualAddress == 0 || export_directory_entry.Size == 0)
{
return;
}
const auto first_section_absolute = reinterpret_cast<uint64_t>(IMAGE_FIRST_SECTION(&nt_headers));
const auto absolute_base = reinterpret_cast<uint64_t>(&nt_headers);
return nt_headers_offset + (first_section_absolute - absolute_base);
}

std::vector<uint8_t> read_mapped_memory(emulator& emu, const mapped_module& binary)
{
std::vector<uint8_t> memory{};
memory.resize(binary.size_of_image);
emu.read_memory(binary.image_base, memory.data(), memory.size());

const uint8_t* ptr = memory.data();
return memory;
}

const auto* export_directory = reinterpret_cast<const IMAGE_EXPORT_DIRECTORY*>(ptr + export_directory_entry.
VirtualAddress);
void collect_exports(mapped_module& binary, const utils::safe_buffer_accessor<const uint8_t> buffer,
const IMAGE_OPTIONAL_HEADER& optional_header)
{
auto& export_directory_entry = optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
if (export_directory_entry.VirtualAddress == 0 || export_directory_entry.Size == 0)
{
return;
}

const auto export_directory = buffer.as<IMAGE_EXPORT_DIRECTORY>(export_directory_entry.
VirtualAddress).get();

//const auto function_count = export_directory->NumberOfFunctions;
const auto names_count = export_directory->NumberOfNames;
const auto names_count = export_directory.NumberOfNames;

const auto* names = reinterpret_cast<const DWORD*>(ptr + export_directory->AddressOfNames);
const auto* ordinals = reinterpret_cast<const WORD*>(ptr + export_directory->AddressOfNameOrdinals);
const auto* functions = reinterpret_cast<const DWORD*>(ptr + export_directory->AddressOfFunctions);
const auto names = buffer.as<DWORD>(export_directory.AddressOfNames);
const auto ordinals = buffer.as<WORD>(export_directory.AddressOfNameOrdinals);
const auto functions = buffer.as<DWORD>(export_directory.AddressOfFunctions);

for (DWORD i = 0; i < names_count; i++)
{
exported_symbol symbol{};
symbol.ordinal = ordinals[i];
symbol.name = reinterpret_cast<const char*>(ptr + names[i]);
symbol.rva = functions[symbol.ordinal];
symbol.ordinal = ordinals.get(i);
symbol.name = buffer.as_string(names.get(i));
symbol.rva = functions.get(symbol.ordinal);
symbol.address = binary.image_base + symbol.rva;

binary.exports.push_back(std::move(symbol));
Expand All @@ -45,7 +58,18 @@ namespace
}
}

void apply_relocations(emulator& emu, const mapped_module& binary,
template <typename T>
requires(std::is_integral_v<T>)
void apply_relocation(const utils::safe_buffer_accessor<uint8_t> buffer, const uint64_t offset,
const uint64_t delta)
{
const auto obj = buffer.as<T>(offset);
const auto value = obj.get();
const auto new_value = value + static_cast<T>(delta);
obj.set(new_value);
}

void apply_relocations(const mapped_module& binary, const utils::safe_buffer_accessor<uint8_t> buffer,
const IMAGE_OPTIONAL_HEADER& optional_header)
{
const auto delta = binary.image_base - optional_header.ImageBase;
Expand All @@ -60,73 +84,68 @@ namespace
return;
}

std::vector<uint8_t> memory{};
memory.resize(binary.size_of_image);
emu.read_memory(binary.image_base, memory.data(), memory.size());

const auto start = memory.data() + directory->VirtualAddress;
const auto end = start + directory->Size;

const auto* relocation = reinterpret_cast<const IMAGE_BASE_RELOCATION*>(start);
auto relocation_offset = directory->VirtualAddress;

while (reinterpret_cast<const uint8_t*>(relocation) < end)
while (relocation_offset < directory->Size)
{
if (relocation->VirtualAddress <= 0 || relocation->SizeOfBlock <= 0)
const auto relocation = buffer.as<IMAGE_BASE_RELOCATION>(relocation_offset).get();

if (relocation.VirtualAddress <= 0 || relocation.SizeOfBlock <= 0)
{
break;
}

const auto dest = memory.data() + relocation->VirtualAddress;

const auto data_size = relocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION);
const auto data_size = relocation.SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION);
const auto entry_count = data_size / sizeof(uint16_t);

const auto entry_start = offset_pointer<uint16_t>(relocation, sizeof(IMAGE_BASE_RELOCATION));
const auto entries = std::span(entry_start, entry_count);
const auto entries = buffer.as<uint16_t>(relocation_offset + sizeof(IMAGE_BASE_RELOCATION));

for (const auto entry : entries)
relocation_offset += relocation.SizeOfBlock;

for (size_t i = 0; i < entry_count; ++i)
{
const auto entry = entries.get(i);

const int type = entry >> 12;
const int offset = entry & 0xfff;
const auto total_offset = relocation.VirtualAddress + offset;

switch (type)
{
case IMAGE_REL_BASED_ABSOLUTE:
break;

case IMAGE_REL_BASED_HIGHLOW:
*reinterpret_cast<DWORD*>(dest + offset) += static_cast<DWORD>(delta);
apply_relocation<DWORD>(buffer, total_offset, delta);
break;

case IMAGE_REL_BASED_DIR64:
*reinterpret_cast<ULONGLONG*>(dest + offset) += delta;
apply_relocation<ULONGLONG>(buffer, total_offset, delta);
break;

default:
throw std::runtime_error("Unknown relocation type: " + std::to_string(type));
}
}

relocation = offset_pointer<IMAGE_BASE_RELOCATION>(relocation, relocation->SizeOfBlock);
}

emu.write_memory(binary.image_base, memory.data(), memory.size());
}

void map_sections(emulator& emu, const mapped_module& binary, const unsigned char* ptr,
const IMAGE_NT_HEADERS& nt_headers)
void map_sections(emulator& emu, const mapped_module& binary,
const utils::safe_buffer_accessor<const uint8_t> buffer,
const IMAGE_NT_HEADERS& nt_headers, const uint64_t nt_headers_offset)
{
const std::span sections(IMAGE_FIRST_SECTION(&nt_headers), nt_headers.FileHeader.NumberOfSections);
const auto first_section_offset = get_first_section_offset(nt_headers, nt_headers_offset);
const auto sections = buffer.as<IMAGE_SECTION_HEADER>(first_section_offset);

for (const auto& section : sections)
for (size_t i = 0; i < nt_headers.FileHeader.NumberOfSections; ++i)
{
const auto section = sections.get(i);
const auto target_ptr = binary.image_base + section.VirtualAddress;

if (section.SizeOfRawData > 0)
{
const void* source_ptr = ptr + section.PointerToRawData;

const auto size_of_data = std::min(section.SizeOfRawData, section.Misc.VirtualSize);
const auto* source_ptr = buffer.get_pointer_for_range(section.PointerToRawData, size_of_data);
emu.write_memory(target_ptr, source_ptr, size_of_data);
}

Expand Down Expand Up @@ -160,21 +179,24 @@ namespace
}
}

std::optional<mapped_module> map_module_from_data(emulator& emu, const std::vector<uint8_t>& data,

std::optional<mapped_module> map_module_from_data(emulator& emu, const std::span<const uint8_t> data,
std::filesystem::path file)
{
mapped_module binary{};
binary.path = std::move(file);
binary.name = binary.path.filename().string();

// TODO: Range checks
auto* ptr = data.data();
auto* dos_header = reinterpret_cast<const IMAGE_DOS_HEADER*>(ptr);
auto* nt_headers = reinterpret_cast<const IMAGE_NT_HEADERS*>(ptr + dos_header->e_lfanew);
auto& optional_header = nt_headers->OptionalHeader;
utils::safe_buffer_accessor buffer{data};

const auto dos_header = buffer.as<IMAGE_DOS_HEADER>(0).get();
const auto nt_headers_offset = dos_header.e_lfanew;

const auto nt_headers = buffer.as<IMAGE_NT_HEADERS>(nt_headers_offset).get();
auto& optional_header = nt_headers.OptionalHeader;

binary.image_base = optional_header.ImageBase;
binary.size_of_image = optional_header.SizeOfImage;
binary.size_of_image = optional_header.SizeOfImage; // TODO: Sanitize

if (!emu.allocate_memory(binary.image_base, binary.size_of_image, memory_permission::read))
{
Expand All @@ -192,11 +214,19 @@ std::optional<mapped_module> map_module_from_data(emulator& emu, const std::vect

printf("Mapping %s at %llX\n", binary.path.generic_string().c_str(), binary.image_base);

emu.write_memory(binary.image_base, ptr, optional_header.SizeOfHeaders);
const auto* header_buffer = buffer.get_pointer_for_range(0, optional_header.SizeOfHeaders);
emu.write_memory(binary.image_base, header_buffer,
optional_header.SizeOfHeaders);

map_sections(emu, binary, buffer, nt_headers, nt_headers_offset);

auto mapped_memory = read_mapped_memory(emu, binary);
utils::safe_buffer_accessor<uint8_t> mapped_buffer{mapped_memory};

apply_relocations(binary, mapped_buffer, optional_header);
collect_exports(binary, mapped_buffer, optional_header);

map_sections(emu, binary, ptr, *nt_headers);
apply_relocations(emu, binary, optional_header);
collect_exports(emu, binary, optional_header);
emu.write_memory(binary.image_base, mapped_memory.data(), mapped_memory.size());

return binary;
}
Expand Down
2 changes: 1 addition & 1 deletion src/windows_emulator/module/module_mapping.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
#include <x64_emulator.hpp>
#include "mapped_module.hpp"

std::optional<mapped_module> map_module_from_data(emulator& emu, const std::vector<uint8_t>& data,
std::optional<mapped_module> map_module_from_data(emulator& emu, std::span<const uint8_t> data,
std::filesystem::path file);
std::optional<mapped_module> map_module_from_file(emulator& emu, std::filesystem::path file);

Expand Down

0 comments on commit 9d21fd8

Please sign in to comment.