Indie Swizzling

Method swizzling has recently hit the mainstream. It’s a technique we use at Code School quite a bit. For example, if you look inside the downloaded projects for Core iOS 7, you’ll see plenty of swizzling going on and most of the time the technique described in mattt’s post is sufficient. But every once in awhile you’ll try to sizzle a method and for some strange reason it doesn’t work. It’s like the class is immune to your swizzling ways.

I ran into one of these immune classes when trying to swizzle the downloadTaskWithURL: method on NSURLSession using the traditional category-based technique. I wanted to make sure the Code School student called downloadTaskWithURL: with the correct parameters1.

Here what the calling-code looked like:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
- (BOOL)application:(UIApplication *)application didFinishLaunchingWithOptions:(NSDictionary *)launchOptions
{
    NSURLSessionConfiguration *default = [NSURLSessionConfiguration defaultSessionConfiguration];
    NSURLSession *session = [NSURLSession sessionWithConfiguration:default
                                                          delegate:self
                                                     delegateQueue:[NSOperationQueue mainQueue]];

    NSURL *URL = [NSURL URLWithString:@"http://initwithfunk.com"];

    NSURLSessionDownloadTask *task = [session downloadTaskWithURL:URL];
    [task resume];

    return YES;
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didCompleteWithError:(NSError *)error
{
    NSHTTPURLResponse *response = (NSHTTPURLResponse *)task.response;

    NSLog(@"Task did complete with status: %d", response.statusCode);
}

And the swizzling category (this isn’t exactly how it worked but essentially it did the same thing):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
#import "NSURLSession+Swizzle.h"
#import <objc/runtime.h>

@implementation NSURLSession (Swizzle)

+ (void)load{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        Class class = self.class;

        SEL originalSelector = @selector(downloadTaskWithURL:);
        SEL swizzledSelector = @selector(xxx_downloadTaskWithURL:);

        Method originalMethod = class_getInstanceMethod(class, originalSelector);
        Method swizzledMethod = class_getInstanceMethod(class, swizzledSelector);

        IMP swizzledImp = method_getImplementation(swizzledMethod);
        IMP originalImp = method_getImplementation(originalMethod);

        BOOL didAddMethod =
        class_addMethod(class,
                        originalSelector,
                        swizzledImp,
                        method_getTypeEncoding(swizzledMethod));

        if (didAddMethod) {
            class_replaceMethod(class,
                                swizzledSelector,
                                originalImp,
                                method_getTypeEncoding(originalMethod));
        } else {
            method_exchangeImplementations(originalMethod, swizzledMethod);
        }

    });
}

- (NSURLSessionDownloadTask *)xxx_downloadTaskWithURL:(NSURL *)url
{
    NSLog(@"%s", __PRETTY_FUNCTION__);

    return [self xxx_downloadTaskWithURL:url];
}

@end

When this ran, the xxx_ method would never get called. I stepped through the +load and everything seemed to be working just fine. But when I logged the pointers of the class I was swizzling in +load and the class of the NSURLSession instance I was using in the app delegate, I discovered that they were pointing to different locations in memory.

So I tried updating the +load method. Instead of using self.class, I would grab the class from an actual instance of NSURLSession, like so:

1
Class class = [NSURLSession sharedSession].class;

Unfortunately, that led to the class_getInstanceMethod(class, swizzledSelector) call returning nil since it was no longer looking in the class that included the xxx_downloadTaskWithURL: method. Replacing that code with class_getInstanceMethod(self.class, swizzledSelector) won’t work either, leading to a -[__NSCFURLSession xxx_downloadTaskWithURL:]: unrecognized selector error when the implementation of the swizzle method tries to call back to the original method that doesn’t exist.

The “Indie” Swizzle2

It turns out there is a more direct way to swizzle methods without first defining the replacement method on the same class as the original method. Mike Ash details this technique in a fantastic swizzling Friday Q&A:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
// Create a function pointer for the original downloadTaskWithURL: to be stored in
static NSURLSessionDownloadTask * (*OriginalDownloadedTaskWithURL)(id, SEL, NSURL *);

