Do not swizzle a subclass of an already swizzled class

Adding unit test for swizzling subclass

Adding subClass swizzling check to UNUserNotificationCenter

Don't swizzle subclass if we have already swizzled its super

Author: emawby

iOS_SDK/OneSignalSDK/OneSignalNotifications/Categories/UIApplicationDelegate+OneSignalNotifications.m

@@ -34,6 +34,7 @@
 #import "OneSignalSelectorHelpers.h"
 #import "SwizzlingForwarder.h"
 #import "OSNotificationsManager.h"
+#import <objc/runtime.h>
 
 // This class hooks into the UIApplicationDelegate selectors to receive iOS 9 and older events.
 //   - UNUserNotificationCenter is used for iOS 10
@@ -58,7 +59,7 @@ - (void) setOneSignalDelegate:(id<UIApplicationDelegate>)delegate {
 
     Class delegateClass = [delegate class];
 
-    if (delegate == nil || [swizzledClasses containsObject:delegateClass]) {
+    if (delegate == nil || [OneSignalNotificationsAppDelegate swizzledClassInHeirarchy:delegateClass]) {
         [self setOneSignalDelegate:delegate];
         return;
     }
@@ -93,6 +94,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)oneSignalDidRegisterForRemoteNotifications:(UIApplication*)app deviceToken:(NSData*)inDeviceToken {
     [OneSignalNotificationsAppDelegate traceCall:@"oneSignalDidRegisterForRemoteNotifications:deviceToken:"];
 

iOS_SDK/OneSignalSDK/OneSignalNotifications/Categories/UNUserNotificationCenter+OneSignalNotifications.m

@@ -33,6 +33,7 @@
 #import "UIApplicationDelegate+OneSignalNotifications.h"
 #import "OSNotificationsManager.h"
 #import <OneSignalCore/OneSignalCore.h>
+#import <objc/runtime.h>
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wundeclared-selector"
 
@@ -182,7 +183,7 @@ - (void) setOneSignalUNDelegate:(id)delegate {
 
     Class delegateClass = [delegate class];
 
-    if (delegate == nil || [swizzledClasses containsObject:delegateClass]) {
+    if (delegate == nil || [OneSignalNotificationsUNUserNotificationCenter swizzledClassInHeirarchy:delegateClass]) {
         [self setOneSignalUNDelegate:delegate];
         return;
     }
@@ -196,6 +197,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/Source/UIApplicationDelegate+OneSignal.m

@@ -34,6 +34,7 @@
 #import "OneSignalTracker.h"
 #import "OneSignalSelectorHelpers.h"
 #import "SwizzlingForwarder.h"
+#import <objc/runtime.h>
 
 @interface OneSignal (UN_extra)
 + (NSString*) appId;
@@ -61,7 +62,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;
     }
@@ -79,6 +80,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)oneSignalApplicationWillTerminate:(UIApplication *)application {
     [OneSignalAppDelegate traceCall:@"oneSignalApplicationWillTerminate:"];
 

iOS_SDK/OneSignalSDK/UnitTests/InAppMessagingIntegrationTests.m

@@ -26,7 +26,7 @@ of this software and associated documentation files (the "Software"), to deal
  */
 
 #import <XCTest/XCTest.h>
-#import "OneSignal.h"
+#import "OneSignalFramework.h"
 #import "OneSignalUserDefaults.h"
 #import "OneSignalHelper.h"
 #import "OSInAppMessageInternal.h"

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];
@@ -579,4 +640,5 @@ - (void)testAppDelegateInheritsFromBaseWhereBothHaveSelectorsButSuperIsNotCalled
     XCTAssertFalse(myAppDelegate.selectorCalledOnParent);
     XCTAssertEqual([OneSignalAppDelegateOverrider callCountForSelector:@"oneSignalReceiveRemoteNotification:UserInfo:fetchCompletionHandler:"], 1);
 }
+
 @end