Implement basic memory manager

src/emulator/address_utils.hpp

@@ -0,0 +1,43 @@
+#pragma once
+#include <cstdint>
+
+inline bool is_within_start_and_end(const uint64_t value, const uint64_t start, const uint64_t end)
+{
+	return value >= start && value < end;
+}
+
+inline bool is_within_start_and_length(const uint64_t value, const uint64_t start, const uint64_t length)
+{
+	return is_within_start_and_end(value, start, start + length);
+}
+
+inline bool regions_intersect(const uint64_t start1, const uint64_t end1, const uint64_t start2, const uint64_t end2)
+{
+	return start1 < end2 && start2 < end1;
+}
+
+inline bool regions_with_length_intersect(const uint64_t start1, const uint64_t length1, const uint64_t start2,
+                                          const uint64_t length2)
+{
+	return regions_intersect(start1, start1 + length1, start2, start2 + length2);
+}
+
+inline uint64_t align_down(const uint64_t value, const uint64_t alignment)
+{
+	return value & ~(alignment - 1);
+}
+
+inline uint64_t align_up(const uint64_t value, const uint64_t alignment)
+{
+	return align_down(value + (alignment - 1), alignment);
+}
+
+inline uint64_t page_align_down(const uint64_t value, const uint64_t page_size = 0x1000)
+{
+	return align_down(value, page_size);
+}
+
+inline uint64_t page_align_up(const uint64_t value, const uint64_t page_size = 0x1000)
+{
+	return align_up(value, page_size);
+}

src/emulator/emulator.hpp

@@ -1,10 +1,9 @@
 #pragma once
 #include <chrono>
-#include <vector>
 #include <functional>
 #include <cassert>
 
-#include "memory_region.hpp"
+#include "memory_manager.hpp"
 
 struct emulator_hook;
 
@@ -15,7 +14,7 @@ using hook_callback = std::function<void()>;
 using simple_memory_hook_callback = std::function<void(uint64_t address, size_t size)>;
 using complex_memory_hook_callback = std::function<void(uint64_t address, size_t size, memory_operation operation)>;
 
