Fix/Unprotected memory safety (#54)

include/safetyhook/inline_hook.hpp

@@ -26,6 +26,8 @@ class InlineHook final {
             SHORT_JUMP_IN_TRAMPOLINE,              ///< The trampoline contains a short jump.
             IP_RELATIVE_INSTRUCTION_OUT_OF_RANGE,  ///< An IP-relative instruction is out of range.
             UNSUPPORTED_INSTRUCTION_IN_TRAMPOLINE, ///< An unsupported instruction was found in the trampoline.
+            FAILED_TO_UNPROTECT,                   ///< Failed to unprotect memory.
+            NOT_ENOUGH_SPACE,                      ///< Not enough space to create the hook.
         } type;
 
         /// @brief Extra information about the error.
@@ -68,6 +70,16 @@ class InlineHook final {
         [[nodiscard]] static Error unsupported_instruction_in_trampoline(uint8_t* ip) {
             return {.type = UNSUPPORTED_INSTRUCTION_IN_TRAMPOLINE, .ip = ip};
         }
+
+        /// @brief Create a FAILED_TO_UNPROTECT error.
+        /// @param ip The IP of the problematic instruction.
+        /// @return The new FAILED_TO_UNPROTECT error.
+        [[nodiscard]] static Error failed_to_unprotect(uint8_t* ip) { return {.type = FAILED_TO_UNPROTECT, .ip = ip}; }
+
+        /// @brief Create a NOT_ENOUGH_SPACE error.
+        /// @param ip The IP of the problematic instruction.
+        /// @return The new NOT_ENOUGH_SPACE error.
+        [[nodiscard]] static Error not_enough_space(uint8_t* ip) { return {.type = NOT_ENOUGH_SPACE, .ip = ip}; }
     };
 
     /// @brief Create an inline hook.

include/safetyhook/utility.hpp

@@ -13,4 +13,26 @@ template <typename T>
 concept FnPtr = requires(T f) { std::is_pointer_v<T>&& std::is_function_v<std::remove_pointer_t<T>>; };
 
 bool is_executable(uint8_t* address);
+
+class UnprotectMemory {
+public:
+    UnprotectMemory() = delete;
+    ~UnprotectMemory();
+    UnprotectMemory(const UnprotectMemory&) = delete;
+    UnprotectMemory(UnprotectMemory&& other) noexcept;
+    UnprotectMemory& operator=(const UnprotectMemory&) = delete;
+    UnprotectMemory& operator=(UnprotectMemory&& other) noexcept;
+
+private:
+    friend std::optional<UnprotectMemory> unprotect(uint8_t*, size_t);
+
+    UnprotectMemory(uint8_t* address, size_t size, uint32_t original_protection)
+        : m_address{address}, m_size{size}, m_original_protection{original_protection} {}
+
+    uint8_t* m_address{};
+    size_t m_size{};
+    uint32_t m_original_protection{};
+};
+
+[[nodiscard]] std::optional<UnprotectMemory> unprotect(uint8_t* address, size_t size);
 } // namespace safetyhook

src/inline_hook.cpp

