Restore correct implementation to NSURLSession

Nocilla/Hooks/NSURLSession/LSNSURLSessionHook.m

@@ -10,31 +10,63 @@
 #import "LSHTTPStubURLProtocol.h"
 #import <objc/runtime.h>
 
-@implementation LSNSURLSessionHook
+@implementation LSNSURLSessionHook {
+    IMP ls_protocolClassesImplementation; // Implementation of `-(nullable NSArray *)protocolClasses` for this class
+    IMP ns_protocolClassesImplementation; // Implementation of `-(nullable NSArray *)protocolClasses` for `NSURLSessionConfiguration` or its underlying class
+}
 
-- (void)load {
-    Class cls = NSClassFromString(@"__NSCFURLSessionConfiguration") ?: NSClassFromString(@"NSURLSessionConfiguration");
-    [self swizzleSelector:@selector(protocolClasses) fromClass:cls toClass:[self class]];
+- (Class)sessionConfigurationClass
+{
+    Class retClass = NSClassFromString(@"__NSCFURLSessionConfiguration");
+    return retClass ?: NSClassFromString(@"NSURLSessionConfiguration");
 }
 
-- (void)unload {
-    Class cls = NSClassFromString(@"__NSCFURLSessionConfiguration") ?: NSClassFromString(@"NSURLSessionConfiguration");
-    [self swizzleSelector:@selector(protocolClasses) fromClass:cls toClass:[self class]];
+- (char const *)typeEncodingForSelector:(SEL)selector
+{
+    return method_getTypeEncoding((class_getInstanceMethod([self class], selector)));
 }
 
-- (void)swizzleSelector:(SEL)selector fromClass:(Class)original toClass:(Class)stub {
+- (instancetype)init
+{
+    self = [super init];
+    if (self)
+    {
+        SEL selector = @selector(protocolClasses);
+        ns_protocolClassesImplementation = class_getMethodImplementation(self.sessionConfigurationClass, selector);
+        ls_protocolClassesImplementation = class_getMethodImplementation([self class], selector);
+    }
 
-    Method originalMethod = class_getInstanceMethod(original, selector);
-    Method stubMethod = class_getInstanceMethod(stub, selector);
-    if (!originalMethod || !stubMethod) {
+    return self;
+}
+
+- (void)load
+{
+    if (!ns_protocolClassesImplementation || !ls_protocolClassesImplementation)
+    {
         [NSException raise:NSInternalInconsistencyException format:@"Couldn't load NSURLSession hook."];
     }
-    method_exchangeImplementations(originalMethod, stubMethod);
+
+    SEL selector              = @selector(protocolClasses);
+    const char *type_encoding = [self typeEncodingForSelector:selector];
+
+    //Swap the implementations
+    class_replaceMethod(self.class, selector, ns_protocolClassesImplementation, type_encoding);
+    class_replaceMethod(self.sessionConfigurationClass, selector, ls_protocolClassesImplementation, type_encoding);
 }
 
-- (NSArray *)protocolClasses {
-    return @[[LSHTTPStubURLProtocol class]];
+- (void)unload
+{
+    SEL selector              = @selector(protocolClasses);
+    const char *type_encoding = [self typeEncodingForSelector:selector];
+
+    //Restore the implementations
+    class_replaceMethod(self.class, selector, ls_protocolClassesImplementation, type_encoding);
+    class_replaceMethod(self.sessionConfigurationClass, selector, ns_protocolClassesImplementation, type_encoding);
 }
 
+- (NSArray *)protocolClasses
+{
+    return @[[LSHTTPStubURLProtocol class]];
+}
 
 @end