[Messaging] - Fix crash for strongSelf dereference (#10740)

Fix crash for strongSelf dereference

Author: aashishpatil-g

FirebaseMessaging/Sources/NSError+FIRMessaging.h

@@ -21,6 +21,7 @@ NS_ASSUME_NONNULL_BEGIN
 // FIRMessaging Internal Error Code
 typedef NS_ENUM(NSUInteger, FIRMessagingErrorCode) {
   kFIRMessagingErrorCodeUnknown = 0,
+  kFIRMessagingErrorCodeInternal = 1,
 
   kFIRMessagingErrorCodeNetwork = 4,
 

FirebaseMessaging/Sources/Token/FIRMessagingTokenManager.m

@@ -239,17 +239,25 @@ - (void)tokenWithAuthorizedEntity:(NSString *)authorizedEntity
   }
 
   FIRMessaging_WEAKIFY(self);
-  [_authService
-      fetchCheckinInfoWithHandler:^(FIRMessagingCheckinPreferences *preferences, NSError *error) {
-        FIRMessaging_STRONGIFY(self);
-        if (error) {
-          newHandler(nil, error);
-          return;
-        }
+  [_authService fetchCheckinInfoWithHandler:^(FIRMessagingCheckinPreferences *preferences,
+                                              NSError *error) {
+    FIRMessaging_STRONGIFY(self);
+    if (error) {
+      newHandler(nil, error);
+      return;
+    }
+
+    if (!self) {
+      NSError *derefErr =
+          [NSError messagingErrorWithCode:kFIRMessagingErrorCodeInternal
+                            failureReason:@"Unable to fetch token. Lost Reference to TokenManager"];
+      handler(nil, derefErr);
+      return;
+    }
 
-        FIRMessaging_WEAKIFY(self);
-        [self->_installations installationIDWithCompletion:^(NSString *_Nullable identifier,
-                                                             NSError *_Nullable error) {
+    FIRMessaging_WEAKIFY(self);
+    [self->_installations
+        installationIDWithCompletion:^(NSString *_Nullable identifier, NSError *_Nullable error) {
           FIRMessaging_STRONGIFY(self);
 
           if (error) {
@@ -275,7 +283,7 @@ - (void)tokenWithAuthorizedEntity:(NSString *)authorizedEntity
                                             handler:newHandler];
           }
         }];
-      }];
+  }];
 }
 
 - (void)fetchNewTokenWithAuthorizedEntity:(NSString *)authorizedEntity
@@ -292,52 +300,60 @@ - (void)fetchNewTokenWithAuthorizedEntity:(NSString *)authorizedEntity
                                              options:options
                                           instanceID:instanceID];
   FIRMessaging_WEAKIFY(self);