@@ -23,19 +23,6 @@
 #include <safetyhook/inline_hook.hpp>
 
 namespace safetyhook {
-class UnprotectMemory {
-public:
-    UnprotectMemory(uint8_t* address, size_t size) : m_address{address}, m_size{size} {
-        VirtualProtect(m_address, m_size, PAGE_EXECUTE_READWRITE, &m_protect);
-    }
-
-    ~UnprotectMemory() { VirtualProtect(m_address, m_size, m_protect, &m_protect); }
-
-private:
-    uint8_t* m_address{};
-    size_t m_size{};
-    DWORD m_protect{};
-};
 
 #pragma pack(push, 1)
 struct JmpE9 {
@@ -78,18 +65,25 @@ static auto make_jmp_ff(uint8_t* src, uint8_t* dst, uint8_t* data) {
     return jmp;
 }
 
-static void emit_jmp_ff(uint8_t* src, uint8_t* dst, uint8_t* data, size_t size = sizeof(JmpFF)) {
+[[nodiscard]] static std::expected<void, InlineHook::Error> emit_jmp_ff(
+    uint8_t* src, uint8_t* dst, uint8_t* data, size_t size = sizeof(JmpFF)) {
     if (size < sizeof(JmpFF)) {
-        return;
+        return std::unexpected{InlineHook::Error::not_enough_space(dst)};
     }
 
-    UnprotectMemory unprotect{src, size};
+    auto um = unprotect(src, size);
+
+    if (!um) {
+        return std::unexpected{InlineHook::Error::failed_to_unprotect(src)};
+    }
 
     if (size > sizeof(JmpFF)) {
         std::fill_n(src, size, static_cast<uint8_t>(0x90));
     }
 
     store(src, make_jmp_ff(src, dst, data));
+
+    return {};
 }
 #endif
 
@@ -101,18 +95,25 @@ constexpr auto make_jmp_e9(uint8_t* src, uint8_t* dst) {
     return jmp;
 }
 
-static void emit_jmp_e9(uint8_t* src, uint8_t* dst, size_t size = sizeof(JmpE9)) {
+[[nodiscard]] static std::expected<void, InlineHook::Error> emit_jmp_e9(
+    uint8_t* src, uint8_t* dst, size_t size = sizeof(JmpE9)) {
     if (size < sizeof(JmpE9)) {
-        return;
+        return std::unexpected{InlineHook::Error::not_enough_space(dst)};
     }
 
-    UnprotectMemory unprotect{src, size};
+    auto um = unprotect(src, size);
+
+    if (!um) {
+        return std::unexpected{InlineHook::Error::failed_to_unprotect(src)};
+    }
 
     if (size > sizeof(JmpE9)) {
         std::fill_n(src, size, static_cast<uint8_t>(0x90));
     }
 
     store(src, make_jmp_e9(src, dst));
+
+    return {};
 }
 
 static bool decode(ZydisDecodedInstruction* ix, uint8_t* ip) {
@@ -167,8 +168,8 @@ InlineHook& InlineHook::operator=(InlineHook&& other) noexcept {
         m_trampoline_size = other.m_trampoline_size;
         m_original_bytes = std::move(other.m_original_bytes);
 
-        other.m_target = 0;
-        other.m_destination = 0;
+        other.m_target = nullptr;
+        other.m_destination = nullptr;
         other.m_trampoline_size = 0;
     }
 
@@ -304,32 +305,48 @@ std::expected<void, InlineHook::Error> InlineHook::e9_hook(const std::shared_ptr
     // jmp from trampoline to original.
     auto src = reinterpret_cast<uint8_t*>(&trampoline_epilogue->jmp_to_original);
     auto dst = m_target + m_original_bytes.size();
-    emit_jmp_e9(src, dst);
+
+    if (auto result = emit_jmp_e9(src, dst); !result) {
+        return std::unexpected{result.error()};
+    }
 
     // jmp from trampoline to destination.
     src = reinterpret_cast<uint8_t*>(&trampoline_epilogue->jmp_to_destination);
     dst = m_destination;
 
 #ifdef _M_X64
     auto data = reinterpret_cast<uint8_t*>(&trampoline_epilogue->destination_address);
-    emit_jmp_ff(src, dst, data);
+
+    if (auto result = emit_jmp_ff(src, dst, data); !result) {
+        return std::unexpected{result.error()};
+    }
 #else
-    emit_jmp_e9(src, dst);
+    if (auto result = emit_jmp_e9(src, dst); !result) {
+        return std::unexpected{result.error()};
+    }
 #endif
 
+    std::optional<Error> error;
+
     // jmp from original to trampoline.
     execute_while_frozen(
-        [this, &trampoline_epilogue] {
-            const auto src = m_target;
-            const auto dst = reinterpret_cast<uint8_t*>(&trampoline_epilogue->jmp_to_destination);
-            emit_jmp_e9(src, dst, m_original_bytes.size());
+        [this, &trampoline_epilogue, &error] {
+            if (auto result = emit_jmp_e9(m_target,
+                    reinterpret_cast<uint8_t*>(&trampoline_epilogue->jmp_to_destination), m_original_bytes.size());
+                !result) {
+                error = result.error();
+            }
         },
         [this](uint32_t, HANDLE, CONTEXT& ctx) {
             for (size_t i = 0; i < m_original_bytes.size(); ++i) {
                 fix_ip(ctx, m_target + i, m_trampoline.data() + i);
             }
         });
 
+    if (error) {
+        return std::unexpected{*error};
+    }
+
     return {};
 }
 
@@ -372,22 +389,31 @@ std::expected<void, InlineHook::Error> InlineHook::ff_hook(const std::shared_ptr
     auto src = reinterpret_cast<uint8_t*>(&trampoline_epilogue->jmp_to_original);
     auto dst = m_target + m_original_bytes.size();
     auto data = reinterpret_cast<uint8_t*>(&trampoline_epilogue->original_address);
-    emit_jmp_ff(src, dst, data);
+
+    if (auto result = emit_jmp_ff(src, dst, data); !result) {
+        return std::unexpected{result.error()};
+    }
+
+    std::optional<Error> error;
 
     // jmp from original to trampoline.
     execute_while_frozen(
-        [this] {
-            const auto src = m_target;
-            const auto dst = m_destination;
-            const auto data = src + sizeof(JmpFF);
-            emit_jmp_ff(src, dst, data, m_original_bytes.size());
+        [this, &error] {
+            if (auto result = emit_jmp_ff(m_target, m_destination, m_target + sizeof(JmpFF), m_original_bytes.size());
+                !result) {
+                error = result.error();
+            }
         },
         [this](uint32_t, HANDLE, CONTEXT& ctx) {
             for (size_t i = 0; i < m_original_bytes.size(); ++i) {
                 fix_ip(ctx, m_target + i, m_trampoline.data() + i);
             }
         });
 
+    if (error) {
+        return std::unexpected{*error};
+    }
+
     return {};
 }
 #endif
@@ -401,8 +427,9 @@ void InlineHook::destroy() {
 
     execute_while_frozen(
         [this] {
-            UnprotectMemory unprotect{m_target, m_original_bytes.size()};
-            std::copy(m_original_bytes.begin(), m_original_bytes.end(), m_target);
+            if (auto um = unprotect(m_target, m_original_bytes.size())) {
+                std::copy(m_original_bytes.begin(), m_original_bytes.end(), m_target);
+            }
         },
         [this](uint32_t, HANDLE, CONTEXT& ctx) {
             for (size_t i = 0; i < m_original_bytes.size(); ++i) {

src/utility.cpp

@@ -53,4 +53,39 @@ bool is_executable(uint8_t* address) {
 
     return is_page_executable(address);
 }
+
+UnprotectMemory::~UnprotectMemory() {
+    if (m_address != nullptr) {
+        DWORD old_protection;
+        VirtualProtect(m_address, m_size, m_original_protection, &old_protection);
+    }
+}
+
+UnprotectMemory::UnprotectMemory(UnprotectMemory&& other) noexcept {
+    *this = std::move(other);
+}
+
+UnprotectMemory& UnprotectMemory::operator=(UnprotectMemory&& other) noexcept {
+    if (this != &other) {
+        m_address = other.m_address;
+        m_size = other.m_size;
+        m_original_protection = other.m_original_protection;
+        other.m_address = nullptr;
+        other.m_size = 0;
+        other.m_original_protection = 0;
+    }
+
+    return *this;
+}
+
+std::optional<UnprotectMemory> unprotect(uint8_t* address, size_t size) {
+    DWORD old_protection;
+
+    if (!VirtualProtect(address, size, PAGE_EXECUTE_READWRITE, &old_protection)) {
+        return {};
+    }
+
+    return UnprotectMemory{address, size, old_protection};
+}
+
 } // namespace safetyhook