// Define our custom implementation of downloadTaskWithURL:, calling OriginalDownloadedTaskWithURL
// to call the original method
static NSURLSessionDownloadTask * `(id self, SEL _cmd, NSURL *url){
    NSLog(@"%s", __PRETTY_FUNCTION__);

    NSURLSessionDownloadTask *task = OriginalDownloadedTaskWithURL(self, _cmd, url);

    return task;
}

+ (void)load
{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{

        // Replace the method on the same class that's used
        // In the calling code
        Class class = [NSURLSession sharedSession].class;

        SEL originalSelector = @selector(downloadTaskWithURL:);

        // The replacemethod method implementation
        IMP replacement = (IMP)CustomDownloadTaskWithURL;

        // This will eventually hold the original downloadTaskWithURL: method
        IMP *store = (IMP *)&OriginalDownloadedTaskWithURL;

        IMP originalImp = NULL;
        Method method = class_getInstanceMethod(class, originalSelector);
        if (method) {
            const char *type = method_getTypeEncoding(method);
            // Replace the original method with the XXXDownloadTaskWithURL IMP function
            originalImp = class_replaceMethod(class, originalSelector, replacement, type);
            if (!originalImp) {
                originalImp = method_getImplementation(method);
            }
        }

        // Put the original method IMP into the store pointer
        if (originalImp && store) { *store = originalImp; }
    });
}

Since all objective-c methods eventually just become a C function with the first two arguments of id self and SEL _cmd, we can just go ahead and define the replacement method as a static function called XXXDownloadTaskWithURL that accepts self, _cmd, and an NSURL * as the “first” argument.

We also create a place to store the original method, a pointer to an IMP pointer called OriginalDownloadedTaskWithURL which will allow our replacement method to call back out to the original. Of course the original description of this approach is a better so definitely check it out.

Why doesn’t category-based swizzling work for NSURLSession?

I don’t know what makes NSURLSession immune to the category based swizzle, but I’d love to (my Objective-C runtime skills are just dangerous enough to put this together). I think it has something to do with it being a class-cluster, or maybe it has to do with it being a toll-free bridged class to __NSCFURLSession. I’m not sure.

If you know the answer, or have a suggestion for further research, I’d love to hear about it and update this post. Message me on twitter @eallam.

UPDATE

Thanks to a tip from Daniel Haight it’s possible to shorten this code a bit by using imp_implementationWithBlock instead of creating a C function for the replacement method (saving us from having to name it):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
// Create a function pointer for the original downloadTaskWithURL: to be stored in
static NSURLSessionDownloadTask * (*OriginalDownloadedTaskWithURL)(id, SEL, NSURL *);

+ (void)load
{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{

        Class class = [NSURLSession sharedSession].class;

        SEL originalSelector = @selector(downloadTaskWithURL:);

        // Use this to create the replacement IMP inline
        IMP replacement = imp_implementationWithBlock(^NSURLSessionDownloadTask * (id _self, NSURL *url){

            NSLog(@"%s", __PRETTY_FUNCTION__);

            return OriginalDownloadedTaskWithURL(_self, _cmd, url);
        });

        IMP *store = (IMP *)&OriginalDownloadedTaskWithURL;

        IMP originalImp = NULL;
        Method method = class_getInstanceMethod(class, originalSelector);
        if (method) {
            const char *type = method_getTypeEncoding(method);
            originalImp = class_replaceMethod(class, originalSelector, replacement, type);
            if (!originalImp) {
                originalImp = method_getImplementation(method);
            }
        }
        if (originalImp && store) { *store = originalImp; }
    });
}
Icon attribution:

  1. These are downloadable throwaway projects to help our students practice and learn. So it was safe to swizzle as long as it worked consistently. I wouldn’t suggest every putting this in shipping code.

  2. “Indie” because the replacement method implementation is independent of the class being swizzled