Don't swizzle missing and optional delegate methods

rsattar · web-flow · commit 5496eff88f3d · 2017-10-04T18:39:30.000-07:00
FCM's swizzling of the user notification center currently swizzles only one of the two optional delegate methods (userNotificationCenter:willPresentNotification:withCompletionHandler:), but not the other (userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:).

The didReceiveNotificationResponse, if implemented by the delegate, is the sole receiver of all user action on a notification, including simply tapping on the notification itself. Prior to this change, if the developer had implemented didReceiveNotificationResponse, then FCM would not be able to collect this event for analytics.

Additionally, I changed the logic in FIRMessagingRemoteNotificationsProxy to check whether these methods are actually implemented before swizzling them. It was always swizzling, which meant it was adding an implementation if the method didn't exist. This would confuse iOS into thinking the developer did implement these methods and NOT fall back to delivering the notifications to the application delegate.

With this change, if the developer did not implement these methods, then FCM will not swizzle those methods. That keeps the behavior true to what the developer intended.
diff --git a/Example/Messaging/App/iOS/NotificationsController.swift b/Example/Messaging/App/iOS/NotificationsController.swift
@@ -121,7 +121,6 @@ class NotificationsController: NSObject {
 // MARK: - UNUserNotificationCenterDelegate
 @available(iOS 10.0, *)
 extension NotificationsController: UNUserNotificationCenterDelegate {
-
   func userNotificationCenter(_ center: UNUserNotificationCenter,
                               willPresent notification: UNNotification,
                               withCompletionHandler completionHandler:
@@ -132,4 +131,11 @@ extension NotificationsController: UNUserNotificationCenterDelegate {
     print("\(jsonString)")
     completionHandler([.alert, .badge, .sound])
   }
+
+  func userNotificationCenter(_ center: UNUserNotificationCenter, didReceive response: UNNotificationResponse, withCompletionHandler completionHandler: @escaping () -> Void) {
+    print("Received notification response")
+    let jsonString = response.notification.request.content.userInfo.jsonString ?? "{}"
+    print("\(jsonString)")
+    completionHandler()
+  }
 }
diff --git a/Example/Messaging/Tests/FIRMessagingRemoteNotificationsProxyTest.m b/Example/Messaging/Tests/FIRMessagingRemoteNotificationsProxyTest.m
@@ -53,6 +53,8 @@ void FCM_swizzle_appDidReceiveRemoteNotificationWithHandler(
     void (^handler)(UIBackgroundFetchResult));
 void FCM_swizzle_willPresentNotificationWithHandler(
     id self, SEL _cmd, id center, id notification, void (^handler)(NSUInteger));
+void FCM_swizzle_didReceiveNotificationResponseWithHandler(
+    id self, SEL _cmd, id center, id response, void (^handler)());
 
 @end
 
@@ -90,6 +92,7 @@ @implementation IncompleteUserNotificationCenterDelegate
 #if __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
 @interface FakeUserNotificationCenterDelegate : NSObject <UNUserNotificationCenterDelegate>
 @property(nonatomic) BOOL willPresentWasCalled;
+@property(nonatomic) BOOL didReceiveResponseWasCalled;
 @end
 @implementation FakeUserNotificationCenterDelegate
 - (void)userNotificationCenter:(UNUserNotificationCenter *)center
@@ -98,6 +101,9 @@ - (void)userNotificationCenter:(UNUserNotificationCenter *)center
     completionHandler {
   self.willPresentWasCalled = YES;
 }
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center didReceiveNotificationResponse:(UNNotificationResponse *)response withCompletionHandler:(void (^)(void))completionHandler {
+  self.didReceiveResponseWasCalled = YES;
+}
 @end
 #endif // __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
 
@@ -219,7 +225,7 @@ - (void)testSwizzlingUserNotificationsCenterDelegate {
 }
 
 // Use a fake delegate that doesn't actually implement the needed delegate method.
-// Our swizzled method should still be called.
+// Our swizzled method should not be called.
 
 - (void)testIncompleteUserNotificationCenterDelegateMethod {
   // Early exit if running on pre iOS 10
@@ -229,20 +235,17 @@ - (void)testIncompleteUserNotificationCenterDelegateMethod {
   IncompleteUserNotificationCenterDelegate *delegate =
       [[IncompleteUserNotificationCenterDelegate alloc] init];
   [self.mockProxy swizzleUserNotificationCenterDelegate:delegate];
-  SEL selector = @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:);
-  XCTAssertTrue([delegate respondsToSelector:selector]);
-  // Invoking delegate method should also invoke our swizzled method
-  // The swizzled method uses the +sharedProxy, which should be
-  // returning our mocked proxy.
-  // Use non-nil, proper classes, otherwise our SDK bails out.
-  [delegate userNotificationCenter:OCMClassMock([UNUserNotificationCenter class])
-           willPresentNotification:[self generateMockNotification]
-             withCompletionHandler:^(NSUInteger options) {}];
-  // Verify our swizzled method was called
-  OCMVerify(FCM_swizzle_willPresentNotificationWithHandler);
+  // Because the incomplete delete does not implement either of the optional delegate methods, we
+  // should swizzle nothing. If we had swizzled them, then respondsToSelector: would return YES
+  // even though the delegate does not implement the methods.
+  SEL willPresentSelector = @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:);
+  XCTAssertFalse([delegate respondsToSelector:willPresentSelector]);
+  SEL didReceiveResponseSelector =
+      @selector(userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:);
+  XCTAssertFalse([delegate respondsToSelector:didReceiveResponseSelector]);
 }
 
-// Use an object that does actually implement the needed method. Both should be called.
+// Use an object that does actually implement the optional methods. Both should be called.
 - (void)testSwizzledUserNotificationsCenterDelegate {
   // Early exit if running on pre iOS 10
   if (![UNNotification class]) {
@@ -261,6 +264,14 @@ - (void)testSwizzledUserNotificationsCenterDelegate {
   OCMVerify(FCM_swizzle_willPresentNotificationWithHandler);
   // Verify our original method was called
   XCTAssertTrue(delegate.willPresentWasCalled);
+
+  [delegate userNotificationCenter:OCMClassMock([UNUserNotificationCenter class])
+    didReceiveNotificationResponse:[self generateMockNotificationResponse]
+             withCompletionHandler:^{}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_appDidReceiveRemoteNotificationWithHandler);
+  // Verify our original method was called
+  XCTAssertTrue(delegate.didReceiveResponseWasCalled);
 }
 
 - (id)generateMockNotification {
@@ -274,6 +285,14 @@ - (id)generateMockNotification {
   return mockNotification;
 }
 
+- (id)generateMockNotificationResponse {
+  // Stub out: response.[mock notification above]
+  id mockNotificationResponse = OCMClassMock([UNNotificationResponse class]);
+  id mockNotification = [self generateMockNotification];
+  OCMStub([mockNotificationResponse notification]).andReturn(mockNotification);
+  return mockNotificationResponse;
+}
+
 #endif // __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
 
 @end
diff --git a/Firebase/Messaging/FIRMessagingRemoteNotificationsProxy.m b/Firebase/Messaging/FIRMessagingRemoteNotificationsProxy.m
@@ -28,6 +28,8 @@
 
 static NSString *kUserNotificationWillPresentSelectorString =
     @"userNotificationCenter:willPresentNotification:withCompletionHandler:";
+static NSString *kUserNotificationDidReceiveResponseSelectorString =
+    @"userNotificationCenter:didReceiveNotificationResponse:withCompletionHandler:";
 
 @interface FIRMessagingRemoteNotificationsProxy ()
 
@@ -219,10 +221,26 @@ - (void)swizzleUserNotificationCenterDelegate:(id)delegate {
   if ([delegate conformsToProtocol:userNotificationCenterProtocol]) {
     SEL willPresentNotificationSelector =
         NSSelectorFromString(kUserNotificationWillPresentSelectorString);
-    [self swizzleSelector:willPresentNotificationSelector
-                  inClass:[delegate class]
-       withImplementation:(IMP)FCM_swizzle_willPresentNotificationWithHandler
-               inProtocol:userNotificationCenterProtocol];
+    // Swizzle the optional method
+    // "userNotificationCenter:willPresentNotification:withCompletionHandler:", if it is
+    // implemented. Do not swizzle otherwise, as an implementation *will* be created, which will
+    // fool iOS into thinking that this method is implemented, and therefore not send notifications
+    // to the fallback method in the app delegate
+    // "application:didReceiveRemoteNotification:fetchCompletionHandler:".
+    if ([delegate respondsToSelector:willPresentNotificationSelector]) {
+      [self swizzleSelector:willPresentNotificationSelector
+                    inClass:[delegate class]
+         withImplementation:(IMP)FCM_swizzle_willPresentNotificationWithHandler
+                 inProtocol:userNotificationCenterProtocol];
+    }
+    SEL didReceiveNotificationResponseSelector =
+        NSSelectorFromString(kUserNotificationDidReceiveResponseSelectorString);
+    if ([delegate respondsToSelector:didReceiveNotificationResponseSelector]) {
+      [self swizzleSelector:didReceiveNotificationResponseSelector
+                    inClass:[delegate class]
+         withImplementation:(IMP)FCM_swizzle_didReceiveNotificationResponseWithHandler
+                 inProtocol:userNotificationCenterProtocol];
+    }
     self.currentUserNotificationCenterDelegate = delegate;
     self.hasSwizzledUserNotificationDelegate = YES;
   }
@@ -235,6 +253,7 @@ - (void)unswizzleUserNotificationCenterDelegate:(id)delegate {
   }
   SEL willPresentNotificationSelector =
       NSSelectorFromString(kUserNotificationWillPresentSelectorString);
+  // Call unswizzle methods, even if the method was not implemented (it will fail gracefully).
   [self unswizzleSelector:willPresentNotificationSelector
                   inClass:[self.currentUserNotificationCenterDelegate class]];
   self.currentUserNotificationCenterDelegate = nil;
@@ -526,25 +545,106 @@ void FCM_swizzle_willPresentNotificationWithHandler(
     return;
   }
 
-  // Valid original method signature, go ahead to swizzle.
+  // Attempt to access the user info
+  id notificationUserInfo = userInfoFromNotification(notification);
+
+  if (!notificationUserInfo) {
+    // Could not access notification.request.content.userInfo.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+
+  [[FIRMessaging messaging] appDidReceiveMessage:notificationUserInfo];
+  // Execute the original implementation.
+  callOriginalMethodIfAvailable();
+}
+
+/**
+ * Swizzle the notification handler for iOS 10+ devices.
+ * Signature of original handler is as below:
+ * - (void)userNotificationCenter:(UNUserNotificationCenter *)center
+ *     didReceiveNotificationResponse:(UNNotificationResponse *)response
+ *     withCompletionHandler:(void (^)(void))completionHandler
+ * In order to make FCM SDK compile and compatible with iOS SDKs before iOS 10, hide the
+ * parameter types from the swizzling implementation.
+ */
+void FCM_swizzle_didReceiveNotificationResponseWithHandler(
+    id self, SEL _cmd, id center, id response, void (^handler)()) {
+
+  FIRMessagingRemoteNotificationsProxy *proxy = [FIRMessagingRemoteNotificationsProxy sharedProxy];
+  IMP original_imp = [proxy originalImplementationForSelector:_cmd];
+
+  void (^callOriginalMethodIfAvailable)() = ^{
+    if (original_imp) {
+      ((void (*)(id, SEL, id, id, void (^)(void)))original_imp)(
+          self, _cmd, center, response, handler);
+    }
+    return;
+  };
+
+  Class notificationCenterClass = NSClassFromString(@"UNUserNotificationCenter");
+  Class responseClass = NSClassFromString(@"UNNotificationResponse");
+  if (!center || ![center isKindOfClass:[notificationCenterClass class]]) {
+    // Invalid parameter type from the original method.
+    // Do not swizzle, just execute the original method.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+
+  if (!response || ![response isKindOfClass:[responseClass class]]) {
+    // Invalid parameter type from the original method.
+    // Do not swizzle, just execute the original method.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+
+  if (!handler) {
+    // Invalid parameter type from the original method.
+    // Do not swizzle, just execute the original method.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+
+  // Try to access the response.notification property
+  SEL notificationSelector = NSSelectorFromString(@"notification");
+  if (![response respondsToSelector:notificationSelector]) {
+    // Cannot access the .notification property.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+  id notificationClass = NSClassFromString(@"UNNotification");
+  id notification = getNamedPropertyFromObject(response, @"notification", notificationClass);
+
+  // With a notification object, use the common code to reach deep into notification
+  // (notification.request.content.userInfo)
+  id notificationUserInfo = userInfoFromNotification(notification);
+  if (!notificationUserInfo) {
+    // Could not access notification.request.content.userInfo.
+    callOriginalMethodIfAvailable();
+    return;
+  }
+
+  [[FIRMessaging messaging] appDidReceiveMessage:notificationUserInfo];
+  // Execute the original implementation.
+  callOriginalMethodIfAvailable();
+}
+
+id userInfoFromNotification(id notification) {
+
   // Select the userInfo field from UNNotification.request.content.userInfo.
   SEL requestSelector = NSSelectorFromString(@"request");
   if (![notification respondsToSelector:requestSelector]) {
-    // This is not the expected notification handler. Do not swizzle, just execute the original
-    // method.
-    callOriginalMethodIfAvailable();
-    return;
+    // Cannot access the request property.
+    return nil;
   }
   Class requestClass = NSClassFromString(@"UNNotificationRequest");
   id notificationRequest = getNamedPropertyFromObject(notification, @"request", requestClass);
 
   SEL notificationContentSelector = NSSelectorFromString(@"content");
   if (!notificationRequest
       || ![notificationRequest respondsToSelector:notificationContentSelector]) {
-    // This is not the expected notification handler. Do not swizzle, just execute the original
-    // method.
-    callOriginalMethodIfAvailable();
-    return;
+    // Cannot access the content property.
+    return nil;
   }
   Class contentClass = NSClassFromString(@"UNNotificationContent");
   id notificationContent = getNamedPropertyFromObject(notificationRequest,
@@ -554,25 +654,19 @@ void FCM_swizzle_willPresentNotificationWithHandler(
   SEL notificationUserInfoSelector = NSSelectorFromString(@"userInfo");
   if (!notificationContent
       || ![notificationContent respondsToSelector:notificationUserInfoSelector]) {
-    // This is not the expected notification handler. Do not swizzle, just execute the original
-    // method.
-    callOriginalMethodIfAvailable();
-    return;
+    // Cannot access the userInfo property.
+    return nil;
   }
   id notificationUserInfo = getNamedPropertyFromObject(notificationContent,
                                                        @"userInfo",
                                                        [NSDictionary class]);
 
   if (!notificationUserInfo) {
-    // This is not the expected notification handler. Do not swizzle, just execute the original
-    // method.
-    callOriginalMethodIfAvailable();
-    return;
+    // This is not the expected notification handler.
+    return nil;
   }
 
-  [[FIRMessaging messaging] appDidReceiveMessage:notificationUserInfo];
-  // Execute the original implementation.
-  callOriginalMethodIfAvailable();
+  return notificationUserInfo;
 }
 
 void FCM_swizzle_applicationReceivedRemoteMessage(
