General hat::memory_protector improvements

BasedInc · Oct 7, 2024 · 82f4466 · 82f4466
1 parent ab8dba3
commit 82f4466
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/include/libhat/MemoryProtector.hpp b/include/libhat/MemoryProtector.hpp
@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <type_traits>
+#include <utility>
 
 namespace hat {
 
@@ -25,10 +26,34 @@ namespace hat {
     class memory_protector {
     public:
         memory_protector(uintptr_t address, size_t size, protection flags);
-        ~memory_protector();
+
+        ~memory_protector() {
+            if (this->set) {
+                this->restore();
+            }
+        }
+
+        memory_protector(memory_protector&& o) noexcept :
+            address(o.address),
+            size(o.size),
+            oldProtection(o.oldProtection),
+            set(std::exchange(o.set, false)) {}
+
+        memory_protector& operator=(memory_protector&& o) noexcept = delete;
+        memory_protector(const memory_protector&) = delete;
+        memory_protector& operator=(const memory_protector&) = delete;
+
+        /// Returns true if the memory protect operation was successful
+        [[nodiscard]] bool is_set() const {
+            return this->set;
+        }
+
     private:
+        void restore();
+
         uintptr_t address;
         size_t size;
         uint32_t oldProtection{}; // Memory protection flags native to Operating System
+        bool set{};
     };
 }
diff --git a/src/os/win32/MemoryProtector.cpp b/src/os/win32/MemoryProtector.cpp
@@ -20,12 +20,11 @@ namespace hat {
     }
 
     memory_protector::memory_protector(const uintptr_t address, const size_t size, const protection flags) : address(address), size(size) {
-        VirtualProtect(reinterpret_cast<LPVOID>(this->address), this->size, ToWinProt(flags), reinterpret_cast<PDWORD>(&this->oldProtection));
+        this->set = 0 != VirtualProtect(reinterpret_cast<LPVOID>(this->address), this->size, ToWinProt(flags), reinterpret_cast<PDWORD>(&this->oldProtection));
     }
 
-    memory_protector::~memory_protector() {
-        DWORD temp;
-        VirtualProtect(reinterpret_cast<LPVOID>(this->address), this->size, this->oldProtection, &temp);
+    void memory_protector::restore() {
+        VirtualProtect(reinterpret_cast<LPVOID>(this->address), this->size, this->oldProtection, reinterpret_cast<PDWORD>(&this->oldProtection));
     }
 }
 #endif