-class emulator
+class emulator : public memory_manager
 {
 public:
 	emulator() = default;
@@ -26,25 +25,12 @@ public:
 	emulator(emulator&&) = delete;
 	emulator& operator=(emulator&&) = delete;
 
-	virtual ~emulator() = default;
-
 	virtual void start(uint64_t start, uint64_t end = 0, std::chrono::microseconds timeout = {}, size_t count = 0) = 0;
 	virtual void stop() = 0;
 
 	virtual void read_raw_register(int reg, void* value, size_t size) = 0;
 	virtual void write_raw_register(int reg, const void* value, size_t size) = 0;
 
-	virtual bool try_map_memory(uint64_t address, size_t size, memory_permission permissions) = 0;
-	virtual void map_memory(uint64_t address, size_t size, memory_permission permissions) = 0;
-	virtual void unmap_memory(uint64_t address, size_t size) = 0;
-
-	virtual void read_memory(uint64_t address, void* data, size_t size) = 0;
-	virtual void write_memory(uint64_t address, const void* data, size_t size) = 0;
-
-	virtual void protect_memory(uint64_t address, size_t size, memory_permission permissions) = 0;
-
-	virtual std::vector<memory_region> get_memory_regions() = 0;
-
 	virtual emulator_hook* hook_memory_access(uint64_t address, size_t size, memory_operation filter,
 	                                          complex_memory_hook_callback callback) = 0;
 	virtual emulator_hook* hook_instruction(int instruction_type, hook_callback callback) = 0;

src/emulator/memory_manager.hpp

@@ -0,0 +1,367 @@
+#pragma once
+#include <map>
+#include <vector>
+#include <optional>
+
+#include "memory_region.hpp"
+#include "address_utils.hpp"
+
+class memory_manager
+{
+public:
+	virtual ~memory_manager() = default;
+
+	virtual void read_memory(uint64_t address, void* data, size_t size) = 0;
+	virtual void write_memory(uint64_t address, const void* data, size_t size) = 0;
+
+	bool protect_memory(const uint64_t address, const size_t size, const memory_permission permissions,
+	                    memory_permission* old_permissions)
+	{
+		const auto entry = this->find_reserved_region(address);
+		if (entry == this->reserved_regions_.end())
+		{
+			return false;
+		}
+
+		const auto end = address + size;
+		const auto region_end = entry->first + entry->second.length;
+
+		if (region_end < end)
+		{
+			throw std::runtime_error("Cross region protect not supported yet!");
+		}
+
+		std::optional<memory_permission> old_first_permissions{};
+
+		auto& committed_regions = entry->second.committed_regions;
+		split_regions(committed_regions, {address, end});
+
+		for (auto& sub_region : committed_regions)
+		{
+			if (sub_region.first >= end)
+			{
+				break;
+			}
+
+			const auto sub_region_end = sub_region.first + sub_region.second.length;
+			if (sub_region.first >= address && sub_region_end <= end)
+			{
+				if (!old_first_permissions.has_value())
+				{
+					old_first_permissions = sub_region.second.pemissions;
+				}
+
+				this->apply_memory_protection(sub_region.first, sub_region.second.length, permissions);
+				sub_region.second.pemissions = permissions;
+			}
+		}
+
+		if (old_permissions)
+		{
+			*old_permissions = old_first_permissions.value_or(memory_permission::none);
+		}
+
+		merge_regions(committed_regions);
+		return true;
+	}
+
+	bool allocate_memory(const uint64_t address, const size_t size, const memory_permission permissions,
+	                     const bool reserve_only = false)
+	{
+		if (this->overlaps_reserved_region(address, size))
+		{
+			return false;
+		}
+
+		const auto entry = this->reserved_regions_.try_emplace(address, size).first;
+
+		if (!reserve_only)
+		{
+			this->map_memory(address, size, permissions);
+			entry->second.committed_regions[address] = committed_region{size, permissions};
+		}
+
+		return true;
+	}
+
+	bool commit_memory(const uint64_t address, const size_t size, const memory_permission permissions)
+	{
+		const auto entry = this->find_reserved_region(address);
+		if (entry == this->reserved_regions_.end())
+		{
+			return false;
+		}
+
+		const auto end = address + size;
+		const auto region_end = entry->first + entry->second.length;
+
+		if (region_end < end)
+		{
+			throw std::runtime_error("Cross region commit not supported yet!");
+		}
+
+		auto& committed_regions = entry->second.committed_regions;
+		split_regions(committed_regions, {address, end});
+
+		uint64_t last_region_start{};
+		const committed_region* last_region{nullptr};
+
+		for (auto& sub_region : committed_regions)
+		{
+			if (sub_region.first >= end)
+			{
+				break;
+			}
+
+			const auto sub_region_end = sub_region.first + sub_region.second.length;
+			if (sub_region.first >= address && sub_region_end <= end)
+			{
+				const auto map_start = last_region ? (last_region_start + last_region->length) : address;
+				const auto map_length = sub_region.first - map_start;
+
+				if (map_length > 0)
+				{
+					this->map_memory(map_start, map_length, permissions);
+					committed_regions[map_start] = committed_region{map_length, permissions};
+				}
+
+				last_region_start = sub_region.first;
+				last_region = &sub_region.second;
+			}
+		}
+
+		if (!last_region || (last_region_start + last_region->length) < end)
+		{
+			const auto map_start = last_region ? (last_region_start + last_region->length) : address;
+			const auto map_length = end - map_start;
+
+			this->map_memory(map_start, map_length, permissions);
+			committed_regions[map_start] = committed_region{map_length, permissions};
+		}
+
+		merge_regions(committed_regions);
+		return true;
+	}
+
+	bool decommit_memory(const uint64_t address, const size_t size)
+	{
+		const auto entry = this->find_reserved_region(address);
+		if (entry == this->reserved_regions_.end())
+		{
+			return false;
+		}
+
+		const auto end = address + size;
+		const auto region_end = entry->first + entry->second.length;
+
+		if (region_end < end)
+		{
+			throw std::runtime_error("Cross region decommit not supported yet!");
+		}
+
+		auto& committed_regions = entry->second.committed_regions;
+
+		split_regions(committed_regions, {address, end});
+
+		for (auto i = committed_regions.begin(); i != committed_regions.end();)
+		{
+			if (i->first >= end)
+			{
+				break;
+			}
+
+			const auto sub_region_end = i->first + i->second.length;
+			if (i->first >= address && sub_region_end <= end)
+			{
+				this->unmap_memory(i->first, i->second.length);
+				i = committed_regions.erase(i);
+				continue;
+			}
+
+			++i;
+		}
+
+		return true;
+	}
+
+	bool release_memory(const uint64_t address, size_t size)
+	{
+		const auto entry = this->reserved_regions_.find(address);
+		if (entry == this->reserved_regions_.end())
+		{
+			return false;
+		}
+
+		if (!size)
+		{
+			size = entry->second.length;
+		}
+
+		if (size > entry->second.length)
+		{
+			throw std::runtime_error("Cross region release not supported yet!");
+		}
+
+		const auto end = address + size;
+		auto& committed_regions = entry->second.committed_regions;
+
+		split_regions(committed_regions, {end});
+
+		for (auto i = committed_regions.begin(); i != committed_regions.end();)
+		{
+			if (i->first >= end)
+			{
+				break;
+			}
+
+			const auto sub_region_end = i->first + i->second.length;
+			if (i->first >= address && sub_region_end <= end)
+			{
+				this->unmap_memory(i->first, i->second.length);
+				i = committed_regions.erase(i);
+			}
+			else
+			{
+				++i;
+			}
+		}
+
+		entry->second.length -= size;
+		if (entry->second.length > 0)
+		{
+			this->reserved_regions_[address + size] = std::move(entry->second);
+		}
+
+		this->reserved_regions_.erase(entry);
+		return true;
+	}
+
+	uint64_t find_free_allocation_base(const size_t size) const
+	{
+		uint64_t start_address = 0x0000000000010000;
+
+		for (const auto& region : this->reserved_regions_)
+		{
+			if (!regions_with_length_intersect(start_address, size, region.first, region.second.length))
+			{
+				return start_address;
+			}
+
+			start_address = region.first + region.second.length;
+		}
+
+		if (start_address + size <= 0x00007ffffffeffff)
+		{
+			return start_address;
+		}
+
+		return 0;
+	}
+
+private:
+	struct committed_region
+	{
+		size_t length{};
+		memory_permission pemissions{};
+	};
+
+	using committed_region_map = std::map<uint64_t, committed_region>;
+
+	struct reserved_region
+	{
+		size_t length{};
+		committed_region_map committed_regions{};
+	};
+
+	using reserved_region_map = std::map<uint64_t, reserved_region>;
+	reserved_region_map reserved_regions_{};
+
+	reserved_region_map::iterator find_reserved_region(const uint64_t address)
+	{
+		if (this->reserved_regions_.empty())
+		{
+			return this->reserved_regions_.end();
+		}
+
+		auto upper_bound = this->reserved_regions_.upper_bound(address);
+		if (upper_bound == this->reserved_regions_.begin())
+		{
+			return this->reserved_regions_.end();
+		}
+
+		const auto entry = --upper_bound;
+		if (entry->first + entry->second.length <= address)
+		{
+			return this->reserved_regions_.end();
+		}
+
+		return entry;
+	}
+
+	bool overlaps_reserved_region(const uint64_t address, const size_t size) const
+	{
+		for (const auto& region : this->reserved_regions_)
+		{
+			if (regions_with_length_intersect(address, size, region.first, region.second.length))
+			{
+				return true;
+			}
+		}
+
+		return false;
+	}
+
+	static void split_regions(committed_region_map& regions, const std::vector<uint64_t>& split_points)
+	{
+		for (auto i = regions.begin(); i != regions.end(); ++i)
+		{
+			for (const auto split_point : split_points)
+			{
+				if (is_within_start_and_length(split_point, i->first, i->second.length) && i->first != split_point)
+				{
+					const auto first_length = split_point - i->first;
+					const auto second_length = i->second.length - first_length;
+
+					i->second.length = first_length;
+
+					regions[split_point] = committed_region{second_length, i->second.pemissions};
+				}
+			}
+		}
+	}
+
+	static void merge_regions(committed_region_map& regions)
+	{
+		for (auto i = regions.begin(); i != regions.end();)
+		{
+			assert(i->second.length > 0);
+
+			auto next = i;
+			std::advance(next, 1);
+
+			if (next == regions.end())
+			{
+				break;
+			}
+
+			assert(next->second.length > 0);
+
+			const auto end = i->first + i->second.length;
+			assert(end <= next->first);
+
+			if (end != next->first || i->second.pemissions != next->second.pemissions)
+			{
+				++i;
+				continue;
+			}
+
+			i->second.length += next->second.length;
+			regions.erase(next);
+		}
+	}
+
+	virtual void map_memory(uint64_t address, size_t size, memory_permission permissions) = 0;
+	virtual void unmap_memory(uint64_t address, size_t size) = 0;
+
+	virtual void apply_memory_protection(uint64_t address, size_t size, memory_permission permissions) = 0;
+};

src/emulator/memory_region.hpp

@@ -1,9 +1,14 @@
 #pragma once
 #include "memory_permission.hpp"
 
-struct memory_region
+struct basic_memory_region
 {
-	uint64_t start;
-	size_t length;
-	memory_permission pemissions;
+	uint64_t start{};
+	size_t length{};
+	memory_permission pemissions{};
+};
+
+struct memory_region : basic_memory_region
+{
+	bool committed{};
 };

src/unicorn_emulator/unicorn_x64_emulator.cpp

@@ -197,11 +197,6 @@ namespace unicorn
 				uce(uc_mem_map(*this, address, size, static_cast<uint32_t>(permissions)));
 			}
 
