Better handle abstraction

src/windows_emulator/handles.hpp

@@ -0,0 +1,174 @@
+#pragma once
+
+struct handle_types
+{
+	enum type : uint16_t
+	{
+		file,
+		event,
+		section,
+		symlink,
+		directory,
+	};
+};
+
+#pragma pack(push)
+#pragma pack(1)
+struct handle_value
+{
+	uint64_t id : 32;
+	uint64_t type : 16;
+	uint64_t padding : 15;
+	uint64_t is_pseudo : 1;
+};
+#pragma pack(pop)
+
+static_assert(sizeof(handle_value) == 8);
+
+union handle
+{
+	handle_value value;
+	uint64_t bits;
+	HANDLE h;
+};
+
+inline bool operator==(const handle& h1, const handle& h2)
+{
+	return h1.bits == h2.bits;
+}
+
+inline bool operator==(const handle& h1, const uint64_t& h2)
+{
+	return h1.bits == h2;
+}
+
+inline handle_value get_handle_value(const uint64_t h)
+{
+	handle hh{};
+	hh.bits = h;
+	return hh.value;
+}
+
+constexpr handle make_handle(const uint32_t id, const handle_types::type type, const bool is_pseudo)
+{
+	handle_value value{};
+
+	value.padding = 0;
+	value.id = id;
+	value.type = type;
+	value.is_pseudo = is_pseudo;
+
+	return {value};
+}
+
+constexpr handle make_pseudo_handle(const uint32_t id, const handle_types::type type)
+{
+	return make_handle(id, type, true);
+}
+
+template <handle_types::type Type, typename T>
+class handle_store
+{
+public:
+	handle store(T value)
+	{
+		auto index = this->find_free_index();
+		this->store_[index] = std::move(value);
+
+		handle h{};
+		h.bits = 0;
+		h.value.is_pseudo = false;
+		h.value.type = Type;
+		h.value.id = index;
+
+		return h;
+	}
+
+	T* get(const handle_value h)
+	{
+		const auto entry = this->get_iterator(h);
+		if (entry == this->store_.end())
+		{
+			return nullptr;
+		}
+
+		return &entry->second;
+	}
+
+	T* get(const handle h)
+	{
+		return this->get(h.value);
+	}
+
+	T* get(const uint64_t h)
+	{
+		handle hh{};
+		hh.bits = h;
+
+		return this->get(hh);
+	}
+
+	bool erase(const handle_value h)
+	{
+		const auto entry = this->get_iterator(h);
+		if (entry == this->store_.end())
+		{
+			return false;
+		}
+
+		this->store_.erase(entry);
+		return true;
+	}
+
+	bool erase(const handle h)
+	{
+		return this->erase(h.value);
+	}
+
+	bool erase(const uint64_t h)
+	{
+		handle hh{};
+		hh.bits = h;
+
+		return this->erase(hh);
+	}
+
+private:
+	using value_map = std::map<uint32_t, T>;
+
+	typename value_map::iterator get_iterator(const handle_value h)
+	{
+		if (h.type != Type || h.is_pseudo)
+		{
+			return this->store_.end();
+		}
+
+		return this->store_.find(h.id);
+	}
+
+	uint32_t find_free_index()
+	{
+		uint32_t index = 1;
+		for (; index > 0; ++index)
+		{
+			if (!this->store_.contains(index))
+			{
+				break;
+			}
+		}
+
+		return index;
+	}
+
+
+	value_map store_{};
+};
+
+constexpr auto KNOWN_DLLS_DIRECTORY = make_pseudo_handle(0x1337, handle_types::directory);
+constexpr auto KNOWN_DLLS_SYMLINK = make_pseudo_handle(0x1337, handle_types::symlink);
+constexpr auto SHARED_SECTION = make_pseudo_handle(0x1337, handle_types::section);
+constexpr auto CONSOLE_SERVER = make_pseudo_handle(0x1337, handle_types::section);
+
+constexpr auto CONSOLE_HANDLE = make_pseudo_handle(0x1, handle_types::file);
+constexpr auto STDOUT_HANDLE = make_pseudo_handle(0x2, handle_types::file);
+constexpr auto STDIN_HANDLE = make_pseudo_handle(0x3, handle_types::file);

src/windows_emulator/process_context.hpp

@@ -1,5 +1,6 @@
 #pragma once
 #include "emulator_utils.hpp"
+#include "handles.hpp"
 
 struct exported_symbol
 {
@@ -37,6 +38,11 @@ struct event
 	}
 };
 
