Properly guard swizzled methods to avoid duplicate sniffing

Classes/Network/PonyDebugger/FLEXNetworkObserver.m

@@ -132,6 +132,24 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     return NSSelectorFromString([NSString stringWithFormat:@"_flex_swizzle_%x_%@", arc4random(), NSStringFromSelector(selector)]);
 }
 
+/// All swizzled delegate methods should make use of this guard.
+/// This will prevent duplicated sniffing when the original implementation calls up to a superclass implementation which we've also swizzled.
+/// The superclass implementation (and implementations in classes above that) will be executed without inteference if called from the original implementation.
++ (void)sniffWithoutDuplicationForObject:(NSObject *)object selector:(SEL)selector sniffingBlock:(void (^)(void))sniffingBlock originalImplementationBlock:(void (^)(void))originalImplementationBlock
+{
+    const void *key = selector;
+
+    // Don't run the sniffing block if we're inside a nested call
+    if (!objc_getAssociatedObject(object, key)) {
+        sniffingBlock();
+    }
+
+    // Mark that we're calling through to the original so we can detect nested calls
+    objc_setAssociatedObject(object, key, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+    originalImplementationBlock();
+    objc_setAssociatedObject(object, key, nil, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+}
+
 + (BOOL)instanceRespondsButDoesNotImplementSelector:(SEL)selector class:(Class)cls
 {
     if ([cls instancesRespondToSelector:selector]) {
@@ -275,8 +293,12 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLConnectionWillSendRequestBlock implementationBlock = ^NSURLRequest *(id <NSURLConnectionDelegate> slf, NSURLConnection *connection, NSURLRequest *request, NSURLResponse *response) {
-        NSURLRequest *returnValue = ((id(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, connection, request, response);
-        undefinedBlock(slf, connection, request, response);
+        __block NSURLRequest *returnValue = nil;
+        [self sniffWithoutDuplicationForObject:connection selector:selector sniffingBlock:^{
+            undefinedBlock(slf, connection, request, response);
+        } originalImplementationBlock:^{
+            returnValue = ((id(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, connection, request, response);
+        }];
         return returnValue;
     };
     
@@ -302,8 +324,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLConnectionDidReceiveResponseBlock implementationBlock = ^(id <NSURLConnectionDelegate> slf, NSURLConnection *connection, NSURLResponse *response) {
-        undefinedBlock(slf, connection, response);
-        ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, response);
+        [self sniffWithoutDuplicationForObject:connection selector:selector sniffingBlock:^{
+            undefinedBlock(slf, connection, response);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, response);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -328,8 +353,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLConnectionDidReceiveDataBlock implementationBlock = ^(id <NSURLConnectionDelegate> slf, NSURLConnection *connection, NSData *data) {
-        undefinedBlock(slf, connection, data);
-        ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, data);
+        [self sniffWithoutDuplicationForObject:connection selector:selector sniffingBlock:^{
+            undefinedBlock(slf, connection, data);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, data);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -354,8 +382,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLConnectionDidFinishLoadingBlock implementationBlock = ^(id <NSURLConnectionDelegate> slf, NSURLConnection *connection) {
-        undefinedBlock(slf, connection);
-        ((void(*)(id, SEL, id))objc_msgSend)(slf, swizzledSelector, connection);
+        [self sniffWithoutDuplicationForObject:connection selector:selector sniffingBlock:^{
+            undefinedBlock(slf, connection);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id))objc_msgSend)(slf, swizzledSelector, connection);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -376,8 +407,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLConnectionDidFailWithErrorBlock implementationBlock = ^(id <NSURLConnectionDelegate> slf, NSURLConnection *connection, NSError *error) {
-        undefinedBlock(slf, connection, error);
-        ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, error);
+        [self sniffWithoutDuplicationForObject:connection selector:selector sniffingBlock:^{
+            undefinedBlock(slf, connection, error);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id))objc_msgSend)(slf, swizzledSelector, connection, error);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -399,8 +433,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
 
     NSURLSessionWillPerformHTTPRedirectionBlock implementationBlock = ^(id <NSURLSessionTaskDelegate> slf, NSURLSession *session, NSURLSessionTask *task, NSHTTPURLResponse *response, NSURLRequest *newRequest, void(^completionHandler)(NSURLRequest *)) {
-        ((id(*)(id, SEL, id, id, id, id, void(^)()))objc_msgSend)(slf, swizzledSelector, session, task, response, newRequest, completionHandler);
-        undefinedBlock(slf, session, task, response, newRequest, completionHandler);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, task, response, newRequest, completionHandler);
+        } originalImplementationBlock:^{
+            ((id(*)(id, SEL, id, id, id, id, void(^)()))objc_msgSend)(slf, swizzledSelector, session, task, response, newRequest, completionHandler);
+        }];
     };
 
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -423,8 +460,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLSessionDidReceiveDataBlock implementationBlock = ^(id <NSURLSessionDataDelegate> slf, NSURLSession *session, NSURLSessionDataTask *dataTask, NSData *data) {
-        undefinedBlock(slf, session, dataTask, data);
-        ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, dataTask, data);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, dataTask, data);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, dataTask, data);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -447,8 +487,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
     
     NSURLSessionDidReceiveResponseBlock implementationBlock = ^(id <NSURLSessionDelegate> slf, NSURLSession *session, NSURLSessionDataTask *dataTask, NSURLResponse *response, void(^completionHandler)(NSURLSessionResponseDisposition disposition)) {
-        undefinedBlock(slf, session, dataTask, response, completionHandler);
-        ((void(*)(id, SEL, id, id, id, void(^)()))objc_msgSend)(slf, swizzledSelector, session, dataTask, response, completionHandler);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, dataTask, response, completionHandler);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id, id, void(^)()))objc_msgSend)(slf, swizzledSelector, session, dataTask, response, completionHandler);
+        }];
     };
     
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -470,8 +513,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
 
     NSURLSessionTaskDidCompleteWithErrorBlock implementationBlock = ^(id <NSURLSessionTaskDelegate> slf, NSURLSession *session, NSURLSessionTask *task, NSError *error) {
-        undefinedBlock(slf, session, task, error);
-        ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, task, error);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, task, error);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, task, error);
+        }];
     };
 
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -520,8 +566,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
 
     NSURLSessionDownloadTaskDidFinishDownloadingBlock implementationBlock = ^(id <NSURLSessionTaskDelegate> slf, NSURLSession *session, NSURLSessionDownloadTask *task, NSURL *location) {
-        undefinedBlock(slf, session, task, location);
-        ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, task, location);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, task, location);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id, id))objc_msgSend)(slf, swizzledSelector, session, task, location);
+        }];
     };
 
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];
@@ -542,8 +591,11 @@ NSString *const kFLEXNetworkObserverEnabledStateChangedNotification = @"kFLEXNet
     };
 
     NSURLSessionDownloadTaskDidWriteDataBlock implementationBlock = ^(id <NSURLSessionTaskDelegate> slf, NSURLSession *session, NSURLSessionDownloadTask *task, int64_t bytesWritten, int64_t totalBytesWritten, int64_t totalBytesExpectedToWrite) {
-        undefinedBlock(slf, session, task, bytesWritten, totalBytesWritten, totalBytesExpectedToWrite);
-        ((void(*)(id, SEL, id, id, int64_t, int64_t, int64_t))objc_msgSend)(slf, swizzledSelector, session, task, bytesWritten, totalBytesWritten, totalBytesExpectedToWrite);
+        [self sniffWithoutDuplicationForObject:session selector:selector sniffingBlock:^{
+            undefinedBlock(slf, session, task, bytesWritten, totalBytesWritten, totalBytesExpectedToWrite);
+        } originalImplementationBlock:^{
+            ((void(*)(id, SEL, id, id, int64_t, int64_t, int64_t))objc_msgSend)(slf, swizzledSelector, session, task, bytesWritten, totalBytesWritten, totalBytesExpectedToWrite);
+        }];
     };
 
     [self replaceImplementationOfSelector:selector withSelector:swizzledSelector forClass:cls withMethodDescription:methodDescription implementationBlock:implementationBlock undefinedBlock:undefinedBlock];