From 1bbd9e9c737cad9e3dbff72dcd5e26b994c1ef1c Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sun, 27 Mar 2022 11:57:26 +0200 Subject: [PATCH] Better memory management --- src/driver/hypervisor.cpp | 49 +++++++------------------- src/driver/memory.cpp | 74 +++++++++++++++++++++++++++++++++++++++ src/driver/memory.hpp | 18 ++++++++++ src/driver/new.cpp | 27 +++++--------- src/driver/new.hpp | 2 -- 5 files changed, 114 insertions(+), 56 deletions(-) create mode 100644 src/driver/memory.cpp create mode 100644 src/driver/memory.hpp diff --git a/src/driver/hypervisor.cpp b/src/driver/hypervisor.cpp index 2775a1f..e70f96b 100644 --- a/src/driver/hypervisor.cpp +++ b/src/driver/hypervisor.cpp @@ -4,6 +4,7 @@ #include "exception.hpp" #include "logging.hpp" #include "finally.hpp" +#include "memory.hpp" #include "thread.hpp" #define IA32_FEATURE_CONTROL_MSR 0x3A @@ -34,34 +35,6 @@ namespace { return is_vmx_supported() && is_vmx_available(); } - - _IRQL_requires_max_(DISPATCH_LEVEL) - - void free_aligned_memory(void* memory) - { - MmFreeContiguousMemory(memory); - } - - _Must_inspect_result_ - _IRQL_requires_max_(DISPATCH_LEVEL) - - void* allocate_aligned_memory(const size_t size) - { - PHYSICAL_ADDRESS lowest{}, highest{}; - lowest.QuadPart = 0; - highest.QuadPart = lowest.QuadPart - 1; - -#if (NTDDI_VERSION >= NTDDI_VISTA) - return MmAllocateContiguousNodeMemory(size, - lowest, - highest, - lowest, - PAGE_READWRITE, - KeGetCurrentNodeNumber()); -#else - return MmAllocateContiguousMemory(size, highest); -#endif - } } hypervisor::hypervisor() @@ -125,23 +98,27 @@ void hypervisor::disable_core() void hypervisor::allocate_vm_states() { + if (this->vm_states_) + { + throw std::runtime_error("VM states are still in use"); + } + const auto core_count = thread::get_processor_count(); const auto allocation_size = sizeof(vmx::vm_state) * core_count; - this->vm_states_ = static_cast(allocate_aligned_memory(allocation_size)); - if(!this->vm_states_) + this->vm_states_ = static_cast(memory::allocate_aligned_memory(allocation_size)); + if (!this->vm_states_) { - throw std::runtime_error("Failed to allocate vm states"); + throw std::runtime_error("Failed to allocate VM states"); } + + RtlSecureZeroMemory(this->vm_states_, allocation_size); } void hypervisor::free_vm_states() { - if(this->vm_states_) - { - free_aligned_memory(this->vm_states_); - this->vm_states_ = nullptr; - } + memory::free_aligned_memory(this->vm_states_); + this->vm_states_ = nullptr; } vmx::vm_state* hypervisor::get_current_vm_state() const diff --git a/src/driver/memory.cpp b/src/driver/memory.cpp new file mode 100644 index 0000000..b8bc8f6 --- /dev/null +++ b/src/driver/memory.cpp @@ -0,0 +1,74 @@ +#include "std_include.hpp" +#include "memory.hpp" + +namespace memory +{ + namespace + { + void* allocate_aligned_memory_internal(const size_t size) + { + PHYSICAL_ADDRESS lowest{}, highest{}; + lowest.QuadPart = 0; + highest.QuadPart = lowest.QuadPart - 1; + +#if (NTDDI_VERSION >= NTDDI_VISTA) + return MmAllocateContiguousNodeMemory(size, + lowest, + highest, + lowest, + PAGE_READWRITE, + KeGetCurrentNodeNumber()); +#else + return MmAllocateContiguousMemory(size, highest); +#endif + } + } + + _IRQL_requires_max_(DISPATCH_LEVEL) + + void free_aligned_memory(void* memory) + { + if (memory) + { + MmFreeContiguousMemory(memory); + } + } + + _Must_inspect_result_ + _IRQL_requires_max_(DISPATCH_LEVEL) + + void* allocate_aligned_memory(const size_t size) + { + void* memory = allocate_aligned_memory_internal(size); + if (memory) + { + RtlSecureZeroMemory(memory, size); + } + + return memory; + } + + _Must_inspect_result_ + _IRQL_requires_max_(DISPATCH_LEVEL) + + void* allocate_non_paged_memory(const size_t size) + { + void* memory = ExAllocatePoolWithTag(NonPagedPool, size, 'MOMO'); + if (memory) + { + RtlSecureZeroMemory(memory, size); + } + + return memory; + } + + _IRQL_requires_max_(DISPATCH_LEVEL) + + void free_non_paged_memory(void* memory) + { + if (memory) + { + ExFreePool(memory); + } + } +} diff --git a/src/driver/memory.hpp b/src/driver/memory.hpp new file mode 100644 index 0000000..dda4acd --- /dev/null +++ b/src/driver/memory.hpp @@ -0,0 +1,18 @@ +#pragma once + +namespace memory +{ + _IRQL_requires_max_(DISPATCH_LEVEL) + void free_aligned_memory(void* memory); + + _Must_inspect_result_ + _IRQL_requires_max_(DISPATCH_LEVEL) + void* allocate_aligned_memory(size_t size); + + _Must_inspect_result_ + _IRQL_requires_max_(DISPATCH_LEVEL) + void* allocate_non_paged_memory(size_t size); + + _IRQL_requires_max_(DISPATCH_LEVEL) + void free_non_paged_memory(void* memory); +} diff --git a/src/driver/new.cpp b/src/driver/new.cpp index 0623bdb..0f2a1be 100644 --- a/src/driver/new.cpp +++ b/src/driver/new.cpp @@ -1,12 +1,13 @@ #include "std_include.hpp" #include "new.hpp" #include "exception.hpp" +#include "memory.hpp" namespace { - void* perform_allocation(const size_t size, const POOL_TYPE pool, const unsigned long tag) + void* perform_checked_non_paged_allocation(const size_t size) { - auto* memory = ExAllocatePoolWithTag(pool, size, tag); + auto* memory = memory::allocate_non_paged_memory(size); if (!memory) { throw std::runtime_error("Memory allocation failed"); @@ -16,24 +17,14 @@ namespace } } -void* operator new(const size_t size, const POOL_TYPE pool, const unsigned long tag) -{ - return perform_allocation(size, pool, tag); -} - -void* operator new[](const size_t size, const POOL_TYPE pool, const unsigned long tag) -{ - return perform_allocation(size, pool, tag); -} - void* operator new(const size_t size) { - return operator new(size, NonPagedPool); + return perform_checked_non_paged_allocation(size); } void* operator new[](const size_t size) { - return operator new[](size, NonPagedPool); + return perform_checked_non_paged_allocation(size); } // Placement new @@ -44,22 +35,22 @@ void* operator new(size_t, void* where) void operator delete(void* ptr, size_t) { - ExFreePool(ptr); + memory::free_non_paged_memory(ptr); } void operator delete(void* ptr) { - ExFreePool(ptr); + memory::free_non_paged_memory(ptr); } void operator delete[](void* ptr, size_t) { - ExFreePool(ptr); + memory::free_non_paged_memory(ptr); } void operator delete[](void* ptr) { - ExFreePool(ptr); + memory::free_non_paged_memory(ptr); } extern "C" void __std_terminate() diff --git a/src/driver/new.hpp b/src/driver/new.hpp index 860dfa3..ba3206f 100644 --- a/src/driver/new.hpp +++ b/src/driver/new.hpp @@ -1,7 +1,5 @@ #pragma once -void* operator new(size_t size, POOL_TYPE pool, unsigned long tag = 'momo'); -void* operator new[](size_t size, POOL_TYPE pool, unsigned long tag = 'momo'); void* operator new(size_t size); void* operator new[](size_t size);