stopping other threads.

6 files changed, 234 insertions, 23 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
+
diff --git a/lib/hook-functions.c b/lib/hook-functions.c
new file mode 100644
index 0000000..975520e
--- /dev/null
+++ b/lib/hook-functions.c
@@ -0,0 +1,32 @@
+#include "substitute.h"
+#include "substitute-internal.h"
+
+static uintptr_t patch_callback(UNUSED void *ctx, uintptr_t pc) {
+    printf("patch_callback: pc=%llx\n", (long long) pc);
+    return pc;
+}
+
+EXPORT
+int substitute_hook_functions(const struct substitute_function_hook *hooks,
+                              size_t nhooks,
+                              int options) {
+    (void) hooks; (void) nhooks;
+    int ret = SUBSTITUTE_OK;
+    void *stop_token;
+    if (!(options & SUBSTITUTE_DONT_STOP_THREADS)) {
+        if ((ret = stop_other_threads(&stop_token)))
+            return ret;
+    }
+    if (!(options & SUBSTITUTE_DONT_STOP_THREADS)) {
+        if ((ret = apply_pc_patch_callback(stop_token, patch_callback, NULL)))
+            goto fail;
+    }
+
+fail:
+    if (!(options & SUBSTITUTE_DONT_STOP_THREADS)) {
+        int r2 = resume_other_threads(stop_token);
+        if (!ret)
+            ret = r2;
+    }
+    return ret;
+}
diff --git a/lib/substitute-internal.h b/lib/substitute-internal.h
index f078688..80052f9 100644
--- a/lib/substitute-internal.h
+++ b/lib/substitute-internal.h
@@ -64,9 +64,15 @@ enum {
      * is to be root */
     SUBSTITUTE_ERR_TASK_FOR_PID = 1000,
 
-    /* substitute_dlopen_in_pid: something didn't work */
     SUBSTITUTE_ERR_MISC,
 };
 
 int substitute_dlopen_in_pid(int pid, const char *filename, int options, char **error);
 #endif
+
+/* Stop the world; return token to be used for applying PC patches and resuming. */
+int stop_other_threads(void **token_ptr);
+int apply_pc_patch_callback(void *token,
+                            uintptr_t (*pc_patch_callback)(void *ctx, uintptr_t pc),
+                            void *ctx);
+int resume_other_threads(void *token);
diff --git a/lib/substitute.h b/lib/substitute.h
index 3adb39a..ada57f2 100644
--- a/lib/substitute.h
+++ b/lib/substitute.h
@@ -44,6 +44,12 @@ enum {
      * preventing pages from being marked executable. */
     SUBSTITUTE_ERR_VM,
 
+    /* substitute_hook_functions: not on the main thread (so stopping all other
+     * threads would be unsafe, as concurrent attempts to do the same from
+     * other threads would result in deadlock), and you did not pass
+     * SUBSTITUTE_DONT_STOP_THREADS */
+    SUBSTITUTE_ERR_NOT_ON_MAIN_THREAD,
+
     /* substitute_interpose_imports: couldn't redo relocation for an import
      * because the type was unknown */
     SUBSTITUTE_ERR_UNKNOWN_RELOCATION_TYPE,
@@ -51,6 +57,9 @@ enum {
     /* substitute_hook_objc_message: no such selector existed in the class's
      * inheritance tree */
     SUBSTITUTE_ERR_NO_SUCH_SELECTOR,
+
+    /* substitute_hook_functions: OS error suspending other threads */
+    SUBSTITUTE_ERR_ADJUSTING_THREADS,
 };
 
 struct substitute_function_hook {
@@ -62,6 +71,11 @@ struct substitute_function_hook {
 /* Get a string representation for a SUBSTITUTE_* error code. */
 const char *substitute_strerror(int err);
 
+/* substitute_hook_functions options */
+enum {
+    SUBSTITUTE_DONT_STOP_THREADS = 1,
+};
+
 /* TODO doc */
 int substitute_hook_functions(const struct substitute_function_hook *hooks,
                               size_t nhooks,