-			bool try_map_memory(const uint64_t address, const size_t size, memory_permission permissions) override
-			{
-				return uc_mem_map(*this, address, size, static_cast<uint32_t>(permissions)) == UC_ERR_OK;
-			}
-
 			void unmap_memory(const uint64_t address, const size_t size) override
 			{
 				uce(uc_mem_unmap(*this, address, size));
@@ -217,12 +212,12 @@ namespace unicorn
 				uce(uc_mem_write(*this, address, data, size));
 			}
 
-			void protect_memory(const uint64_t address, const size_t size, memory_permission permissions) override
+			void apply_memory_protection(const uint64_t address, const size_t size, memory_permission permissions) override
 			{
 				uce(uc_mem_protect(*this, address, size, static_cast<uint32_t>(permissions)));
 			}
 
-			std::vector<memory_region> get_memory_regions() override
+			/*std::vector<memory_region> get_memory_regions() override
 			{
 				const unicorn_memory_regions regions{*this};
 				const auto region_span = regions.get_span();
@@ -235,13 +230,14 @@ namespace unicorn
 					memory_region reg{};
 					reg.start = region.begin;
 					reg.length = region.end - region.begin;
+					reg.committed = true;
 					reg.pemissions = static_cast<memory_permission>(region.perms) & memory_permission::all;
 
 					result.push_back(reg);
 				}
 
 				return result;
-			}
+			}*/
 
 			emulator_hook* hook_instruction(int instruction_type,
 			                                hook_callback callback)

