Allow FirstPartyAuth to specify a token factory func. (#4773)

* Allow firstparty credentials to specify an authToken factory that is used in lieu of direct GAPI.

* Remove stray import

* Add return type to private method

Co-authored-by: wu-hui <[email protected]>

Author: richieforeman

packages/firestore/src/api/credentials.ts

@@ -37,12 +37,15 @@ import { Deferred } from '../util/promise';
 // TODO(mikelehen): This should be split into multiple files and probably
 // moved to an auth/ folder to match other platforms.
 
+export type AuthTokenFactory = () => string;
+
 export interface FirstPartyCredentialsSettings {
   // These are external types. Prevent minification.
   ['type']: 'gapi';
   ['client']: unknown;
   ['sessionIndex']: string;
   ['iamToken']: string | null;
+  ['authTokenFactory']: AuthTokenFactory | null;
 }
 
 export interface ProviderCredentialsSettings {
@@ -397,11 +400,40 @@ export class FirstPartyToken implements Token {
   user = User.FIRST_PARTY;
   headers = new Map();
 
-  constructor(gapi: Gapi, sessionIndex: string, iamToken: string | null) {
-    this.headers.set('X-Goog-AuthUser', sessionIndex);
-    const authHeader = gapi['auth']['getAuthHeaderValueForFirstParty']([]);
-    if (authHeader) {
-      this.headers.set('Authorization', authHeader);
+  constructor(
+    private readonly gapi: Gapi,
+    private readonly sessionIndex: string,
+    private readonly iamToken: string | null,
+    private readonly authTokenFactory: AuthTokenFactory | null
+  ) {}
+
+  /** Gets an authorization token, using a provided factory function, or falling back to First Party GAPI. */
+  private getAuthToken(): string | null {
+    if (this.authTokenFactory) {
+      return this.authTokenFactory();
+    } else {
+      // Make sure this really is a Gapi client.
+      hardAssert(
+        !!(
+          typeof this.gapi === 'object' &&
+          this.gapi !== null &&
+          this.gapi['auth'] &&
+          this.gapi['auth']['getAuthHeaderValueForFirstParty']
+        ),
+        'unexpected gapi interface'
+      );
+      return this.gapi['auth']['getAuthHeaderValueForFirstParty']([]);
+    }
+  }
+
+  get authHeaders(): { [header: string]: string } {
+    const headers: { [header: string]: string } = {
+      'X-Goog-AuthUser': this.sessionIndex
+    };
+    // Use array notation to prevent minification
+    const authHeaderTokenValue = this.getAuthToken();
+    if (authHeaderTokenValue) {
+      headers['Authorization'] = authHeaderTokenValue;
     }
     if (iamToken) {
       this.headers.set('X-Goog-Iam-Authorization-Token', iamToken);
@@ -420,12 +452,18 @@ export class FirstPartyAuthCredentialsProvider
   constructor(
     private gapi: Gapi,
     private sessionIndex: string,
-    private iamToken: string | null
+    private iamToken: string | null,
+    private authTokenFactory: AuthTokenFactory | null
   ) {}
 
   getToken(): Promise<Token | null> {
     return Promise.resolve(
-      new FirstPartyToken(this.gapi, this.sessionIndex, this.iamToken)
+      new FirstPartyToken(
+        this.gapi,
+        this.sessionIndex,
+        this.iamToken,
+        this.authTokenFactory
+      )
     );
   }
 
@@ -634,20 +672,11 @@ export function makeAuthCredentialsProvider(
   switch (credentials['type']) {
     case 'gapi':
       const client = credentials['client'] as Gapi;
-      // Make sure this really is a Gapi client.
-      hardAssert(
-        !!(
-          typeof client === 'object' &&
-          client !== null &&
-          client['auth'] &&
-          client['auth']['getAuthHeaderValueForFirstParty']
-        ),
-        'unexpected gapi interface'
-      );
-      return new FirstPartyAuthCredentialsProvider(
+      return new FirstPartyCredentialsProvider(
         client,
         credentials['sessionIndex'] || '0',
-        credentials['iamToken'] || null
+        credentials['iamToken'] || null,
+        credentials['authTokenFactory'] || null
       );
 
     case 'provider':