fixes and test - both tests work on all archs

6 files changed, 100 insertions, 35 deletions

diff --git a/Makefile b/Makefile
index 33c22e2..fbc04af 100644
--- a/Makefile
+++ b/Makefile
@@ -44,7 +44,7 @@ out/libsubstitute.dylib: $(LIB_OBJS)
 
 define define_test
 out/test-$(1): test/test-$(2).[cm]* $(HEADERS) $(GENERATED) Makefile out/libsubstitute.dylib
-	$(3) -o $$@ $$< -Ilib -Isubstrate -Lout -lsubstitute
+	$(3) -g -o $$@ $$< -Ilib -Isubstrate -Lout -lsubstitute
 all: out/test-$(1)
 endef
 $(eval $(call define_test,tdarm-simple,td-simple,$(CC) -std=c11 -DHDR='"dis-arm.inc.h"' -Dxdis=dis_arm))
@@ -57,6 +57,7 @@ $(eval $(call define_test,find-syms-cpp,find-syms,$(CXX) -x c++ -std=c++98))
 $(eval $(call define_test,substrate,substrate,$(CXX) -std=c++98))
 $(eval $(call define_test,jump-dis,jump-dis,$(CC) -std=c11))
 $(eval $(call define_test,imp-forwarding,imp-forwarding,$(CC) -std=c11 -framework Foundation -lobjc))
+$(eval $(call define_test,objc-hook,objc-hook,$(CC) -std=c11 -framework Foundation -lsubstitute))
 
 out/insns-arm.o: test/insns-arm.S Makefile
 	clang -arch armv7 -c -o $@ $<
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
diff --git a/test/test-objc-hook.m b/test/test-objc-hook.m
new file mode 100644
index 0000000..ae7d3ff
--- /dev/null
+++ b/test/test-objc-hook.m
@@ -0,0 +1,46 @@
+#include "substitute.h"
+#import <Foundation/Foundation.h>
+
+@interface Base : NSObject {
+}
+- (void)foo:(NSString *)str;
+- (void)bar:(NSString *)str;
+@end
+
+@implementation Base
+- (void)foo:(NSString *)str {
+    NSLog(@"base foo: %@", str);
+}
+- (void)bar:(NSString *)str {
+    NSLog(@"base bar: %@", str);
+}
+@end
+
+@interface Derived : Base {
+}
+@end
+
+@implementation Derived
+- (void)foo:(NSString *)str {
+    NSLog(@"derived foo: %@", str);
+}
+@end
+
+static void (*old_foo)(id, SEL, NSString *);
+static void new_foo(id self, SEL sel, NSString *str) {
+    NSLog(@"new foo: %@; calling orig", str);
+    return old_foo(self, sel, str);
+}
+static void (*old_bar)(id, SEL, NSString *);
+static void new_bar(id self, SEL sel, NSString *str) {
+    NSLog(@"new bar: %@; calling orig", str);
+    return old_bar(self, sel, str);
+}
+
+int main() {
+    assert(!substitute_hook_objc_message([Derived class], @selector(foo:), new_foo, &old_foo, NULL));
+    assert(!substitute_hook_objc_message([Derived class], @selector(bar:), new_bar, &old_bar, NULL));
+    Derived *d = [[Derived alloc] init];
+    [d foo:@"hi!"];
+    [d bar:@"hello!"];
+}