stopping other threads.

3 files changed, 181 insertions, 22 deletions

diff --git a/lib/darwin/inject.c b/lib/darwin/inject.c
index 252c6bf..a515ada 100644
--- a/lib/darwin/inject.c
+++ b/lib/darwin/inject.c
@@ -2,6 +2,7 @@
 #include "substitute.h"
 #include "substitute-internal.h"
 #include "darwin/read.h"
+#include "darwin/thread-state.h"
 #include <mach/mach.h>
 #include <mach-o/dyld_images.h>
 #include <dlfcn.h>
@@ -379,23 +380,6 @@ got_symbol:;
     return true;
 }
 
-struct _x86_thread_state_32 {
-    uint32_t eax, ebx, ecx, edx, edi, esi, ebp, esp;
-    uint32_t ss, eflags, eip, cs, ds, es, fs, gs;
-};
-struct _x86_thread_state_64 {
-    uint64_t rax, rbx, rcx, rdx, rdi, rsi, rbp, rsp;
-    uint64_t r8, r9, r10, r11, r12, r13, r14, r15;
-    uint64_t rip, rflags, cs, fs, gs;
-};
-struct _arm_thread_state_32 {
-    uint32_t r[13], sp, lr, pc, cpsr;
-};
-struct _arm_thread_state_64 {
-    uint64_t x[29], fp, lr, sp, pc;
-    uint32_t cpsr, pad;
-};
-
 int substitute_dlopen_in_pid(int pid, const char *filename, int options, char **error) {
     mach_port_t task;
     mach_vm_address_t target_stack = 0;
@@ -457,7 +441,7 @@ int substitute_dlopen_in_pid(int pid, const char *filename, int options, char **
         pthread_create_addr = libs[1].symaddr;
     }
 
-    __attribute__((unused))
+    UNUSED
     extern char inject_page_start[],
                 inject_start_x86_64[],
                 inject_start_i386[],
@@ -542,14 +526,14 @@ int substitute_dlopen_in_pid(int pid, const char *filename, int options, char **
         u.x64.rdi = target_stack_top;
         u.x64.rip = target_code_page + (inject_start_x86_64 - inject_page_start);
         state_size = sizeof(u.x64);
-        flavor = 4;
+        flavor = _x86_thread_state_64_flavor;
         break;
     case CPU_TYPE_I386:
         u.x32.esp = target_stack_top;
         u.x32.ecx = target_stack_top;
         u.x32.eip = target_code_page + (inject_start_i386 - inject_page_start);
         state_size = sizeof(u.x32);
-        flavor = 1;
+        flavor = _x86_thread_state_32_flavor;
         break;
 #endif
 #if defined(__arm__) || defined(__arm64__)
@@ -558,14 +542,14 @@ int substitute_dlopen_in_pid(int pid, const char *filename, int options, char **
         u.a32.r[0] = target_stack_top;
         u.a32.pc = target_code_page + (inject_start_arm - inject_page_start);
         state_size = sizeof(u.a32);
-        flavor = 9;
+        flavor = _arm_thread_state_32_flavor;
         break;
     case CPU_TYPE_ARM64:
         u.a64.sp = target_stack_top;
         u.a64.x[0] = target_stack_top;
         u.a64.pc = target_code_page + (inject_start_arm64 - inject_page_start);
         state_size = sizeof(u.a64);
-        flavor = 6;
+        flavor = _arm_thread_state_64_flavor;
         break;
 #endif
     default:
diff --git a/lib/darwin/stop-other-threads.c b/lib/darwin/stop-other-threads.c
new file mode 100644
index 0000000..dd00a2a
--- /dev/null
+++ b/lib/darwin/stop-other-threads.c
@@ -0,0 +1,151 @@
+#include "substitute.h"
+#include "substitute-internal.h"
+#include "darwin/thread-state.h"
+#include <pthread.h>
+#include <mach/mach.h>
+#include <CoreFoundation/CoreFoundation.h>
+
+static void release_port(UNUSED CFAllocatorRef allocator, const void *value) {
+    mach_port_t thread = (mach_port_t) value;
+    thread_resume(thread);
+    mach_port_deallocate(mach_task_self(), thread);
+}
+static CFSetCallBacks suspend_port_callbacks = {
+    .version = 0,
+    .release = release_port,
+};
+
+static bool apply_one_pcp(mach_port_t thread,
+                          uintptr_t (*callback)(void *ctx, uintptr_t pc),
+                          void *ctx) {
+    int flavor;
+#if defined(__x86_64__)
+    struct _x86_thread_state_64 state;
+    flavor = _x86_thread_state_64_flavor;
+#elif defined(__i386__)
+    struct _x86_thread_state_32 state;
+    flavor = _x86_thread_state_32_flavor;
+#elif defined(__arm__)
+    struct _arm_thread_state_32 state;
+    flavor = _arm_thread_state_32_flavor;
+#elif defined(__arm64__)
+    struct _arm_thread_state_64 state;
+    flavor = _arm_thread_state_64_flavor;
+#else
+    #error ?
+#endif
+
+    mach_msg_type_number_t real_cnt = sizeof(state) / sizeof(int);
+    mach_msg_type_number_t cnt = real_cnt;
+    kern_return_t kr = thread_get_state(thread, flavor, (thread_state_t) &state, &cnt);
+    if (kr || cnt != real_cnt)
+        return false;
+
+    uintptr_t *pcp;
+#if defined(__x86_64__)
+    pcp = (uintptr_t *) &state.rip;
+#elif defined(__i386__)
+    pcp = (uintptr_t *) &state.eip;
+#elif defined(__arm__) || defined(__arm64__)
+    pcp = (uintptr_t *) &state.pc;
+#endif
+    uintptr_t new = callback(ctx, *pcp);
+    if (new != *pcp) {
+        *pcp = new;
+        kr = thread_set_state(thread, flavor, (thread_state_t) &state, real_cnt);
+        if (kr)
+            return false;
+    }
+    return true;
+}
+
+int apply_pc_patch_callback(void *token,
+                            uintptr_t (*pc_patch_callback)(void *ctx, uintptr_t pc),
+                            void *ctx) {
+    CFMutableSetRef suspended_set = token;
+    CFIndex count = CFSetGetCount(suspended_set);
+    if (!count)
+        return SUBSTITUTE_OK;
+    /* great API there CF */
+    const void **ports = malloc(sizeof(*ports) * count);
+    CFSetGetValues(suspended_set, ports);
+    int ret = SUBSTITUTE_OK;
+    for (CFIndex i = 0; i < count; i++) {
+        if (!apply_one_pcp((mach_port_t) ports[i], pc_patch_callback, ctx)) {
+            ret = SUBSTITUTE_ERR_ADJUSTING_THREADS;
+            break;
+        }
+    }
+    free(ports);
+    return ret;
+}
+
+int stop_other_threads(void **token_ptr) {
+    if (!pthread_main_np())
+        return SUBSTITUTE_ERR_NOT_ON_MAIN_THREAD;
+
+    int ret;
+    mach_port_t self = mach_thread_self();
+
+    /* The following shenanigans are for catching any new threads that are
+     * created while we're looping, without suspending anything twice.  Keep
+     * looping until only threads we already suspended before this loop are
+     * there. */
+    CFMutableSetRef suspended_set = CFSetCreateMutable(NULL, 0, &suspend_port_callbacks);
+
+    thread_act_array_t ports = 0;
+    mach_msg_type_number_t nports = 0;
+
+    bool got_new = true;
+    while (got_new) {
+        got_new = false;
+
+        kern_return_t kr = task_threads(mach_task_self(), &ports, &nports);
+        if (kr) { /* ouch */
+            ret = SUBSTITUTE_ERR_ADJUSTING_THREADS;
+            goto fail;
+        }
+
+        for (mach_msg_type_number_t i = 0; i < nports; i++) {
+            mach_port_t port = ports[i];
+            void *casted_port = (void *) (uintptr_t) port;
+            if (port == self ||
+                CFSetContainsValue(suspended_set, casted_port)) {
+                /* already suspended, ignore */
+                mach_port_deallocate(mach_task_self(), port);
+            } else {
+                got_new = true;
+                printf("suspending %d (self=%d)\n", port, self);
+                kr = thread_suspend(port);
+                if (kr == KERN_TERMINATED) {
+                    /* too late */
+                    mach_port_deallocate(mach_task_self(), port);
+                } else if (kr) {
+                    ret = SUBSTITUTE_ERR_ADJUSTING_THREADS;
+                    for (; i < nports; i++)
+                        mach_port_deallocate(mach_task_self(), ports[i]);
+                    vm_deallocate(mach_task_self(), (vm_address_t) ports,
+                                  nports * sizeof(*ports));
+                    goto fail;
+                }
+                CFSetAddValue(suspended_set, casted_port);
+            }
+        }
+        vm_deallocate(mach_task_self(), (vm_address_t) ports,
+                      nports * sizeof(*ports));
+    }
+
+    /* Success - keep the set around for when we're done. */
+    *token_ptr = suspended_set;
+    return SUBSTITUTE_OK;
+
+fail:
+    CFRelease(suspended_set);
+    return ret;
+}
+
+int resume_other_threads(void *token) {
+    CFMutableSetRef suspended_set = token;
+    CFRelease(suspended_set);
+    return SUBSTITUTE_OK; /* eh */
+}
diff --git a/lib/darwin/thread-state.h b/lib/darwin/thread-state.h
new file mode 100644
index 0000000..0bae7be
--- /dev/null
+++ b/lib/darwin/thread-state.h
@@ -0,0 +1,24 @@
+#pragma once
+#include <stdint.h>
+
+struct _x86_thread_state_32 {
+    uint32_t eax, ebx, ecx, edx, edi, esi, ebp, esp;
+    uint32_t ss, eflags, eip, cs, ds, es, fs, gs;
+};
+#define _x86_thread_state_32_flavor 1
+struct _x86_thread_state_64 {
+    uint64_t rax, rbx, rcx, rdx, rdi, rsi, rbp, rsp;
+    uint64_t r8, r9, r10, r11, r12, r13, r14, r15;
+    uint64_t rip, rflags, cs, fs, gs;
+};
+#define _x86_thread_state_64_flavor 4
+struct _arm_thread_state_32 {
+    uint32_t r[13], sp, lr, pc, cpsr;
+};
+#define _arm_thread_state_32_flavor 9
+struct _arm_thread_state_64 {
+    uint64_t x[29], fp, lr, sp, pc;
+    uint32_t cpsr, pad;
+};
+#define _arm_thread_state_64_flavor 6
+