Validate replacement block signatures

CHANGELOG.md

@@ -15,6 +15,7 @@
 * Preserve floating-point arguments when object hooks invoke original methods on arm64. Thanks to [@ishutinvv](https://github.com/ishutinvv).
 * Run class-availability hooks after Objective-C loads a new image, fixing https://github.com/steipete/InterposeKit/issues/26.
 * Reject Core Foundation-backed object hooks before dynamic subclassing, preventing crashes such as https://github.com/steipete/InterposeKit/issues/23.
+* Reject replacement blocks whose Objective-C signatures do not match the hooked method, fixing https://github.com/steipete/InterposeKit/issues/27.
 
 ## 0.01
 

README.md

@@ -77,7 +77,7 @@ Hi there 👋 and Interpose
 - Interpose works on classes and individual objects.
 - Hooks can easily be undone via calling `revert()`. This also checks and errors if someone else changed stuff in between.
 - Mostly Swift, no `NSInvocation`, which requires boxing and can be slow.
-- No Type checking. If you have a typo or forget a `convention` part, this will crash at runtime.
+- Replacement block signatures are checked when hooks are created. The `methodSignature` used to call the original implementation remains caller-typed.
 - Yes, you have to type the resulting type twice This is a tradeoff, else we need `NSInvocation`.
 - Delayed Interposing helps when a class is loaded at runtime. This is useful for [Mac Catalyst](https://steipete.com/posts/mac-catalyst-crash-hunt/).
 
@@ -171,7 +171,6 @@ Add `github "steipete/InterposeKit"` to your `Cartfile`.
 ## Improvement Ideas
 
 - Write proposal to allow to [convert the calling convention of existing types](https://twitter.com/steipete/status/1266799174563041282?s=21).
-- Use the C block struct to perform type checking between Method type and C type (I do that in  [Aspects library](https://github.com/steipete/Aspects)), it's still a runtime crash but could be at hook time, not when we call it.
 - Add a way to get all current hooks from an object/class.
 - Add a way to revert hooks without super helper.
 - Add a way to apply multiple hooks to classes

Sources/InterposeKit/AnyHook.swift

@@ -1,5 +1,9 @@
 import Foundation
 
+#if SWIFT_PACKAGE && !os(Linux)
+import SuperBuilder
+#endif
+
 /// Base class, represents a hook to exactly one method.
 public class AnyHook {
     /// The class this hook is based on.
@@ -66,6 +70,20 @@ public class AnyHook {
         return method
     }
 
+    func validateImplementationBlock(_ block: AnyObject) throws {
+        #if os(Linux)
+        _ = block
+        #else
+        let method = try validate()
+        guard IKTBlockSignatureMatchesMethod(block, method) else {
+            let methodSignature = method_getTypeEncoding(method).map(String.init(cString:)) ?? "<missing>"
+            let hookSignature = IKTBlockGetTypeEncoding(block).map(String.init(cString:))
+            throw InterposeError.incompatibleHookSignature(
+                `class`, selector, methodSignature: methodSignature, hookSignature: hookSignature)
+        }
+        #endif
+    }
+
     private func execute(newState: State, task: () throws -> Void) throws {
         do {
             try task()

Sources/InterposeKit/ClassHook.swift

@@ -9,7 +9,9 @@ extension Interpose {
         public init(`class`: AnyClass, selector: Selector,
                     implementation: (ClassHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
             try super.init(class: `class`, selector: selector)
-            replacementIMP = imp_implementationWithBlock(implementation(self) as Any)
+            let block = implementation(self) as AnyObject
+            try validateImplementationBlock(block)
+            replacementIMP = imp_implementationWithBlock(block)
         }
 
         override func replaceImplementation() throws {

Sources/InterposeKit/InterposeError.swift

@@ -18,6 +18,10 @@ public enum InterposeError: LocalizedError {
     /// Unable to add method  for object-based interposing.
     case unableToAddMethod(AnyClass, Selector)
 
+    /// The replacement block uses a different Objective-C calling convention than the method.
+    case incompatibleHookSignature(
+        AnyClass, Selector, methodSignature: String, hookSignature: String?)
+
     /// Object-based hooking does not work if an object is using KVO.
     /// The KVO mechanism also uses subclasses created at runtime but doesn't check for additional overrides.
     /// Adding a hook eventually crashes the KVO management code so we reject hooking altogether in this case.
@@ -62,6 +66,9 @@ extension InterposeError: Equatable {
             return "Failed to allocate class pair: \(klass), \(subclassName)"
         case .unableToAddMethod(let klass, let selector):
             return "Unable to add method: -[\(klass) \(selector)]"
+        case let .incompatibleHookSignature(klass, selector, methodSignature, hookSignature):
+            return "Hook signature \(hookSignature ?? "<missing>") does not match method signature " +
+                "\(methodSignature): -[\(klass) \(selector)]"
         case .keyValueObservationDetected(let obj):
             return "Unable to hook object that uses Key Value Observing: \(obj)"
         case .coreFoundationObjectDetected(let obj):

Sources/InterposeKit/ObjectHook.swift

@@ -22,6 +22,7 @@ extension Interpose {
             self.object = object
             try super.init(class: type(of: object), selector: selector)
             let block = implementation(self) as AnyObject
+            try validateImplementationBlock(block)
             replacementIMP = imp_implementationWithBlock(block)
             guard replacementIMP != nil else {
                 throw InterposeError.unknownError("imp_implementationWithBlock failed for \(block) - slots exceeded?")

Sources/SuperBuilder/include/ITKSuperBuilder.h

@@ -1,5 +1,6 @@
 #if __APPLE__
 #import <Foundation/Foundation.h>
+#import <objc/runtime.h>
 #endif
 
 NS_ASSUME_NONNULL_BEGIN
@@ -9,6 +10,14 @@ typedef void (*ITKImageDidLoadCallback)(void);
 /// Schedules `callback` asynchronously when an image load can be inspected safely.
 FOUNDATION_EXPORT void IKTRegisterImageDidLoadCallback(ITKImageDidLoadCallback callback);
 
+#if __APPLE__
+/// Returns the Objective-C type encoding embedded in `block`, or NULL if unavailable.
+FOUNDATION_EXPORT const char *_Nullable IKTBlockGetTypeEncoding(id block);
+
+/// Checks whether `block` can replace `method` without changing its Objective-C calling convention.
+FOUNDATION_EXPORT BOOL IKTBlockSignatureMatchesMethod(id block, Method method);
+#endif
+
 /**
 Adds an empty super implementation instance method to originalClass.
 If a method already exists, this will return NO and a descriptive error message.

Sources/SuperBuilder/src/ITKSuperBuilder.m

@@ -21,6 +21,80 @@
 
 NSString *const SuperBuilderErrorDomain = @"com.steipete.superbuilder";
 
+typedef NS_OPTIONS(int32_t, IKTBlockFlags) {
+    IKTBlockFlagsHasCopyDisposeHelpers = 1 << 25,
+    IKTBlockFlagsHasSignature = 1 << 30
+};
+
+typedef struct {
+    __unsafe_unretained Class isa;
+    IKTBlockFlags flags;
+    int32_t reserved;
+    void (*invoke)(void *, ...);
+    void *descriptor;
+} IKTBlockLiteral;
+
+const char *IKTBlockGetTypeEncoding(id block) {
+    IKTBlockLiteral *layout = (__bridge void *)block;
+    if (!(layout->flags & IKTBlockFlagsHasSignature) || layout->descriptor == NULL) {
+        return NULL;
+    }
+
+    uint8_t *descriptor = layout->descriptor;
+    descriptor += 2 * sizeof(uintptr_t);
+    if (layout->flags & IKTBlockFlagsHasCopyDisposeHelpers) {
+        descriptor += 2 * sizeof(void *);
+    }
+    return *(const char **)descriptor;
+}
+
+static const char *IKTSkipTypeQualifiers(const char *typeEncoding) {
+    while (strchr("rnNoORV", typeEncoding[0]) != NULL) {
+        typeEncoding++;
+    }
+    return typeEncoding;
+}
+
+static BOOL IKTTypeEncodingsMatch(const char *first, const char *second) {
+    if (first == NULL || second == NULL) {
+        return NO;
+    }
+    first = IKTSkipTypeQualifiers(first);
+    second = IKTSkipTypeQualifiers(second);
+    if (first[0] == '@' && second[0] == '@') {
+        return YES;
+    }
+    return strcmp(first, second) == 0;
+}
+
+BOOL IKTBlockSignatureMatchesMethod(id block, Method method) {
+    const char *blockTypeEncoding = IKTBlockGetTypeEncoding(block);
+    const char *methodTypeEncoding = method_getTypeEncoding(method);
+    if (blockTypeEncoding == NULL || methodTypeEncoding == NULL) {
+        return NO;
+    }
+
+    NSMethodSignature *blockSignature = [NSMethodSignature signatureWithObjCTypes:blockTypeEncoding];
+    NSMethodSignature *methodSignature = [NSMethodSignature signatureWithObjCTypes:methodTypeEncoding];
+    if (blockSignature == nil || methodSignature == nil ||
+        blockSignature.numberOfArguments != methodSignature.numberOfArguments ||
+        blockSignature.numberOfArguments < 2 ||
+        !IKTTypeEncodingsMatch(blockSignature.methodReturnType, methodSignature.methodReturnType) ||
+        !IKTTypeEncodingsMatch([blockSignature getArgumentTypeAtIndex:1],
+                               [methodSignature getArgumentTypeAtIndex:0])) {
+        return NO;
+    }
+
+    // Block argument 0 is the block itself; method argument 1 is _cmd.
+    for (NSUInteger index = 2; index < blockSignature.numberOfArguments; index++) {
+        if (!IKTTypeEncodingsMatch([blockSignature getArgumentTypeAtIndex:index],
+                                   [methodSignature getArgumentTypeAtIndex:index])) {
+            return NO;
+        }
+    }
+    return YES;
+}
+
 static os_unfair_lock _imageDidLoadLock = OS_UNFAIR_LOCK_INIT;
 static ITKImageDidLoadCallback *_imageDidLoadCallbacks;
 static size_t _imageDidLoadCallbackCount;

Tests/InterposeKitTests/InterposeKitTests.swift

@@ -3,6 +3,20 @@
 
 final class InterposeKitTests: InterposeKitTestCase {
 
+    func testRejectsMismatchedClassHookSignature() {
+        XCTAssertThrowsError(try Interpose(TestClass.self).hook(
+            #selector(TestClass.returnInt),
+            methodSignature: (@convention(c) (AnyObject, Selector) -> Int).self,
+            hookSignature: (@convention(block) (AnyObject) -> Double).self
+        ) { _ in
+            { _ in 1.0 }
+        }) { error in
+            guard case InterposeError.incompatibleHookSignature = error else {
+                return XCTFail("Unexpected error: \(error)")
+            }
+        }
+    }
+
     override func setUpWithError() throws {
         Interpose.isLoggingEnabled = true
     }

Tests/InterposeKitTests/ObjectInterposeTests.swift

@@ -4,6 +4,38 @@
 
 final class ObjectInterposeTests: InterposeKitTestCase {
 
+    func testRejectsMismatchedHookReturnType() {
+        let testObj = TestClass()
+
+        XCTAssertThrowsError(try testObj.hook(
+            #selector(TestClass.returnInt),
+            methodSignature: (@convention(c) (AnyObject, Selector) -> Int).self,
+            hookSignature: (@convention(block) (AnyObject) -> Double).self
+        ) { _ in
+            { _ in 1.0 }
+        }) { error in
+            guard case InterposeError.incompatibleHookSignature = error else {
+                return XCTFail("Unexpected error: \(error)")
+            }
+        }
+    }
+
+    func testRejectsMismatchedHookArgumentType() {
+        let testObj = TestClass()
+
+        XCTAssertThrowsError(try testObj.hook(
+            #selector(TestClass.doubleString),
+            methodSignature: (@convention(c) (AnyObject, Selector, String) -> String).self,
+            hookSignature: (@convention(block) (AnyObject, Int) -> String).self
+        ) { _ in
+            { _, _ in "" }
+        }) { error in
+            guard case InterposeError.incompatibleHookSignature = error else {
+                return XCTFail("Unexpected error: \(error)")
+            }
+        }
+    }
+
     func testRejectsCoreFoundationBackedObject() throws {
         let url = try XCTUnwrap(NSURL(string: "https://www.google.com"))
         let expectedError = InterposeError.coreFoundationObjectDetected(url)
@@ -16,7 +48,7 @@
             methodSignature: (@convention(c) (AnyObject, Selector) -> String).self,
             hookSignature: (@convention(block) (AnyObject) -> String).self
         ) { _ in
             { _ in "www.facebook.com" }
         }) { error in
             XCTAssertEqual(error as? InterposeError, expectedError)
         }