imp forwarding works

6 files changed, 283 insertions, 114 deletions

diff --git a/Makefile b/Makefile
index fdf108a..33c22e2 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ CXX := clang++
 CFLAGS := -O3 -Wall -Wextra -Werror -arch x86_64
 override CC := $(CC) $(CFLAGS)
 override CXX := $(CXX) $(CFLAGS) -fno-exceptions -fno-asynchronous-unwind-tables
-
+LIB_LDFLAGS := -lobjc -dynamiclib -fvisibility=hidden
 IMAON2 := /Users/comex/c/imaon2
 GEN_JS := node --harmony --harmony_arrow_functions $(IMAON2)/tables/gen.js
 
@@ -40,7 +40,7 @@ LIB_OBJS := \
 	out/substrate-compat.o \
 	out/jump-dis-arm-multi.o
 out/libsubstitute.dylib: $(LIB_OBJS)
-	$(CC) -dynamiclib -fvisibility=hidden -o $@ $(LIB_OBJS) -lobjc
+	$(CC) -o $@ $(LIB_OBJS) $(LIB_LDFLAGS)
 
 define define_test
 out/test-$(1): test/test-$(2).[cm]* $(HEADERS) $(GENERATED) Makefile out/libsubstitute.dylib
diff --git a/lib/objc-asm.S b/lib/objc-asm.S
index 94d9de9..7bfe159 100644
--- a/lib/objc-asm.S
+++ b/lib/objc-asm.S
@@ -1,122 +1,71 @@
-/* These all forward to the IMP obtained from block->get_imp(block); this
- * signature is just so that I don't have to have two different assembly
- * implementations for the temporary and fake-old blocks.  To make things
- * confusing, we don't know whether it will be called as stret or not, so we
- * have to guess.
- * (Could avoid the guessing by generating custom assembly instead, like
- * Substrate does, but meh.)
- */
-.globl _temp_block_invoke
-.section __TEXT,__text,regular,pure_instructions
-#if defined(__x86_64__)
-_temp_block_invoke:
-    push %rcx
-    mov 0x10(%rsi), %rax /* block if stret, else self */
-    lea _temp_block_invoke(%rip), %rcx
-    cmp %rax, %rcx
-    pop %rcx
-    je 2f
+#include "objc.h"
+.text
+.align _PAGE_SHIFT
+#ifdef __arm__
+.thumb_func _remap_start
+.thumb
+#endif
+.globl _remap_start
+_remap_start:
 
-    mov 0x20(%rdi), %rax /* get_imp */
-    push %rdi; push %rsi; push %rdx; push %rcx; push %r8; push %r9
-    call *%rax
-    pop %r9; pop %r8; pop %rcx; pop %rdx; pop %rdi; pop %rsi
-    mov 0x28(%rsi), %rsi /* selector */
-    jmp *%rax
-2: /* stret */
-    mov 0x20(%rsi), %rax /* get_imp */
+.set i, 0
+#define my_rpe (0b + (_PAGE_SIZE - i * TRAMPOLINE_SIZE + i * REMAP_PAGE_ENTRY_SIZE))
+.rept TRAMPOLINES_PER_PAGE
+0:
+#if defined(__x86_64__)
     push %rdi; push %rsi; push %rdx; push %rcx; push %r8; push %r9
-    mov %rsi, %rdi
-    call *%rax
-    pop %r9; pop %r8; pop %rcx; pop %rsi; pop %rdx; pop %rdi
-    mov 0x28(%rdx), %rdx /* selector */
+    lea my_rpe(%rip), %rdx
+    mov 8(%rdx), %rdi
+    mov 16(%rdx), %rsi
+    call *(%rdx)
+    pop %r9; pop %r8; pop %rcx; pop %rdx; pop %rsi; pop %rdi
     jmp *%rax
 #elif defined(__i386__)
-_temp_block_invoke:
     call 1f
 1:
     pop %edx
