Method replacement - dynamic subclassing in objective c

NSObject+MethodBlockReplacement.h

//
//  NSObject+MethodBlockReplacement.h
//
//  Created by Martin Andreas Ullrich on 10.10.13.
//  Copyright (c) 2013 CSS Computer-Systems-Support GmbH. All rights reserved.
//

#import <Foundation/Foundation.h>

#import <objc/message.h>

/*!
 use this macro to declare a typed version of objc_msgSendSuper.
 supply method name, return type and paramter list (without SEL)
 */
#define RKTypedSuperMethod(methodName, returnType, ...) returnType (*methodName)( struct objc_super *, SEL, ##__VA_ARGS__ ) = (returnType (*) ( struct objc_super *, SEL, ##__VA_ARGS__ ))objc_msgSendSuper;

/*!
 use this macro to declare a typed version of objc_msgSendSuper that returns a struct value.
 supply method name, return type and paramter list (without SEL)
 */
#define RKTypedSuperMethod_StructReturn(methodName, returnType, ...) returnType (*methodName)( struct objc_super *, SEL, ##__VA_ARGS__ ) = (returnType (*) ( struct objc_super *, SEL, ##__VA_ARGS__ ))objc_msgSendSuper_stret;

/// use this macro for generating a default objc_super struct pointer for an object to pass to objc_msgSendSuper()
#define RKSuperStructPointer(obj) (&(struct objc_super){ .receiver = obj, .super_class = class_getSuperclass(object_getClass(obj))})

@interface NSObject (MethodBlockReplacement)

/*!
Intercepts calls to selector by calling a specified block instead. this uses dynamic subclassing and is not usable with CF-bridged objects as they use a custom dispatching mechanism.
 
The block's parameters must be a self-pointer followed by the exact paramter list of the correspoinding method signature. The SEL argument (_cmd) is not passed to the block.
@param selector A selector to intercept
@param block A block to invoke when a message with selector is sent to the reciever. the block's first argument is a pointer to the reciever (= self-pointer), followed by the method parameters.
*/
- (void)replaceSelector:(SEL)selector withBlock:(void *)block;

/// Clears all block intercepted methods (if present)
- (void)removeAllBlocksReplacingSelectors;

@end

NSObject+MethodBlockReplacement.m

//
//  NSObject+MethodBlockReplacement.m
//
//  Created by Martin Andreas Ullrich on 10.10.13.
//  Copyright (c) 2013 CSS Computer-Systems-Support GmbH. All rights reserved.
//

#import "NSObject+MethodBlockReplacement.h"
#import <objc/runtime.h>
#import <objc/message.h>

static NSString * const InterceptClassPrefix = @"RKBlockReplaced";

static BOOL isCrazyCFClass(__unsafe_unretained Class cls)
{
    NSString *className = NSStringFromClass(cls);
    if ([className hasPrefix:@"__NSCF"] || [className hasPrefix:@"NSCF"]) {
        return YES;
    }
    return NO;
}

// tests if a class (not any superclass!!) implements a specific selector
static BOOL classImplementsSelector(__unsafe_unretained Class class, SEL selector)
{
    unsigned int count = 0;
    Method *methods = class_copyMethodList(class, &count);
    BOOL implementsSelector = NO;
    for (unsigned int i = 0; i < count; i++) {
        if (sel_isEqual(selector, method_getName(methods[i]))) {
            implementsSelector = YES;
            break;
        }
    }
    free(methods);
    return implementsSelector;
}

static Class interceptedClassForObject(id object)
{
    Class cls = object_getClass(object);
    if ([NSStringFromClass(cls) hasPrefix:InterceptClassPrefix]) {
        return cls;
    }
    return nil;
}

static void interceptedImp_dealloc(__unsafe_unretained id object, SEL cmd); // define to use..

static Class dynamicallySubclassObject(id object)
{
    Class originalClass = object_getClass(object);
    NSString *className = [NSString stringWithFormat:@"%@_%@_%@", InterceptClassPrefix, NSStringFromClass(originalClass), [[NSUUID UUID] UUIDString]];
    Class cls = objc_allocateClassPair(object_getClass(object), [className cStringUsingEncoding:NSASCIIStringEncoding], 0);
    if (cls) {
        objc_registerClassPair(cls);
        class_addMethod(cls, NSSelectorFromString(@"dealloc"), (IMP)interceptedImp_dealloc, "v@:");
        object_setClass(object, cls);
    } else {
        @throw [NSException exceptionWithName:NSGenericException reason:@"could not allocate class pair for method interception" userInfo:nil];
    }
    return cls;
}

static void removeInterceptedMethods(__unsafe_unretained Class cls)
{
    unsigned int methodCount = 0;
    Method *methods = class_copyMethodList(cls, &methodCount);
    
    for (unsigned int i = 0; i < methodCount; i++) {
        Method m = methods[i];
        IMP methodImp = method_getImplementation(m);
        if (methodImp != (IMP)interceptedImp_dealloc) {
            imp_removeBlock(methodImp); // this may fail for IMPs that aren't block trampolines
        }
    }
    
    free(methods);
}

