InlineHook/MidHook: Add enable() and disable()

include/safetyhook/easy.hpp

@@ -12,29 +12,34 @@ namespace safetyhook {
 /// @brief Easy to use API for creating an InlineHook.
 /// @param target The address of the function to hook.
 /// @param destination The address of the destination function.
+/// @param flags The flags to use.
 /// @return The InlineHook object.
-[[nodiscard]] InlineHook create_inline(void* target, void* destination);
+[[nodiscard]] InlineHook create_inline(void* target, void* destination, InlineHook::Flags flags = InlineHook::Default);
 
 /// @brief Easy to use API for creating an InlineHook.
 /// @param target The address of the function to hook.
 /// @param destination The address of the destination function.
+/// @param flags The flags to use.
 /// @return The InlineHook object.
-[[nodiscard]] InlineHook create_inline(FnPtr auto target, FnPtr auto destination) {
- return create_inline(reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination));
+[[nodiscard]] InlineHook create_inline(
+ FnPtr auto target, FnPtr auto destination, InlineHook::Flags flags = InlineHook::Default) {
+ return create_inline(reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination), flags);
 }
 
 /// @brief Easy to use API for creating a MidHook.
 /// @param target the address of the function to hook.
 /// @param destination The destination function.
+/// @param flags The flags to use.
 /// @return The MidHook object.
-[[nodiscard]] MidHook create_mid(void* target, MidHookFn destination);
+[[nodiscard]] MidHook create_mid(void* target, MidHookFn destination, MidHook::Flags = MidHook::Default);
 
 /// @brief Easy to use API for creating a MidHook.
 /// @param target the address of the function to hook.
 /// @param destination The destination function.
+/// @param flags The flags to use.
 /// @return The MidHook object.
-[[nodiscard]] MidHook create_mid(FnPtr auto target, MidHookFn destination) {
- return create_mid(reinterpret_cast<void*>(target), destination);
+[[nodiscard]] MidHook create_mid(FnPtr auto target, MidHookFn destination, MidHook::Flags flags = MidHook::Default) {
+ return create_mid(reinterpret_cast<void*>(target), destination, flags);
 }
 
 /// @brief Easy to use API for creating a VmtHook.

include/safetyhook/inline_hook.hpp

@@ -87,42 +87,54 @@ class InlineHook final {
  [[nodiscard]] static Error not_enough_space(uint8_t* ip) { return {.type = NOT_ENOUGH_SPACE, .ip = ip}; }
  };
 
+ /// @brief Flags for InlineHook.
+ enum Flags : int {
+ Default = 0, ///< Default flags.
+ StartDisabled = 1 << 0, ///< Start the hook disabled.
+ };
+
  /// @brief Create an inline hook.
  /// @param target The address of the function to hook.
  /// @param destination The destination address.
+ /// @param flags The flags to use.
  /// @return The InlineHook or an InlineHook::Error if an error occurred.
  /// @note This will use the default global Allocator.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_inline).
- [[nodiscard]] static std::expected<InlineHook, Error> create(void* target, void* destination);
+ [[nodiscard]] static std::expected<InlineHook, Error> create(
+ void* target, void* destination, Flags flags = Default);
 
  /// @brief Create an inline hook.
  /// @param target The address of the function to hook.
  /// @param destination The destination address.
+ /// @param flags The flags to use.
  /// @return The InlineHook or an InlineHook::Error if an error occurred.
  /// @note This will use the default global Allocator.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_inline).
- [[nodiscard]] static std::expected<InlineHook, Error> create(FnPtr auto target, FnPtr auto destination) {
- return create(reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination));
+ [[nodiscard]] static std::expected<InlineHook, Error> create(
+ FnPtr auto target, FnPtr auto destination, Flags flags = Default) {
+ return create(reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination), flags);
  }
 
  /// @brief Create an inline hook with a given Allocator.
  /// @param allocator The allocator to use.
  /// @param target The address of the function to hook.
  /// @param destination The destination address.
+ /// @param flags The flags to use.
  /// @return The InlineHook or an InlineHook::Error if an error occurred.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_inline).
  [[nodiscard]] static std::expected<InlineHook, Error> create(
- const std::shared_ptr<Allocator>& allocator, void* target, void* destination);
+ const std::shared_ptr<Allocator>& allocator, void* target, void* destination, Flags flags = Default);
 
  /// @brief Create an inline hook with a given Allocator.
  /// @param allocator The allocator to use.
  /// @param target The address of the function to hook.
  /// @param destination The destination address.
+ /// @param flags The flags to use.
  /// @return The InlineHook or an InlineHook::Error if an error occurred.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_inline).
  [[nodiscard]] static std::expected<InlineHook, Error> create(
- const std::shared_ptr<Allocator>& allocator, FnPtr auto target, FnPtr auto destination) {
- return create(allocator, reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination));
+ const std::shared_ptr<Allocator>& allocator, FnPtr auto target, FnPtr auto destination, Flags flags = Default) {
+ return create(allocator, reinterpret_cast<void*>(target), reinterpret_cast<void*>(destination), flags);
  }
 
  InlineHook() = default;
