Add cached token changes to v9 appcheck (#5226)

Author: hsubox76

packages-exp/app-check-exp/src/api.ts

@@ -31,8 +31,10 @@ import { AppCheckProvider, ListenerType } from './types';
 import {
   getToken as getTokenInternal,
   addTokenListener,
-  removeTokenListener
+  removeTokenListener,
+  isValid
 } from './internal-api';
+import { readTokenFromStorage } from './storage';
 
 declare module '@firebase/component' {
   interface NameServiceMapping {
@@ -95,7 +97,13 @@ function _activate(
   const state = getState(app);
 
   const newState: AppCheckState = { ...state, activated: true };
-  newState.provider = provider;
+  newState.provider = provider; // Read cached token from storage if it exists and store it in memory.
+  newState.cachedTokenPromise = readTokenFromStorage(app).then(cachedToken => {
+    if (cachedToken && isValid(cachedToken)) {
+      setState(app, { ...getState(app), token: cachedToken });
+    }
+    return cachedToken;
+  });
 
   // Use value of global `automaticDataCollectionEnabled` (which
   // itself defaults to false if not specified in config) if

packages-exp/app-check-exp/src/internal-api.test.ts

@@ -142,13 +142,13 @@ describe('internal api', () => {
     });
 
     it('notifies listeners using cached token', async () => {
+      storageReadStub.resolves(fakeCachedAppCheckToken);
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
         isTokenAutoRefreshEnabled: false
       });
 
       const clock = useFakeTimers();
-      storageReadStub.returns(Promise.resolve(fakeCachedAppCheckToken));
 
       const listener1 = spy();
       const listener2 = spy();
@@ -271,12 +271,12 @@ describe('internal api', () => {
 
     it('loads persisted token to memory and returns it', async () => {
       const clock = useFakeTimers();
+
+      storageReadStub.resolves(fakeCachedAppCheckToken);
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY)
       });
 
-      storageReadStub.returns(Promise.resolve(fakeCachedAppCheckToken));
-
       const clientStub = stub(client, 'exchangeToken');
 
       expect(getState(app).token).to.equal(undefined);
@@ -367,6 +367,10 @@ describe('internal api', () => {
   describe('addTokenListener', () => {
     it('adds token listeners', () => {
       const listener = (): void => {};
+      setState(app, {
+        ...getState(app),
+        cachedTokenPromise: Promise.resolve(undefined)
+      });
 
       addTokenListener(
         { app } as AppCheckService,
@@ -379,7 +383,11 @@ describe('internal api', () => {
 
     it('starts proactively refreshing token after adding the first listener', () => {
       const listener = (): void => {};
-      setState(app, { ...getState(app), isTokenAutoRefreshEnabled: true });
+      setState(app, {
+        ...getState(app),
+        isTokenAutoRefreshEnabled: true,
+        cachedTokenPromise: Promise.resolve(undefined)
+      });
       expect(getState(app).tokenObservers.length).to.equal(0);
       expect(getState(app).tokenRefresher).to.equal(undefined);
 
@@ -419,17 +427,15 @@ describe('internal api', () => {
     });
 
     it('notifies the listener with the valid token in storage', done => {
+      storageReadStub.resolves({
+        token: `fake-cached-app-check-token`,
+        expireTimeMillis: Date.now() + 60000,
+        issuedAtTimeMillis: 0
+      });
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
         isTokenAutoRefreshEnabled: true
       });
-      storageReadStub.returns(
-        Promise.resolve({
-          token: `fake-cached-app-check-token`,
-          expireTimeMillis: Date.now() + 60000,
-          issuedAtTimeMillis: 0
-        })
-      );
 
       const fakeListener: AppCheckTokenListener = token => {
         expect(token).to.deep.equal({
@@ -449,6 +455,10 @@ describe('internal api', () => {
   describe('removeTokenListener', () => {
     it('should remove token listeners', () => {
       const listener = (): void => {};
+      setState(app, {
+        ...getState(app),
+        cachedTokenPromise: Promise.resolve(undefined)
+      });
       addTokenListener(
         { app } as AppCheckService,
         ListenerType.INTERNAL,
@@ -463,6 +473,10 @@ describe('internal api', () => {
     it('should stop proactively refreshing token after deleting the last listener', () => {
       const listener = (): void => {};
       setState(app, { ...getState(app), isTokenAutoRefreshEnabled: true });
+      setState(app, {
+        ...getState(app),
+        cachedTokenPromise: Promise.resolve(undefined)
+      });
 
       addTokenListener(
         { app } as AppCheckService,

packages-exp/app-check-exp/src/internal-api.ts

@@ -28,7 +28,7 @@ import { TOKEN_REFRESH_TIME } from './constants';
 import { Refresher } from './proactive-refresh';
 import { ensureActivated } from './util';
 import { exchangeToken, getExchangeDebugTokenRequest } from './client';
-import { writeTokenToStorage, readTokenFromStorage } from './storage';
+import { writeTokenToStorage } from './storage';
 import { getDebugToken, isDebugMode } from './debug';
 import { base64 } from '@firebase/util';
 import { logger } from './logger';
@@ -76,8 +76,8 @@ export async function getToken(
    * If there is no token in memory, try to load token from indexedDB.
    */
   if (!token) {
-    // readTokenFromStorage() always resolves. In case of an error, it resolves with `undefined`.
-    const cachedToken = await readTokenFromStorage(app);
+    // cachedTokenPromise contains the token found in IndexedDB or undefined if not found.
+    const cachedToken = await state.cachedTokenPromise;
     if (cachedToken && isValid(cachedToken)) {
       token = cachedToken;
 
@@ -175,14 +175,28 @@ export function addTokenListener(
     newState.tokenRefresher.start();
   }
 
-  // invoke the listener async immediately if there is a valid token
+  // Invoke the listener async immediately if there is a valid token
+  // in memory.
   if (state.token && isValid(state.token)) {
     const validToken = state.token;
     Promise.resolve()
       .then(() => listener({ token: validToken.token }))
       .catch(() => {
         /* we don't care about exceptions thrown in listeners */
       });
+  } else if (state.token == null) {
+    // Only check cache if there was no token. If the token was invalid,
+    // skip this and rely on exchange endpoint.
+    void state
+      .cachedTokenPromise! // Storage token promise. Always populated in `activate()`.
+      .then(cachedToken => {
+        if (cachedToken && isValid(cachedToken)) {
+          listener({ token: cachedToken.token });
+        }
+      })
+      .catch(() => {
+        /** Ignore errors in listeners. */
+      });
   }
 
   setState(app, newState);
@@ -288,7 +302,7 @@ function notifyTokenListeners(
   }
 }
 
-function isValid(token: AppCheckTokenInternal): boolean {
+export function isValid(token: AppCheckTokenInternal): boolean {
   return token.expireTimeMillis - Date.now() > 0;
 }
 

packages-exp/app-check-exp/src/state.ts

@@ -29,6 +29,7 @@ export interface AppCheckState {
   tokenObservers: AppCheckTokenObserver[];
   provider?: AppCheckProvider;
   token?: AppCheckTokenInternal;
+  cachedTokenPromise?: Promise<AppCheckTokenInternal | undefined>;
   tokenRefresher?: Refresher;
   reCAPTCHAState?: ReCAPTCHAState;
   isTokenAutoRefreshEnabled?: boolean;