Force refresh token if RPC fails with "Unauthenticated" error (#1373)

"Unauthenticated" is presumed to mean that token is expired (which might happen if local clock is wrong) and retried, subject to the usual backoff logic.

2 files changed, 33 insertions, 30 deletions

diff --git a/Firestore/Source/Remote/FSTDatastore.mm b/Firestore/Source/Remote/FSTDatastore.mm
index 5f79122..fdbeea3 100644
--- a/Firestore/Source/Remote/FSTDatastore.mm
+++ b/Firestore/Source/Remote/FSTDatastore.mm
@@ -157,10 +157,9 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
     case FIRFirestoreErrorCodeResourceExhausted:
     case FIRFirestoreErrorCodeInternal:
     case FIRFirestoreErrorCodeUnavailable:
+    // Unauthenticated means something went wrong with our token and we need to retry with new
+    // credentials which will happen automatically.
     case FIRFirestoreErrorCodeUnauthenticated:
-      // Unauthenticated means something went wrong with our token and we need
-      // to retry with new credentials which will happen automatically.
-      // TODO(b/37325376): Give up after second unauthenticated error.
       return NO;
     case FIRFirestoreErrorCodeInvalidArgument:
     case FIRFirestoreErrorCodeNotFound:
@@ -174,6 +173,7 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
     case FIRFirestoreErrorCodeOutOfRange:
     case FIRFirestoreErrorCodeUnimplemented:
     case FIRFirestoreErrorCodeDataLoss:
+      return YES;
     default:
       return YES;
   }
@@ -239,6 +239,9 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
                        handler:^(GCFSCommitResponse *response, NSError *_Nullable error) {
                          error = [FSTDatastore firestoreErrorForError:error];
                          [self.workerDispatchQueue dispatchAsync:^{
+                           if (error != nil && error.code == FIRFirestoreErrorCodeUnauthenticated) {
+                             _credentials->InvalidateToken();
+                           }
                            LOG_DEBUG("RPC CommitRequest completed. Error: %s", error);
                            [FSTDatastore logHeadersForRPC:rpc RPCName:@"CommitRequest"];
                            completion(error);
@@ -273,6 +276,9 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
                                [self.workerDispatchQueue dispatchAsync:^{
                                  if (error) {
                                    LOG_DEBUG("RPC BatchGetDocuments completed. Error: %s", error);
+                                   if (error.code == FIRFirestoreErrorCodeUnauthenticated) {
+                                     _credentials->InvalidateToken();
+                                   }
                                    [FSTDatastore logHeadersForRPC:rpc RPCName:@"BatchGetDocuments"];
                                    completion(nil, error);
                                    return;
@@ -310,25 +316,21 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
 
 - (void)invokeRPCWithFactory:(GRPCProtoCall * (^)(void))rpcFactory
                 errorHandler:(FSTVoidErrorBlock)errorHandler {
-  // TODO(mikelehen): We should force a refresh if the previous RPC failed due to an expired token,
-  // but I'm not sure how to detect that right now. http://b/32762461
-  _credentials->GetToken(
-      /*force_refresh=*/false, [self, rpcFactory, errorHandler](util::StatusOr<Token> result) {
-        [self.workerDispatchQueue dispatchAsyncAllowingSameQueue:^{
-          if (!result.ok()) {
-            errorHandler(util::MakeNSError(result.status()));
-          } else {
-            GRPCProtoCall *rpc = rpcFactory();
-            const Token &token = result.ValueOrDie();
-            [FSTDatastore
-                prepareHeadersForRPC:rpc
-                          databaseID:&self.databaseInfo->database_id()
-                               token:(token.user().is_authenticated() ? token.token()
-                                                                      : absl::string_view())];
-            [rpc start];
-          }
-        }];
-      });
+  _credentials->GetToken([self, rpcFactory, errorHandler](util::StatusOr<Token> result) {
+    [self.workerDispatchQueue dispatchAsyncAllowingSameQueue:^{
+      if (!result.ok()) {
+        errorHandler(util::MakeNSError(result.status()));
+      } else {
+        GRPCProtoCall *rpc = rpcFactory();
+        const Token &token = result.ValueOrDie();
+        [FSTDatastore prepareHeadersForRPC:rpc
+                                databaseID:&self.databaseInfo->database_id()
+                                     token:(token.user().is_authenticated() ? token.token()
+                                                                            : absl::string_view())];
+        [rpc start];
+      }
+    }];
+  });
 }
 
 - (FSTWatchStream *)createWatchStream {
diff --git a/Firestore/Source/Remote/FSTStream.mm b/Firestore/Source/Remote/FSTStream.mm
index 3a6c035..2c9c8d9 100644
--- a/Firestore/Source/Remote/FSTStream.mm
+++ b/Firestore/Source/Remote/FSTStream.mm
@@ -265,12 +265,11 @@ static const NSTimeInterval kIdleTimeout = 60.0;
   HARD_ASSERT(_delegate == nil, "Delegate must be nil");
   _delegate = delegate;
 
-  _credentials->GetToken(
-      /*force_refresh=*/false, [self](util::StatusOr<Token> result) {
-        [self.workerDispatchQueue dispatchAsyncAllowingSameQueue:^{
-          [self resumeStartWithToken:result];
-        }];
-      });
+  _credentials->GetToken([self](util::StatusOr<Token> result) {
+    [self.workerDispatchQueue dispatchAsyncAllowingSameQueue:^{
+      [self resumeStartWithToken:result];
+    }];
+  });
 }
 
 /** Add an access token to our RPC, after obtaining one from the credentials provider. */
@@ -283,8 +282,6 @@ static const NSTimeInterval kIdleTimeout = 60.0;
   }
   HARD_ASSERT(self.state == FSTStreamStateAuth, "State should still be auth (was %s)", self.state);
 
-  // TODO(mikelehen): We should force a refresh if the previous RPC failed due to an expired token,
-  // but I'm not sure how to detect that right now. http://b/32762461
   if (!result.ok()) {
     // RPC has not been started yet, so just invoke higher-level close handler.
     [self handleStreamClose:util::MakeNSError(result.status())];
@@ -383,6 +380,10 @@ static const NSTimeInterval kIdleTimeout = 60.0;
     LOG_DEBUG("%s %s Using maximum backoff delay to prevent overloading the backend.", [self class],
               (__bridge void *)self);
     [self.backoff resetToMax];
+  } else if (error != nil && error.code == FIRFirestoreErrorCodeUnauthenticated) {
+    // "unauthenticated" error means the token was rejected. Try force refreshing it in case it just
+    // expired.
+    _credentials->InvalidateToken();
   }
 
   if (finalState != FSTStreamStateError) {