@@ -285,15 +297,32 @@ class InlineHook final {
  return original<RetT(SAFETYHOOK_FASTCALL*)(Args...)>()(args...);
  }
 
+ /// @brief Enable the hook.
+ [[nodiscard]] std::expected<void, Error> enable();
+
+ /// @brief Disable the hook.
+ [[nodiscard]] std::expected<void, Error> disable();
+
+ /// @brief Check if the hook is enabled.
+ [[nodiscard]] bool enabled() const { return m_enabled; }
+
 private:
  friend class MidHook;
 
+ enum class Type {
+ Unset,
+ E9,
+ FF,
+ };
+
  uint8_t* m_target{};
  uint8_t* m_destination{};
  Allocation m_trampoline{};
  std::vector<uint8_t> m_original_bytes{};
  uintptr_t m_trampoline_size{};
  std::recursive_mutex m_mutex{};
+ bool m_enabled{};
+ Type m_type{Type::Unset};
 
  std::expected<void, Error> setup(
  const std::shared_ptr<Allocator>& allocator, uint8_t* target, uint8_t* destination);

include/safetyhook/mid_hook.hpp

@@ -52,43 +52,55 @@ class MidHook final {
  }
  };
 
+ /// @brief Flags for MidHook.
+ enum Flags : int {
+ Default = 0, ///< Default flags.
+ StartDisabled = 1, ///< Start the hook disabled.
+ };
+
  /// @brief Creates a new MidHook object.
  /// @param target The address of the function to hook.
  /// @param destination_fn The destination function.
+ /// @param flags The flags to use.
  /// @return The MidHook object or a MidHook::Error if an error occurred.
  /// @note This will use the default global Allocator.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_mid).
- [[nodiscard]] static std::expected<MidHook, Error> create(void* target, MidHookFn destination_fn);
+ [[nodiscard]] static std::expected<MidHook, Error> create(
+ void* target, MidHookFn destination_fn, Flags flags = Default);
 
  /// @brief Creates a new MidHook object.
  /// @param target The address of the function to hook.
  /// @param destination_fn The destination function.
+ /// @param flags The flags to use.
  /// @return The MidHook object or a MidHook::Error if an error occurred.
  /// @note This will use the default global Allocator.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_mid).
- [[nodiscard]] static std::expected<MidHook, Error> create(FnPtr auto target, MidHookFn destination_fn) {
- return create(reinterpret_cast<void*>(target), destination_fn);
+ [[nodiscard]] static std::expected<MidHook, Error> create(
+ FnPtr auto target, MidHookFn destination_fn, Flags flags = Default) {
+ return create(reinterpret_cast<void*>(target), destination_fn, flags);
  }
 
  /// @brief Creates a new MidHook object with a given Allocator.
  /// @param allocator The Allocator to use.
  /// @param target The address of the function to hook.
  /// @param destination_fn The destination function.
+ /// @param flags The flags to use.
  /// @return The MidHook object or a MidHook::Error if an error occurred.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_mid).
  [[nodiscard]] static std::expected<MidHook, Error> create(
- const std::shared_ptr<Allocator>& allocator, void* target, MidHookFn destination_fn);
+ const std::shared_ptr<Allocator>& allocator, void* target, MidHookFn destination_fn, Flags flags = Default);
 
  /// @brief Creates a new MidHook object with a given Allocator.
  /// @tparam T The type of the function to hook.
  /// @param allocator The Allocator to use.
  /// @param target The address of the function to hook.
  /// @param destination_fn The destination function.
+ /// @param flags The flags to use.
  /// @return The MidHook object or a MidHook::Error if an error occurred.
  /// @note If you don't care about error handling, use the easy API (safetyhook::create_mid).
