Add debug token caching to app-check-exp

Author: hsubox76

packages-exp/app-check-exp/src/client.test.ts

@@ -52,7 +52,8 @@ describe('client', () => {
   });
 
   it('returns a AppCheck token', async () => {
-    useFakeTimers();
+    // To get a consistent expireTime/issuedAtTime.
+    const clock = useFakeTimers();
     fetchStub.returns(
       Promise.resolve({
         status: 200,
@@ -77,6 +78,7 @@ describe('client', () => {
       expireTimeMillis: 3600,
       issuedAtTimeMillis: 0
     });
+    clock.restore();
   });
 
   it('throws when there is a network error', async () => {

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

@@ -24,6 +24,7 @@ import {
   removeTokenListener
 } from './internal-api';
 import { Provider } from '@firebase/component';
+import { getState } from './state';
 
 /**
  * AppCheck Service class.
@@ -34,6 +35,10 @@ export class AppCheckService implements AppCheck, _FirebaseService {
     public platformLoggerProvider: Provider<'platform-logger'>
   ) {}
   _delete(): Promise<void> {
+    const { tokenObservers } = getState(this.app);
+    for (const tokenObserver of tokenObservers) {
+      removeTokenListener(this.app, tokenObserver.next);
+    }
     return Promise.resolve();
   }
 }

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

@@ -113,10 +113,6 @@ describe('internal api', () => {
         fakeRecaptchaToken
       );
       expect(token).to.deep.equal({ token: fakeRecaptchaAppCheckToken.token });
-      // TODO: Permanently fix.
-      // Small delay to prevent common test flakiness where this test runs
-      // into afterEach() sometimes
-      await new Promise(resolve => setTimeout(resolve, 50));
     });
 
     it('resolves with a dummy token and an error if failed to get a token', async () => {
@@ -148,7 +144,7 @@ describe('internal api', () => {
     it('notifies listeners using cached token', async () => {
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
-        isTokenAutoRefreshEnabled: true
+        isTokenAutoRefreshEnabled: false
       });
 
       const clock = useFakeTimers();
@@ -216,6 +212,7 @@ describe('internal api', () => {
     });
 
     it('calls 3P error handler if there is an error getting a token', async () => {
+      stub(console, 'error');
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
         isTokenAutoRefreshEnabled: true
@@ -239,6 +236,7 @@ describe('internal api', () => {
     });
 
     it('ignores listeners that throw', async () => {
+      stub(console, 'error');
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
         isTokenAutoRefreshEnabled: true
@@ -247,9 +245,7 @@ describe('internal api', () => {
       stub(client, 'exchangeToken').returns(
         Promise.resolve(fakeRecaptchaAppCheckToken)
       );
-      const listener1 = (): void => {
-        throw new Error();
-      };
+      const listener1 = stub().throws(new Error());
       const listener2 = spy();
 
       addTokenListener(
@@ -265,6 +261,9 @@ describe('internal api', () => {
 
       await getToken(appCheck as AppCheckService);
 
+      expect(listener1).to.be.calledWith({
+        token: fakeRecaptchaAppCheckToken.token
+      });
       expect(listener2).to.be.calledWith({
         token: fakeRecaptchaAppCheckToken.token
       });
@@ -344,7 +343,7 @@ describe('internal api', () => {
       });
     });
 
-    it('exchanges debug token if in debug mode', async () => {
+    it('exchanges debug token if in debug mode and there is no cached token', async () => {
       const exchangeTokenStub: SinonStub = stub(
         client,
         'exchangeToken'
@@ -393,15 +392,10 @@ describe('internal api', () => {
       expect(getState(app).tokenRefresher?.isRunning()).to.be.true;
     });
 
-    it('notifies the listener with the valid token in memory immediately', done => {
+    it('notifies the listener with the valid token in memory immediately', async () => {
       const clock = useFakeTimers();
-      const fakeListener: AppCheckTokenListener = token => {
-        expect(token).to.deep.equal({
-          token: `fake-memory-app-check-token`
-        });
-        clock.restore();
-        done();
-      };
+
+      const listener = stub();
 
       setState(app, {
         ...getState(app),
@@ -415,12 +409,16 @@ describe('internal api', () => {
       addTokenListener(
         { app } as AppCheckService,
         ListenerType.INTERNAL,
-        fakeListener
+        listener
       );
+      await clock.runAllAsync();
+      expect(listener).to.be.calledWith({
+        token: 'fake-memory-app-check-token'
+      });
+      clock.restore();
     });
 
     it('notifies the listener with the valid token in storage', done => {
-      const clock = useFakeTimers();
       const appCheck = initializeAppCheck(app, {
         provider: new ReCaptchaV3Provider(FAKE_SITE_KEY),
         isTokenAutoRefreshEnabled: true
@@ -437,60 +435,15 @@ describe('internal api', () => {
         expect(token).to.deep.equal({
           token: `fake-cached-app-check-token`
         });
-        clock.restore();
-        done();
-      };
-
-      addTokenListener(
-        appCheck as AppCheckService,
-        ListenerType.INTERNAL,
-        fakeListener
-      );
-
-      clock.tick(1);
-    });
-
-    it('notifies the listener with the debug token immediately', done => {
-      const fakeListener: AppCheckTokenListener = token => {
-        expect(token).to.deep.equal({
-          token: `my-debug-token`
-        });
         done();
       };
 
-      const debugState = getDebugState();
-      debugState.enabled = true;
-      debugState.token = new Deferred();
-      debugState.token.resolve('my-debug-token');
-
-      const appCheck = initializeAppCheck(app, {
-        provider: new ReCaptchaV3Provider(FAKE_SITE_KEY)
-      });
       addTokenListener(
         appCheck as AppCheckService,
         ListenerType.INTERNAL,
         fakeListener
       );
     });
-
-    it('does NOT start token refresher in debug mode', () => {
-      const debugState = getDebugState();
-      debugState.enabled = true;
-      debugState.token = new Deferred();
-      debugState.token.resolve('my-debug-token');
-
-      const appCheck = initializeAppCheck(app, {
-        provider: new ReCaptchaV3Provider(FAKE_SITE_KEY)
-      });
-      addTokenListener(
-        appCheck as AppCheckService,
-        ListenerType.INTERNAL,
-        () => {}
-      );
-
-      const state = getState(app);
-      expect(state.tokenRefresher).is.undefined;
-    });
   });
 
   describe('removeTokenListener', () => {

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

@@ -23,7 +23,7 @@ import {
   ListenerType
 } from './types';
 import { AppCheckTokenListener } from './public-types';
-import { getDebugState, getState, setState } from './state';
+import { getState, setState } from './state';
 import { TOKEN_REFRESH_TIME } from './constants';
 import { Refresher } from './proactive-refresh';
 import { ensureActivated } from './util';
@@ -63,25 +63,17 @@ export async function getToken(
 ): Promise<AppCheckTokenResult> {
   const app = appCheck.app;
   ensureActivated(app);
-  /**
-   * DEBUG MODE
-   * return the debug token directly
-   */
-  if (isDebugMode()) {
-    const tokenFromDebugExchange: AppCheckTokenInternal = await exchangeToken(
-      getExchangeDebugTokenRequest(app, await getDebugToken()),
-      appCheck.platformLoggerProvider
-    );
-    return { token: tokenFromDebugExchange.token };
-  }
 
   const state = getState(app);
 
+  /**
+   * First check if there is a token in memory from a previous `getToken()` call.
+   */
   let token: AppCheckTokenInternal | undefined = state.token;
   let error: Error | undefined = undefined;
 
   /**
-   * try to load token from indexedDB if it's the first time this function is called
+   * 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`.
@@ -95,13 +87,30 @@ export async function getToken(
     }
   }
 
-  // return the cached token if it's valid
+  // Return the cached token (from either memory or indexedDB) if it's valid
   if (!forceRefresh && token && isValid(token)) {
     return {
       token: token.token
     };
   }
 
+  /**
+   * DEBUG MODE
+   * If debug mode is set, and there is no cached token, fetch a new App
+   * Check token using the debug token, and return it directly.
+   */
+  if (isDebugMode()) {
+    const tokenFromDebugExchange: AppCheckTokenInternal = await exchangeToken(
+      getExchangeDebugTokenRequest(app, await getDebugToken()),
+      appCheck.platformLoggerProvider
+    );
+    // Write debug token to indexedDB.
+    await writeTokenToStorage(app, tokenFromDebugExchange);
+    // Write debug token to state.
+    setState(app, { ...state, token: tokenFromDebugExchange });
+    return { token: tokenFromDebugExchange.token };
+  }
+
   /**
    * request a new token
    */
@@ -125,7 +134,7 @@ export async function getToken(
     interopTokenResult = {
       token: token.token
     };
-    // write the new token to the memory state as well ashe persistent storage.
+    // write the new token to the memory state as well as the persistent storage.
     // Only do it if we got a valid new token
     setState(app, { ...state, token });
     await writeTokenToStorage(app, token);
@@ -152,50 +161,31 @@ export function addTokenListener(
     ...state,
     tokenObservers: [...state.tokenObservers, tokenObserver]
   };
-
   /**
-   * DEBUG MODE
-   *
-   * invoke the listener once with the debug token.
+   * Invoke the listener with the valid token, then start the token refresher
    */
-  if (isDebugMode()) {
-    const debugState = getDebugState();
-    if (debugState.enabled && debugState.token) {
-      debugState.token.promise
-        .then(token => listener({ token }))
-        .catch(() => {
-          /* we don't care about exceptions thrown in listeners */
-        });
-    }
-  } else {
-    /**
-     * PROD MODE
-     *
-     * invoke the listener with the valid token, then start the token refresher
-     */
-    if (!newState.tokenRefresher) {
-      const tokenRefresher = createTokenRefresher(appCheck);
-      newState.tokenRefresher = tokenRefresher;
-    }
+  if (!newState.tokenRefresher) {
+    const tokenRefresher = createTokenRefresher(appCheck);
+    newState.tokenRefresher = tokenRefresher;
+  }
 
-    // Create the refresher but don't start it if `isTokenAutoRefreshEnabled`
-    // is not true.
-    if (
-      !newState.tokenRefresher.isRunning() &&
-      state.isTokenAutoRefreshEnabled === true
-    ) {
-      newState.tokenRefresher.start();
-    }
+  // Create the refresher but don't start it if `isTokenAutoRefreshEnabled`
+  // is not true.
+  if (
+    !newState.tokenRefresher.isRunning() &&
+    state.isTokenAutoRefreshEnabled === true
+  ) {
+    newState.tokenRefresher.start();
+  }
 
-    // invoke the listener async immediately if there is a valid token
-    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 */
-        });
-    }
+  // invoke the listener async immediately if there is a valid token
+  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 */
+      });
   }
 
   setState(app, newState);