Test/Thread trapping (#63)

Co-authored-by: bottiger1 <55270538+bottiger1@users.noreply.github.com>

CMakeLists.txt

@@ -190,7 +190,7 @@ if(SAFETYHOOK_AMALGAMATE) # amalgamate
 	add_custom_command(
 	    OUTPUT ${AMALGAMATED_FILE} ${AMALGAMATED_HEADER}
 	    DEPENDS ${HEADER_FILES} ${SOURCE_FILES} ${AMALGAMATE_SCRIPT}
-	    COMMAND ${Python3_EXECUTABLE} ${AMALGAMATE_SCRIPT} ${AMALGAMATED_FILE} ${AMALGAMATED_HEADER}
+	    COMMAND ${Python3_EXECUTABLE} ${AMALGAMATE_SCRIPT}
 	    MAIN_DEPENDENCY ${AMALGAMATE_SCRIPT}
 	    COMMENT "Amalgamating"
 	)

amalgamate.py

@@ -4,9 +4,12 @@
 from typing import List, Set
 from glob import glob
 from shutil import rmtree
+from textwrap import dedent
 
 import os
 import re
+import sys
+import argparse
 
 SAFETYHOOK_ROOT = Path(__file__).resolve().parent
 PUBLIC_INCLUDE_PATHS = [
@@ -18,6 +21,10 @@
 OUTPUT_DIR = SAFETYHOOK_ROOT / 'amalgamated-dist'
 FILE_HEADER = ['// DO NOT EDIT. This file is auto-generated by `amalgamate.py`.', '']
 
+parser = argparse.ArgumentParser(description='bundles cpp and hpp files together')
+parser.add_argument('--polyfill', action='store_true',
+                    help='replace std::except with a polyfill so it can be compiled on C++20 or older. https://raw.githubusercontent.com/TartanLlama/expected/master/include/tl/expected.hpp')
+
 
 # Python versions before 3.10 don't have the root_dir argument for glob, so we 
 # crudely emulate it here.
@@ -155,7 +162,25 @@ def merge_sources(*, source_dir: Path, covered_headers: Set[Path]):
     return output
 
 
+def do_polyfill(content):
+    return content.replace('#include <expected>',
+                           dedent('''
+                           #if __has_include("tl/expected.hpp")
+                               #include "tl/expected.hpp"
+                           #elif __has_include("expected.hpp")
+                               #include "expected.hpp"
+                           #else
+                               #error "No <expected> polyfill found"
+                           #endif
+                           ''')) \
+        .replace('std::expected', 'tl::expected') \
+        .replace('std::unexpected', 'tl::unexpected')
+
+
 def main():
+    args = parser.parse_args()
+    polyfill = args.polyfill is True
+
     if OUTPUT_DIR.exists():
         print('Output directory exists. Deleting.')
         rmtree(OUTPUT_DIR)
@@ -164,20 +189,26 @@ def main():
 
     covered_headers = set()
     with open(OUTPUT_DIR / 'safetyhook.hpp', 'w') as f:
-        f.write('\n'.join(FILE_HEADER + merge_headers(
+        content = '\n'.join(FILE_HEADER + merge_headers(
             header='safetyhook.hpp',
             search_paths=PUBLIC_INCLUDE_PATHS,
             covered_headers=covered_headers,
             stack=[],
-        )))
+        ))
+        if polyfill:
+            content = do_polyfill(content)
+        f.write(content)
 
     print(covered_headers)
 
     with open(OUTPUT_DIR / 'safetyhook.cpp', 'w') as f:
-        f.write('\n'.join(FILE_HEADER + merge_sources(
+        content = '\n'.join(FILE_HEADER + merge_sources(
             source_dir=SAFETYHOOK_ROOT / 'src',
             covered_headers=covered_headers,
-        )))
+        ))
+        if polyfill:
+            content = do_polyfill(content)
+        f.write(content)
 
 
 if __name__ == '__main__':

cmake.toml

@@ -97,7 +97,7 @@ set(AMALGAMATE_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/amalgamate.py)
 add_custom_command(
     OUTPUT ${AMALGAMATED_FILE} ${AMALGAMATED_HEADER}
     DEPENDS ${HEADER_FILES} ${SOURCE_FILES} ${AMALGAMATE_SCRIPT}
-    COMMAND ${Python3_EXECUTABLE} ${AMALGAMATE_SCRIPT} ${AMALGAMATED_FILE} ${AMALGAMATED_HEADER}
+    COMMAND ${Python3_EXECUTABLE} ${AMALGAMATE_SCRIPT}
     MAIN_DEPENDENCY ${AMALGAMATE_SCRIPT}
     COMMENT "Amalgamating"
 )

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{};

include/safetyhook/os.hpp

@@ -63,19 +63,9 @@ struct SystemInfo {
 
 SystemInfo system_info();
 
-using ThreadId = uint32_t;
-using ThreadHandle = void*;
 using ThreadContext = void*;
 
-/// @brief Executes a function while all other threads are frozen. Also allows for visiting each frozen thread and
-/// modifying it's context.
-/// @param run_fn The function to run while all other threads are frozen.
-/// @param visit_fn The function that will be called for each frozen thread.
-/// @note The visit function will be called in the order that the threads were frozen.
-/// @note The visit function will be called before the run function.
-/// @note Keep the logic inside run_fn and visit_fn as simple as possible to avoid deadlocks.
-void execute_while_frozen(const std::function<void()>& run_fn,
-    const std::function<void(ThreadId, ThreadHandle, ThreadContext)>& visit_fn = {});
+void trap_threads(uint8_t* from, uint8_t* to, size_t len, const std::function<void()>& run_fn);
 
 /// @brief Will modify the context of a thread's IP to point to a new address if its IP is at the old address.
 /// @param ctx The thread context to modify.

src/allocator.cpp

@@ -5,6 +5,8 @@
 #include "safetyhook/os.hpp"
 #include "safetyhook/utility.hpp"
 
+#include "safetyhook/utility.hpp"
+
 #include "safetyhook/allocator.hpp"
 
 namespace safetyhook {

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 {};