diff --git a/include/mach_detours.h b/include/mach_detours.h index 002e880..382a022 100644 --- a/include/mach_detours.h +++ b/include/mach_detours.h @@ -11,6 +11,7 @@ typedef void* detour_func_t; #define detour_err_in_progress (err_local | 1) +#define detour_err_wrong_thread (err_local | 2) mach_error_t detour_transaction_begin(); mach_error_t detour_transaction_abort(); diff --git a/src/detours_internal.h b/src/detours_internal.h index e3a734b..5d204e5 100644 --- a/src/detours_internal.h +++ b/src/detours_internal.h @@ -9,11 +9,19 @@ #if DETOUR_DEBUG #define DETOUR_TRACE(x) printf x #define DETOUR_BREAK() raise(SIGTRAP) +#define DETOUR_CHECK(x) \ + { \ + kern_return_t check_error = (x); \ + if (check_error != err_none) { \ + DETOUR_BREAK(); \ + } \ + } #else #include #include #define DETOUR_TRACE(x) #define DETOUR_BREAK() +#define DETOUR_CHECK(x) (x) #endif #endif diff --git a/src/mach_detours.c b/src/mach_detours.c index 840deb1..d3f444e 100644 --- a/src/mach_detours.c +++ b/src/mach_detours.c @@ -8,8 +8,9 @@ #include #include +#include -#include +#include #include //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -53,6 +54,15 @@ static void internal_detour_runnable_trampoline_regions() } } +static void internal_detour_free_trampoline(detour_trampoline* trampoline) +{ + detour_region* region = (detour_region*)((uintptr_t)trampoline & ~(uintptr_t)0xffff); + + memset(trampoline, 0, sizeof(*trampoline)); + trampoline->ptr_remain = (uint8_t*)region->free_list_head; + region->free_list_head = trampoline; +} + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // Transactions @@ -102,3 +112,45 @@ mach_error_t detour_transaction_begin() return s_pending_error; } + +mach_error_t detour_transaction_abort() +{ + if (s_transaction_thread != mach_thread_self()) { + return detour_err_wrong_thread; + } + + // Restore all the page permissions. + const mach_port_t port = mach_task_self(); + for (detour_operation* operation = s_pending_operations_head; operation != nullptr;) { + DETOUR_CHECK(mach_vm_protect(port, (mach_vm_address_t)operation->target, operation->trampoline->restore_code_size, false, operation->perm)); + + if (operation->kind == detour_operation_kind_attach) { + if (operation->trampoline) { + internal_detour_free_trampoline(operation->trampoline); + operation->trampoline = nullptr; + } + } + + detour_operation* next = operation->next; + free(operation); + operation = next; + } + s_pending_operations_head = nullptr; + + // Make sure the trampoline pages are no longer writable. + internal_detour_runnable_trampoline_regions(); + + // Resume any suspended threads. + for (detour_pending_thread* thread = s_pending_threads_head; thread != nullptr;) { + // There is nothing we can do if this fails. + DETOUR_CHECK(thread_resume(thread->thread)); + + detour_pending_thread* next = thread->next; + free(thread); + thread = next; + } + s_pending_threads_head = nullptr; + s_transaction_thread = THREAD_NULL; + + return err_none; +}