Adding subClass swizzling check to UNUserNotificationCenter

Author: emawby

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

@@ -466,7 +466,7 @@ - (void)testAppDelegateSubClassIsNotSwizzledTwice {
     XCTAssertTrue([OtherLibraryASwizzler selectorCalled]);
 
     // 7. Ensure other library subclass selector is still called too.
-    XCTAssertTrue([OtherLibraryASwizzler selectorCalled]);
+    XCTAssertTrue([OtherLibraryBSwizzlerSubClass selectorCalled]);
 }
 
 - (void)testCompatibleWithOtherSwizzlerWhenSwapingBetweenNil {