-    lea _temp_block_invoke-1b(%edx), %edx
-    mov 8(%esp), %ecx /* block if stret, else self */
-    mov 0xc(%ecx), %eax
-    cmp %eax, %edx
-    je 2f
-
-    mov 4(%esp), %eax /* block */
-    push %eax
-    mov 0x14(%eax), %eax /* get_imp */
-    call *%eax
-    add $4, %esp
-    mov 8(%esp), %edx /* self */
-    mov 4(%esp), %ecx /* block */
-    mov %edx, 4(%esp)
-    mov 0x18(%ecx), %ecx
-    mov %ecx, 8(%esp)
-    jmp *%eax
-2: /* stret */
-    int3
-    mov 8(%esp), %eax /* block */
-    push %eax
-    mov 0x14(%eax), %eax /* get_imp */
-    call *%eax
-    add $4, %esp
-    int3
-    mov 12(%esp), %ecx /* self */
-    mov %ecx, 8(%esp)
-    mov 8(%esp), %ecx /* block */
-    mov 0x18(%ecx), %ecx
-    mov %ecx, 12(%esp)
+    lea my_rpe-1b(%edx), %edx
+    push 8(%edx)
+    push 4(%edx)
+    call *(%edx)
+    add 8, %esp
     jmp *%eax
 #elif defined(__arm__)