src/windows_emulator/main.cpp

@@ -23,7 +23,7 @@ namespace
 {
 	void setup_stack(x64_emulator& emu, const uint64_t stack_base, const size_t stack_size)
 	{
-		emu.map_memory(stack_base, stack_size, memory_permission::read_write);
+		emu.allocate_memory(stack_base, stack_size, memory_permission::read_write);
 
 		const uint64_t stack_end = stack_base + stack_size;
 		emu.reg(x64_register::rsp, stack_end);
@@ -43,14 +43,14 @@ namespace
 		};
 
 		emu.write_register(x64_register::msr, &value, sizeof(value));
-		emu.map_memory(segment_base, size, memory_permission::read_write);
+		emu.allocate_memory(segment_base, size, memory_permission::read_write);
 
 		return {emu, segment_base, size};
 	}
 
 	emulator_object<KUSER_SHARED_DATA> setup_kusd(x64_emulator& emu)
 	{
-		emu.map_memory(KUSD_ADDRESS, page_align_up(sizeof(KUSER_SHARED_DATA)), memory_permission::read);
+		emu.allocate_memory(KUSD_ADDRESS, page_align_up(sizeof(KUSER_SHARED_DATA)), memory_permission::read);
 
 		const emulator_object<KUSER_SHARED_DATA> kusd_object{emu, KUSD_ADDRESS};
 		kusd_object.access([](KUSER_SHARED_DATA& kusd)
@@ -84,22 +84,15 @@ namespace
 		binary.image_base = optional_header.ImageBase;
 		binary.size_of_image = optional_header.SizeOfImage;
 
-		while (true)
+		if (!emu.allocate_memory(binary.image_base, binary.size_of_image, memory_permission::read))
 		{
-			try
+			binary.image_base = emu.find_free_allocation_base(binary.size_of_image);
+			if ((optional_header.DllCharacteristics &
+					IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) == 0 || //
+				!emu.allocate_memory(
+					binary.image_base, binary.size_of_image, memory_permission::read))
 			{
-				emu.map_memory(binary.image_base, binary.size_of_image, memory_permission::read);
-				break;
-			}
-			catch (...)
-			{
-				binary.image_base += 0x10000;
-
-				if (binary.image_base < optional_header.ImageBase || (optional_header.DllCharacteristics &
-					IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) == 0)
-				{
-					throw std::runtime_error("Failed to map range");
-				}
+				throw std::runtime_error("Failed to map binary");
 			}
 		}
 
@@ -139,7 +132,7 @@ namespace
 
 			const auto size_of_section = page_align_up(std::max(section.SizeOfRawData, section.Misc.VirtualSize));
 
-			emu.protect_memory(target_ptr, size_of_section, permissions);
+			emu.protect_memory(target_ptr, size_of_section, permissions, nullptr);
 		}
 
 		auto& export_directory_entry = optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];

