diff --git a/loader/Pe.hpp b/loader/Pe.hpp index 34fcb99..b120010 100644 --- a/loader/Pe.hpp +++ b/loader/Pe.hpp @@ -447,6 +447,10 @@ class Pe return *this; } + Pe(Pe&& other) noexcept : m_base(other.m_base), m_type(other.m_type) { + other.m_base = nullptr; + } + static Pe fromFile(const void* const buffer) noexcept { return Pe(ImgType::file, buffer); diff --git a/loader/mappedmodule.cpp b/loader/mappedmodule.cpp index 256909b..8d207f9 100644 --- a/loader/mappedmodule.cpp +++ b/loader/mappedmodule.cpp @@ -1,9 +1,11 @@ +#include #include #include #include #include #include +#include #include "mappedmodule.h" @@ -132,30 +134,6 @@ static void ProcessRelocations(const Pe::PeNative& pe, uintptr_t delta) { } } -static void UpdateSectionPermissions(const Pe::PeNative& pe) { - - for (const auto& section : pe.sections()) { - - auto pagePermissions = PAGE_READONLY; - auto oldPermissions = 1ul; - - if (section.Characteristics & IMAGE_SCN_CNT_CODE || section.Characteristics & IMAGE_SCN_MEM_EXECUTE) { - if (section.Characteristics & IMAGE_SCN_MEM_WRITE) { - pagePermissions = PAGE_EXECUTE_WRITECOPY; - } - else { - pagePermissions = PAGE_EXECUTE_READ; - } - } - else if (section.Characteristics & IMAGE_SCN_MEM_WRITE) { - pagePermissions = PAGE_READWRITE; - } - - if(pagePermissions != PAGE_READWRITE) - VirtualProtect(LPVOID(pe.byRva(section.VirtualAddress)), section.SizeOfRawData, pagePermissions, &oldPermissions); - } -} - static const Pe::GenericTypes::SecHeader* FindSection(const std::string& name, const Pe::PeNative& mappedPe) { for (auto& sec : mappedPe.sections()) { @@ -401,13 +379,8 @@ MappedModule::MappedModule(Logger logger, const std::vector& peBytes) throw std::runtime_error("PE file is not valid"); } - //We try to allocate at the preferred base, this will save relocation later - _mappedImage = VirtualAlloc((LPVOID)(pe.headers().opt()->ImageBase), pe.imageSize(), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); - - //Preferred base was not available, so just allocate at any available address - if (_mappedImage == nullptr) { - _mappedImage = VirtualAlloc(nullptr, pe.imageSize(), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); - } + _alloc = std::make_unique>(pe.imageSize()); + _mappedImage = _alloc->data(); _logger("Allocated image @ 0x%p\n", _mappedImage); @@ -453,7 +426,10 @@ MappedModule::MappedModule(Logger logger, const std::vector& peBytes) //Update the section permissions to expected values, //for example RX for the .text section - UpdateSectionPermissions(_mappedPe); + DWORD old=0; + // found some weird crashes in TLS callbacks when applying proper per-section protection to the unique_ptr-managed allocation. + // They were referring to non-existant stack regions + VirtualProtect(PVOID(pe.imageBase()), pe.imageSize(), PAGE_EXECUTE_READWRITE, &old); _logger("Set section permissions\n"); AddExceptionSupport(_mappedPe); @@ -474,15 +450,24 @@ MappedModule::MappedModule(Logger logger, const std::vector& peBytes) } } +MappedModule::MappedModule(MappedModule&& other) noexcept : + _mappedPe(std::move(other._mappedPe)), + _mappedImage(other._mappedImage), + _loaded(other._loaded), + _logger(other._logger), + _alloc(std::move(other._alloc)), + _refs(other._refs.fetch_add(1)){} + MappedModule::~MappedModule() { - - if (_mappedPe.headers().nt()->FileHeader.Characteristics & IMAGE_FILE_DLL) { - auto dllmain = dllmain_ptr(uintptr_t(_mappedPe.base()) + _mappedPe.headers().opt()->AddressOfEntryPoint); - dllmain(HINSTANCE(_mappedPe.base()), DLL_PROCESS_DETACH, 0); - } - RemoveExceptionSupport(_mappedPe); - VirtualFree(_mappedImage, 0, MEM_RELEASE); + if(_refs.fetch_sub(1) == 1) { + if (_mappedPe.headers().nt()->FileHeader.Characteristics & IMAGE_FILE_DLL) { + auto dllmain = dllmain_ptr(uintptr_t(_mappedPe.base()) + _mappedPe.headers().opt()->AddressOfEntryPoint); + dllmain(HINSTANCE(_mappedPe.base()), DLL_PROCESS_DETACH, 0); + } + + RemoveExceptionSupport(_mappedPe); + } } FARPROC MappedModule::GetProcAddress(const char* name) const { diff --git a/loader/mappedmodule.h b/loader/mappedmodule.h index ce2d952..2d1152b 100644 --- a/loader/mappedmodule.h +++ b/loader/mappedmodule.h @@ -1,6 +1,8 @@ #pragma once #include "Pe.hpp" +#include +#include #include #include #include @@ -21,9 +23,22 @@ class MappedModule { LPVOID _mappedImage; bool _loaded; Logger _logger; + std::unique_ptr> _alloc; + std::atomic _refs = 1; public: MappedModule(Logger logger, const std::vector& peBytes); + MappedModule(MappedModule&& other) noexcept; + MappedModule& operator=(MappedModule&& other) noexcept { + + _mappedPe = std::move(other._mappedPe); + _mappedImage = other._mappedImage; + _loaded = other._loaded; + _logger = other._logger; + _alloc = std::move(other._alloc); + _refs = other._refs.fetch_add(1); + return *this; + } ~MappedModule(); FARPROC GetProcAddress(const char* name) const; const Pe::PeNative& GetModule() const;