-  FIRMessagingTokenOperationCompletion completion =
-      ^(FIRMessagingTokenOperationResult result, NSString *_Nullable token,
-        NSError *_Nullable error) {
-        FIRMessaging_STRONGIFY(self);
-        if (error) {
-          handler(nil, error);
-          return;
-        }
-        if ([self isDefaultTokenWithAuthorizedEntity:authorizedEntity scope:scope]) {
-          [self postTokenRefreshNotificationWithDefaultFCMToken:token];
-        }
-        NSString *firebaseAppID = options[kFIRMessagingTokenOptionsFirebaseAppIDKey];
-        FIRMessagingTokenInfo *tokenInfo =
-            [[FIRMessagingTokenInfo alloc] initWithAuthorizedEntity:authorizedEntity
-                                                              scope:scope
-                                                              token:token
-                                                         appVersion:FIRMessagingCurrentAppVersion()
-                                                      firebaseAppID:firebaseAppID];
-        tokenInfo.APNSInfo = [[FIRMessagingAPNSInfo alloc] initWithTokenOptionsDictionary:options];
-
-        [self->_tokenStore
-            saveTokenInfo:tokenInfo
-                  handler:^(NSError *error) {
-                    if (!error) {
-                      // Do not send the token back in case the save was unsuccessful. Since with
-                      // the new asychronous fetch mechanism this can lead to infinite loops, for
-                      // example, we will return a valid token even though we weren't able to store
-                      // it in our cache. The first token will lead to a onTokenRefresh callback
-                      // wherein the user again calls `getToken` but since we weren't able to save
-                      // it we won't hit the cache but hit the server again leading to an infinite
-                      // loop.
-                      FIRMessagingLoggerDebug(
-                          kFIRMessagingMessageCodeTokenManager001,
-                          @"Token fetch successful, token: %@, authorizedEntity: %@, scope:%@",
-                          token, authorizedEntity, scope);
-
-                      if (handler) {
-                        handler(token, nil);
-                      }
-                    } else {
-                      if (handler) {
-                        handler(nil, error);
-                      }
-                    }
-                  }];
-      };
+  FIRMessagingTokenOperationCompletion completion = ^(FIRMessagingTokenOperationResult result,
+                                                      NSString *_Nullable token,
+                                                      NSError *_Nullable error) {
+    FIRMessaging_STRONGIFY(self);
+    if (error) {
+      handler(nil, error);
+      return;
+    }
+
+    if (!self) {
+      NSError *lostRefError = [NSError messagingErrorWithCode:kFIRMessagingErrorCodeInternal
+                                                failureReason:@"Lost Reference to TokenManager"];
+      handler(nil, lostRefError);
+      return;
+    }
+
+    if ([self isDefaultTokenWithAuthorizedEntity:authorizedEntity scope:scope]) {
+      [self postTokenRefreshNotificationWithDefaultFCMToken:token];
+    }
+    NSString *firebaseAppID = options[kFIRMessagingTokenOptionsFirebaseAppIDKey];
+    FIRMessagingTokenInfo *tokenInfo =
+        [[FIRMessagingTokenInfo alloc] initWithAuthorizedEntity:authorizedEntity
+                                                          scope:scope
+                                                          token:token
+                                                     appVersion:FIRMessagingCurrentAppVersion()
+                                                  firebaseAppID:firebaseAppID];
+    tokenInfo.APNSInfo = [[FIRMessagingAPNSInfo alloc] initWithTokenOptionsDictionary:options];
+
+    [self->_tokenStore
+        saveTokenInfo:tokenInfo
+              handler:^(NSError *error) {
+                if (!error) {
+                  // Do not send the token back in case the save was unsuccessful. Since with
+                  // the new asychronous fetch mechanism this can lead to infinite loops, for
+                  // example, we will return a valid token even though we weren't able to store
+                  // it in our cache. The first token will lead to a onTokenRefresh callback
+                  // wherein the user again calls `getToken` but since we weren't able to save
+                  // it we won't hit the cache but hit the server again leading to an infinite
+                  // loop.
+                  FIRMessagingLoggerDebug(
+                      kFIRMessagingMessageCodeTokenManager001,
+                      @"Token fetch successful, token: %@, authorizedEntity: %@, scope:%@", token,
+                      authorizedEntity, scope);
+
+                  if (handler) {
+                    handler(token, nil);
+                  }
+                } else {
+                  if (handler) {
+                    handler(nil, error);
+                  }
+                }
+              }];
+  };
   // Add completion handler, and ensure it's called on the main queue
   [operation addCompletionHandler:^(FIRMessagingTokenOperationResult result,
                                     NSString *_Nullable token, NSError *_Nullable error) {
@@ -456,6 +472,15 @@ - (void)deleteWithHandler:(void (^)(NSError *))handler {
       handler(error);
       return;
     }
+
+    if (!self) {
+      NSError *lostRefError =
+          [NSError messagingErrorWithCode:kFIRMessagingErrorCodeInternal
+                            failureReason:@"Cannot delete token. Lost reference to TokenManager"];
+      handler(lostRefError);
+      return;
+    }
+
     [self deleteAllTokensLocallyWithHandler:^(NSError *localError) {
       [self postTokenRefreshNotificationWithDefaultFCMToken:nil];
       self->_defaultFCMToken = nil;