Safely access buffer when mapping modules

src/common/utils/buffer_accessor.hpp

@@ -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_{};
+	};
+}

src/windows_emulator/module/module_mapping.cpp

@@ -2,9 +2,28 @@
 #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)
+	{
+		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());
+
+		return memory;
+	}
+
+	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)
@@ -12,28 +31,22 @@ namespace
 			return;
 		}
 
-		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();
-
-		const auto* export_directory = reinterpret_cast<const IMAGE_EXPORT_DIRECTORY*>(ptr + export_directory_entry.
-			VirtualAddress);
+		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));
@@ -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;
@@ -60,34 +84,31 @@ namespace
 			return;
 		}
 
-		std::vector<uint8_t> memory{};
-		memory.resize(binary.size_of_image);
-		emu.read_memory(binary.image_base, memory.data(), memory.size());
+		auto relocation_offset = directory->VirtualAddress;
 
-		const auto start = memory.data() + directory->VirtualAddress;
-		const auto end = start + directory->Size;
-
-		const auto* relocation = reinterpret_cast<const IMAGE_BASE_RELOCATION*>(start);
-
-		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)
 				{
@@ -95,38 +116,36 @@ namespace
 					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);
 			}
 
@@ -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))
 	{
@@ -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, ptr, *nt_headers);
-	apply_relocations(emu, binary, optional_header);
-	collect_exports(emu, binary, optional_header);
+	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);
+
+	emu.write_memory(binary.image_base, mapped_memory.data(), mapped_memory.size());
 
 	return binary;
 }

src/windows_emulator/module/module_mapping.hpp

@@ -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);