thread tests

include/mach_detours.h

@@ -20,11 +20,13 @@ typedef void* detour_func_t;
 #define	detour_err_too_large	(err_local | 4)
 
 mach_error_t detour_transaction_begin();
+mach_error_t detour_transaction_begin_managed();
 mach_error_t detour_transaction_abort();
 mach_error_t detour_transaction_commit();
 mach_error_t detour_transaction_commit_ex(detour_func_t** out_failed_target);
 
 mach_error_t detour_manage_thread(thread_t thread);
+mach_error_t detour_manage_all_threads();
 
 mach_error_t detour_attach(detour_func_t* inout_pointer, detour_func_t detour);
 mach_error_t detour_attach_ex(detour_func_t* inout_pointer, detour_func_t detour, detour_func_t* out_real_trampoline, detour_func_t* out_real_target, detour_func_t* out_real_detour);

src/mach_detours.c

@@ -55,7 +55,7 @@ static mach_error_t internal_detour_writable_trampoline_regions()
             DETOUR_BREAK();
             result = error;
         }
-        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d)\n", (void*)pRegion, DETOUR_REGION_SIZE, VM_PROT_READ | VM_PROT_WRITE));
+        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d) - internal_detour_writable_trampoline_regions\n", (void*)pRegion, DETOUR_REGION_SIZE, VM_PROT_READ | VM_PROT_WRITE));
     }
     return result;
 }
@@ -70,7 +70,7 @@ static void internal_detour_runnable_trampoline_regions()
         if (error != err_none) {
             DETOUR_BREAK();
         }
-        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d)\n", (void*)pRegion, DETOUR_REGION_SIZE, VM_PROT_READ | VM_PROT_EXECUTE));
+        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d) - internal_detour_runnable_trampoline_regions\n", (void*)pRegion, DETOUR_REGION_SIZE, VM_PROT_READ | VM_PROT_EXECUTE));
     }
 }
 
@@ -323,7 +323,7 @@ static detour_pending_thread* s_pending_threads_head = nullptr;
 static mach_error_t s_pending_error = err_none;
 static detour_func_t* s_pending_error_pointer = nullptr;
 
-mach_error_t detour_transaction_begin()
+static mach_error_t internal_detour_transaction_begin_managed_impl(bool manage_all_threads)
 {
     // Make sure only one thread can start a transaction.
     thread_t expected = THREAD_NULL;
@@ -336,12 +336,28 @@ mach_error_t detour_transaction_begin()
     s_pending_threads_head = nullptr;
     s_pending_error_pointer = nullptr;
 
+    if (manage_all_threads) {
+        s_pending_error = detour_manage_all_threads();
+        if (s_pending_error != err_none) {
+            return s_pending_error;
+        }
+    }
+
     // Make sure the trampoline pages are writable.
     s_pending_error = internal_detour_writable_trampoline_regions();
-
     return s_pending_error;
 }
 
+mach_error_t detour_transaction_begin()
+{
+    return internal_detour_transaction_begin_managed_impl(false);
+}
+
+mach_error_t detour_transaction_begin_managed()
+{
+    return internal_detour_transaction_begin_managed_impl(true);
+}
+
 mach_error_t detour_transaction_abort()
 {
     if (s_transaction_thread == THREAD_NULL) {
@@ -357,7 +373,7 @@ mach_error_t detour_transaction_abort()
         DETOUR_CHECK(
             mach_vm_protect(port, (mach_vm_address_t)operation->target, operation->trampoline->restore_code_size, false,
                 operation->perm));
-        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu8 ", %d)\n", (void*)operation->target, operation->trampoline->restore_code_size, operation->perm));
+        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu8 ", %d) - detour_transaction_abort\n", (void*)operation->target, operation->trampoline->restore_code_size, operation->perm));
 
         if (operation->kind == detour_operation_kind_attach) {
             if (operation->trampoline) {
@@ -476,7 +492,7 @@ mach_error_t detour_transaction_commit_ex(detour_func_t** out_failed_target)
         DETOUR_CHECK(
             mach_vm_protect(port, (mach_vm_address_t)operation->target, operation->trampoline->restore_code_size, false,
                 operation->perm));
-        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu8 ", %d)\n", (void*)operation->target, operation->trampoline->restore_code_size, operation->perm));
+        DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu8 ", %d) - detour_transaction_commit_ex\n", (void*)operation->target, operation->trampoline->restore_code_size, operation->perm));
 
         if (operation->kind == detour_operation_kind_detach && operation->trampoline) {
             internal_detour_free_trampoline(operation->trampoline);
@@ -561,6 +577,23 @@ mach_error_t detour_manage_thread(thread_t thread)
     return err_none;
 }
 