- [[nodiscard]] static std::expected<MidHook, Error> create(
- const std::shared_ptr<Allocator>& allocator, FnPtr auto target, MidHookFn destination_fn) {
- return create(allocator, reinterpret_cast<void*>(target), destination_fn);
+ [[nodiscard]] static std::expected<MidHook, Error> create(const std::shared_ptr<Allocator>& allocator,
+ FnPtr auto target, MidHookFn destination_fn, Flags flags = Default) {
+ return create(allocator, reinterpret_cast<void*>(target), destination_fn, flags);
  }
 
  MidHook() = default;
@@ -123,6 +135,15 @@ class MidHook final {
  /// @return true if the hook is valid, false otherwise.
  explicit operator bool() const { return static_cast<bool>(m_stub); }
 
+ /// @brief Enable the hook.
+ [[nodiscard]] std::expected<void, Error> enable();
+
+ /// @brief Disable the hook.
+ [[nodiscard]] std::expected<void, Error> disable();
+
+ /// @brief Check if the hook is enabled.
+ [[nodiscard]] bool enabled() const { return m_hook.enabled(); }
+
 private:
  InlineHook m_hook{};
  uint8_t* m_target{};

src/easy.cpp

@@ -1,16 +1,16 @@
 #include "safetyhook/easy.hpp"
 
 namespace safetyhook {
-InlineHook create_inline(void* target, void* destination) {
- if (auto hook = InlineHook::create(target, destination)) {
+InlineHook create_inline(void* target, void* destination, InlineHook::Flags flags) {
+ if (auto hook = InlineHook::create(target, destination, flags)) {
  return std::move(*hook);
  } else {
  return {};
  }
 }
 
-MidHook create_mid(void* target, MidHookFn destination) {
- if (auto hook = MidHook::create(target, destination)) {
+MidHook create_mid(void* target, MidHookFn destination, MidHook::Flags flags) {
+ if (auto hook = MidHook::create(target, destination, flags)) {
  return std::move(*hook);
  } else {
  return {};

src/inline_hook.cpp

@@ -114,12 +114,12 @@ static bool decode(ZydisDecodedInstruction* ix, uint8_t* ip) {
  return ZYAN_SUCCESS(ZydisDecoderDecodeInstruction(&decoder, nullptr, ip, 15, ix));
 }
 
-std::expected<InlineHook, InlineHook::Error> InlineHook::create(void* target, void* destination) {
- return create(Allocator::global(), target, destination);
+std::expected<InlineHook, InlineHook::Error> InlineHook::create(void* target, void* destination, Flags flags) {
+ return create(Allocator::global(), target, destination, flags);
 }
 
 std::expected<InlineHook, InlineHook::Error> InlineHook::create(
- const std::shared_ptr<Allocator>& allocator, void* target, void* destination) {
+ const std::shared_ptr<Allocator>& allocator, void* target, void* destination, Flags flags) {
  InlineHook hook{};
 
  if (const auto setup_result =
@@ -128,6 +128,12 @@ std::expected<InlineHook, InlineHook::Error> InlineHook::create(
  return std::unexpected{setup_result.error()};
  }
 
+ if (!(flags & StartDisabled)) {
+ if (auto enable_result = hook.enable(); !enable_result) {
+ return std::unexpected{enable_result.error()};
+ }
+ }
+
  return hook;
 }
 
@@ -146,10 +152,14 @@ InlineHook& InlineHook::operator=(InlineHook&& other) noexcept {
  m_trampoline = std::move(other.m_trampoline);
  m_trampoline_size = other.m_trampoline_size;
  m_original_bytes = std::move(other.m_original_bytes);
+ m_enabled = other.m_enabled;
+ m_type = other.m_type;
 
  other.m_target = nullptr;
  other.m_destination = nullptr;
  other.m_trampoline_size = 0;
+ other.m_enabled = false;
+ other.m_type = Type::Unset;
  }
 
  return *this;
@@ -305,20 +315,7 @@ std::expected<void, InlineHook::Error> InlineHook::e9_hook(const std::shared_ptr
  }
 #endif
 
- std::optional<Error> error;
-
- // jmp from original to trampoline.
- trap_threads(m_target, m_trampoline.data(), 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();
- }
- });
-
- if (error) {
- return std::unexpected{*error};
- }
+ m_type = Type::E9;
 
  return {};
 }
@@ -367,34 +364,77 @@ std::expected<void, InlineHook::Error> InlineHook::ff_hook(const std::shared_ptr
  return std::unexpected{result.error()};
  }
 
+ m_type = Type::FF;
+
+ return {};
+}
+#endif
+
+std::expected<void, InlineHook::Error> InlineHook::enable() {
+ std::scoped_lock lock{m_mutex};
+
+ if (m_enabled) {
+ return {};
+ }
+
  std::optional<Error> error;
 
  // jmp from original to trampoline.
  trap_threads(m_target, m_trampoline.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();
+ if (m_type == Type::E9) {
+ auto trampoline_epilogue = reinterpret_cast<TrampolineEpilogueE9*>(
+ m_trampoline.address() + m_trampoline_size - sizeof(TrampolineEpilogueE9));
+
+ 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();
+ }
  }
+
+#if SAFETYHOOK_ARCH_X86_64
+ if (m_type == Type::FF) {
+ if (auto result = emit_jmp_ff(m_target, m_destination, m_target + sizeof(JmpFF), m_original_bytes.size());
+ !result) {
+ error = result.error();
+ }
+ }
+#endif
  });
 
  if (error) {
  return std::unexpected{*error};
  }
 
+ m_enabled = true;
+
  return {};
 }
-#endif
 
-void InlineHook::destroy() {
+std::expected<void, InlineHook::Error> InlineHook::disable() {
  std::scoped_lock lock{m_mutex};
 
- if (!m_trampoline) {
- return;
+ if (!m_enabled) {
+ return {};
  }
 
  trap_threads(m_trampoline.data(), m_target, m_original_bytes.size(),
  [this] { std::copy(m_original_bytes.begin(), m_original_bytes.end(), m_target); });
 
+ m_enabled = false;
+
+ return {};
+}
+
+void InlineHook::destroy() {
+ [[maybe_unused]] auto disable_result = disable();
+
+ std::scoped_lock lock{m_mutex};
+
+ if (!m_trampoline) {
+ return;
+ }
+
  m_trampoline.free();
 }
 } // namespace safetyhook