Merge pull request #926 from OneSignal/fix/pre-existing_notification_…

…delegate_for_master

Fix pre-existing notification delegate

iOS_SDK/OneSignalSDK/Source/OneSignal.m

@@ -2945,10 +2945,7 @@ +(void)setupUNUserNotificationCenterDelegate {
     if (!NSClassFromString(@"UNUserNotificationCenter"))
         return;
 
-    [OneSignalUNUserNotificationCenter swizzleSelectors];
-
-    // Set our own delegate if one hasn't been set already from something else.
-    [OneSignalHelper registerAsUNNotificationCenterDelegate];
+    [OneSignalUNUserNotificationCenter setup];
 }
 
 +(BOOL) shouldDisableBasedOnProcessArguments {

iOS_SDK/OneSignalSDK/Source/OneSignalHelper.h

@@ -51,7 +51,6 @@
 + (BOOL)handleIAMPreview:(OSNotification *)notification;
 
 // - iOS 10
-+ (void)registerAsUNNotificationCenterDelegate;
 + (void)clearCachedMedia;
 + (UNNotificationRequest*)prepareUNNotificationRequest:(OSNotification*)notification;
 + (void)addNotificationRequest:(OSNotification*)notification completionHandler:(void (^)(UIBackgroundFetchResult))completionHandler;

iOS_SDK/OneSignalSDK/Source/OneSignalHelper.m

@@ -518,7 +518,7 @@ + (UILocalNotification*)prepareUILocalNotification:(OSNotification *)osNotificat
     return notification;
 }
 
