Swizzle UNUserNotificationCenterDelegate once

Ensure we only swizzle an UNUserNotificationCenterDelegate class once.
This was added to prevent an infinite loops when another library that
swizzles the same selectors is present in the app. See the previous
commit where this same logic was added for AppDelegate for more details.

iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.h

@@ -49,6 +49,7 @@
 - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
          didReceiveNotificationResponse:(UNNotificationResponse *)response
                   withCompletionHandler:(void(^)())completionHandler;
++ (void)traceCall:(NSString*)selector;
 @end
 
 

iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.m

@@ -83,16 +83,6 @@ + (void)setup {
     [OneSignalUNUserNotificationCenter registerDelegate];
 }
 
-static Class delegateUNClass = nil;
-
-// Store an array of all UIAppDelegate subclasses to iterate over in cases where UIAppDelegate swizzled methods are not overriden in main AppDelegate
-// But rather in one of the subclasses
-static NSArray* delegateUNSubclasses = nil;
-
-//ensures setDelegate: swizzles will never get executed twice for the same delegate object
-//captures a weak reference to avoid retain cycles
-__weak static id previousDelegate;
-
 + (void)swizzleSelectors {
     injectSelector(
         [UNUserNotificationCenter class],
@@ -186,16 +176,27 @@ - (void)onesignalGetNotificationSettingsWithCompletionHandler:(void(^)(UNNotific
     [self onesignalGetNotificationSettingsWithCompletionHandler:wrapperBlock];
 }
 
+
+// A Set to keep track of which classes we have already swizzled so we only
+// swizzle each one once. If we swizzled more than once then this will create
+// an infinite loop, this includes swizzling with ourselves but also with
+// another SDK that swizzles.
+static NSMutableSet<Class>* swizzledClasses;
+
 // Take the received delegate and swizzle in our own hooks.
 //  - Selector will be called once if developer does not set a UNUserNotificationCenter delegate.
 //  - Selector will be called a 2nd time if the developer does set one.
 - (void) setOneSignalUNDelegate:(id)delegate {
-    if (previousDelegate == delegate) {
+    if (swizzledClasses == nil)
+        swizzledClasses = [NSMutableSet new];
+
+    Class delegateClass = [delegate class];
+
+    if (delegate == nil || [swizzledClasses containsObject:delegateClass]) {
         [self setOneSignalUNDelegate:delegate];
         return;
     }
-
-    previousDelegate = delegate;
+    [swizzledClasses addObject:delegateClass];
 
     [OneSignal onesignalLog:ONE_S_LL_VERBOSE message:@"OneSignalUNUserNotificationCenter setOneSignalUNDelegate Fired!"];
 
@@ -206,7 +207,7 @@ - (void) setOneSignalUNDelegate:(id)delegate {
 }
 
 + (void)swizzleSelectorsOnDelegate:(id)delegate {
-    delegateUNClass = [delegate class];
+    Class delegateUNClass = [delegate class];
     injectSelector(
         delegateUNClass,
         @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:),
@@ -256,6 +257,8 @@ - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
                 willPresentNotification:(UNNotification *)notification
                   withCompletionHandler:(void (^)(UNNotificationPresentationOptions options))completionHandler {
 
+    [OneSignalUNUserNotificationCenter traceCall:@"onesignalUserNotificationCenter:willPresentNotification:withCompletionHandler:"];
+
     // return if the user has not granted privacy permissions or if not a OneSignal payload
     if ([OSPrivacyConsentController shouldLogMissingPrivacyConsentErrorWithMethodName:nil] || ![OneSignalHelper isOneSignalPayload:notification.request.content.userInfo]) {
         BOOL hasReceiver = [OneSignalUNUserNotificationCenter forwardNotificationWithCenter:center notification:notification OneSignalCenter:self completionHandler:completionHandler];
@@ -302,6 +305,7 @@ void finishProcessingNotification(UNNotification *notification,
 - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
          didReceiveNotificationResponse:(UNNotificationResponse *)response
                   withCompletionHandler:(void(^)())completionHandler {
+    [OneSignalUNUserNotificationCenter traceCall:@"onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:"];
     // return if the user has not granted privacy permissions or if not a OneSignal payload
     if ([OSPrivacyConsentController shouldLogMissingPrivacyConsentErrorWithMethodName:nil] || ![OneSignalHelper isOneSignalPayload:response.notification.request.content.userInfo]) {
         SwizzlingForwarder *forwarder = [[SwizzlingForwarder alloc]
@@ -434,6 +438,12 @@ The iOS SDK used to call some local notification selectors (such as didReceiveLo
         completionHandler();
 }
 
+// Used to log all calls, also used in unit tests to observer
+// the OneSignalUserNotificationCenter selectors get called.
++(void) traceCall:(NSString*)selector {
+    [OneSignal onesignalLog:ONE_S_LL_VERBOSE message:selector];
+}
+
 @end
 
 #pragma clang diagnostic pop

iOS_SDK/OneSignalSDK/UnitTests/OneSignalUNUserNotificationCenterSwizzlingTest.m

@@ -9,6 +9,8 @@
 #import "OneSignalHelperOverrider.h"
 #import "OneSignalHelper.h"
 #import "OneSignalUNUserNotificationCenterHelper.h"
+#import "TestHelperFunctions.h"
+#import "OneSignalUNUserNotificationCenterOverrider.h"
 
 @interface DummyNotificationCenterDelegateForDoesSwizzleTest : NSObject<UNUserNotificationCenterDelegate>
 @end
@@ -123,6 +125,41 @@ @interface UNUserNotificationCenterDelegateForInfiniteLoopTest : UIResponder<UNU
 @implementation UNUserNotificationCenterDelegateForInfiniteLoopTest
 @end
 
+
+@interface UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTest : UIResponder<UNUserNotificationCenterDelegate>
+@end
+@implementation UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTest
+@end
+@interface OtherUNNotificationLibraryASwizzler : NSObject
++(void)swizzleAppDelegate;
++(BOOL)selectorCalled;
+@end
+@implementation OtherUNNotificationLibraryASwizzler
+static BOOL selectorCalled = false;
++(BOOL)selectorCalled {
+    return selectorCalled;
+}
+
++(void)swizzleUNUserNotificationCenterDelegate
+{
+    swizzleExistingSelector(
+        [UNUserNotificationCenter.currentNotificationCenter.delegate class],
+        @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:),
+        [self class],
+        @selector(userNotificationCenterLibraryA:willPresentNotification:withCompletionHandler:)
+    );
+}
+-(void)userNotificationCenterLibraryA:(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 userNotificationCenterLibraryA:center willPresentNotification:notification withCompletionHandler:completionHandler];
+}
+@end
+
+
+
 @interface OneSignalUNUserNotificationCenterSwizzlingTest : XCTestCase
 @end
 
@@ -282,6 +319,35 @@ - (void)testDoubleSwizzleInfiniteLoop {
     [localOrignalDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter willPresentNotification:[self createBasiciOSNotification] withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
 }
 
+- (void)testCompatibleWithOtherSwizzlerWhenSwapingBetweenNil {
+    // 1. Create a new delegate and assign it
+    id myAppDelegate = [UNUserNotificationCenterDelegateForInfiniteLoopWithAnotherSwizzlerTest new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myAppDelegate;
+
+    // 2. Other library swizzles
+    [OtherUNNotificationLibraryASwizzler swizzleUNUserNotificationCenterDelegate];
+
+    // 3. Nil and set it again to trigger OneSignal swizzling again.
+    UNUserNotificationCenter.currentNotificationCenter.delegate = nil;
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myAppDelegate;
+
+    // 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]);
+}
+
 - (void)testSwizzleExistingSelectors {
     UNUserNotificationCenterDelegateForExistingSelectorsTest* myNotifCenterDelegate = [UNUserNotificationCenterDelegateForExistingSelectorsTest new];
     UNUserNotificationCenter.currentNotificationCenter.delegate = myNotifCenterDelegate;

iOS_SDK/OneSignalSDK/UnitTests/Shadows/OneSignalUNUserNotificationCenterOverrider.h

@@ -0,0 +1,10 @@
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface OneSignalUNUserNotificationCenterOverrider : NSObject
++ (int)callCountForSelector:(NSString*)selector;
++ (void)reset;
+@end
+
+NS_ASSUME_NONNULL_END

iOS_SDK/OneSignalSDK/UnitTests/Shadows/OneSignalUNUserNotificationCenterOverrider.m

@@ -0,0 +1,32 @@
+#import "OneSignalUNUserNotificationCenterOverrider.h"
+
+#import "TestHelperFunctions.h"
+#import "UNUserNotificationCenter+OneSignal.h"
+
+@implementation OneSignalUNUserNotificationCenterOverrider
+static NSMutableDictionary* callCount;
+
++ (void)load {
+    callCount = [NSMutableDictionary new];
+
+    injectStaticSelector(
+        [OneSignalUNUserNotificationCenterOverrider class],
+        @selector(overrideTraceCall:),
+        [OneSignalUNUserNotificationCenter class],
+        @selector(traceCall:)
+    );
+}
+
++ (void)reset {
+    callCount = [NSMutableDictionary new];
+}
+
++ (void)overrideTraceCall:(NSString*)selector {
+    NSNumber *value = callCount[selector];
+    callCount[selector] = [NSNumber numberWithInt:[value intValue] + 1];
+}
+
++ (int)callCountForSelector:(NSString*)selector {
+    return [callCount[selector] intValue];
+}
+@end

iOS_SDK/OneSignalSDK/UnitTests/UnitTestCommonMethods.m

@@ -61,6 +61,7 @@
 #import "OneSignalUserDefaults.h"
 #import "OneSignalLog.h"
 #import "OneSignalAppDelegateOverrider.h"
+#import "OneSignalUNUserNotificationCenterOverrider.h"
 
 NSString * serverUrlWithPath(NSString *path) {
     return [OS_API_SERVER_URL stringByAppendingString:path];
@@ -235,6 +236,7 @@ + (void)clearStateForAppRestart:(XCTestCase *)testCase {
 
     [OneSignalLifecycleObserver removeObserver];
     [OneSignalAppDelegateOverrider reset];
+    [OneSignalUNUserNotificationCenterOverrider reset];
 }
 
 + (void)beforeAllTest:(XCTestCase *)testCase {