Merge pull request #1284 from OneSignal/Fix/swizzling_subclass_of_already_swizzled_class

Fix swizzling subclass of already swizzled class

Author: emawby

iOS_SDK/OneSignalSDK/Source/UIApplicationDelegate+OneSignal.m

@@ -38,6 +38,7 @@
 #import "OneSignalHelper.h"
 #import "OSMessagingController.h"
 #import "SwizzlingForwarder.h"
+#import <objc/runtime.h>
 
 @interface OneSignal (UN_extra)
 + (void) didRegisterForRemoteNotifications:(UIApplication*)app deviceToken:(NSData*)inDeviceToken;
@@ -75,7 +76,7 @@ - (void) setOneSignalDelegate:(id<UIApplicationDelegate>)delegate {
 
     Class delegateClass = [delegate class];
 
-    if (delegate == nil || [swizzledClasses containsObject:delegateClass]) {
+    if (delegate == nil || [OneSignalAppDelegate swizzledClassInHeirarchy:delegateClass]) {
         [self setOneSignalDelegate:delegate];
         return;
     }
@@ -120,6 +121,22 @@ - (void) setOneSignalDelegate:(id<UIApplicationDelegate>)delegate {
     [self setOneSignalDelegate:delegate];
 }
 
++ (BOOL)swizzledClassInHeirarchy:(Class)delegateClass {
+    if ([swizzledClasses containsObject:delegateClass]) {
+        [OneSignalLog onesignalLog:ONE_S_LL_VERBOSE message:[NSString stringWithFormat:@"OneSignal already swizzled %@", NSStringFromClass(delegateClass)]];
+        return true;
+    }
+    Class superClass = class_getSuperclass(delegateClass);
+    while(superClass) {
+        if ([swizzledClasses containsObject:superClass]) {
+            [OneSignalLog onesignalLog:ONE_S_LL_VERBOSE message:[NSString stringWithFormat:@"OneSignal already swizzled %@ in super class: %@", NSStringFromClass(delegateClass), NSStringFromClass(superClass)]];
+            return true;
+        }
+        superClass = class_getSuperclass(superClass);
+    }
+    return false;
+}
+
 + (void)swizzlePreiOS10Methods:(Class)delegateClass {
     if ([OneSignalHelper isIOSVersionGreaterThanOrEqual:@"10.0"])
         return;

iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.m

@@ -37,6 +37,7 @@
 #import "UIApplicationDelegate+OneSignal.h"
 #import "OneSignalCommonDefines.h"
 #import "SwizzlingForwarder.h"
+#import <objc/runtime.h>
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wundeclared-selector"
 
@@ -192,7 +193,7 @@ - (void) setOneSignalUNDelegate:(id)delegate {
 
     Class delegateClass = [delegate class];
 
-    if (delegate == nil || [swizzledClasses containsObject:delegateClass]) {
+    if (delegate == nil || [OneSignalUNUserNotificationCenter swizzledClassInHeirarchy:delegateClass]) {
         [self setOneSignalUNDelegate:delegate];
         return;
     }
@@ -206,6 +207,22 @@ - (void) setOneSignalUNDelegate:(id)delegate {
     [self setOneSignalUNDelegate:delegate];
 }
 
++ (BOOL)swizzledClassInHeirarchy:(Class)delegateClass {
+    if ([swizzledClasses containsObject:delegateClass]) {
+        [OneSignalLog onesignalLog:ONE_S_LL_VERBOSE message:[NSString stringWithFormat:@"OneSignal already swizzled %@", NSStringFromClass(delegateClass)]];
+        return true;
+    }
+    Class superClass = class_getSuperclass(delegateClass);
+    while(superClass) {
+        if ([swizzledClasses containsObject:superClass]) {
+            [OneSignalLog onesignalLog:ONE_S_LL_VERBOSE message:[NSString stringWithFormat:@"OneSignal already swizzled %@ in super class: %@", NSStringFromClass(delegateClass), NSStringFromClass(superClass)]];
+            return true;
+        }
+        superClass = class_getSuperclass(superClass);
+    }
+    return false;
+}
+
 + (void)swizzleSelectorsOnDelegate:(id)delegate {
     Class delegateUNClass = [delegate class];
     injectSelector(

iOS_SDK/OneSignalSDK/UnitTests/OneSignalUNUserNotificationCenterSwizzlingTest.m

@@ -130,7 +130,7 @@ @interface UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTes
 @end
 @implementation UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTest
 @end
-@interface OtherUNNotificationLibraryASwizzler : NSObject
+@interface OtherUNNotificationLibraryASwizzler : UIResponder<UNUserNotificationCenterDelegate>
 +(void)swizzleUNUserNotificationCenterDelegate;
 +(BOOL)selectorCalled;
 @end
@@ -157,7 +157,33 @@ -(void)userNotificationCenterLibraryA:(UNUserNotificationCenter *)center willPre
         [self userNotificationCenterLibraryA:center willPresentNotification:notification withCompletionHandler:completionHandler];
 }
 @end
+@interface OtherUNNotificationLibraryBSubClassSwizzler : OtherUNNotificationLibraryASwizzler
++(void)swizzleUNUserNotificationCenterDelegate;
++(BOOL)selectorCalled;
+@end
+@implementation OtherUNNotificationLibraryBSubClassSwizzler
 
++(BOOL)selectorCalled {
+    return selectorCalled;
+}
+
++(void)swizzleUNUserNotificationCenterDelegate
+{
+    swizzleExistingSelector(
+        [UNUserNotificationCenter.currentNotificationCenter.delegate class],
+        @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:),
+        [self class],
+        @selector(userNotificationCenterLibraryB:willPresentNotification:withCompletionHandler:)
+    );
+}
+-(void)userNotificationCenterLibraryB:(UNUserNotificationCenter *)center willPresentNotification:(UNNotification *)notification withCompletionHandler:(void (^)(UNNotificationPresentationOptions))completionHandler
+{
+    selectorCalled = true;
+    // Standard basic swizzling forwarder another library may have.
+    if ([self respondsToSelector:@selector(userNotificationCenterLibraryA:willPresentNotification:withCompletionHandler:)])
+        [self userNotificationCenterLibraryB:center willPresentNotification:notification withCompletionHandler:completionHandler];
+}
+@end
 
 
 @interface OneSignalUNUserNotificationCenterSwizzlingTest : XCTestCase
@@ -319,6 +345,41 @@ - (void)testDoubleSwizzleInfiniteLoop {
     [localOrignalDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter willPresentNotification:[self createBasiciOSNotification] withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
 }
 
+- (void)testNotificationCenterSubClassIsNotSwizzledTwice {
+    // 1. Create a new delegate and assign it
+    id myDelegate = [UNUserNotificationCenterDelegateForInfiniteLoopTest new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myDelegate;
+    
+    // 2. Create another Library's app delegate and assign it then swizzle
+    id thierDelegate = [OtherUNNotificationLibraryASwizzler new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = thierDelegate;
+    [OtherUNNotificationLibraryASwizzler swizzleUNUserNotificationCenterDelegate];
+    
+    // 3. Create another Library's app delegate subclass and assign it then swizzle
+    id thierDelegateSubClass = [OtherUNNotificationLibraryBSubClassSwizzler new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = thierDelegateSubClass;
+    [OtherUNNotificationLibraryBSubClassSwizzler swizzleUNUserNotificationCenterDelegate];
+    
+    // 4. Call something to confirm we don't get stuck in an infinite call loop
+    id<UNUserNotificationCenterDelegate> delegate =
+        UNUserNotificationCenter.currentNotificationCenter.delegate;
+    [delegate
+        userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+        willPresentNotification:[self createBasiciOSNotification]
+        withCompletionHandler:^(UNNotificationPresentationOptions options) {}
+    ];
+    
+    // 5. Ensure OneSignal's selector is called.
+    XCTAssertEqual([OneSignalUNUserNotificationCenterOverrider
+        callCountForSelector:@"onesignalUserNotificationCenter:willPresentNotification:withCompletionHandler:"], 1);
+    
+    // 6. Ensure other library selector is still called too.
+    XCTAssertTrue([OtherUNNotificationLibraryASwizzler selectorCalled]);
+    
+    // 7. Ensure other library subclass selector is still called too.
+    XCTAssertTrue([OtherUNNotificationLibraryBSubClassSwizzler selectorCalled]);
+}
+
 - (void)testCompatibleWithOtherSwizzlerWhenSwapingBetweenNil {
     // 1. Create a new delegate and assign it
     id myAppDelegate = [UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTest new];

iOS_SDK/OneSignalSDK/UnitTests/UIApplicationDelegateSwizzlingTests.m

@@ -7,6 +7,8 @@
 #import "TestHelperFunctions.h"
 #import "OneSignalAppDelegateOverrider.h"
 
+#define ONESIGNALApplicationDelegate UIApplicationDelegate
+
 @interface AppDelegateForAddsMissingSelectorsTest : UIResponder<UIApplicationDelegate>
 @end
 @implementation AppDelegateForAddsMissingSelectorsTest
@@ -90,7 +92,8 @@ @interface AppDelegateForInfiniteLoopWithAnotherSwizzlerTest : UIResponder<UIApp
 @end
 @implementation AppDelegateForInfiniteLoopWithAnotherSwizzlerTest
 @end
-@interface OtherLibraryASwizzler : NSObject
+
+@interface OtherLibraryASwizzler : UIResponder<ONESIGNALApplicationDelegate>
 +(void)swizzleAppDelegate;
 +(BOOL)selectorCalled;
 @end
@@ -109,6 +112,7 @@ +(void)swizzleAppDelegate
         @selector(applicationWillTerminateLibraryA:)
     );
 }
+
 - (void)applicationWillTerminateLibraryA:(UIApplication *)application
 {
     selectorCalled = true;
@@ -118,6 +122,33 @@ - (void)applicationWillTerminateLibraryA:(UIApplication *)application
 }
 @end
 
+
+@interface OtherLibraryBSwizzlerSubClass : OtherLibraryASwizzler
+@end
+@implementation OtherLibraryBSwizzlerSubClass
++(BOOL)selectorCalled {
+    return selectorCalled;
+}
+
++(void)swizzleAppDelegate
+{
+    swizzleExistingSelector(
+        [UIApplication.sharedApplication.delegate class],
+        @selector(applicationWillTerminate:),
+        [self class],
+        @selector(applicationWillTerminateLibraryB:)
+    );
+}
+
+- (void)applicationWillTerminateLibraryB:(UIApplication *)application
+{
+    selectorCalled = true;
+    // Standard basic swizzling forwarder another library may have.
+    if ([self respondsToSelector:@selector(applicationWillTerminateLibraryB:)])
+        [self applicationWillTerminateLibraryB:application];
+}
+@end
+
 @interface AppDelegateForExistingSelectorsTest : UIResponder<UIApplicationDelegate> {
     @public NSMutableDictionary *selectorCallsDict;
 }
@@ -408,6 +439,36 @@ - (void)testDoubleSwizzleInfiniteLoop {
     [localOrignalDelegate applicationWillTerminate:UIApplication.sharedApplication];
 }
 
+- (void)testAppDelegateSubClassIsNotSwizzledTwice {
+    // 1. Create a new delegate and assign it
+    id myAppDelegate = [AppDelegateForInfiniteLoopWithAnotherSwizzlerTest new];
+    UIApplication.sharedApplication.delegate = myAppDelegate;
+    
+    // 2. Create another Library's app delegate and assign it then swizzle
+    id thierAppDelegate = [OtherLibraryASwizzler new];
+    UIApplication.sharedApplication.delegate = thierAppDelegate;
+    [OtherLibraryASwizzler swizzleAppDelegate];
+    
+    // 3. Create another Library's app delegate subclass and assign it then swizzle
+    id thierAppDelegateSubClass = [OtherLibraryBSwizzlerSubClass new];
+    UIApplication.sharedApplication.delegate = thierAppDelegateSubClass;
+    [OtherLibraryBSwizzlerSubClass swizzleAppDelegate];
+    
+    // 4. Call something to confirm we don't get stuck in an infinite call loop
+    id<UIApplicationDelegate> delegate = UIApplication.sharedApplication.delegate;
+    [delegate applicationWillTerminate:UIApplication.sharedApplication];
+    
+    // 5. Ensure OneSignal's selector is called.
+    XCTAssertEqual([OneSignalAppDelegateOverrider
+        callCountForSelector:@"oneSignalApplicationWillTerminate:"], 1);
+    
+    // 6. Ensure other library selector is still called too.
+    XCTAssertTrue([OtherLibraryASwizzler selectorCalled]);
+    
+    // 7. Ensure other library subclass selector is still called too.
+    XCTAssertTrue([OtherLibraryBSwizzlerSubClass selectorCalled]);
+}
+
 - (void)testCompatibleWithOtherSwizzlerWhenSwapingBetweenNil {
     // 1. Create a new delegate and assign it
     id myAppDelegate = [AppDelegateForInfiniteLoopWithAnotherSwizzlerTest new];
@@ -624,4 +685,5 @@ - (void)testAppDelegateInheritsFromBaseWhereBothHaveSelectorsButSuperIsNotCalled
     XCTAssertFalse(myAppDelegate.selectorCalledOnParent);
     XCTAssertEqual([OneSignalAppDelegateOverrider callCountForSelector:@"oneSignalReceiveRemoteNotification:UserInfo:fetchCompletionHandler:"], 1);
 }
+
 @end