Fixed appcheck implementation

Author: maneesht

common/api-review/data-connect.api.md

@@ -4,6 +4,9 @@
 
 ```ts
 
+import { AppCheckInternalComponentName } from '@firebase/app-check-interop-types';
+import { AppCheckTokenListener } from '@firebase/app-check-interop-types';
+import { AppCheckTokenResult } from '@firebase/app-check-interop-types';
 import { FirebaseApp } from '@firebase/app';
 import { FirebaseAuthInternalName } from '@firebase/auth-interop-types';
 import { FirebaseAuthTokenData } from '@firebase/auth-interop-types';
@@ -12,6 +15,15 @@ import { FirebaseOptions } from '@firebase/app-types';
 import { LogLevelString } from '@firebase/logger';
 import { Provider } from '@firebase/component';
 
+// @public
+export class AppCheckTokenProvider {
+    constructor(appName_: string, appCheckProvider?: Provider<AppCheckInternalComponentName>);
+    // (undocumented)
+    addTokenChangeListener(listener: AppCheckTokenListener): void;
+    // (undocumented)
+    getToken(forceRefresh?: boolean): Promise<AppCheckTokenResult>;
+}
+
 // @public (undocumented)
 export type AuthTokenListener = (token: string | null) => void;
 
@@ -46,7 +58,7 @@ export interface ConnectorConfig {
 
 // @public
 export class DataConnect {
-    constructor(app: FirebaseApp, dataConnectOptions: DataConnectOptions, _authProvider: Provider<FirebaseAuthInternalName>);
+    constructor(app: FirebaseApp, dataConnectOptions: DataConnectOptions, _authProvider: Provider<FirebaseAuthInternalName>, _appCheckProvider: Provider<AppCheckInternalComponentName>);
     // (undocumented)
     readonly app: FirebaseApp;
     // (undocumented)
@@ -283,7 +295,7 @@ export function terminate(dataConnect: DataConnect): Promise<void>;
 export function toQueryRef<Data, Variables>(serializedRef: SerializedRef<Data, Variables>): QueryRef<Data, Variables>;
 
 // @public (undocumented)
-export type TransportClass = new (options: DataConnectOptions, apiKey?: string, authProvider?: AuthTokenProvider, transportOptions?: TransportOptions, _isUsingGen?: boolean) => DataConnectTransport;
+export type TransportClass = new (options: DataConnectOptions, apiKey?: string, authProvider?: AuthTokenProvider, appCheckProvider?: AppCheckTokenProvider, transportOptions?: TransportOptions, _isUsingGen?: boolean) => DataConnectTransport;
 
 // @public
 export interface TransportOptions {

packages/data-connect/src/api/index.ts

@@ -20,5 +20,6 @@ export * from './DataConnect';
 export * from './Reference';
 export * from './Mutation';
 export * from './query';
+export * from '../core/AppCheckTokenProvider';
 export { setLogLevel } from '../logger';
 export { validateArgs } from '../util/validateArgs';

packages/data-connect/src/network/transport/rest.ts

@@ -35,7 +35,6 @@ export class RESTTransport implements DataConnectTransport {
   private _serviceName: string;
   private _accessToken: string | null = null;
   private _appCheckToken: string | null = null;
-  private _authInitialized = false;
   private _lastToken: string | null = null;
   constructor(
     options: DataConnectOptions,
@@ -103,24 +102,25 @@ export class RESTTransport implements DataConnectTransport {
     this._accessToken = newToken;
   }
 
-  getWithAuth(forceToken = false): Promise<string> {
+  async getWithAuth(forceToken = false): Promise<string> {
     let starterPromise: Promise<string | null> = new Promise(resolve =>
       resolve(this._accessToken)
     );
-    if (!this._authInitialized) {
-      if (this.authProvider) {
-        starterPromise = this.authProvider
-          .getToken(/*forceToken=*/ forceToken)
-          .then(data => {
-            if (!data) {
-              return null;
-            }
-            this._accessToken = data.accessToken;
-            return this._accessToken;
-          });
-      } else {
-        starterPromise = new Promise(resolve => resolve(''));
-      }
+    if (this.appCheckProvider) {
+      this._appCheckToken = (await this.appCheckProvider.getToken()).token;
+    }
+    if (this.authProvider) {
+      starterPromise = this.authProvider
+        .getToken(/*forceToken=*/ forceToken)
+        .then(data => {
+          if (!data) {
+            return null;
+          }
+          this._accessToken = data.accessToken;
+          return this._accessToken;
+        });
+    } else {
+      starterPromise = new Promise(resolve => resolve(''));
     }
     return starterPromise;
   }