Swizzle an unknown selector of a specific class at runtime - objective-c

Question

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?

Answer 1

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;
}