+mach_error_t detour_manage_all_threads()
+{
+    const mach_port_t port = mach_task_self();
+    thread_act_array_t threads = nullptr;
+    mach_msg_type_number_t numThreads = 0;
+    const kern_return_t result = task_threads(port, &threads, &numThreads);
+    if (result != err_none) {
+        return result;
+    }
+    for (mach_msg_type_number_t i = 0; i < numThreads; i++) {
+        DETOUR_CHECK( detour_manage_thread(threads[i]) );
+    }
+
+    DETOUR_CHECK( vm_deallocate(port, (vm_address_t)threads, numThreads * sizeof(*threads)) );
+    return err_none;
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 // Attach/Detach
 
@@ -794,7 +827,7 @@ mach_error_t detour_attach_ex(detour_func_t* inout_pointer, detour_func_t detour
         DETOUR_BREAK();
         goto fail;
     }
-    DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d)\n", (void*)target, target_override_len, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY));
+    DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d) - detour_attach_ex\n", (void*)target, target_override_len, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY));
 
     DETOUR_TRACE(("detours: target=%p: "
         "%02x %02x %02x %02x "
@@ -916,7 +949,7 @@ mach_error_t detour_detach(detour_func_t* inout_pointer, detour_func_t detour)
         DETOUR_BREAK();
         goto fail;
     }
-    DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d)\n", (void*)target, restore_code_size, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY));
+    DETOUR_TRACE(("mach_vm_protect(%p, %" PRIu32 ", %d) - detour_detach\n", (void*)target, restore_code_size, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY));
 
     op->kind = detour_operation_kind_detach;
     op->pointer = (uint8_t**)inout_pointer;

tests/CMakeLists.txt

@@ -4,6 +4,7 @@ add_executable(mach_detours_tests
         test_dylib_function.cpp
         test_local_function.cpp
         test_system_function.cpp
+        test_threads.cpp
         test_transaction.cpp
 )
 

tests/test_threads.cpp

@@ -0,0 +1,123 @@
+// Copyright (c) Lysann Tranvouez.  All rights reserved.
+
+#include <catch2/catch_test_macros.hpp>
+
+#include <mach_detours.h>
+#include <pthread.h>
+#include <unistd.h>
+
+static const char* callee(int a, int b, int c)
+{
+    return "localFunction";
+}
+static int g_b = 6;
+static int g_c = 3;
+
+__attribute__((target("branch-protection=pac-ret+bti")))
+static const char* localFunction()
+{
+    // This is just some "complicated" code to ensure localFunction gets the branch protection instructions (needs a callee)
+    int a = 53;
+    int b = g_b++;
+    auto ret = callee(a, b, 17);
+    g_c = g_b;
+    return ret;
+}
+static const char* localFunctionDetour()
+{
+    return "localFunctionDetour";
+}
+static const char* (*realLocalFunction)() = localFunction;
+
+static void* threadFunc(void* pUserArg)
+{
+    const bool* pCancel = static_cast<const bool*>(pUserArg);
+    while (!*pCancel) {
+        localFunction();
+        localFunction();
+        localFunction();
+        localFunction();
+        localFunction();
+    }
+    return nullptr;
+}
+
+TEST_CASE( "Handling other threads", "[attach][local][threads]" )
+{
+    bool cancelThreads = false;
+
+    constexpr auto numThreads = 150;
+    pthread_t threads[numThreads];
+    for (auto i = 0; i < numThreads; i++) {
+        pthread_create(&threads[i], nullptr, threadFunc, (void*)&cancelThreads);
+    }
+
+    SECTION( "running threads modify a value" )
+    {
+        const int saved_b = g_b;
+        usleep( 100 ); // let the threads update g_b a bit
+        CHECK( saved_b != g_b );
+    }
+
+    SECTION( "registered threads are suspended during transaction and resumed afterwards" )
+    {
+        SECTION( "manual thread managing" )
+        {
+            CHECK( detour_transaction_begin() == err_none );
+            for (auto&& thread : threads) {
+                CHECK( detour_manage_thread( pthread_mach_thread_np(thread) ) == err_none );
+            }
+        }
+        SECTION( "manage_all_threads" )
+        {
+            CHECK( detour_transaction_begin() == err_none );
+            CHECK( detour_manage_all_threads() == err_none );
+        }
+        SECTION( "transaction_begin_managed" )
+        {
+            CHECK( detour_transaction_begin_managed() == err_none );
+        }
+
+        int saved_b = g_b;
+        usleep( 300 ); // let any running thread update g_b a bit
+        CHECK( saved_b == g_b );
+
+        CHECK( detour_transaction_commit() == err_none );
+
+        saved_b = g_b;
+        usleep( 100 ); // let the threads update g_b a bit
+        CHECK( saved_b != g_b );
+    }
+
+    SECTION( "aborting a transaction resumes too" )
+    {
+        CHECK( detour_transaction_begin_managed() == err_none );
+        CHECK( detour_transaction_abort() == err_none );
+
+        const int saved_b = g_b;
+        usleep( 100 ); // let the threads update g_b a bit
+        CHECK( saved_b != g_b );
+    }
+
+    SECTION( "thread's PC gets redirected on commit, if it was in an overridden area" )
+    {
+        // Note that this test case doesn't guarantee that any thread is in the section of code that gets moved into the
+        // trampoline. So if buggy it could be flaky. You can try using more threads to get a more reliable repro.
+
+        CHECK( detour_transaction_begin_managed() == err_none );
+        CHECK( detour_attach_and_commit(reinterpret_cast<detour_func_t*>(&realLocalFunction), reinterpret_cast<detour_func_t>(localFunctionDetour)) == err_none );
+
+        const int saved_b = g_b;
+        usleep( 300 ); // let the threads update g_b a bit
+        CHECK( saved_b != g_b );
+
+        // clean up
+        CHECK( detour_transaction_begin_managed() == err_none );
+        CHECK( detour_detach_and_commit(reinterpret_cast<detour_func_t*>(&realLocalFunction), reinterpret_cast<detour_func_t>(localFunctionDetour)) == err_none );
+    }
+
+    cancelThreads = true;
+    for (auto i = 0; i < numThreads; i++) {
+        pthread_join(threads[i], nullptr);
+    }
+}