src/windows_emulator/memory_utils.hpp

@@ -2,37 +2,7 @@
 #include <cstdint>
 #include <emulator.hpp>
 
-inline bool is_within_start_and_end(const uint64_t value, const uint64_t start, const uint64_t end)
-{
-	return value >= start && value < end;
-}
-
-inline bool is_within_start_and_length(const uint64_t value, const uint64_t start, const uint64_t length)
-{
-	return is_within_start_and_end(value, start, start + length);
-}
-
-inline uint64_t align_down(const uint64_t value, const uint64_t alignment)
-{
-	return value & ~(alignment - 1);
-}
-
-inline uint64_t align_up(const uint64_t value, const uint64_t alignment)
-{
-	return align_down(value + (alignment - 1), alignment);
-}
-
-inline uint64_t page_align_down(const uint64_t value)
-{
-	return align_down(value, 0x1000);
-}
-
-inline uint64_t page_align_up(const uint64_t value)
-{
-	return align_up(value, 0x1000);
-}
-
-
+/*
 inline memory_permission get_memory_protection(emulator& emu, const uint64_t address)
 {
 	for (const auto& region : emu.get_memory_regions())
@@ -58,6 +28,7 @@ inline bool is_memory_allocated(emulator& emu, const uint64_t address)
 
 	return false;
 }
+*/
 
 inline memory_permission map_nt_to_emulator_protection(const uint32_t nt_protection)
 {

src/windows_emulator/syscalls.cpp

@@ -67,7 +67,7 @@ namespace
 		};
 
 		const auto ret = std::apply(handler, std::move(func_args));
-		c.emu.reg<uint64_t>(x64_register::rax, static_cast<uint64_t>(ret));
+		c.emu.reg<int64_t>(x64_register::rax, ret);
 	}
 
 	NTSTATUS handle_NtQueryPerformanceCounter(const syscall_context&,
@@ -115,7 +115,8 @@ namespace
 		return STATUS_NOT_SUPPORTED;
 	}
 
