Merge pull request #1097 from OneSignal/swizzle_forward_target_userno…

…tificationcenter Use Forward Target Swizzling check for UNUserNotificationCenterDelegate methods
OneSignal · May 25, 2022 · 2879e22 · 2879e22
2 parents 0d4bc5c + 7cb0c50
commit 2879e22
Show file tree

Hide file tree

Showing 2 changed files with 265 additions and 14 deletions.
diff --git a/iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.m b/iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.m
@@ -36,7 +36,7 @@
 #import "OneSignalSelectorHelpers.h"
 #import "UIApplicationDelegate+OneSignal.h"
 #import "OneSignalCommonDefines.h"
-
+#import "SwizzlingForwarder.h"
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wundeclared-selector"
 
@@ -205,21 +205,31 @@ + (void)swizzleSelectorsOnDelegate:(id)delegate {
                         @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:), delegateUNSubclasses, [OneSignalUNUserNotificationCenter class], delegateUNClass);
 }
 
-+ (void)forwardNotificationWithCenter:(UNUserNotificationCenter *)center
++ (BOOL)forwardNotificationWithCenter:(UNUserNotificationCenter *)center
                          notification:(UNNotification *)notification
                       OneSignalCenter:(id)instance
                     completionHandler:(void (^)(UNNotificationPresentationOptions options))completionHandler {
-    // Call orginal selector if one was set.
-    if ([instance respondsToSelector:@selector(onesignalUserNotificationCenter:willPresentNotification:withCompletionHandler:)])
-        [instance onesignalUserNotificationCenter:center willPresentNotification:notification withCompletionHandler:completionHandler];
-    // Or call a legacy AppDelegate selector
-    else {
+    SwizzlingForwarder *forwarder = [[SwizzlingForwarder alloc]
+        initWithTarget:instance
+        withYourSelector:@selector(
+            onesignalUserNotificationCenter:willPresentNotification:withCompletionHandler:
+        )
+        withOriginalSelector:@selector(
+            userNotificationCenter:willPresentNotification:withCompletionHandler:
+        )
+    ];
+    if (forwarder.hasReceiver) {
+        [forwarder invokeWithArgs:@[center, notification, completionHandler]];
+        return true;
+    } else {
+        // call a legacy AppDelegate selector
         [OneSignalUNUserNotificationCenter callLegacyAppDeletegateSelector:notification
                                                 isTextReply:false
                                            actionIdentifier:nil
                                                    userText:nil
                                     fromPresentNotification:true
                                       withCompletionHandler:^() {}];
+        return false;
     }
 }
 
@@ -232,8 +242,8 @@ - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
 
     // 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]) {
-        [OneSignalUNUserNotificationCenter forwardNotificationWithCenter:center notification:notification OneSignalCenter:self completionHandler:completionHandler];
-        if (![self respondsToSelector:@selector(onesignalUserNotificationCenter:willPresentNotification:withCompletionHandler:)]) {
+        BOOL hasReceiver = [OneSignalUNUserNotificationCenter forwardNotificationWithCenter:center notification:notification OneSignalCenter:self completionHandler:completionHandler];
+        if (!hasReceiver) {
             completionHandler(7);
         }
         return;
@@ -278,10 +288,20 @@ - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
                   withCompletionHandler:(void(^)())completionHandler {
     // 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]) {
-        if ([self respondsToSelector:@selector(onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:)])
-            [self onesignalUserNotificationCenter:center didReceiveNotificationResponse:response withCompletionHandler:completionHandler];
-        else
+        SwizzlingForwarder *forwarder = [[SwizzlingForwarder alloc]
+            initWithTarget:self
+            withYourSelector:@selector(
+                onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:
+            )
+            withOriginalSelector:@selector(
+                userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:
+            )
+        ];
+        if (forwarder.hasReceiver) {
+            [forwarder invokeWithArgs:@[center, response, completionHandler]];
+        } else {
             completionHandler();
+        }
         return;
     }
 
@@ -290,8 +310,18 @@ - (void)onesignalUserNotificationCenter:(UNUserNotificationCenter *)center
     [OneSignalUNUserNotificationCenter processiOS10Open:response];
 
     // Call orginal selector if one was set.
-    if ([self respondsToSelector:@selector(onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:)])
-        [self onesignalUserNotificationCenter:center didReceiveNotificationResponse:response withCompletionHandler:completionHandler];
+    SwizzlingForwarder *forwarder = [[SwizzlingForwarder alloc]
+        initWithTarget:self
+        withYourSelector:@selector(
+            onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:
+        )
+        withOriginalSelector:@selector(
+            userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:
+        )
+    ];
+    if (forwarder.hasReceiver) {
+        [forwarder invokeWithArgs:@[center, response, completionHandler]];
+    }
     // Or call a legacy AppDelegate selector
     //  - If not a dismiss event as their isn't a iOS 9 selector for it.
     else if (![OneSignalUNUserNotificationCenter isDismissEvent:response]) {

diff --git a/iOS_SDK/OneSignalSDK/UnitTests/OneSignalUNUserNotificationCenterSwizzlingTest.m b/iOS_SDK/OneSignalSDK/UnitTests/OneSignalUNUserNotificationCenterSwizzlingTest.m
@@ -30,6 +30,98 @@ -(void)userNotificationCenter:(UNUserNotificationCenter *)center didReceiveNotif
 }
 @end
 
+@interface UNUserNotificationCenterDelegateForwardingTargetForSelectorTest : UIResponder<UNUserNotificationCenterDelegate>
+@end
+@implementation UNUserNotificationCenterDelegateForwardingTargetForSelectorTest {
+    id forwardingInstance;
+}
+- (instancetype)initForwardingTarget:(id)forwardingTarget {
+    self = [super init];
+    forwardingInstance = forwardingTarget;
+    return self;
+}
+
+- (id)forwardingTargetForSelector:(SEL)selector {
+    return forwardingInstance;
+}
+@end
+
+@interface UNUserNotificationCenterDelegateForwardReceiver : UIResponder<UNUserNotificationCenterDelegate> {
+    @public NSMutableDictionary *selectorCallsDict;
+}
+@end
+@implementation UNUserNotificationCenterDelegateForwardReceiver
+
+- (instancetype)init {
+    self = [super init];
+    selectorCallsDict = [NSMutableDictionary new];
+    return self;
+}
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center willPresentNotification:(UNNotification *)notification withCompletionHandler:(void (^)(UNNotificationPresentationOptions))completionHandler {
+    SEL thisSelector = @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center didReceiveNotificationResponse:(UNNotificationResponse *)response withCompletionHandler:(void (^)(void))completionHandler {
+    SEL thisSelector = @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center openSettingsForNotification:(UNNotification *)notification {
+    SEL thisSelector = @selector(userNotificationCenter:openSettingsForNotification:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+@end
+
+@interface UNUserNotificationCenterDelegateForExistingSelectorsTest : UIResponder<UNUserNotificationCenterDelegate> {
+    @public NSMutableDictionary *selectorCallsDict;
+}
+@end
+@implementation UNUserNotificationCenterDelegateForExistingSelectorsTest
+
+- (instancetype)init {
+    self = [super init];
+    selectorCallsDict = [NSMutableDictionary new];
+    return self;
+}
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center willPresentNotification:(UNNotification *)notification withCompletionHandler:(void (^)(UNNotificationPresentationOptions))completionHandler {
+    SEL thisSelector = @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center didReceiveNotificationResponse:(UNNotificationResponse *)response withCompletionHandler:(void (^)(void))completionHandler {
+    SEL thisSelector = @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center openSettingsForNotification:(UNNotification *)notification {
+    SEL thisSelector = @selector(userNotificationCenter:openSettingsForNotification:);
+    [selectorCallsDict
+       setObject:@(true)
+       forKey:NSStringFromSelector(thisSelector)
+    ];
+}
+@end
+
+@interface UNUserNotificationCenterDelegateForInfiniteLoopTest : UIResponder<UNUserNotificationCenterDelegate>
+@end
+@implementation UNUserNotificationCenterDelegateForInfiniteLoopTest
+@end
 
 @interface OneSignalUNUserNotificationCenterSwizzlingTest : XCTestCase
 @end
@@ -50,6 +142,37 @@ - (void)tearDown {
     [OneSignalUNUserNotificationCenterHelper restoreDelegateAsOneSignal];
 }
 
+- (UNNotificationResponse *)createBasiciOSNotificationResponse {
+  id userInfo = @{@"custom":
+                       @{ @"i": @"b2f7f966-d8cc-11e4-bed1-df8f05be55ba" }
+                };
+
+  return [UnitTestCommonMethods createBasiciOSNotificationResponseWithPayload:userInfo];
+}
+
+- (UNNotification *)createNonOneSignaliOSNotification {
+    id userInfo = @{@"aps": @{
+                            @"mutable-content": @1,
+                            @"alert": @"Message Body"
+                            }
+                    };
+
+    return [UnitTestCommonMethods createBasiciOSNotificationWithPayload:userInfo];
+}
+
+- (UNNotification *)createBasiciOSNotification {
+    id userInfo = @{@"aps": @{
+                            @"mutable-content": @1,
+                            @"alert": @"Message Body"
+                            },
+                    @"os_data": @{
+                            @"i": @"b2f7f966-d8cc-11e4-bed1-df8f05be55ba",
+                            @"buttons": @[@{@"i": @"id1", @"n": @"text1"}],
+                            }};
+
+    return [UnitTestCommonMethods createBasiciOSNotificationWithPayload:userInfo];
+}
+
 // Tests to make sure that UNNotificationCenter setDelegate: duplicate calls don't double-swizzle for the same object
 - (void)testAUNUserNotificationCenterDelegateAssigningDoesSwizzle {
     UNUserNotificationCenter *center = [UNUserNotificationCenter currentNotificationCenter];
@@ -93,4 +216,102 @@ - (void)testUNUserNotificationCenterDelegateAssignedBeforeOneSignal {
     XCTAssertNotEqual(originalDummyImp, swizzledDummyImp);
 }
 
+- (void)testForwardingTargetForSelector {
+    UNUserNotificationCenterDelegateForwardReceiver *receiver = [UNUserNotificationCenterDelegateForwardReceiver new];
+    id myNotifCenterDelegate = [[UNUserNotificationCenterDelegateForwardingTargetForSelectorTest alloc]
+                        initForwardingTarget:receiver];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myNotifCenterDelegate;
+    id<UNUserNotificationCenterDelegate> notifCenterDelegate = UNUserNotificationCenter.currentNotificationCenter.delegate;
+
+    [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                        willPresentNotification:[self createBasiciOSNotification]
+                          withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
+    XCTAssertTrue([receiver->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:)
+        )
+    ]);
+    [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                 didReceiveNotificationResponse:[self createBasiciOSNotificationResponse]
+                          withCompletionHandler:^{}];
+    XCTAssertTrue([receiver->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:)
+        )
+    ]);
+    if (@available(iOS 12.0, *)) {
+        [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                        openSettingsForNotification:[self createBasiciOSNotification]];
+    }
+    XCTAssertTrue([receiver->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:openSettingsForNotification:)
+        )
+    ]);
+}
+
+- (void)testForwardingTargetForNonOneSignalNotification {
+    UNUserNotificationCenterDelegateForwardReceiver *receiver = [UNUserNotificationCenterDelegateForwardReceiver new];
+    id myNotifCenterDelegate = [[UNUserNotificationCenterDelegateForwardingTargetForSelectorTest alloc]
+                        initForwardingTarget:receiver];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myNotifCenterDelegate;
+    id<UNUserNotificationCenterDelegate> notifCenterDelegate = UNUserNotificationCenter.currentNotificationCenter.delegate;
+
+    [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                        willPresentNotification:[self createNonOneSignaliOSNotification]
+                          withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
+    XCTAssertTrue([receiver->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:)
+        )
+    ]);
+}
+
+- (void)testDoubleSwizzleInfiniteLoop {
+    // 1. Save original delegate
+    id<UNUserNotificationCenterDelegate> localOrignalDelegate = UNUserNotificationCenter.currentNotificationCenter.delegate;
+
+    // 2. Create a new delegate and assign it
+    id myDelegate = [UNUserNotificationCenterDelegateForInfiniteLoopTest new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myDelegate;
+
+    // 3. Put the original delegate back
+    UNUserNotificationCenter.currentNotificationCenter.delegate = localOrignalDelegate;
+
+    // 4. Call something to confirm we don't get stuck in an infinite call loop
+    [localOrignalDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter willPresentNotification:[self createBasiciOSNotification] withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
+}
+
+- (void)testSwizzleExistingSelectors {
+    UNUserNotificationCenterDelegateForExistingSelectorsTest* myNotifCenterDelegate = [UNUserNotificationCenterDelegateForExistingSelectorsTest new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = myNotifCenterDelegate;
+    id<UNUserNotificationCenterDelegate> notifCenterDelegate = UNUserNotificationCenter.currentNotificationCenter.delegate;
+
+    [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                        willPresentNotification:[self createBasiciOSNotification]
+                          withCompletionHandler:^(UNNotificationPresentationOptions options) {}];
+    XCTAssertTrue([myNotifCenterDelegate->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:)
+        )
+    ]);
+    [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                 didReceiveNotificationResponse:[self createBasiciOSNotificationResponse]
+                          withCompletionHandler:^{}];
+    XCTAssertTrue([myNotifCenterDelegate->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:)
+        )
+    ]);
+    if (@available(iOS 12.0, *)) {
+        [notifCenterDelegate userNotificationCenter:UNUserNotificationCenter.currentNotificationCenter
+                        openSettingsForNotification:[self createBasiciOSNotification]];
+    }
+    XCTAssertTrue([myNotifCenterDelegate->selectorCallsDict
+        objectForKey:NSStringFromSelector(
+                                          @selector(userNotificationCenter:openSettingsForNotification:)
+        )
+    ]);
+}
+
 @end