Start using custom containers

src/driver/ept.cpp

@@ -168,7 +168,7 @@ namespace vmx
 	}
 
 	void ept::install_page_hook(void* destination, const void* source, const size_t length,
-	                            ept_translation_hint* translation_hint)
+	                            const ept_translation_hint* translation_hint)
 	{
 		auto* hook = this->get_or_create_ept_hook(destination, translation_hint);
 
@@ -194,7 +194,7 @@ namespace vmx
 	}
 
 	void ept::install_hook(const void* destination, const void* source, const size_t length,
-	                       ept_translation_hint* translation_hint)
+		const utils::list<ept_translation_hint>& hints)
 	{
 		auto current_destination = reinterpret_cast<uint64_t>(destination);
 		auto current_source = reinterpret_cast<uint64_t>(source);
@@ -207,17 +207,16 @@ namespace vmx
 			const auto page_remaining = PAGE_SIZE - page_offset;
 			const auto data_to_write = min(page_remaining, current_length);
 
-			ept_translation_hint* relevant_hint = nullptr;
+
-			ept_translation_hint* current_hint = translation_hint;
+
-			while (current_hint)
+			const ept_translation_hint* relevant_hint = nullptr;
+			for(const auto& hint : hints)
 			{
-				if (current_hint->virtual_base_address == aligned_destination)
+				if (hint.virtual_base_address == aligned_destination)
 				{
-					relevant_hint = current_hint;
+					relevant_hint = &hint;
 					break;
 				}
-
-				current_hint = current_hint->next_hint;
 			}
 
 			this->install_page_hook(reinterpret_cast<void*>(current_destination),
@@ -539,7 +538,7 @@ namespace vmx
 		return nullptr;
 	}
 
-	ept_hook* ept::get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint)
+	ept_hook* ept::get_or_create_ept_hook(void* destination, const ept_translation_hint* translation_hint)
 	{
 		const auto virtual_target = PAGE_ALIGN(destination);
 
@@ -654,18 +653,13 @@ namespace vmx
 		target_entry->flags = new_pointer.flags;
 	}
 
-	ept_translation_hint* ept::generate_translation_hints(const void* destination, const size_t length)
+	utils::list<ept_translation_hint> ept::generate_translation_hints(const void* destination, const size_t length)
 	{
+		utils::list<ept_translation_hint> hints{};
+
 		auto current_destination = reinterpret_cast<uint64_t>(destination);
 		auto current_length = length;
 
-		ept_translation_hint* current_hints = nullptr;
-
-		auto destructor = utils::finally([&current_hints]()
-		{
-			ept::free_translation_hints(current_hints);
-		});
-
 		while (current_length != 0)
 		{
 			const auto aligned_destination = PAGE_ALIGN(current_destination);
@@ -673,42 +667,25 @@ namespace vmx
 			const auto page_remaining = PAGE_SIZE - page_offset;
 			const auto data_to_write = min(page_remaining, current_length);
 
-			auto* new_hint = memory::allocate_non_paged_object<ept_translation_hint>();
+			ept_translation_hint current_hint{};
-			if (!new_hint)
-			{
-				throw std::runtime_error("Failed to allocate hint");
-			}
 
-			new_hint->next_hint = current_hints;
+			current_hint.virtual_base_address = aligned_destination;
-			current_hints = new_hint;
+			current_hint.physical_base_address = memory::get_physical_address(aligned_destination);
-			current_hints->virtual_base_address = aligned_destination;
-			current_hints->physical_base_address = memory::get_physical_address(aligned_destination);
 
-			if (!current_hints->physical_base_address)
+			if (!current_hint.physical_base_address)
 			{
 				throw std::runtime_error("Failed to resolve physical address");
 			}
 
-			memcpy(&current_hints->page[0], aligned_destination, PAGE_SIZE);
+			memcpy(&current_hint.page[0], aligned_destination, PAGE_SIZE);
+
+			hints.push_back(current_hint);
 
 			current_length -= data_to_write;
 			current_destination += data_to_write;
 		}
 
-		destructor.cancel();
+		return hints;
-
-		return current_hints;
-	}
-
-	void ept::free_translation_hints(ept_translation_hint* hints)
-	{
-		auto* hint = hints;
-		while (hint)
-		{
-			auto* current_hint = hint;
-			hint = hint->next_hint;
-			memory::free_non_paged_object(current_hint);
-		}
 	}
 
 	uint64_t* ept::get_access_records(size_t* count)

src/driver/ept.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #define DECLSPEC_PAGE_ALIGN DECLSPEC_ALIGN(PAGE_SIZE)
+#include "list.hpp"
 
 namespace vmx
 {
@@ -55,8 +56,6 @@ namespace vmx
 
 		uint64_t physical_base_address{};
 		const void* virtual_base_address{};
-
-		ept_translation_hint* next_hint{nullptr};
 	};
 
 	struct guest_context;
