diff --git a/src/driver/ept.cpp b/src/driver/ept.cpp index f333211..1ed057b 100644 --- a/src/driver/ept.cpp +++ b/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& hints) { auto current_destination = reinterpret_cast(destination); auto current_source = reinterpret_cast(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(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::generate_translation_hints(const void* destination, const size_t length) { + utils::list hints{}; + auto current_destination = reinterpret_cast(destination); auto current_length = length; - ept_translation_hint* current_hints = nullptr; - - auto destructor = utils::finally([¤t_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(); - if (!new_hint) - { - throw std::runtime_error("Failed to allocate hint"); - } + ept_translation_hint current_hint{}; - new_hint->next_hint = current_hints; - current_hints = new_hint; - current_hints->virtual_base_address = aligned_destination; - current_hints->physical_base_address = memory::get_physical_address(aligned_destination); + current_hint.virtual_base_address = aligned_destination; + current_hint.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(¤t_hints->page[0], aligned_destination, PAGE_SIZE); + memcpy(¤t_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 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); - } + return hints; } uint64_t* ept::get_access_records(size_t* count) diff --git a/src/driver/ept.hpp b/src/driver/ept.hpp index 30483e8..8e36379 100644 --- a/src/driver/ept.hpp +++ b/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& 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 void free_translation_hints(ept_translation_hint* hints); + static utils::list generate_translation_hints(const void* destination, size_t length); 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); }; diff --git a/src/driver/hypervisor.cpp b/src/driver/hypervisor.cpp index 19c5566..88318df 100644 --- a/src/driver/hypervisor.cpp +++ b/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& 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& 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& 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; - //vm_state->ept->install_hook(destination, source, length, translation_hint); + (void)hints; + //vm_state->ept->install_hook(destination, source, length, hints); if (this->is_enabled()) { diff --git a/src/driver/hypervisor.hpp b/src/driver/hypervisor.hpp index 7ae7e93..f9ba928 100644 --- a/src/driver/hypervisor.hpp +++ b/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& 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& 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& hints = {}); vmx::state* get_current_vm_state() const; }; diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp index 0f8be63..1227648 100644 --- a/src/driver/irp.cpp +++ b/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 generate_translation_hints(uint32_t process_id, const void* target_address, + size_t size) { - vmx::ept_translation_hint* translation_hints{nullptr}; + utils::list 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 (name) + if (const auto name = process_handle.get_image_filename()) { 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); - }); - - t.join(); + }).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; diff --git a/src/driver/list.hpp b/src/driver/list.hpp index 3db8fe8..d353d4b 100644 --- a/src/driver/list.hpp +++ b/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_; - while (*inseration_point) - { - inseration_point = &(*inseration_point)->next; - } + auto& entry = this->add_uninitialized_entry(); - auto* list_base = this->list_allocator_.allocate(sizeof(ListEntry) + alignof(ListEntry)); - auto* entry_base = this->object_allocator_.allocate(sizeof(T) + alignof(T)); + new(&entry) T(obj); - auto* entry = align_pointer(entry_base); - auto* list_entry = align_pointer(list_base); + return entry; + } - list_entry->this_base = list_base; - list_entry->entry_base = entry_base; - list_entry->next = nullptr; - list_entry->entry = entry; + T& push_back(T&& obj) + { + auto& entry = this->add_uninitialized_entry(); - *inseration_point = list_entry; + new(&entry) T(std::move(obj)); - new(entry) T(std::move(obj)); + return entry; + } - return *entry; + template + T& emplace_back(Args&&... args) + { + auto& entry = this->add_uninitialized_entry(); + + new(&entry) T(std::forward(args)...); + + 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()); } @@ -287,12 +289,17 @@ namespace utils this->object_allocator_.free(list_entry->entry_base); this->list_allocator_.free(list_entry->this_base); - return { *inseration_point }; + return {*inseration_point}; } 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(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(entry_base); + auto* list_entry = align_pointer(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; + } }; } diff --git a/src/driver/type_traits.hpp b/src/driver/type_traits.hpp index 395155b..563dd28 100644 --- a/src/driver/type_traits.hpp +++ b/src/driver/type_traits.hpp @@ -27,7 +27,7 @@ namespace std }; template - typename remove_reference::type&& move(T&& arg) + typename remove_reference::type&& move(T&& arg) noexcept { return static_cast::type&&>(arg); } diff --git a/src/driver/vector.hpp b/src/driver/vector.hpp index 8df0e80..023b8b8 100644 --- a/src/driver/vector.hpp +++ b/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_) - { - this->reserve(max(this->capacity_, 5) * 2); - } + auto& entry = this->add_uninitialized_entry(); - auto* data = this->data() + this->size_++; - new(data) T(std::move(obj)); + new(&entry) T(obj); - return *data; + return entry; + } + + T& push_back(T&& obj) + { + auto& entry = this->add_uninitialized_entry(); + + new(&entry) T(std::move(obj)); + + return entry; + } + + template + T& emplace_back(Args&&... args) + { + auto& entry = this->add_uninitialized_entry(); + + new(&entry) T(std::forward(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);