From 4cbbaed72fb510555b5775a4da48264892c1c231 Mon Sep 17 00:00:00 2001 From: Maurice Heumann Date: Tue, 27 Dec 2022 13:30:20 +0100 Subject: [PATCH] Cleanup on process termination --- src/driver/driver_main.cpp | 11 ++--- src/driver/ept.cpp | 67 ++++++++++++++++++++------ src/driver/ept.hpp | 18 +++++-- src/driver/hypervisor.cpp | 85 ++++++++++++--------------------- src/driver/hypervisor.hpp | 17 +++---- src/driver/irp.cpp | 10 ++-- src/driver/list.hpp | 8 ++-- src/driver/process.cpp | 30 ++++++++---- src/driver/process.hpp | 9 +++- src/driver/process_callback.cpp | 6 ++- src/driver/process_callback.hpp | 2 +- src/driver/std_include.hpp | 3 ++ 12 files changed, 155 insertions(+), 111 deletions(-) diff --git a/src/driver/driver_main.cpp b/src/driver/driver_main.cpp index 6853e2f..19143a9 100644 --- a/src/driver/driver_main.cpp +++ b/src/driver/driver_main.cpp @@ -19,7 +19,7 @@ public: this->sleep_notification(type); }), process_callback_( - [this](const HANDLE parent_id, const HANDLE process_id, const process_callback::type type) + [this](const process_id parent_id, const process_id process_id, const process_callback::type type) { this->process_notification(parent_id, process_id, type); }), @@ -66,16 +66,11 @@ private: } } - void process_notification(HANDLE /*parent_id*/, const HANDLE process_id, const process_callback::type type) + void process_notification(process_id /*parent_id*/, const process_id process_id, const process_callback::type type) { - if (type == process_callback::type::create) - { - debug_log("Created process: %X\n", process_id); - } - if (type == process_callback::type::destroy) { - debug_log("Destroyed process: %X\n", process_id); + this->hypervisor_.handle_process_termination(process_id); } } }; diff --git a/src/driver/ept.cpp b/src/driver/ept.cpp index fdd5c3d..d45ba06 100644 --- a/src/driver/ept.cpp +++ b/src/driver/ept.cpp @@ -97,7 +97,7 @@ namespace vmx void reset_all_watch_point_pages(utils::list& watch_points) { - for(const auto& watch_point : watch_points) + for (const auto& watch_point : watch_points) { if (watch_point.target_page) { @@ -121,6 +121,8 @@ namespace vmx ept_hook::~ept_hook() { + this->target_page->flags = this->original_entry.flags; + if (mapped_virtual_address) { memory::unmap_physical_memory(mapped_virtual_address, PAGE_SIZE); @@ -143,10 +145,13 @@ namespace vmx this->disable_all_hooks(); } - void ept::install_page_hook(void* destination, const void* source, const size_t length, + void ept::install_page_hook(void* destination, const void* source, const size_t length, const process_id source_pid, + const process_id target_pid, const ept_translation_hint* translation_hint) { auto* hook = this->get_or_create_ept_hook(destination, translation_hint); + hook->source_pid = source_pid; + hook->target_pid = target_pid; const auto page_offset = ADDRMASK_EPT_PML1_OFFSET(reinterpret_cast(destination)); memcpy(hook->fake_page + page_offset, source, length); @@ -170,7 +175,8 @@ namespace vmx } void ept::install_hook(const void* destination, const void* source, const size_t length, - const utils::list& hints) + const process_id source_pid, const process_id target_pid, + const utils::list& hints) { auto current_destination = reinterpret_cast(destination); auto current_source = reinterpret_cast(source); @@ -184,9 +190,8 @@ namespace vmx const auto data_to_write = min(page_remaining, current_length); - const ept_translation_hint* relevant_hint = nullptr; - for(const auto& hint : hints) + for (const auto& hint : hints) { if (hint.virtual_base_address == aligned_destination) { @@ -196,7 +201,8 @@ namespace vmx } this->install_page_hook(reinterpret_cast(current_destination), - reinterpret_cast(current_source), data_to_write, relevant_hint); + reinterpret_cast(current_source), data_to_write, source_pid, + target_pid, relevant_hint); current_length -= data_to_write; current_destination += data_to_write; @@ -204,12 +210,9 @@ namespace vmx } } - void ept::disable_all_hooks() const + void ept::disable_all_hooks() { - for(auto& hook : this->ept_hooks) - { - hook.target_page->flags = hook.original_entry.flags; - } + this->ept_hooks.clear(); } void ept::handle_violation(guest_context& guest_context) @@ -333,7 +336,8 @@ namespace vmx } } - void ept::install_code_watch_point(const uint64_t physical_page) + void ept::install_code_watch_point(const uint64_t physical_page, const process_id source_pid, + const process_id target_pid) { const auto physical_base_address = reinterpret_cast(PAGE_ALIGN(physical_page)); @@ -343,6 +347,8 @@ namespace vmx } auto& watch_point = this->allocate_ept_code_watch_point(); + watch_point.source_pid = source_pid; + watch_point.target_pid = target_pid; this->split_large_page(physical_base_address); @@ -420,7 +426,7 @@ namespace vmx pml1* ept::find_pml1_table(const uint64_t physical_address) { - for(auto& split : this->ept_splits) + for (auto& split : this->ept_splits) { if (memory::get_physical_address(&split.pml1[0]) == physical_address) { @@ -461,7 +467,7 @@ namespace vmx ept_code_watch_point* ept::find_ept_code_watch_point(const uint64_t physical_address) { - for(auto& watch_point : this->ept_code_watch_points) + for (auto& watch_point : this->ept_code_watch_points) { if (watch_point.physical_base_address == physical_address) { @@ -628,4 +634,37 @@ namespace vmx *count = i; return this->access_records; } + + bool ept::handle_process_termination(const process_id process) + { + bool changed = false; + + for (auto i = this->ept_hooks.begin(); i != this->ept_hooks.end();) + { + if (i->source_pid == process || i->target_pid == process) + { + i = this->ept_hooks.erase(i); + changed = true; + } + else + { + ++i; + } + } + + for (auto i = this->ept_code_watch_points.begin(); i != this->ept_code_watch_points.end();) + { + if (i->source_pid == process || i->target_pid == process) + { + i = this->ept_code_watch_points.erase(i); + changed = true; + } + else + { + ++i; + } + } + + return changed; + } } diff --git a/src/driver/ept.hpp b/src/driver/ept.hpp index 5feef1f..a109130 100644 --- a/src/driver/ept.hpp +++ b/src/driver/ept.hpp @@ -26,6 +26,8 @@ namespace vmx { uint64_t physical_base_address{}; pml1* target_page{}; + process_id source_pid{0}; + process_id target_pid{0}; }; struct ept_hook @@ -43,6 +45,9 @@ namespace vmx pml1 original_entry{}; pml1 execute_entry{}; pml1 readwrite_entry{}; + + process_id source_pid{0}; + process_id target_pid{0}; }; struct ept_translation_hint @@ -68,11 +73,12 @@ namespace vmx void initialize(); - void install_code_watch_point(uint64_t physical_page); + void install_code_watch_point(uint64_t physical_page, process_id source_pid, process_id target_pid); - void install_hook(const void* destination, const void* source, size_t length, + void install_hook(const void* destination, const void* source, size_t length, process_id source_pid, + process_id target_pid, const utils::list& hints = {}); - void disable_all_hooks() const; + void disable_all_hooks(); void handle_violation(guest_context& guest_context); void handle_misconfiguration(guest_context& guest_context) const; @@ -84,6 +90,8 @@ namespace vmx uint64_t* get_access_records(size_t* count); + bool handle_process_termination(process_id process); + private: DECLSPEC_PAGE_ALIGN pml4 epml4[EPT_PML4E_ENTRY_COUNT]; DECLSPEC_PAGE_ALIGN pml3 epdpt[EPT_PDPTE_ENTRY_COUNT]; @@ -110,8 +118,8 @@ namespace vmx void split_large_page(uint64_t physical_address); - void install_page_hook(void* destination, const void* source, size_t length, - const ept_translation_hint* translation_hint = nullptr); + void install_page_hook(void* destination, const void* source, size_t length, process_id source_pid, + process_id target_pid, 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 23ce220..7762b71 100644 --- a/src/driver/hypervisor.cpp +++ b/src/driver/hypervisor.cpp @@ -164,11 +164,12 @@ bool hypervisor::is_enabled() const } bool hypervisor::install_ept_hook(const void* destination, const void* source, const size_t length, + const process_id source_pid, const process_id target_pid, const utils::list& hints) { try { - this->ept_->install_hook(destination, source, length, hints); + this->ept_->install_hook(destination, source, length, source_pid, target_pid, hints); } catch (std::exception& e) { @@ -181,23 +182,16 @@ bool hypervisor::install_ept_hook(const void* destination, const void* source, c return false; } - volatile long failures = 0; - thread::dispatch_on_all_cores([&] - { - if (!this->try_install_ept_hook_on_core(destination, source, length, hints)) - { - InterlockedIncrement(&failures); - } - }); - - return failures == 0; + this->invalidate_cores(); + return true; } -bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, const bool invalidate) const +bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, const process_id source_pid, + const process_id target_pid, const bool invalidate) const { try { - this->ept_->install_code_watch_point(physical_page); + this->ept_->install_code_watch_point(physical_page, source_pid, target_pid); } catch (std::exception& e) { @@ -221,12 +215,13 @@ bool hypervisor::install_ept_code_watch_point(const uint64_t physical_page, cons return true; } -bool hypervisor::install_ept_code_watch_points(const uint64_t* physical_pages, const size_t count) const +bool hypervisor::install_ept_code_watch_points(const uint64_t* physical_pages, const size_t count, + const process_id source_pid, const process_id target_pid) const { bool success = true; for (size_t i = 0; i < count; ++i) { - success &= this->install_ept_code_watch_point(physical_pages[i], false); + success &= this->install_ept_code_watch_point(physical_pages[i], source_pid, target_pid, false); } thread::dispatch_on_all_cores([&] @@ -243,7 +238,7 @@ void hypervisor::disable_all_ept_hooks() const thread::dispatch_on_all_cores([&] { - auto* vm_state = this->get_current_vm_state(); + const auto* vm_state = this->get_current_vm_state(); if (!vm_state) { return; @@ -266,6 +261,17 @@ hypervisor* hypervisor::get_instance() return instance; } +void hypervisor::handle_process_termination(const process_id process) +{ + if (!this->ept_->handle_process_termination(process)) + { + return; + } + + debug_log("Handled termination of %X\n", process); + this->invalidate_cores(); +} + void hypervisor::enable() { const auto cr3 = __readcr3(); @@ -430,7 +436,7 @@ vmx::gdt_entry convert_gdt_entry(const uint64_t gdt_base, const uint16_t selecto result.access_rights.granularity = gdt_entry->granularity; result.access_rights.reserved1 = 0; - result.access_rights.unusable = !gdt_entry->present; + result.access_rights.unusable = ~gdt_entry->present; return result; } @@ -826,45 +832,16 @@ void hypervisor::free_vm_states() } } -bool hypervisor::try_install_ept_hook_on_core(const void* destination, const void* source, const size_t length, - const utils::list& hints) +void hypervisor::invalidate_cores() const { - try + thread::dispatch_on_all_cores([&] { - this->install_ept_hook_on_core(destination, source, length, hints); - return true; - } - catch (std::exception& e) - { - debug_log("Failed to install ept hook on core %d: %s\n", thread::get_processor_index(), e.what()); - return false; - } - catch (...) - { - debug_log("Failed to install ept hook on core %d.\n", thread::get_processor_index()); - return false; - } -} - -void hypervisor::install_ept_hook_on_core(const void* destination, const void* source, const size_t length, - const utils::list& hints) -{ - auto* vm_state = this->get_current_vm_state(); - if (!vm_state) - { - throw std::runtime_error("No vm state available"); - } - - (void)destination; - (void)source; - (void)length; - (void)hints; - //vm_state->ept->install_hook(destination, source, length, hints); - - if (this->is_enabled()) - { - vm_state->ept->invalidate(); - } + const auto* vm_state = this->get_current_vm_state(); + if (vm_state && this->is_enabled()) + { + vm_state->ept->invalidate(); + } + }); } vmx::state* hypervisor::get_current_vm_state() const diff --git a/src/driver/hypervisor.hpp b/src/driver/hypervisor.hpp index f9ba928..4ba6499 100644 --- a/src/driver/hypervisor.hpp +++ b/src/driver/hypervisor.hpp @@ -19,11 +19,13 @@ public: bool is_enabled() const; - bool install_ept_hook(const void* destination, const void* source, size_t length, - const utils::list& hints = {}); + bool install_ept_hook(const void* destination, const void* source, size_t length, process_id source_pid, + process_id target_pid, 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; + bool install_ept_code_watch_point(uint64_t physical_page, process_id source_pid, process_id target_pid, + bool invalidate = true) const; + bool install_ept_code_watch_points(const uint64_t* physical_pages, size_t count, process_id source_pid, + process_id target_pid) const; void disable_all_ept_hooks() const; @@ -31,6 +33,8 @@ public: static hypervisor* get_instance(); + void handle_process_termination(process_id process); + private: uint32_t vm_state_count_{0}; vmx::state** vm_states_{nullptr}; @@ -43,10 +47,7 @@ private: void allocate_vm_states(); void free_vm_states(); - bool try_install_ept_hook_on_core(const void* destination, const void* source, size_t length, - const utils::list& hints = {}); - void install_ept_hook_on_core(const void* destination, const void* source, size_t length, - const utils::list& hints = {}); + void invalidate_cores() const; vmx::state* get_current_vm_state() const; }; diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp index 0c4571f..e0250f7 100644 --- a/src/driver/irp.cpp +++ b/src/driver/irp.cpp @@ -92,7 +92,7 @@ namespace } hypervisor->install_ept_hook(request.target_address, buffer.get(), request.source_data_size, - translation_hints); + process::get_current_process_id(), request.process_id, translation_hints); } void unhook() @@ -160,7 +160,7 @@ namespace throw std::runtime_error("Failed to copy buffer"); } - thread::kernel_thread t([watch_request_copy, hypervisor, &index, &page_buffer] + thread::kernel_thread t([watch_request_copy, &index, &page_buffer] { debug_log("Looking up process: %d\n", watch_request_copy.process_id); @@ -210,7 +210,8 @@ namespace t.join(); debug_log("Installing watch points...\n"); - (void)hypervisor->install_ept_code_watch_points(page_buffer.get(), index); + (void)hypervisor->install_ept_code_watch_points(page_buffer.get(), index, process::get_current_process_id(), + watch_request_copy.process_id); debug_log("Watch points installed\n"); } @@ -232,7 +233,7 @@ namespace void get_records(const PIRP irp, const PIO_STACK_LOCATION irp_sp) { - auto* hypervisor = hypervisor::get_instance(); + const auto* hypervisor = hypervisor::get_instance(); if (!hypervisor) { throw std::runtime_error("Hypervisor not installed"); @@ -251,7 +252,6 @@ namespace irp->IoStatus.Status = STATUS_SUCCESS; const auto irp_sp = IoGetCurrentIrpStackLocation(irp); - if (irp_sp) { const auto ioctr_code = irp_sp->Parameters.DeviceIoControl.IoControlCode; diff --git a/src/driver/list.hpp b/src/driver/list.hpp index 2e39f75..5cb09b7 100644 --- a/src/driver/list.hpp +++ b/src/driver/list.hpp @@ -36,9 +36,9 @@ namespace utils return *this->entry_->entry; } - T& operator->() const + T* operator->() const { - return *this->entry_->entry; + return this->entry_->entry; } bool operator==(const iterator& i) const @@ -94,9 +94,9 @@ namespace utils 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 diff --git a/src/driver/process.cpp b/src/driver/process.cpp index d8f9624..deedd65 100644 --- a/src/driver/process.cpp +++ b/src/driver/process.cpp @@ -37,7 +37,7 @@ namespace process process_handle::process_handle(const process_handle& obj) { - this->operator=(std::move(obj)); + this->operator=(obj); } process_handle& process_handle::operator=(const process_handle& obj) @@ -80,14 +80,14 @@ namespace process return KeWaitForSingleObject(this->handle_, Executive, KernelMode, FALSE, &zero_time) != STATUS_WAIT_0; } - uint32_t process_handle::get_id() const + process_id process_handle::get_id() const { if (!this->handle_) { return 0; } - return uint32_t(uint64_t(PsGetProcessId(this->handle_))); + return process_id_from_handle(PsGetProcessId(this->handle_)); } const char* process_handle::get_image_filename() const @@ -111,16 +111,25 @@ namespace process this->own_ = false; } - process_handle find_process_by_id(const uint32_t process_id) + process_id process_id_from_handle(HANDLE handle) { - PEPROCESS process{}; - const uint64_t process_id_long = process_id; - if (PsLookupProcessByProcessId(HANDLE(process_id_long), &process) != STATUS_SUCCESS) + return process_id(size_t(handle)); + } + + HANDLE handle_from_process_id(const process_id process) + { + return HANDLE(size_t(process)); + } + + process_handle find_process_by_id(const process_id process) + { + PEPROCESS process_obj{}; + if (PsLookupProcessByProcessId(handle_from_process_id(process), &process_obj) != STATUS_SUCCESS) { return {}; } - return process_handle{process, true}; + return process_handle{process_obj, true}; } process_handle get_current_process() @@ -128,6 +137,11 @@ namespace process return process_handle{PsGetCurrentProcess(), false}; } + process_id get_current_process_id() + { + return get_current_process().get_id(); + } + scoped_process_attacher::scoped_process_attacher(const process_handle& process) { if (!process || !process.is_alive()) diff --git a/src/driver/process.hpp b/src/driver/process.hpp index adfc5d8..eef929d 100644 --- a/src/driver/process.hpp +++ b/src/driver/process.hpp @@ -19,7 +19,7 @@ namespace process operator PEPROCESS() const; bool is_alive() const; - uint32_t get_id() const; + process_id get_id() const; const char* get_image_filename() const; @@ -30,9 +30,14 @@ namespace process void release(); }; - process_handle find_process_by_id(uint32_t process_id); + process_id process_id_from_handle(HANDLE handle); + HANDLE handle_from_process_id(process_id process); + + process_handle find_process_by_id(process_id process); process_handle get_current_process(); + process_id get_current_process_id(); + class scoped_process_attacher { public: diff --git a/src/driver/process_callback.cpp b/src/driver/process_callback.cpp index 0a00628..50f8fda 100644 --- a/src/driver/process_callback.cpp +++ b/src/driver/process_callback.cpp @@ -1,5 +1,6 @@ #include "std_include.hpp" #include "process_callback.hpp" +#include "process.hpp" #include "list.hpp" #include "logging.hpp" @@ -13,12 +14,13 @@ namespace process_callback return list; } - void process_notification_callback(const HANDLE parent_id, const HANDLE process_id, const BOOLEAN create) + void process_notification_callback(const HANDLE parent, const HANDLE process, const BOOLEAN create) { const auto& list = get_callback_list(); for (const auto& callback : list) { - callback(parent_id, process_id, create == FALSE ? type::destroy : type::create); + callback(process::process_id_from_handle(parent), process::process_id_from_handle(process), + create == FALSE ? type::destroy : type::create); } } diff --git a/src/driver/process_callback.hpp b/src/driver/process_callback.hpp index 737ccb1..16d16ab 100644 --- a/src/driver/process_callback.hpp +++ b/src/driver/process_callback.hpp @@ -9,7 +9,7 @@ namespace process_callback destroy, }; - using callback = void(HANDLE parent_id, HANDLE process_id, type type); + using callback = void(process_id parent_id, process_id process_id, type type); using callback_function = std::function; void* add(callback_function callback); diff --git a/src/driver/std_include.hpp b/src/driver/std_include.hpp index 011dde1..9c7ef99 100644 --- a/src/driver/std_include.hpp +++ b/src/driver/std_include.hpp @@ -16,3 +16,6 @@ #include "nt_ext.hpp" #include "new.hpp" #include "exception.hpp" + +// Not sure if this is good, but fuck it. +using process_id = uint32_t;