@@ -77,7 +76,7 @@ namespace vmx
 		void install_code_watch_point(uint64_t physical_page);
 
 		void install_hook(const void* destination, const void* source, size_t length,
-		                  ept_translation_hint* translation_hint = nullptr);
+		                 const utils::list<ept_translation_hint>& hints = {});
 		void disable_all_hooks() const;
 
 		void handle_violation(guest_context& guest_context);
@@ -86,8 +85,7 @@ namespace vmx
 		ept_pointer get_ept_pointer() const;
 		void invalidate() const;
 
-		static ept_translation_hint* generate_translation_hints(const void* destination, size_t length);
+		static utils::list<ept_translation_hint> generate_translation_hints(const void* destination, size_t length);
-		static void free_translation_hints(ept_translation_hint* hints);
 
 		uint64_t* get_access_records(size_t* count);
 
@@ -113,12 +111,12 @@ namespace vmx
 		ept_code_watch_point* allocate_ept_code_watch_point();
 		ept_code_watch_point* find_ept_code_watch_point(uint64_t physical_address) const;
 
-		ept_hook* get_or_create_ept_hook(void* destination, ept_translation_hint* translation_hint = nullptr);
+		ept_hook* get_or_create_ept_hook(void* destination, const ept_translation_hint* translation_hint = nullptr);
 
 		void split_large_page(uint64_t physical_address);
 
 		void install_page_hook(void* destination, const void* source, size_t length,
-		                       ept_translation_hint* translation_hint = nullptr);
+		                       const ept_translation_hint* translation_hint = nullptr);
 
 		void record_access(uint64_t rip);
 	};

src/driver/hypervisor.cpp

@@ -163,11 +163,11 @@ bool hypervisor::is_enabled() const
 }
 
 bool hypervisor::install_ept_hook(const void* destination, const void* source, const size_t length,
-                                  vmx::ept_translation_hint* translation_hint)
+                                  const utils::list<vmx::ept_translation_hint>& hints)
 {
 	try
 	{
-		this->ept_->install_hook(destination, source, length, translation_hint);
+		this->ept_->install_hook(destination, source, length, hints);
 	}
 	catch (std::exception& e)
 	{
@@ -183,7 +183,7 @@ bool hypervisor::install_ept_hook(const void* destination, const void* source, c
 	volatile long failures = 0;
 	thread::dispatch_on_all_cores([&]
 	{
-		if (!this->try_install_ept_hook_on_core(destination, source, length, translation_hint))
+		if (!this->try_install_ept_hook_on_core(destination, source, length, hints))
 		{
 			InterlockedIncrement(&failures);
 		}
@@ -192,7 +192,7 @@ bool hypervisor::install_ept_hook(const void* destination, const void* source, c
 	return failures == 0;
 }
 
-bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, bool invalidate) const
+bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, const bool invalidate) const
 {
 	try
 	{
@@ -826,11 +826,11 @@ void hypervisor::free_vm_states()
 }
 
 bool hypervisor::try_install_ept_hook_on_core(const void* destination, const void* source, const size_t length,
-                                              vmx::ept_translation_hint* translation_hint)
+                                              const utils::list<vmx::ept_translation_hint>& hints)
 {
 	try
 	{
-		this->install_ept_hook_on_core(destination, source, length, translation_hint);
+		this->install_ept_hook_on_core(destination, source, length, hints);
 		return true;
 	}
 	catch (std::exception& e)
@@ -846,7 +846,7 @@ bool hypervisor::try_install_ept_hook_on_core(const void* destination, const voi
 }
 
 void hypervisor::install_ept_hook_on_core(const void* destination, const void* source, const size_t length,
-                                          vmx::ept_translation_hint* translation_hint)
+                                          const utils::list<vmx::ept_translation_hint>& hints)
 {
 	auto* vm_state = this->get_current_vm_state();
 	if (!vm_state)
@@ -857,8 +857,8 @@ void hypervisor::install_ept_hook_on_core(const void* destination, const void* s
 	(void)destination;
 	(void)source;
 	(void)length;
-	(void)translation_hint;
+	(void)hints;
-	//vm_state->ept->install_hook(destination, source, length, translation_hint);
+	//vm_state->ept->install_hook(destination, source, length, hints);
 
 	if (this->is_enabled())
 	{

src/driver/hypervisor.hpp

@@ -20,7 +20,7 @@ public:
 	bool is_enabled() const;
 
 	bool install_ept_hook(const void* destination, const void* source, size_t length,
-	                      vmx::ept_translation_hint* translation_hint = nullptr);
+	                      const utils::list<vmx::ept_translation_hint>& hints = {});
 
 	bool install_ept_code_watch_point(uint64_t physical_page, bool invalidate = true) const;
 	bool install_ept_code_watch_points(const uint64_t* physical_pages, size_t count) const;
@@ -44,9 +44,9 @@ private:
 	void free_vm_states();
 
 	bool try_install_ept_hook_on_core(const void* destination, const void* source, size_t length,
-	                                  vmx::ept_translation_hint* translation_hint = nullptr);
+	                                  const utils::list<vmx::ept_translation_hint>& hints = {});
 	void install_ept_hook_on_core(const void* destination, const void* source, size_t length,
-	                              vmx::ept_translation_hint* translation_hint = nullptr);
+	                              const utils::list<vmx::ept_translation_hint>& hints = {});
 
 	vmx::state* get_current_vm_state() const;
 };

src/driver/irp.cpp

@@ -37,11 +37,12 @@ namespace
 		return STATUS_SUCCESS;
 	}
 
-	vmx::ept_translation_hint* generate_translation_hints(uint32_t process_id, const void* target_address, size_t size)
+	utils::list<vmx::ept_translation_hint> generate_translation_hints(uint32_t process_id, const void* target_address,
+	                                                                  size_t size)
 	{
-		vmx::ept_translation_hint* translation_hints{nullptr};
+		utils::list<vmx::ept_translation_hint> translation_hints{};
 
-		thread::kernel_thread t([&translation_hints, process_id, target_address, size]
+		thread::kernel_thread([&translation_hints, process_id, target_address, size]
 		{
 			debug_log("Looking up process: %d\n", process_id);
 
@@ -52,8 +53,7 @@ namespace
 				return;
 			}
 
-			const auto name = process_handle.get_image_filename();
+			if (const auto name = process_handle.get_image_filename())
-			if (name)
 			{
 				debug_log("Attaching to %s\n", name);
 			}
@@ -62,9 +62,7 @@ namespace
 
 			debug_log("Generating translation hints for address: %p\n", target_address);
 			translation_hints = vmx::ept::generate_translation_hints(target_address, size);
-		});
+		}).join();
-
-		t.join();
 
 		return translation_hints;
 	}
