Fix potential UAF/double-free footgun

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);

loader/mappedmodule.cpp

@@ -1,9 +1,11 @@
 
+#include <memory>
 #include <windows.h>
 #include <vector>
 #include <string>
 #include <fstream>
 #include <format>
+#include <winnt.h>
 
 #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<void>(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<std::byte>& 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<std::vector<std::byte>>(pe.imageSize());
+    _mappedImage = _alloc->data();
 
     _logger("Allocated image @ 0x%p\n", _mappedImage);
 
@@ -453,7 +426,10 @@ MappedModule::MappedModule(Logger logger, const std::vector<std::byte>& 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<std::byte>& 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 {

loader/mappedmodule.h

@@ -1,6 +1,8 @@
 #pragma once
 
 #include "Pe.hpp"
+#include <atomic>
+#include <memory>
 #include <string>
 #include <vector>
 #include <functional>
@@ -21,9 +23,22 @@ class MappedModule {
 	LPVOID _mappedImage;
 	bool _loaded;
 	Logger _logger;
+  	std::unique_ptr<std::vector<std::byte>> _alloc;
+	std::atomic<int32_t> _refs = 1;
 
 public:
 	MappedModule(Logger logger, const std::vector<std::byte>& 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;