+struct file
+{
+	std::wstring name{};
+};
+
 struct process_context
 {
 	emulator_object<TEB> teb{};
@@ -47,9 +53,8 @@ struct process_context
 	mapped_binary executable{};
 	mapped_binary ntdll{};
 
-	std::map<uint32_t, event> events{};
-	std::map<uint32_t, HANDLE> os_handles{};
-	std::map<uint32_t, std::wstring> files{};
+	handle_store<handle_types::event, event> events{};
+	handle_store<handle_types::file, file> files{};
 	emulator_allocator gs_segment{};
 
 	bool verbose{false};

src/windows_emulator/syscalls.cpp

@@ -224,24 +224,18 @@ namespace
 	NTSTATUS handle_NtSetEvent(const syscall_context& c, const uint64_t handle,
 	                           const emulator_object<LONG> previous_state)
 	{
-		const auto value = get_handle_value(handle);
-		if (value.type != handle_types::event)
-		{
-			return STATUS_INVALID_HANDLE;
-		}
-
-		const auto entry = c.proc.events.find(value.id);
-		if (entry == c.proc.events.end())
+		const auto entry = c.proc.events.get(handle);
+		if (!entry)
 		{
 			return STATUS_INVALID_HANDLE;
 		}
 
 		if (previous_state.value())
 		{
-			previous_state.write(entry->second.signaled ? 1ULL : 0ULL);
+			previous_state.write(entry->signaled ? 1ULL : 0ULL);
 		}
 
-		entry->second.signaled = true;
+		entry->signaled = true;
 		return STATUS_SUCCESS;
 	}
 
@@ -253,24 +247,14 @@ namespace
 			return STATUS_SUCCESS;
 		}
 
-		if (value.type == handle_types::event)
+		if (value.type == handle_types::event && c.proc.events.erase(handle))
 		{
-			const auto entry = c.proc.events.find(value.id);
-			if (entry != c.proc.events.end())
-			{
-				c.proc.events.erase(entry);
-				return STATUS_SUCCESS;
-			}
+			return STATUS_SUCCESS;
 		}
 
-		if (value.type == handle_types::file)
+		if (value.type == handle_types::file && c.proc.files.erase(handle))
 		{
-			const auto entry = c.proc.files.find(value.id);
-			if (entry != c.proc.files.end())
-			{
-				c.proc.files.erase(entry);
-				return STATUS_SUCCESS;
-			}
+			return STATUS_SUCCESS;
 		}
 
 		return STATUS_INVALID_HANDLE;
@@ -297,19 +281,9 @@ namespace
 			return STATUS_NOT_SUPPORTED;
 		}
 
-		uint32_t index = 1;
-		for (;; ++index)
-		{
-			if (!c.proc.events.contains(index))
-			{
-				break;
-			}
-		}
-
-		const auto h = make_handle(index, handle_types::event, false);
-		event_handle.write(h.bits);
-
-		c.proc.events.try_emplace(index, initial_state != FALSE, event_type);
+		event e{initial_state != FALSE, event_type};
+		const auto handle = c.proc.events.store(std::move(e));
+		event_handle.write(handle.bits);
 
 		static_assert(sizeof(EVENT_TYPE) == sizeof(uint32_t));
 		static_assert(sizeof(ACCESS_MASK) == sizeof(uint32_t));
@@ -347,33 +321,19 @@ namespace
 	                           const ULONG /*share_access*/,
 	                           const ULONG /*open_options*/)
 	{
-		uint32_t index = 1;
-		for (;; ++index)
+		file f{};
+		const auto attributes = object_attributes.read();
+		f.name = read_unicode_string(c.emu, attributes.ObjectName);
+
+		if (!std::filesystem::exists(f.name))
 		{
-			if (!c.proc.files.contains(index))
-			{
-				break;
-			}
+			return STATUS_FILE_INVALID;
 		}
 
-		const auto h = make_handle(index, handle_types::file, false);
-		file_handle.write(h.bits);
-
-		auto status = STATUS_SUCCESS;
-		object_attributes.access([&](const OBJECT_ATTRIBUTES& attributes)
-		{
-			auto section = read_unicode_string(c.emu, attributes.ObjectName);
-			if (!std::filesystem::exists(section))
-			{
-				status = STATUS_FILE_INVALID;
-			}
-			else
-			{
-				c.proc.files.try_emplace(index, std::move(section));
-			}
-		});
+		const auto handle = c.proc.files.store(std::move(f));
+		file_handle.write(handle.bits);
 
-		return status;
+		return STATUS_SUCCESS;
 	}
 
 	NTSTATUS handle_NtOpenSection(const syscall_context& c, const emulator_object<uint64_t> section_handle,
@@ -412,19 +372,8 @@ namespace
 			return STATUS_FILE_INVALID;
 		}
 
-		uint32_t index = 1;
-		for (;; ++index)
-		{
-			if (!c.proc.files.contains(index))
-			{
-				break;
-			}
-		}
-
-		const auto h = make_handle(index, handle_types::file, false);
-		section_handle.write(h.bits);
-
-		c.proc.files.try_emplace(index, std::move(filename));
+		const auto handle = c.proc.files.store({std::move(filename)});
+		section_handle.write(handle.bits);
 
 		return STATUS_SUCCESS;
 	}
@@ -441,20 +390,13 @@ namespace
 			return STATUS_INVALID_HANDLE;
 		}
 
-		const auto value = get_handle_value(section_handle);
-		if (value.type != handle_types::file)
-		{
-			return STATUS_INVALID_HANDLE;
-		}
-
-		const auto section_entry = c.proc.files.find(value.id);
-		if (section_entry == c.proc.files.end())
+		const auto section_entry = c.proc.files.get(section_handle);
+		if (!section_entry)
 		{
 			return STATUS_INVALID_HANDLE;
 		}
 
-		const auto& section_name = section_entry->second;
-		const auto binary = map_file(c.emu, section_name);
+		const auto binary = map_file(c.emu, section_entry->name);
 		if (!binary.has_value())
 		{
 			return STATUS_FILE_INVALID;