fixes and test - both tests work on all archs

4 files changed, 52 insertions, 34 deletions

diff --git a/lib/objc-asm.S b/lib/objc-asm.S
index 7bfe159..7b874c4 100644
--- a/lib/objc-asm.S
+++ b/lib/objc-asm.S
@@ -13,12 +13,13 @@ _remap_start:
 .rept TRAMPOLINES_PER_PAGE
 0:
 #if defined(__x86_64__)
-    push %rdi; push %rsi; push %rdx; push %rcx; push %r8; push %r9
+    /* double push for align */
+    push %rdi; push %rsi; push %rdx; push %rcx; push %r8; push %r9; push %r9
     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
+    pop %r9; pop %r9; pop %r8; pop %rcx; pop %rdx; pop %rsi; pop %rdi
     jmp *%rax
 #elif defined(__i386__)
     call 1f
@@ -28,11 +29,11 @@ _remap_start:
     push 8(%edx)
     push 4(%edx)
     call *(%edx)
-    add 8, %esp
+    add $$8, %esp
     jmp *%eax
 #elif defined(__arm__)
     push {r0-r4, lr} /* r4 for align */
-    mov r3, #(my_rpe - 1f)
+    mov r3, #(my_rpe - (1f + 2))
     add r3, pc
 1:
     ldr r0, [r3, #4]
@@ -43,23 +44,23 @@ _remap_start:
     pop {r0-r4, lr}
     bx r9
 #elif defined(__arm64__)
+    stp x30, x8, [sp, #-0x10]!
     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 x0, my_rpe+4
-    ldr x1, my_rpe+8
+    ldr x0, my_rpe+8
+    ldr x1, my_rpe+0x10
     ldr x2, my_rpe
     blr x2
     mov x9, x0
 
-    ldr x30, [sp], #0x10
     ldp x1, x0, [sp], #0x10
     ldp x3, x2, [sp], #0x10
     ldp x5, x4, [sp], #0x10
     ldp x7, x6, [sp], #0x10
+    ldp x30, x8, [sp], #0x10
 
     br x9
 #else
diff --git a/lib/objc.c b/lib/objc.c
index 98b84e6..92b4a34 100644
--- a/lib/objc.c
+++ b/lib/objc.c
@@ -35,7 +35,7 @@ LIST_HEAD(tramp_info_page_list, tramp_info_page_header)
 
 extern char remap_start[];
 
-static int get_trampoline(void *func, void *arg1, void *arg2, void *tramp) {
+static int get_trampoline(void *func, void *arg1, void *arg2, void *tramp_ptr) {
     int ret, rerrno = 0;
     pthread_mutex_lock(&tramp_mutex);
 
@@ -89,7 +89,11 @@ static int get_trampoline(void *func, void *arg1, void *arg2, void *tramp) {
     entry->func = func;
     entry->arg1 = arg1;
     entry->arg2 = arg2;
-    *(void **) tramp = (page - PAGE_SIZE) + (entry - entries) * TRAMPOLINE_SIZE;
+    void *tramp = (page - PAGE_SIZE) + (entry - entries) * TRAMPOLINE_SIZE;
+#ifdef __arm__
+    tramp += 1;
+#endif
+    *(void **) tramp_ptr = tramp;
     ret = SUBSTITUTE_OK;
 out:
     pthread_mutex_unlock(&tramp_mutex);
@@ -126,41 +130,52 @@ 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) {
+EXPORT
+int substitute_hook_objc_message(Class class, SEL selector, void *replacement,
+                                 void *old_ptr, bool *created_imp_ptr) {
     int ret;
-    Method meth = class_getClassMethod(class, selector);
+    Method meth = class_getInstanceMethod(class, selector);
     if (meth == NULL)
         return SUBSTITUTE_ERR_NO_SUCH_SELECTOR;
     const char *types = method_getTypeEncoding(meth);
 
+    if (created_imp_ptr)
+        *created_imp_ptr = false;
+
     /* temporary trampoline just tries again */
-    IMP temp;
-    if ((ret = get_trampoline(dereference, old_ptr, NULL, &temp)))
-        return ret;
-    *old_ptr = temp;
+    IMP temp = NULL;
+    if (old_ptr) {
+        if ((ret = get_trampoline(dereference, old_ptr, NULL, &temp)))
+            return ret;
+        *(IMP *) old_ptr = temp;
+    }
 
     IMP old = class_replaceMethod(class, selector, replacement, types);
     if (old) {
-        *old_ptr = old;
-        *created_imp_ptr = false;
+        if (old_ptr)
+            *(IMP *) old_ptr = old;
     } 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__);
+        if (old_ptr) {
+            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);
+            if (created_imp_ptr)
+                *created_imp_ptr = true;
         }
-        ret = get_trampoline(class_getMethodImplementation, super, selector, old_ptr);
-        *created_imp_ptr = true;
     }
 
-    free_trampoline(temp);
+    if (temp)
+        free_trampoline(temp);
     return SUBSTITUTE_OK;
 }
 
+EXPORT
 void substitute_free_created_imp(IMP imp) {
     free_trampoline(imp);
 }
diff --git a/lib/objc.h b/lib/objc.h
index 24427e1..141df11 100644
--- a/lib/objc.h
+++ b/lib/objc.h
@@ -7,9 +7,9 @@
 #endif
 #define _PAGE_SIZE (1 << _PAGE_SHIFT)
 #if defined(__x86_64__)
-#define TRAMPOLINE_SIZE 0x23
+#define TRAMPOLINE_SIZE 0x27
 #elif defined(__i386__)
-#define TRAMPOLINE_SIZE 0x1c
+#define TRAMPOLINE_SIZE 0x19
 #elif defined(__arm__)
 #define TRAMPOLINE_SIZE 0x18
 #elif defined(__arm64__)
diff --git a/lib/substitute.h b/lib/substitute.h
index 4dbfc7c..398b7fd 100644
--- a/lib/substitute.h
+++ b/lib/substitute.h
@@ -187,7 +187,9 @@ int substitute_interpose_imports(const struct substitute_image *handle,
  *
  * @klass            the class
  * @selector         the selector
- * @replacement      the new implementation
+ * @replacement      the new implementation (other APIs would call this an
+ *                   IMP, but that isn't in general the real type of the
+ *                   implementation, so declared as a void * here)
  * @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 can
@@ -198,8 +200,8 @@ int substitute_interpose_imports(const struct substitute_image *handle,
  * @return           SUBSTITUTE_OK
  *                   SUBSTITUTE_ERR_NO_SUCH_SELECTOR
  */
-int substitute_hook_objc_message(Class klass, SEL selector, IMP replacement,
-                                 IMP *old_ptr, bool *created_imp_ptr);
+int substitute_hook_objc_message(Class klass, SEL selector, void *replacement,
+                                 void *old_ptr, bool *created_imp_ptr);
 
 void substitute_free_created_imp(IMP imp);
 #endif