diff --git a/include/safetyhook/inline_hook.hpp b/include/safetyhook/inline_hook.hpp index a9c38da..381b5c4 100644 --- a/include/safetyhook/inline_hook.hpp +++ b/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. diff --git a/include/safetyhook/utility.hpp b/include/safetyhook/utility.hpp index 25b3707..933b338 100644 --- a/include/safetyhook/utility.hpp +++ b/include/safetyhook/utility.hpp @@ -13,4 +13,26 @@ template concept FnPtr = requires(T f) { std::is_pointer_v&& std::is_function_v>; }; 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 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 unprotect(uint8_t* address, size_t size); } // namespace safetyhook diff --git a/src/inline_hook.cpp b/src/inline_hook.cpp index fb98900..4a311bc 100644 --- a/src/inline_hook.cpp +++ b/src/inline_hook.cpp @@ -23,19 +23,6 @@ #include 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 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(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 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(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,7 +305,10 @@ std::expected InlineHook::e9_hook(const std::shared_ptr // jmp from trampoline to original. auto src = reinterpret_cast(&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(&trampoline_epilogue->jmp_to_destination); @@ -312,17 +316,26 @@ std::expected InlineHook::e9_hook(const std::shared_ptr #ifdef _M_X64 auto data = reinterpret_cast(&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; + // jmp from original to trampoline. execute_while_frozen( - [this, &trampoline_epilogue] { - const auto src = m_target; - const auto dst = reinterpret_cast(&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(&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) { @@ -330,6 +343,10 @@ std::expected InlineHook::e9_hook(const std::shared_ptr } }); + if (error) { + return std::unexpected{*error}; + } + return {}; } @@ -372,15 +389,20 @@ std::expected InlineHook::ff_hook(const std::shared_ptr auto src = reinterpret_cast(&trampoline_epilogue->jmp_to_original); auto dst = m_target + m_original_bytes.size(); auto data = reinterpret_cast(&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; // 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) { @@ -388,6 +410,10 @@ std::expected InlineHook::ff_hook(const std::shared_ptr } }); + 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) { diff --git a/src/utility.cpp b/src/utility.cpp index 65d5b57..8fc6cc1 100644 --- a/src/utility.cpp +++ b/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 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 \ No newline at end of file