objective-cruntimeselectormethod-swizzling

Swizzle an unknown selector of a specific class at runtime


I have a class with the following method in one of my library's public headers. A user of my class will pass in the selector and class.

-(void)someMethodWithSelector(SEL)aSelector ofClass:(Class)clazz

At compile time, I don't know what the selector will look like, how many parameters will be passed, etc... but what I want is to be able to swizzle the passed selector at runtime, perform some extra logic, and call the original method afterwards.

I know how to swizzle class and instance methods, but I'm unsure how I would proceed given this scenario.

Has anyone had any experience dealing with a similar approach?


Solution

  • MikeAsh was able to figure out this problem, so all the credit of this answer goes to him

    @import Foundation;
    @import ObjectiveC;
    static NSMutableSet *swizzledClasses;
    static NSMutableDictionary *swizzledBlocks; // Class -> SEL (as string) -> block
    static IMP forwardingIMP;
    static dispatch_once_t once;
    
    void Swizzle(Class c, SEL sel, void (^block)(NSInvocation *)) {
        dispatch_once(&once, ^{
            swizzledClasses = [NSMutableSet set];
            swizzledBlocks = [NSMutableDictionary dictionary];
            forwardingIMP = class_getMethodImplementation([NSObject class], @selector(thisReallyShouldNotExistItWouldBeExtremelyWeirdIfItDid));
        });
        if(![swizzledClasses containsObject: c]) {
            SEL fwdSel = @selector(forwardInvocation:);
            Method m = class_getInstanceMethod(c, fwdSel);
            __block IMP orig;
            IMP imp = imp_implementationWithBlock(^(id self, NSInvocation *invocation) {
                NSString *selStr = NSStringFromSelector([invocation selector]);
                void (^block)(NSInvocation *) = swizzledBlocks[c][selStr];
                if(block != nil) {
                    NSString *originalStr = [@"omniswizzle_" stringByAppendingString: selStr];
                    [invocation setSelector: NSSelectorFromString(originalStr)];
                    block(invocation);
                } else {
                    ((void (*)(id, SEL, NSInvocation *))orig)(self, fwdSel, invocation);
                }
            });
            orig = method_setImplementation(m, imp);
            [swizzledClasses addObject: c];
        }
        NSMutableDictionary *classDict = swizzledBlocks[c];
        if(classDict == nil) {
            classDict = [NSMutableDictionary dictionary];
            swizzledBlocks[(id)c] = classDict;
        }
        classDict[NSStringFromSelector(sel)] = block;
        Method m = class_getInstanceMethod(c, sel);
        NSString *newSelStr = [@"omniswizzle_" stringByAppendingString: NSStringFromSelector(sel)];
        SEL newSel = NSSelectorFromString(newSelStr);
        class_addMethod(c, newSel, method_getImplementation(m), method_getTypeEncoding(m));
        method_setImplementation(m, forwardingIMP);
    }
    

    Here is how we would call the function:

    int main(int argc, const char * argv[]) {
        @autoreleasepool {
            Swizzle([NSBundle class], @selector(objectForInfoDictionaryKey:), ^(NSInvocation *inv) {
                NSLog(@"invocation is %@ - calling now", inv);
                [inv invoke];
                NSLog(@"after");
            });
    
            NSLog(@"%@", [[NSBundle bundleForClass: [NSString class]] objectForInfoDictionaryKey: (__bridge NSString *)kCFBundleVersionKey]);
        }
        return 0;
    }