@@ -83,17 +81,11 @@ namespace
 			throw std::runtime_error("Failed to copy buffer");
 		}
 
-		vmx::ept_translation_hint* translation_hints = nullptr;
-		auto destructor = utils::finally([&translation_hints]()
-		{
-			vmx::ept::free_translation_hints(translation_hints);
-		});
-
 		memcpy(buffer.get(), request.source_data, request.source_data_size);
-		translation_hints = generate_translation_hints(request.process_id, request.target_address,
+		auto translation_hints = generate_translation_hints(request.process_id, request.target_address,
 		                                               request.source_data_size);
 
-		if (!translation_hints)
+		if (translation_hints.empty())
 		{
 			debug_log("Failed to generate tranlsation hints\n");
 			return;

src/driver/list.hpp

@@ -31,14 +31,14 @@ namespace utils
 			{
 			}
 
-			T* operator*() const
+			T& operator*() const
 			{
-				return this->entry_->entry;
+				return *this->entry_->entry;
 			}
 
-			T* operator->() const
+			T& operator->() const
 			{
-				return this->entry_->entry;
+				return *this->entry_->entry;
 			}
 
 			bool operator==(const iterator& i) const
@@ -46,7 +46,7 @@ namespace utils
 				return this->entry_ == i.entry_;
 			}
 
-			iterator operator++() const
+			iterator operator++()
 			{
 				this->entry_ = this->entry_->next;
 				return *this;
@@ -89,14 +89,14 @@ namespace utils
 			{
 			}
 
-			const T* operator*() const
+			const T& operator*() const
 			{
-				return this->entry_->entry;
+				return *this->entry_->entry;
 			}
 
-			const T* operator->() const
+			const T& operator->() const
 			{
-				return this->entry_->entry;
+				return *this->entry_->entry;
 			}
 
 			bool operator==(const const_iterator& i) const
@@ -104,7 +104,7 @@ namespace utils
 				return this->entry_ == i.entry_;
 			}
 
-			const_iterator operator++() const
+			const_iterator operator++()
 			{
 				this->entry_ = this->entry_->next;
 				return *this;
@@ -166,30 +166,32 @@ namespace utils
 			return *this;
 		}
 
-		T& push_back(T obj)
+		T& push_back(const T& obj)
 		{
-			auto** inseration_point = &this->entries_;
+			auto& entry = this->add_uninitialized_entry();
-			while (*inseration_point)
+
-			{
+			new(&entry) T(obj);
-				inseration_point = &(*inseration_point)->next;
+
+			return entry;
 		}
 
-			auto* list_base = this->list_allocator_.allocate(sizeof(ListEntry) + alignof(ListEntry));
+		T& push_back(T&& obj)
-			auto* entry_base = this->object_allocator_.allocate(sizeof(T) + alignof(T));
+		{
+			auto& entry = this->add_uninitialized_entry();
 
-			auto* entry = align_pointer<T>(entry_base);
+			new(&entry) T(std::move(obj));
-			auto* list_entry = align_pointer<ListEntry>(list_base);
 
-			list_entry->this_base = list_base;
+			return entry;
-			list_entry->entry_base = entry_base;
+		}
-			list_entry->next = nullptr;
-			list_entry->entry = entry;
 
-			*inseration_point = list_entry;
+		template <typename... Args>
+		T& emplace_back(Args&&... args)
+		{
+			auto& entry = this->add_uninitialized_entry();
 
-			new(entry) T(std::move(obj));
+			new(&entry) T(std::forward<Args>(args)...);
 
-			return *entry;
+			return entry;
 		}
 
 		T& operator[](const size_t index)
@@ -232,7 +234,7 @@ namespace utils
 
 		void clear()
 		{
-			while (this->entries_)
+			while (!this->empty())
 			{
 				this->erase(this->begin());
 			}
@@ -293,6 +295,11 @@ namespace utils
 			throw std::runtime_error("Bad iterator");
 		}
 
+		bool empty() const
+		{
+			return this->entries_ == nullptr;
+		}
+
 	private:
 		friend iterator;
 
@@ -309,5 +316,29 @@ namespace utils
 
 			return reinterpret_cast<U*>(ptr);
 		}
+
+		T& add_uninitialized_entry()
+		{
+			auto** inseration_point = &this->entries_;
+			while (*inseration_point)
+			{
+				inseration_point = &(*inseration_point)->next;
+			}
+
+			auto* list_base = this->list_allocator_.allocate(sizeof(ListEntry) + alignof(ListEntry));
+			auto* entry_base = this->object_allocator_.allocate(sizeof(T) + alignof(T));
+
+			auto* entry = align_pointer<T>(entry_base);
+			auto* list_entry = align_pointer<ListEntry>(list_base);
+
+			list_entry->this_base = list_base;
+			list_entry->entry_base = entry_base;
+			list_entry->next = nullptr;
+			list_entry->entry = entry;
+
+			*inseration_point = list_entry;
+
+			return *entry;
+		}
 	};
 }

