and now for something completely different: assembly maybestret-IMPL

forwarding functions for atomicity

4 files changed, 147 insertions, 3 deletions

diff --git a/Makefile b/Makefile
index 7d29f12..fdf108a 100644
--- a/Makefile
+++ b/Makefile
@@ -17,7 +17,7 @@ $(shell mkdir -p out generated)
 HEADERS := lib/*.h
 GENERATED := generated/generic-dis-arm.inc.h generated/generic-dis-thumb.inc.h generated/generic-dis-thumb2.inc.h generated/generic-dis-arm64.inc.h
 define do_prefix
-generated/generic-dis-$(1).inc.h: Makefile
+generated/generic-dis-$(1).inc.h:
 	$(GEN_JS) --gen-hook-disassembler $(2) --dis-pattern='P(XXX)' $(IMAON2)/out/out-$(3).json > $$@ || rm -f $$@
 generateds: generated/generic-dis-$(1).inc.h
 endef
@@ -28,18 +28,22 @@ $(eval $(call do_prefix,arm64,,AArch64))
 
 out/%.o: lib/%.c Makefile $(HEADERS)
 	$(CC) -fvisibility=hidden -std=c11 -c -o $@ $<
+out/%.o: lib/%.S Makefile $(HEADERS)
+	$(CC) -fvisibility=hidden -c -o $@ $<
 out/jump-dis-arm-multi.o: generated/generic-dis-arm.inc.h generated/generic-dis-thumb.inc.h generated/generic-dis-thumb2.inc.h
 
 LIB_OBJS := \
 	out/find-syms.o \
 	out/interpose.o \
+	out/objc.o \
+	out/objc-asm.o \
 	out/substrate-compat.o \
 	out/jump-dis-arm-multi.o
 out/libsubstitute.dylib: $(LIB_OBJS)
-	$(CC) -dynamiclib -fvisibility=hidden -o $@ $(LIB_OBJS)
+	$(CC) -dynamiclib -fvisibility=hidden -o $@ $(LIB_OBJS) -lobjc
 
 define define_test
-out/test-$(1): test/test-$(2).c* $(HEADERS) $(GENERATED) Makefile out/libsubstitute.dylib
+out/test-$(1): test/test-$(2).[cm]* $(HEADERS) $(GENERATED) Makefile out/libsubstitute.dylib
 	$(3) -o $$@ $$< -Ilib -Isubstrate -Lout -lsubstitute
 all: out/test-$(1)
 endef
@@ -52,6 +56,7 @@ $(eval $(call define_test,find-syms,find-syms,$(CC) -std=c89))
 $(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))
 
 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
new file mode 100644
index 0000000..d3dae8a
--- /dev/null
+++ b/lib/objc-asm.S
@@ -0,0 +1,81 @@
+/* These all try to re-invoke old_ptr.  To make things worse, we don't know
+ * whether it will be called as stret or not, so we have to guess. */
+.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
+    mov %rdi, %rax
+    mov %rsi, %rdi
+    mov 0x28(%rax), %rsi
+    mov 0x20(%rax), %rax
+    jmp *(%rax)
+2: /* stret */
+    mov %rsi, %rax
+    mov %rdx, %rsi
+    mov 0x28(%rsi), %rdx
+    mov 0x20(%rsi), %rax
+    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 */
+    mov %ecx, 4(%esp)
+    mov 0x18(%eax), %ecx
+    mov %ecx, 8(%esp)
+    mov 0x14(%eax), %eax
+    jmp *(%eax)
+2: /* stret */
+    mov 12(%esp), %ecx; /* self */
+    mov 8(%esp), %eax; /* block */
+    mov %ecx, 4(%esp)
+    mov 0x18(%eax), %ecx
+    mov %ecx, 12(%esp)
+    mov 0x14(%eax), %eax
+    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
+    mov r9, r0
+    mov r0, r1
+    ldr r1, [r9, #0x18]
+    ldr r9, [r9, #0x14]
+    ldr r9, [r9]
+    bx r9
+2: /* stret */
+    mov r9, r1
+    mov r1, r2
+    ldr r2, [r9, #0x18]
+    ldr r9, [r9, #0x14]
+    ldr r9, [r9]
+    bx r9
+#elif defined(__arm64__)
+.align 2
+_temp_block_invoke:
+    mov x9, x0
+    mov x0, x1
+    ldr x1, [x9, #0x28]
+    ldr x9, [x9, #0x20]
+    ldr x9, [x9]
+    br x9
+#else
+#error No forwarding assembly definition for this arch
+#endif
diff --git a/lib/substitute.h b/lib/substitute.h
index 11141bf..181fb7c 100644
--- a/lib/substitute.h
+++ b/lib/substitute.h
@@ -43,6 +43,10 @@ enum {
     /* substitute_interpose_imports: couldn't redo relocation for an import
      * because the type was unknown */
     SUBSTITUTE_ERR_UNKNOWN_RELOCATION_TYPE,
+
+    /* substitute_hook_objc_message: no such selector existed in the class's
+     * inheritance tree */
+    SUBSTITUTE_ERR_NO_SUCH_SELECTOR,
 };
 
 struct substitute_function_hook {
@@ -167,6 +171,34 @@ int substitute_interpose_imports(const struct substitute_image *handle,
 
 #endif /* 1 */
 
+#if defined(__APPLE__) && __BLOCKS__
+#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
+ * runtime call class_replaceMethod, plus the superclass-call generation
+ * functionality described below, and some code to ensure atomicity if the
+ * method is called while the function is in progress.  However, it will race
+ * with code that modifies class methods without using atomic runtime calls,
+ * such as Substrate.
+ *
+ * @klass            the class
+ * @selector         the selector
+ * @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.
+ * @created_imp_ptr  optional - out pointer to whether a fake superclass-call
+ *                   IMP has been placed in <old_ptr>
+ *
+ * @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);
+#endif
+
 #ifdef __cplusplus
 } /* extern */
 #endif
diff --git a/test/test-imp-forwarding.m b/test/test-imp-forwarding.m
new file mode 100644
index 0000000..da2ece2
--- /dev/null
+++ b/test/test-imp-forwarding.m
@@ -0,0 +1,26 @@
+#include "../lib/objc.c"
+#include <objc/runtime.h>
+#include <stdio.h>
+
+static void imp1(id self, SEL sel, int a, int b) {
+    printf("imp1: self=%p sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
+}
+
+struct big {
+    int x[10];
+};
+
+static struct big imp2(id self, SEL sel, int a, int b) {
+    printf("imp2: self=%p sel=%s a=%d b=%d\n", self, sel_getName(sel), a, b);
+    return (struct big) {{0}};
+}
+
+int main() {
+    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 (*)(id, SEL, int, int)) new)(@"bar", sel, 1, 2);
+}