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

"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.

Author: var-const

packages/firestore/src/api/credentials.ts

@@ -72,11 +72,14 @@ export type UserListener = (user: User) => void;
  * listening for changes.
  */
 export interface CredentialsProvider {
+  /** Requests a token for the current user. */
+  getToken(): Promise<Token | null>;
+
   /**
-   * Requests a token for the current user, optionally forcing a refreshed
-   * token to be fetched.
+   * Marks the last retrieved token as invalid, making the next GetToken request
+   * force-refresh the token.
    */
-  getToken(forceRefresh: boolean): Promise<Token | null>;
+  invalidateToken(): void;
 
   /**
    * Specifies a listener to be notified of user changes (sign-in / sign-out).
@@ -99,10 +102,12 @@ export class EmptyCredentialsProvider implements CredentialsProvider {
 
   constructor() {}
 
-  getToken(forceRefresh: boolean): Promise<Token | null> {
+  getToken(): Promise<Token | null> {
     return Promise.resolve<Token | null>(null);
   }
 
+  invalidateToken(): void {}
+
   setUserChangeListener(listener: UserListener): void {
     assert(!this.userListener, 'Can only call setUserChangeListener() once.');
     this.userListener = listener;
@@ -138,6 +143,8 @@ export class FirebaseCredentialsProvider implements CredentialsProvider {
   /** The User listener registered with setUserChangeListener(). */
   private userListener: UserListener | null = null;
 
+  private forceRefresh = false;
+
   constructor(private readonly app: FirebaseApp) {
     // We listen for token changes but all we really care about is knowing when
     // the uid may have changed.
@@ -160,7 +167,7 @@ export class FirebaseCredentialsProvider implements CredentialsProvider {
     );
   }
 
-  getToken(forceRefresh: boolean): Promise<Token | null> {
+  getToken(): Promise<Token | null> {
     assert(
       this.tokenListener != null,
       'getToken cannot be called after listener removed.'
@@ -170,6 +177,8 @@ export class FirebaseCredentialsProvider implements CredentialsProvider {
     // fail (with an ABORTED error) if there is a user change while the request
     // is outstanding.
     const initialUserCounter = this.userCounter;
+    const forceRefresh = this.forceRefresh;
+    this.forceRefresh = false;
     return (this.app as _FirebaseApp).INTERNAL.getToken(forceRefresh).then(
       tokenData => {
         // Cancel the request since the user changed while the request was
@@ -195,6 +204,10 @@ export class FirebaseCredentialsProvider implements CredentialsProvider {
     );
   }
 
+  invalidateToken(): void {
+    this.forceRefresh = true;
+  }
+
   setUserChangeListener(listener: UserListener): void {
     assert(!this.userListener, 'Can only call setUserChangeListener() once.');
     this.userListener = listener;
@@ -285,7 +298,7 @@ export class FirstPartyCredentialsProvider implements CredentialsProvider {
     );
   }
 
-  getToken(forceRefresh: boolean): Promise<Token | null> {
+  getToken(): Promise<Token | null> {
     return Promise.resolve(new FirstPartyToken(this.gapi, this.sessionIndex));
   }
 
@@ -297,6 +310,8 @@ export class FirstPartyCredentialsProvider implements CredentialsProvider {
   }
 
   removeUserChangeListener(): void {}
+
+  invalidateToken(): void {}
 }
 
 /**

packages/firestore/src/remote/datastore.ts

@@ -21,6 +21,7 @@ import { MaybeDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { Mutation, MutationResult } from '../model/mutation';
 import { assert } from '../util/assert';
+import { Code, FirestoreError } from '../util/error';
 import { AsyncQueue } from '../util/async_queue';
 
 import { Connection } from './connection';
@@ -110,24 +111,38 @@ export class Datastore {
 
   /** Gets an auth token and invokes the provided RPC. */
   private invokeRPC<Req, Resp>(rpcName: string, request: Req): Promise<Resp> {
-    // TODO(mikelehen): Retry (with backoff) on token failures?
-    return this.credentials.getToken(/*forceRefresh=*/ false).then(token => {
-      return this.connection.invokeRPC<Req, Resp>(rpcName, request, token);
-    });
+    return this.credentials
+      .getToken()
+      .then(token => {
+        return this.connection.invokeRPC<Req, Resp>(rpcName, request, token);
+      })
+      .catch((error: FirestoreError) => {
+        if (error.code === Code.UNAUTHENTICATED) {
+          this.credentials.invalidateToken();
+        }
+        throw error;
+      });
   }
 
   /** Gets an auth token and invokes the provided RPC with streamed results. */
   private invokeStreamingRPC<Req, Resp>(
     rpcName: string,
     request: Req
   ): Promise<Resp[]> {
-    // TODO(mikelehen): Retry (with backoff) on token failures?
-    return this.credentials.getToken(/*forceRefresh=*/ false).then(token => {
-      return this.connection.invokeStreamingRPC<Req, Resp>(
-        rpcName,
-        request,
-        token
-      );
-    });
+    return this.credentials
+      .getToken()
+      .then(token => {
+        return this.connection.invokeStreamingRPC<Req, Resp>(
+          rpcName,
+          request,
+          token
+        );
+      })
+      .catch((error: FirestoreError) => {
+        if (error.code === Code.UNAUTHENTICATED) {
+          this.credentials.invalidateToken();
+        }
+        throw error;
+      });
   }
 }

packages/firestore/src/remote/persistent_stream.ts

@@ -331,6 +331,10 @@ export abstract class PersistentStream<
         'Using maximum backoff delay to prevent overloading the backend.'
       );
       this.backoff.resetToMax();
+    } else if (error && error.code === Code.UNAUTHENTICATED) {
+      // "unauthenticated" error means the token was rejected. Try force refreshing it in case it
+      // just expired.
+      this.credentialsProvider.invalidateToken();
     }
 
     // Clean up the underlying stream because we are no longer interested in events.
@@ -384,7 +388,7 @@ export abstract class PersistentStream<
 
     this.state = PersistentStreamState.Auth;
 
-    this.credentialsProvider.getToken(/*forceRefresh=*/ false).then(
+    this.credentialsProvider.getToken().then(
       token => {
         // Normally we'd have to schedule the callback on the AsyncQueue.
         // However, the following calls are safe to be called outside the
@@ -478,7 +482,8 @@ export abstract class PersistentStream<
     });
   }
 
-  private handleStreamClose(error?: FirestoreError): Promise<void> {
+  // Visible for tests
+  handleStreamClose(error?: FirestoreError): Promise<void> {
     assert(this.isStarted(), "Can't handle server close on non-started stream");
     log.debug(LOG_TAG, `close with error: ${error}`);
 

packages/firestore/src/remote/rpc_error.ts

@@ -60,7 +60,6 @@ export function isPermanentError(code: Code): boolean {
     case Code.UNAVAILABLE:
     // 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.
     case Code.UNAUTHENTICATED:
       return false;
     case Code.INVALID_ARGUMENT:

packages/firestore/test/integration/remote/stream.test.ts

@@ -15,6 +15,7 @@
  */
 
 import { expect } from 'chai';
+import { EmptyCredentialsProvider, Token } from '../../../src/api/credentials';
 import { SnapshotVersion } from '../../../src/core/snapshot_version';
 import { MutationResult } from '../../../src/model/mutation';
 import {
@@ -29,10 +30,10 @@ import {
   WatchTargetChange
 } from '../../../src/remote/watch_change';
 import { AsyncQueue, TimerId } from '../../../src/util/async_queue';
+import { Code, FirestoreError } from '../../../src/util/error';
 import { Deferred } from '../../../src/util/promise';
 import { setMutation } from '../../util/helpers';
 import { withTestDatastore } from '../util/internal_helpers';
-import { FirestoreError } from '@firebase/firestore-types';
 
 /**
  * StreamEventType combines the events that can be observed by the
@@ -152,6 +153,24 @@ describe('Watch Stream', () => {
   });
 });
 
+class MockCredentialsProvider extends EmptyCredentialsProvider {
+  private states: string[] = [];
+
+  get observedStates(): string[] {
+    return this.states;
+  }
+
+  getToken(): Promise<Token | null> {
+    this.states.push('getToken');
+    return super.getToken();
+  }
+
+  invalidateToken(): void {
+    this.states.push('invalidateToken');
+    super.invalidateToken();
+  }
+}
+
 describe('Write Stream', () => {
   let streamListener: StreamStatusListener;
 
@@ -262,4 +281,51 @@ describe('Write Stream', () => {
         });
     }, queue);
   });
+
+  it('force refreshes auth token on receiving unauthenticated error', () => {
+    const queue = new AsyncQueue();
+    const credentials = new MockCredentialsProvider();
+
+    return withTestDatastore(
+      ds => {
+        const writeStream = ds.newPersistentWriteStream();
+        writeStream.start(streamListener);
+        return streamListener
+          .awaitCallback('open')
+          .then(() => {
+            // Simulate callback from GRPC with an unauthenticated error -- this should invalidate
+            // the token.
+            writeStream.handleStreamClose(
+              new FirestoreError(Code.UNAUTHENTICATED, '')
+            );
+            return streamListener.awaitCallback('close');
+          })
+          .then(() => {
+            writeStream.start(streamListener);
+            return streamListener.awaitCallback('open');
+          })
+          .then(() => {
+            // Simulate a different error -- token should not be invalidated this time.
+            writeStream.handleStreamClose(
+              new FirestoreError(Code.UNAVAILABLE, '')
+            );
+            return streamListener.awaitCallback('close');
+          })
+          .then(() => {
+            writeStream.start(streamListener);
+            return streamListener.awaitCallback('open');
+          })
+          .then(() => {
+            expect(credentials.observedStates).to.deep.equal([
+              'getToken',
+              'invalidateToken',
+              'getToken',
+              'getToken'
+            ]);
+          });
+      },
+      queue,
+      credentials
+    );
+  });
 });

packages/firestore/test/integration/util/internal_helpers.ts

@@ -19,7 +19,10 @@ import * as firestore from '@firebase/firestore-types';
 import { DatabaseId, DatabaseInfo } from '../../../src/core/database_info';
 import { Datastore } from '../../../src/remote/datastore';
 
-import { EmptyCredentialsProvider } from '../../../src/api/credentials';
+import {
+  CredentialsProvider,
+  EmptyCredentialsProvider
+} from '../../../src/api/credentials';
 import { PlatformSupport } from '../../../src/platform/platform';
 import { AsyncQueue } from '../../../src/util/async_queue';
 import { DEFAULT_SETTINGS, DEFAULT_PROJECT_ID } from './helpers';
@@ -41,7 +44,8 @@ export function getDefaultDatabaseInfo(): DatabaseInfo {
 
 export function withTestDatastore(
   fn: (datastore: Datastore) => Promise<void>,
-  queue?: AsyncQueue
+  queue?: AsyncQueue,
+  credentialsProvider?: CredentialsProvider
 ): Promise<void> {
   const databaseInfo = getDefaultDatabaseInfo();
   return PlatformSupport.getPlatform()
@@ -53,7 +57,7 @@ export function withTestDatastore(
       const datastore = new Datastore(
         queue || new AsyncQueue(),
         conn,
-        new EmptyCredentialsProvider(),
+        credentialsProvider || new EmptyCredentialsProvider(),
         serializer
       );