src/driver/type_traits.hpp

@@ -27,7 +27,7 @@ namespace std
 	};
 
 	template <typename T>
-	typename remove_reference<T>::type&& move(T&& arg)
+	typename remove_reference<T>::type&& move(T&& arg) noexcept
 	{
 		return static_cast<typename remove_reference<T>::type&&>(arg);
 	}

src/driver/vector.hpp

@@ -91,17 +91,32 @@ namespace utils
 			}
 		}
 
-		T& push_back(T obj)
+		T& push_back(const T& obj)
 		{
-			if (this->size_ + 1 > this->capacity_)
+			auto& entry = this->add_uninitialized_entry();
-			{
+
-				this->reserve(max(this->capacity_, 5) * 2);
+			new(&entry) T(obj);
+
+			return entry;
 		}
 
-			auto* data = this->data() + this->size_++;
+		T& push_back(T&& obj)
-			new(data) T(std::move(obj));
+		{
+			auto& entry = this->add_uninitialized_entry();
 
-			return *data;
+			new(&entry) T(std::move(obj));
+
+			return entry;
+		}
+
+		template <typename... Args>
+		T& emplace_back(Args&&... args)
+		{
+			auto& entry = this->add_uninitialized_entry();
+
+			new(&entry) T(std::forward<Args>(args)...);
+
+			return entry;
 		}
 
 		T& operator[](const size_t index)
@@ -217,12 +232,28 @@ namespace utils
 			return iterator;
 		}
 
+		bool empty() const
+		{
+			return this->size_ == 0;
+		}
+
 	private:
 		Allocator allocator_{};
 		void* storage_{nullptr};
 		size_t capacity_{0};
 		size_t size_{0};
 
+		T& add_uninitialized_entry()
+		{
+			if (this->size_ + 1 > this->capacity_)
+			{
+				this->reserve(max(this->capacity_, 5) * 2);
+			}
+
+			auto* data = this->data() + this->size_++;
+			return *data;
+		}
+
 		void* allocate_memory_for_capacity(const size_t capacity)
 		{
 			constexpr auto alignment = alignof(T);