-//Shared instance as OneSignal is delegate of UNUserNotificationCenterDelegate and CLLocationManagerDelegate
+//Shared instance as OneSignal is delegate CLLocationManagerDelegate
 static OneSignal* singleInstance = nil;
 + (OneSignal*)sharedInstance {
     @synchronized( singleInstance ) {
@@ -539,20 +539,6 @@ +(NSString*)randomStringWithLength:(int)length {
     return randomString;
 }
 
-+ (void)registerAsUNNotificationCenterDelegate {
-    let curNotifCenter = [UNUserNotificationCenter currentNotificationCenter];
-
-    /*
-        Sets the OneSignal shared instance as a delegate of UNUserNotificationCenter
-        OneSignal does not implement the delegate methods, we simply set it as a delegate
-        in order to swizzle the UNUserNotificationCenter methods in case the developer
-        does not set a UNUserNotificationCenter delegate themselves
-    */
-
-    if (!curNotifCenter.delegate)
-        curNotifCenter.delegate = (id)[self sharedInstance];
-}
-
 + (UNNotificationRequest*)prepareUNNotificationRequest:(OSNotification*)notification {
     let content = [UNMutableNotificationContent new];
 

iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.h

@@ -31,7 +31,10 @@
 #import "OneSignal.h"
 
 @interface OneSignalUNUserNotificationCenter : NSObject
++ (void)setup;
 + (void)swizzleSelectors;
++ (void)swizzleSelectorsOnDelegate:(id)delegate;
++ (void)registerDelegate;
 + (void)setUseiOS10_2_workaround:(BOOL)enable;
 @end
 

iOS_SDK/OneSignalSDK/Source/UNUserNotificationCenter+OneSignal.m

@@ -51,6 +51,21 @@ + (BOOL)shouldLogMissingPrivacyConsentErrorWithMethodName:(NSString *)methodName
 + (void)handleWillPresentNotificationInForegroundWithPayload:(NSDictionary *)payload withCompletion:(OSNotificationDisplayResponse)completionHandler;
 @end
 
+@interface OSUNUserNotificationCenterDelegate : NSObject
++ (OSUNUserNotificationCenterDelegate*)sharedInstance;
+@end
+
+@implementation OSUNUserNotificationCenterDelegate
+static OSUNUserNotificationCenterDelegate* singleInstance = nil;
++ (OSUNUserNotificationCenterDelegate*)sharedInstance {
+    @synchronized(singleInstance) {
+        if (!singleInstance)
+            singleInstance = [OSUNUserNotificationCenterDelegate new];
+    }
+    return singleInstance;
+}
+@end
+
 // This class hooks into the following iSO 10 UNUserNotificationCenterDelegate selectors:
 // - userNotificationCenter:willPresentNotification:withCompletionHandler:
 //   - Reads OneSignal.inFocusDisplayType to respect it's setting.
@@ -64,6 +79,11 @@ + (void)handleWillPresentNotificationInForegroundWithPayload:(NSDictionary *)pay
 
 @implementation OneSignalUNUserNotificationCenter
 
++ (void)setup {
+    [OneSignalUNUserNotificationCenter swizzleSelectors];
+    [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
@@ -87,6 +107,30 @@ + (void)swizzleSelectors {
                         [OneSignalUNUserNotificationCenter class], [UNUserNotificationCenter class]);
 }
 
++ (void)registerDelegate {
+    let curNotifCenter = [UNUserNotificationCenter currentNotificationCenter];
+
+    if (!curNotifCenter.delegate) {
+        /*
+          Set OSUNUserNotificationCenterDelegate.sharedInstance as a
+            UNUserNotificationCenterDelegate.
+          Note that OSUNUserNotificationCenterDelegate does not contain the methods such as
+            "willPresentNotification" as this assigment triggers setOneSignalUNDelegate which
+            will attach the selectors to the class at runtime.
+         */
+        curNotifCenter.delegate = (id)OSUNUserNotificationCenterDelegate.sharedInstance;
+    }
+    else {
+        /*
+         This handles the case where a delegate may have already been assigned before
+           OneSignal is loaded into memory.
+         This re-assignment triggers setOneSignalUNDelegate providing it with the
+           existing delegate instance so OneSignal can swizzle in its logic.
+         */
+        curNotifCenter.delegate = curNotifCenter.delegate;
+    }
+}
+
 static BOOL useiOS10_2_workaround = true;
 + (void)setUseiOS10_2_workaround:(BOOL)enable {
     useiOS10_2_workaround = enable;
@@ -140,11 +184,18 @@ - (void) setOneSignalUNDelegate:(id)delegate {
         [self setOneSignalUNDelegate:delegate];
         return;
     }
-    
+
     previousDelegate = delegate;
-    
+
     [OneSignal onesignal_Log:ONE_S_LL_VERBOSE message:@"OneSignalUNUserNotificationCenter setOneSignalUNDelegate Fired!"];
-
+
+    [OneSignalUNUserNotificationCenter swizzleSelectorsOnDelegate:delegate];
+
+    // Call orignal iOS implemenation
+    [self setOneSignalUNDelegate:delegate];
+}
+
++ (void)swizzleSelectorsOnDelegate:(id)delegate {
     delegateUNClass = getClassWithProtocolInHierarchy([delegate class], @protocol(UNUserNotificationCenterDelegate));
     delegateUNSubclasses = ClassGetSubclasses(delegateUNClass);
 
@@ -153,8 +204,6 @@ - (void) setOneSignalUNDelegate:(id)delegate {
 
     injectToProperClass(@selector(onesignalUserNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:),
                         @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:), delegateUNSubclasses, [OneSignalUNUserNotificationCenter class], delegateUNClass);
-
-    [self setOneSignalUNDelegate:delegate];
 }
 
 + (void)forwardNotificationWithCenter:(UNUserNotificationCenter *)center

iOS_SDK/OneSignalSDK/UnitTests/OneSignalUNUserNotificationCenterSwizzingTest.m

@@ -0,0 +1,79 @@
+#import <objc/runtime.h>
+
+#import <XCTest/XCTest.h>
+
+#import "UnitTestCommonMethods.h"
+#import "OneSignalExtensionBadgeHandler.h"
+#import "UNUserNotificationCenterOverrider.h"
+#import "UNUserNotificationCenter+OneSignal.h"
+#import "OneSignalHelperOverrider.h"
+#import "OneSignalHelper.h"
+#import "DummyNotificationCenterDelegate.h"
+#import "OneSignalUNUserNotificationCenterHelper.h"
+
+@interface OneSignalUNUserNotificationCenterSwizzingTest : XCTestCase
+@end
+
+@implementation OneSignalUNUserNotificationCenterSwizzingTest
+
+// Called BEFORE each test method
+- (void)setUp {
+    [super setUp];
+    [UnitTestCommonMethods beforeEachTest:self];
+
+    [OneSignalUNUserNotificationCenter setUseiOS10_2_workaround:true];
+}
+
+// Called AFTER each test method
+- (void)tearDown {
+    [super tearDown];
+}
+
+// Tests to make sure that UNNotificationCenter setDelegate: duplicate calls don't double-swizzle for the same object
+- (void)testAUNUserNotificationCenterDelegateAssigningDoesSwizzle {
+    UNUserNotificationCenter *center = [UNUserNotificationCenter currentNotificationCenter];
+
+    let dummyDelegate = [[DummyNotificationCenterDelegate alloc] init];
+
+    IMP original = class_getMethodImplementation([dummyDelegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
+
+    // This triggers UNUserNotificationCenter+OneSignal.m setOneSignalUNDelegate which does the implemenation swizzling
+    center.delegate = dummyDelegate;
+
+    IMP swizzled = class_getMethodImplementation([dummyDelegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
+    // Since we swizzled the implemenations should be different.
+    XCTAssertNotEqual(original, swizzled);
+
+    // Calling setDelegate: a second time on the same object should not re-exchange method implementations
+    // thus the new method implementation should still be the same, swizzled == newSwizzled should be true
+    center.delegate = dummyDelegate;
+
+    IMP newSwizzled = class_getMethodImplementation([dummyDelegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
+
+    XCTAssertEqual(swizzled, newSwizzled);
+
+    [OneSignalUNUserNotificationCenterHelper restoreDelegateAsOneSignal];
+}
+
+- (void)testUNUserNotificationCenterDelegateAssignedBeforeOneSignal {
+    [OneSignalUNUserNotificationCenterHelper putIntoPreloadedState];
+
+    // Create and assign a delegate with iOS
+    let dummyDelegate = [DummyNotificationCenterDelegate new];
+    UNUserNotificationCenter.currentNotificationCenter.delegate = dummyDelegate;
+
+    // Save original implemenation reference, before OneSignal is loaded.
+    IMP originalDummyImp = class_getMethodImplementation([dummyDelegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
+
+    // Setup the OneSignal delegate where it will be loaded into memeory
+    [OneSignalUNUserNotificationCenter setup];
+
+    IMP swizzledDummyImp = class_getMethodImplementation([dummyDelegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
+
+    // Since we swizzled the implemenations should be different.
+    XCTAssertNotEqual(originalDummyImp, swizzledDummyImp);
+
+    [OneSignalUNUserNotificationCenterHelper restoreDelegateAsOneSignal];
+}
+
+@end

iOS_SDK/OneSignalSDK/UnitTests/UNNotificationCenter/OneSignalUNUserNotificationCenterHelper.h

@@ -0,0 +1,10 @@
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface OneSignalUNUserNotificationCenterHelper : NSObject
++ (void)restoreDelegateAsOneSignal;
++ (void)putIntoPreloadedState;
+@end
+
+NS_ASSUME_NONNULL_END

iOS_SDK/OneSignalSDK/UnitTests/UNNotificationCenter/OneSignalUNUserNotificationCenterHelper.m

@@ -0,0 +1,37 @@
+#import "OneSignalUNUserNotificationCenterHelper.h"
+
+#import "UNUserNotificationCenter+OneSignal.h"
+#import "OneSignalHelper.h"
+
+@implementation OneSignalUNUserNotificationCenterHelper
+
+// Call this after you assign UNUserNotificationCenter.currentNotificationCenter.delegate to restore
+//   the delegate as well as swizzling effects.
+// In a nutshell this swizzles a 2nd time in reverse to swap method implementations back.
++ (void)restoreDelegateAsOneSignal {
+    // Undo swizzling of notification events by swizzling again to swap method implementations back.
+    [OneSignalUNUserNotificationCenter swizzleSelectorsOnDelegate:UNUserNotificationCenter.currentNotificationCenter.delegate];
+
+    // Temp unswizzle setDelegate, so we can re-assign the delegate below without side-effects.
+    [OneSignalUNUserNotificationCenter swizzleSelectors];
+
+    // Clear the delegate and call registerAsUNNotificationCenterDelegate so it assigns OneSignal
+    //   as the delegate like it does the first time when it is loaded into memory.
+    UNUserNotificationCenter.currentNotificationCenter.delegate = nil;
+    [OneSignalUNUserNotificationCenter registerDelegate];
+
+    // Undo our temp unswizzle by swizzle one more time. Undo our undo :)
+    [OneSignalUNUserNotificationCenter swizzleSelectors];
+}
+
+// OneSignal auto swizzles UNUserNotificationCenterDelegate when the class is loaded into memory.
+// This undoes this affect so we can test setup a pre-loaded OneSignal state.
+// Why? This way we can setup a senario where someone else's delegate is assigned before OneSignal get loaded.
++ (void)putIntoPreloadedState {
+    // Swizzle to undo the swizzle. (yes, swizzling a 2nd time reverses the swizzling)
+    [OneSignalUNUserNotificationCenter swizzleSelectors];
+
+    // Unassign OneSignal as the delegate
+    UNUserNotificationCenter.currentNotificationCenter.delegate = nil;
+}
+@end

iOS_SDK/OneSignalSDK/UnitTests/UnitTests.m

@@ -55,6 +55,7 @@
 #import "OneSignalExtensionBadgeHandler.h"
 #import "OneSignalDialogControllerOverrider.h"
 #import "OneSignalNotificationCategoryController.h"
+#import "OneSignalUNUserNotificationCenterHelper.h"
 
 // Shadows
 #import "NSObjectOverrider.h"
@@ -1998,34 +1999,6 @@ - (void)testPushNotificationToken {
 
     [NSBundleOverrider setPrivacyState:false];
 }
-
-//tests to make sure that UNNotificationCenter setDelegate: duplicate calls don't double-swizzle for the same object
-// TODO: This test causes the UNUserNotificationCenter singleton's Delegate property to get nullified
-// Unfortunately the fix is not as simple as just setting it back to the original when the test is done
-// To avoid breaking other tests, this test should be executed last, and since tests are alphabetical order, adding Z's does this.
-- (void)testZSwizzling {
-    UNUserNotificationCenter *center = [UNUserNotificationCenter currentNotificationCenter];
-
-    DummyNotificationCenterDelegate *delegate = [[DummyNotificationCenterDelegate alloc] init];
-
-    IMP original = class_getMethodImplementation([delegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
-
-    center.delegate = delegate;
-
-    IMP swizzled = class_getMethodImplementation([delegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
-
-    XCTAssertNotEqual(original, swizzled);
-
-    //calling setDelegate: a second time on the same object should not re-exchange method implementations
-    //thus the new method implementation should still be the same, swizzled == newSwizzled should be true
-    center.delegate = delegate;
-
-    IMP newSwizzled = class_getMethodImplementation([delegate class], @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:));
-
-    XCTAssertNotEqual(original, newSwizzled);
-    XCTAssertEqual(swizzled, newSwizzled);
-
-}
 
 - (NSDictionary *)setUpWillShowInForegroundHandlerTestWithBlock:(OSNotificationWillShowInForegroundBlock)willShowInForegroundBlock withNotificationOpenedBlock:(OSNotificationOpenedBlock)openedBlock withPayload: (NSDictionary *)payload {