static void clearAllInterceptions(id object)
{
    __unsafe_unretained Class dynamicSubclass = interceptedClassForObject(object);
    if (dynamicSubclass) {
        Class cls = dynamicSubclass;
        
        do {
            Class superClass = class_getSuperclass(cls);
            object_setClass(object, superClass); // make sure object has no reference to our to-be-removed class
            
            removeInterceptedMethods(cls);
            
            objc_disposeClassPair(cls);
            
            cls = superClass;
        } while (cls && [NSStringFromClass(cls) hasPrefix:InterceptClassPrefix]);
    }
}

static void interceptedImp_dealloc(__unsafe_unretained id object, SEL cmd)
{
    __unsafe_unretained Class currentClass = object_getClass(object);
    removeInterceptedMethods(currentClass);
    objc_msgSendSuper(&((struct objc_super){ .receiver = object, .super_class = class_getSuperclass(object_getClass(object))}), cmd);
    objc_disposeClassPair(currentClass);
}

static void replaceSelectorWithBlock(id object, SEL selector, void *block)
{
    Class dynamicSubclass = interceptedClassForObject(object);
    if (!dynamicSubclass) {
        dynamicSubclass = dynamicallySubclassObject(object);
    }
    
    BOOL alreadyExists = NO;
    unsigned int methodCount = 0;
    Method *methods = class_copyMethodList(dynamicSubclass, &methodCount);
    for (unsigned int i = 0; i < methodCount; i++) {
        if (method_getName(methods[i]) == selector) {
            alreadyExists = YES;
            break;
        }
    }
    free(methods);
    
    if (alreadyExists) {
        dynamicSubclass = dynamicallySubclassObject(object);
    }
    
    Method superMethod = class_getInstanceMethod(class_getSuperclass(dynamicSubclass), selector);
    IMP blockImp = imp_implementationWithBlock((__bridge id)(block));
    class_addMethod(dynamicSubclass, selector, blockImp, method_getTypeEncoding(superMethod));
}

@implementation NSObject (MethodBlockReplacement)

- (void)replaceSelector:(SEL)selector withBlock:(void *)block
{
    if (isCrazyCFClass(object_getClass(self))) {
        @throw [NSException exceptionWithName:NSInternalInconsistencyException reason:@"RootKit block-based method interception is not designed to work with CoreFoundation bridged objects." userInfo:nil];
    }
    replaceSelectorWithBlock(self, selector, block);
}

- (void)removeAllBlocksReplacingSelectors
{
    clearAllInterceptions(self);
}

@end

NSObject+MethodBlockReplacement.Tests.m

//
//  InterceptionTests.m
//
//  Created by Martin Ullrich on 19.07.13.
//  Copyright (c) 2013 CSS Computer-Systems-Support GmbH. All rights reserved.
//

#import "InterceptionTests.h"
#import <RootKit/RootKit.h>
#import <objc/message.h>

@implementation InterceptionTests

- (void)testBlockInterception
{
    { // < this block is to make ARC generate releases within a defined scope
        
        id null = [NSNull null];
        
        id obj = [NSObject new];
        
        STAssertFalse([obj isEqual:null], @"a new NSObject is iequal to NSNull..");
        
        BOOL (^yesBlock)(id this, id other) = ^(id this, id other){
            return YES;
        };
        
        [obj replaceSelector:@selector(isEqual:) withBlock:(__bridge void *)(yesBlock)];
        
        STAssertTrue([obj isEqual:null], @"Method interception unsuccessful");
        
        [obj removeAllBlocksReplacingSelectors];
        
        STAssertFalse([obj isEqual:null], @"clearing method interception unsuccessful");
        
        [obj replaceSelector:@selector(isEqual:) withBlock:(__bridge void *)(yesBlock)];
        
        STAssertTrue([obj isEqual:null], @"Method interception unsuccessful");
        
        [obj replaceSelector:@selector(isEqual:) withBlock:(void*)^BOOL(id this, id other){
            return NO;
        }];
        
        STAssertFalse([obj isEqual:null], @"Method interception unsuccessful");
        
        [obj removeAllBlocksReplacingSelectors];
        
        STAssertFalse([obj isEqual:null], @"clearing method interception unsuccessful");
        
        obj = [NSObject new];
        
        [obj replaceSelector:@selector(isEqual:) withBlock:(void*)^BOOL(id this, id other){
            RKTypedSuperMethod(typedSuper, BOOL, id);
            BOOL value = typedSuper(RKSuperStructPointer(this), @selector(isEqual:), other);
            return value;
        }];
        
        STAssertFalse([obj isEqual:null], @"clearing method interception unsuccessful");
        
    } // ARC should have generated -release calls by now (even without optimizations), so..
    // scan runtime for leftover classes
    
    NSString *prefix = @"RKBlockReplaced";
    
    unsigned int classCount = 0;
    Class *classes = objc_copyClassList(&classCount);
    for (unsigned int i = 0; i < classCount; i++) {
        if ([NSStringFromClass(classes[i]) hasPrefix:prefix]) {
            STFail(@"there still are dynamic subclasses left over from testing block interception");
        }
    }
    free(classes);
}

@end