diff --git a/include/mach_detours.h b/include/mach_detours.h index d875b9f..f551d61 100644 --- a/include/mach_detours.h +++ b/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); diff --git a/src/mach_detours.c b/src/mach_detours.c index 24ef035..ba601e1 100644 --- a/src/mach_detours.c +++ b/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; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 17c8a60..e572dbb 100644 --- a/tests/CMakeLists.txt +++ b/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 ) diff --git a/tests/test_threads.cpp b/tests/test_threads.cpp new file mode 100644 index 0000000..db50fed --- /dev/null +++ b/tests/test_threads.cpp @@ -0,0 +1,123 @@ +// Copyright (c) Lysann Tranvouez. All rights reserved. + +#include + +#include +#include +#include + +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(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(&realLocalFunction), reinterpret_cast(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(&realLocalFunction), reinterpret_cast(localFunctionDetour)) == err_none ); + } + + cancelThreads = true; + for (auto i = 0; i < numThreads; i++) { + pthread_join(threads[i], nullptr); + } +}