-	NTSTATUS handle_NtSetEvent(const syscall_context& c, const uint64_t handle, const emulator_object<LONG> previous_state)
+	NTSTATUS handle_NtSetEvent(const syscall_context& c, const uint64_t handle,
+	                           const emulator_object<LONG> previous_state)
 	{
 		if (handle & EVENT_BIT)
 		{
@@ -452,70 +453,21 @@ namespace
 		const auto size = page_align_up(bytes_to_protect.read());
 		bytes_to_protect.write(static_cast<uint32_t>(size));
 
-		const auto current_uc_protection = get_memory_protection(c.emu, address);
-		const auto current_protection = map_emulator_to_nt_protection(current_uc_protection);
+		const auto requested_protection = map_nt_to_emulator_protection(protection);
+
+		memory_permission old_protection_value{};
+		c.emu.protect_memory(address, size, requested_protection, &old_protection_value);
+
+		const auto current_protection = map_emulator_to_nt_protection(old_protection_value);
 		old_protection.write(current_protection);
 
-		const auto requested_protection = map_nt_to_emulator_protection(protection);
-		c.emu.protect_memory(address, size, requested_protection);
-
 		return STATUS_SUCCESS;
 	}
 
-	NTSTATUS handle_NtAllocateVirtualMemory(const syscall_context& c, const uint64_t process_handle,
-	                                        const emulator_object<uint64_t> base_address,
-	                                        uint64_t /*zero_bits*/,
-	                                        const emulator_object<uint64_t> bytes_to_allocate,
-	                                        const uint32_t /*allocation_type*/, const uint32_t page_protection)
-	{
-		if (process_handle != ~0ULL)
-		{
-			return STATUS_NOT_IMPLEMENTED;
-		}
-
-		constexpr auto allocation_granularity = 0x10000;
-		const auto allocation_bytes = bytes_to_allocate.read();
-		//allocation_bytes = align_up(allocation_bytes, allocation_granularity);
-		//bytes_to_allocate.write(allocation_bytes);
-
-		const auto protection = map_nt_to_emulator_protection(page_protection);
-
-		auto allocate_anywhere = false;
-		auto allocation_base = base_address.read();
-		if (!allocation_base)
-		{
-			allocate_anywhere = true;
-			allocation_base = allocation_granularity;
-		}
-		else if (is_memory_allocated(c.emu, allocation_base))
-		{
-			return STATUS_SUCCESS;
-		}
-
-		bool succeeded = false;
-
-		while (true)
-		{
-			succeeded = c.emu.try_map_memory(allocation_base, allocation_bytes, protection);
-			if (succeeded || !allocate_anywhere)
-			{
-				break;
-			}
-
-			allocation_base += allocation_granularity;
-		}
-
-		base_address.write(allocation_base);
-
-		return succeeded
-			       ? STATUS_SUCCESS
-			       : STATUS_NOT_SUPPORTED; // No idea what the correct code is
-	}
-
 	NTSTATUS handle_NtAllocateVirtualMemoryEx(const syscall_context& c, const uint64_t process_handle,
 	                                          const emulator_object<uint64_t> base_address,
 	                                          const emulator_object<uint64_t> bytes_to_allocate,
-	                                          const uint32_t /*allocation_type*/,
+	                                          const uint32_t allocation_type,
 	                                          const uint32_t page_protection)
 	{
 		if (process_handle != ~0ULL)
@@ -523,48 +475,56 @@ namespace
 			return STATUS_NOT_IMPLEMENTED;
 		}
 
-		constexpr auto allocation_granularity = 0x10000;
 		const auto allocation_bytes = bytes_to_allocate.read();
-		//allocation_bytes = align_up(allocation_bytes, allocation_granularity);
-		//bytes_to_allocate.write(allocation_bytes);
 
 		const auto protection = map_nt_to_emulator_protection(page_protection);
 
-		auto allocate_anywhere = false;
-		auto allocation_base = base_address.read();
-		if (!allocation_base)
+		auto potential_base = base_address.read();
+		if (!potential_base)
 		{
-			allocate_anywhere = true;
-			allocation_base = allocation_granularity;
-		}
-		else if (is_memory_allocated(c.emu, allocation_base))
-		{
-			return STATUS_SUCCESS;
+			potential_base = c.emu.find_free_allocation_base(allocation_bytes);
 		}
 
-		bool succeeded = false;
-
-		while (true)
+		if (!potential_base)
 		{
-			succeeded = c.emu.try_map_memory(allocation_base, allocation_bytes, protection);
-			if (succeeded || !allocate_anywhere)
-			{
-				break;
-			}
-
-			allocation_base += allocation_granularity;
+			return STATUS_MEMORY_NOT_ALLOCATED;
 		}
 
-		base_address.write(allocation_base);
+		base_address.write(potential_base);
 
-		return succeeded
+		const bool reserve = allocation_type & MEM_RESERVE;
+		const bool commit = allocation_type & MEM_COMMIT;
+
+		if ((allocation_type & ~(MEM_RESERVE | MEM_COMMIT | MEM_TOP_DOWN)) || (!commit && !reserve))
+		{
+			throw std::runtime_error("Unsupported allocation type!");
+		}
+
+		if (commit && !reserve)
+		{
+			return c.emu.commit_memory(potential_base, allocation_bytes, protection)
+				       ? STATUS_SUCCESS
+				       : STATUS_MEMORY_NOT_ALLOCATED;
+		}
+
+		return c.emu.allocate_memory(potential_base, allocation_bytes, protection, !commit)
 			       ? STATUS_SUCCESS
-			       : STATUS_NOT_SUPPORTED; // No idea what the correct code is
+			       : STATUS_MEMORY_NOT_ALLOCATED;
+	}
+
+	NTSTATUS handle_NtAllocateVirtualMemory(const syscall_context& c, const uint64_t process_handle,
+	                                        const emulator_object<uint64_t> base_address,
+	                                        uint64_t /*zero_bits*/,
+	                                        const emulator_object<uint64_t> bytes_to_allocate,
+	                                        const uint32_t allocation_type, const uint32_t page_protection)
+	{
+		return handle_NtAllocateVirtualMemoryEx(c, process_handle, base_address, bytes_to_allocate, allocation_type,
+		                                        page_protection);
 	}
 
 	NTSTATUS handle_NtFreeVirtualMemory(const syscall_context& c, const uint64_t process_handle,
 	                                    const emulator_object<uint64_t> base_address,
-	                                    const emulator_object<uint64_t> bytes_to_allocate)
+	                                    const emulator_object<uint64_t> bytes_to_allocate, uint32_t free_type)
 	{
 		if (process_handle != ~0ULL)
 		{
@@ -574,20 +534,21 @@ namespace
 		const auto allocation_base = base_address.read();
 		const auto allocation_size = bytes_to_allocate.read();
 
-		bool succeeded = false;
-		try
+		if (free_type & MEM_RELEASE)
 		{
-			c.emu.unmap_memory(allocation_base, allocation_size);
-			succeeded = true;
-		}
-		catch (...)
-		{
-			succeeded = false;
+			return c.emu.release_memory(allocation_base, allocation_size)
+				       ? STATUS_SUCCESS
+				       : STATUS_MEMORY_NOT_ALLOCATED;
 		}
 
-		return succeeded
-			       ? STATUS_SUCCESS
-			       : STATUS_NOT_SUPPORTED; // No idea what the correct code is
+		if (free_type & MEM_DECOMMIT)
+		{
+			return c.emu.decommit_memory(allocation_base, allocation_size)
+				       ? STATUS_SUCCESS
+				       : STATUS_MEMORY_NOT_ALLOCATED;
+		}
+
+		throw std::runtime_error("Bad free type");
 	}
 
 #define handle(id, handler) \
@@ -648,10 +609,12 @@ void handle_syscall(x64_emulator& emu, process_context& context)
 	{
 		printf("Syscall threw an exception: %X (%llX) - %s\n", syscall_id, address, e.what());
 		emu.reg<uint64_t>(x64_register::rax, STATUS_UNSUCCESSFUL);
+		emu.stop();
 	}
 	catch (...)
 	{
 		printf("Syscall threw an unknown exception: %X (%llX)\n", syscall_id, address);
 		emu.reg<uint64_t>(x64_register::rax, STATUS_UNSUCCESSFUL);
+		emu.stop();
 	}
 }