-.thumb_func _temp_block_invoke
-.thumb
-_temp_block_invoke:
-1:
-    ldr r9, [r1, #0xc]
-    adr r12, 1b
-    cmp r9, r12
-    beq 2f
-
     push {r0-r4, lr} /* r4 for align */
-    ldr r9, [r0, #0x14]
-    blx r9
-    mov r12, r0
-    pop {r0-r4, lr}
+    mov r3, #(my_rpe - 1f)
+    add r3, pc
+1:
+    ldr r0, [r3, #4]
+    ldr r1, [r3, #8]
+    ldr r2, [r3]
+    blx r2
     mov r9, r0
-    mov r0, r1
-    ldr r1, [r9, #0x18]
-    bx r12
-2: /* stret */
-    push {r0-r3, lr}
-    ldr r9, [r1, #0x14]
-    blx r9
-    mov r12, r0
-    pop {r0-r3, lr}
-    mov r9, r1
-    mov r1, r2
-    ldr r2, [r9, #0x18]
-    bx r12
+    pop {r0-r4, lr}
+    bx r9
 #elif defined(__arm64__)
-.align 2
-_temp_block_invoke:
     stp x7, x6, [sp, #-0x10]!
     stp x5, x4, [sp, #-0x10]!
     stp x3, x2, [sp, #-0x10]!
     stp x1, x0, [sp, #-0x10]!
     str x30, [sp, #-0x10]!
 
-    ldr x9, [x0, #0x20]
-    blr x9
-    mov x10, x0
+    ldr x0, my_rpe+4
+    ldr x1, my_rpe+8
+    ldr x2, my_rpe
+    blr x2
+    mov x9, x0
 
     ldr x30, [sp], #0x10
-    ldp x0, x9, [sp], #0x10
+    ldp x1, x0, [sp], #0x10
     ldp x3, x2, [sp], #0x10
     ldp x5, x4, [sp], #0x10
     ldp x7, x6, [sp], #0x10
 
-    ldr x1, [x9, #0x28]
-    br x10
+    br x9
 #else
 #error No forwarding assembly definition for this arch
 #endif
+
+.set i, i + 1
+.endr
+
diff --git a/lib/objc.c b/lib/objc.c
new file mode 100644
index 0000000..98b84e6
--- /dev/null
+++ b/lib/objc.c
@@ -0,0 +1,167 @@
+#if defined(__APPLE__)
+#include "substitute.h"
+#include "substitute-internal.h"
+#include "objc.h"
+#include <stddef.h>
+#include <pthread.h>
+#include <mach/mach.h>
+#include <sys/mman.h>
+#include <sys/queue.h>
+#include <errno.h>
+
+/* These trampolines will call e->func(e->arg1, e->arg2), and jump to there,
+ * preserving all arguments.  imp_implementationWithBlock would be easier and
+ * maybe a bit faster, but it's impossible to avoid throwing away registers
+ * without having to ask whether the selector is stret or not. */
+
+struct tramp_info_page_header {
+    struct tramp_info_page_entry *first_free;
+    size_t nfree;
+    LIST_ENTRY(tramp_info_page_header) free_pages;
+};
+
+struct tramp_info_page_entry {
+    union {
+        struct tramp_info_page_entry *next_free;
+        void *func;
+    };
+    void *arg1;
+    void *arg2;
+};
+
+static pthread_mutex_t tramp_mutex = PTHREAD_MUTEX_INITIALIZER;
+LIST_HEAD(tramp_info_page_list, tramp_info_page_header)
+    tramp_free_page_list = LIST_HEAD_INITIALIZER(tramp_info_page_list);
+
+extern char remap_start[];
+
+static int get_trampoline(void *func, void *arg1, void *arg2, void *tramp) {
+    int ret, rerrno = 0;
+    pthread_mutex_lock(&tramp_mutex);
+
+    struct tramp_info_page_header *header = LIST_FIRST(&tramp_free_page_list);
+    if (!header) {
+        if (PAGE_SIZE > _PAGE_SIZE)
+            panic("%s: strange PAGE_SIZE %x\n", __func__, PAGE_SIZE);
+        void *new_pages = mmap(NULL, _PAGE_SIZE * 2, PROT_READ | PROT_WRITE,
+                               MAP_SHARED | MAP_ANON, -1, 0);
+        if (new_pages == MAP_FAILED) {
+            ret = SUBSTITUTE_ERR_OOM;
+            rerrno = errno;
+            goto out;
+        }
+        vm_address_t tramp_page = (vm_address_t) new_pages;
+        vm_prot_t cur_prot, max_prot;
+        kern_return_t kr = vm_remap(
+            mach_task_self(),
+            &tramp_page,
+            _PAGE_SIZE,
+            _PAGE_SIZE - 1,
+            VM_FLAGS_OVERWRITE | VM_FLAGS_FIXED,
+            mach_task_self(),
+            (vm_address_t) remap_start,
+            FALSE, /* copy */
+            &cur_prot,
+            &max_prot,
+            VM_INHERIT_NONE);
+        if (kr != KERN_SUCCESS || tramp_page != (vm_address_t) new_pages) {
+            ret = SUBSTITUTE_ERR_VM;
+            goto out;
+        }
+        header = new_pages + _PAGE_SIZE * 2 - sizeof(*header);
+        header->first_free = NULL;
+        header->nfree = TRAMPOLINES_PER_PAGE;
+        LIST_INSERT_HEAD(&tramp_free_page_list, header, free_pages);
+    }
+
+    void *page = (void *) (((uintptr_t) header) & ~(_PAGE_SIZE - 1));
+    struct tramp_info_page_entry *entries = page;
+    struct tramp_info_page_entry *entry = header->first_free;
+    if (entry == NULL) {
+        entry = &entries[TRAMPOLINES_PER_PAGE - header->nfree];
+        entry->next_free = NULL;
+    }
+
+    header->first_free = entry->next_free;
+    if (--header->nfree == 0)
+        LIST_REMOVE(header, free_pages);
+
+    entry->func = func;
+    entry->arg1 = arg1;
+    entry->arg2 = arg2;
+    *(void **) tramp = (page - PAGE_SIZE) + (entry - entries) * TRAMPOLINE_SIZE;
+    ret = SUBSTITUTE_OK;
+out:
+    pthread_mutex_unlock(&tramp_mutex);
+    errno = rerrno;
+    return ret;
+}
+
+static void free_trampoline(void *tramp) {
+    pthread_mutex_lock(&tramp_mutex);
+    void *page = (void *) (((uintptr_t) tramp) & ~(_PAGE_SIZE - 1));
+    size_t i = (tramp - page) / TRAMPOLINE_SIZE;
+    struct tramp_info_page_entry *entries = page + _PAGE_SIZE;
+    struct tramp_info_page_entry *entry = &entries[i];
+    struct tramp_info_page_header *header = page + 2 * _PAGE_SIZE - sizeof(*header);
+
+    entry->next_free = header->first_free;
+    header->first_free = entry;
+    header->nfree++;
+    if (header->nfree == 1)
+        LIST_INSERT_HEAD(&tramp_free_page_list, header, free_pages);
+    else if (header->nfree == TRAMPOLINES_PER_PAGE &&
+        /* have others? */
+        (LIST_FIRST(&tramp_free_page_list) != header ||
+         LIST_NEXT(header, free_pages))) {
+        /* free the trampoline and info pages */
+        LIST_REMOVE(header, free_pages);
+        munmap(page, 2 * _PAGE_SIZE);
+    }
+
+    pthread_mutex_unlock(&tramp_mutex);
+}
+
+static IMP dereference(IMP *old_ptr, UNUSED void *_) {
+    return *old_ptr;
+}
+
+int substitute_hook_objc_message(Class class, SEL selector, IMP replacement,
+                                 IMP *old_ptr, bool *created_imp_ptr) {
+    int ret;
+    Method meth = class_getClassMethod(class, selector);
+    if (meth == NULL)
+        return SUBSTITUTE_ERR_NO_SUCH_SELECTOR;
+    const char *types = method_getTypeEncoding(meth);
+
+    /* temporary trampoline just tries again */
+    IMP temp;
+    if ((ret = get_trampoline(dereference, old_ptr, NULL, &temp)))
+        return ret;
+    *old_ptr = temp;
+
+    IMP old = class_replaceMethod(class, selector, replacement, types);
+    if (old) {
+        *old_ptr = old;
+        *created_imp_ptr = false;
+    } else {
+        Class super = class_getSuperclass(class);
+        if (!super) {
+            /* this ought to only be possible if the method was removed in the
+             * meantime, since we found the method above and it couldn't have
+             * been found in a superclass, but the objc2 runtime doesn't allow
+             * removing methods. */
+            panic("%s: no superclass but the method didn't exist\n", __func__);
+        }
+        ret = get_trampoline(class_getMethodImplementation, super, selector, old_ptr);
+        *created_imp_ptr = true;
+    }
+
+    free_trampoline(temp);
+    return SUBSTITUTE_OK;
+}
+
+void substitute_free_created_imp(IMP imp) {
+    free_trampoline(imp);
+}
+#endif
diff --git a/lib/objc.h b/lib/objc.h
new file mode 100644
index 0000000..24427e1
--- /dev/null
+++ b/lib/objc.h
@@ -0,0 +1,23 @@
+#pragma once
+/* PAGE_SIZE is not actually a constant on iOS */
+#if defined(__arm64__)
+#define _PAGE_SHIFT 14
+#else
+#define _PAGE_SHIFT 12
+#endif
+#define _PAGE_SIZE (1 << _PAGE_SHIFT)
+#if defined(__x86_64__)
+#define TRAMPOLINE_SIZE 0x23
+#elif defined(__i386__)
+#define TRAMPOLINE_SIZE 0x1c
+#elif defined(__arm__)
+#define TRAMPOLINE_SIZE 0x18
+#elif defined(__arm64__)
+#define TRAMPOLINE_SIZE 0x40
+#endif
+#ifdef __LP64__
+#define REMAP_PAGE_ENTRY_SIZE 24
+#else
+#define REMAP_PAGE_ENTRY_SIZE 12
+#endif
+#define TRAMPOLINES_PER_PAGE (_PAGE_SIZE / TRAMPOLINE_SIZE)
diff --git a/lib/substitute.h b/lib/substitute.h
index 181fb7c..4dbfc7c 100644
--- a/lib/substitute.h
+++ b/lib/substitute.h
@@ -34,10 +34,14 @@ enum {
      * patch region at the beginning */
     SUBSTITUTE_ERR_FUNC_JUMPS_TO_START,
 
-    /* mmap or mprotect failure other than ENOMEM (preserved in errno on return
-     * from the substitute_* function).  Most likely to come up with
-     * substitute_hook_functions, if the kernel is preventing pages from being
-     * marked executable. */
+    /* out of memory */
+    SUBSTITUTE_ERR_OOM,
+
+    /* substitute_hook_functions:    mmap or mprotect failure other than ENOMEM
+     *                               (preserved in errno on return)
+     * substitute_hook_objc_message: vm_remap failure
+     * Most likely to come up with substitute_hook_functions if the kernel is
+     * preventing pages from being marked executable. */
     SUBSTITUTE_ERR_VM,
 
     /* substitute_interpose_imports: couldn't redo relocation for an import
@@ -171,7 +175,7 @@ int substitute_interpose_imports(const struct substitute_image *handle,
 
 #endif /* 1 */
 
-#if defined(__APPLE__) && __BLOCKS__
+#if defined(__APPLE__)
 #include <objc/runtime.h>
 /* Hook a method implementation for a given Objective-C class.  By itself, this
  * function is thread safe: it is simply a wrapper for the atomic Objective-C
@@ -186,9 +190,8 @@ int substitute_interpose_imports(const struct substitute_image *handle,
  * @replacement      the new implementation
  * @old_ptr          optional - out pointer to the 'old implementation'.
  *                   If there is no old implementation, a custom IMP is
- *                   returned that delegates to the superclass.  This IMP is
- *                   created with imp_implementationWithBlock, so it can be
- *                   freed if desired with imp_removeBlock.
+ *                   returned that delegates to the superclass.  This IMP can
+ *                   be freed if desired with imp_removeBlock.
  * @created_imp_ptr  optional - out pointer to whether a fake superclass-call
  *                   IMP has been placed in <old_ptr>
  *
@@ -197,6 +200,8 @@ int substitute_interpose_imports(const struct substitute_image *handle,
  */
 int substitute_hook_objc_message(Class klass, SEL selector, IMP replacement,
                                  IMP *old_ptr, bool *created_imp_ptr);
+
+void substitute_free_created_imp(IMP imp);
 #endif
 
 #ifdef __cplusplus
diff --git a/test/test-imp-forwarding.m b/test/test-imp-forwarding.m
index e7bff27..0beb24a 100644
--- a/test/test-imp-forwarding.m
+++ b/test/test-imp-forwarding.m
@@ -1,8 +1,14 @@
 #include "../lib/objc.c"
 #include <objc/runtime.h>
 #include <stdio.h>
+#include <assert.h>
 #import <Foundation/Foundation.h>
 
+static void *what_to_call(void *a, void *b) {
+    printf("what_to_call: %p %p\n", a, b);
+    return a;
+}
+
 static void imp1(id self, SEL sel, int a, int b) {
     NSLog(@"imp1: self=%@ sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
 }
@@ -13,7 +19,7 @@ struct big {
 
 static struct big imp2(id self, SEL sel, int a, int b) {
     NSLog(@"imp2: self=%@ sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
-    return (struct big) {{0}};
+    return (struct big) {{4}};
 }
 
 struct big (^test)(id, int) = ^(id self, int a) {
@@ -22,14 +28,33 @@ struct big (^test)(id, int) = ^(id self, int a) {
 };
 
 int main() {
-    IMP testi = imp_implementationWithBlock(test);
-    ((struct big (*)(id, SEL, int)) testi)(@"test", @selector(dumb), 5);
-    IMP old = (IMP) imp1;
     SEL sel = @selector(some);
-    struct temp_block_literal temp_block = get_temp_block(&old, sel);
-    IMP new = imp_implementationWithBlock((id) &temp_block);
-    ((void (*)(id, SEL, int, int)) new)(@"foo", sel, 1, 2);
-    old = (IMP) imp2;
-    struct big big = ((struct big (*)(id, SEL, int, int)) new)(@"bar", sel, 1, 2);
+    IMP imp;
+    assert(!get_trampoline(what_to_call, imp1, (void *) 0x123, &imp));
+    printf("imp = %p\n", imp);
+    ((void (*)(id, SEL, int, int)) imp)(@"foo", sel, 1, 2);
+    free_trampoline(imp);
+    assert(!get_trampoline(what_to_call, imp2, (void *) 0x123, &imp));
+    struct big big = ((struct big (*)(id, SEL, int, int)) imp)(@"bar", sel, 1, 2);
     printf("out? %d\n", big.x[0]);
+
+    /* test alloc/free */
+    enum { n = 10000 };
+    static void *imps[n];
+    for(size_t i = 0; i < n; i++) {
+        assert(!get_trampoline(what_to_call, imp1, (void *) 0x123, &imps[i]));
+    }
+    for(size_t i = 0; i < n; i++) {
+        assert(imps[i]);
+        if (arc4random() & 1) {
+            free_trampoline(imps[i]);
+            imps[i] = NULL;
+        }
+    }
+    for(size_t i = 0; i < n; i++) {
+        if (imps[i]) {
+            free_trampoline(imps[i]);
+            imps[i] = NULL;
+        }
+    }
 }