diff --git a/Nocilla/Hooks/NSURLSession/LSNSURLSessionHook.m b/Nocilla/Hooks/NSURLSession/LSNSURLSessionHook.m index ab584f59..ba7fe1f5 100644 --- a/Nocilla/Hooks/NSURLSession/LSNSURLSessionHook.m +++ b/Nocilla/Hooks/NSURLSession/LSNSURLSessionHook.m @@ -10,31 +10,63 @@ #import "LSHTTPStubURLProtocol.h" #import -@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