Store FirebaseToken to persistence - Required for BYO-CIAM (#9138)

* Maintain persistence for firebaseToken in AuthImpl object

* Adding Unit Test for auth_impl firebaseToken

Author: mansisampat

packages/auth/demo/src/index.js

@@ -1526,7 +1526,7 @@ function onExchangeToken(event) {
 
   exchangeCIAMToken(byoCiamInput.value)
     .then(response => {
-      byoCiamResult.textContent = response.accessToken;
+      byoCiamResult.textContent = response;
       console.log('Token:', response);
     })
     .catch(error => {
@@ -2086,7 +2086,7 @@ function initApp() {
   const regionalApp = initializeApp(config, `${auth.name}-rgcip`);
 
   regionalAuth = initializeAuth(regionalApp, {
-    persistence: inMemoryPersistence,
+    persistence: indexedDBLocalPersistence,
     popupRedirectResolver: browserPopupRedirectResolver,
     tenantConfig: tenantConfig
   });

packages/auth/src/core/auth/auth_impl.test.ts

@@ -46,8 +46,12 @@ import { mockEndpointWithParams } from '../../../test/helpers/api/helper';
 import { Endpoint, RecaptchaClientType, RecaptchaVersion } from '../../api';
 import * as mockFetch from '../../../test/helpers/mock_fetch';
 import { AuthErrorCode } from '../errors';
-import { PasswordValidationStatus } from '../../model/public_types';
+import {
+  FirebaseToken,
+  PasswordValidationStatus
+} from '../../model/public_types';
 import { PasswordPolicyImpl } from './password_policy_impl';
+import { PersistenceUserManager } from '../persistence/persistence_user_manager';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -150,6 +154,153 @@ describe('core/auth/auth_impl', () => {
     });
   });
 
+  describe('#updateFirebaseToken', () => {
+    const token: FirebaseToken = {
+      token: 'test-token',
+      expirationTime: 123456789
+    };
+
+    it('sets the field on the auth object', async () => {
+      await auth._updateFirebaseToken(token);
+      expect((auth as any).firebaseToken).to.eql(token);
+    });
+
+    it('calls persistence._set with correct values', async () => {
+      await auth._updateFirebaseToken(token);
+      expect(persistenceStub._set).to.have.been.calledWith(
+        'firebase:persistence-token:api-key:test-app', // key
+        {
+          token: token.token,
+          expirationTime: token.expirationTime
+        }
+      );
+    });
+
+    it('setting to null triggers persistence._remove', async () => {
+      await auth._updateFirebaseToken(null);
+      expect(persistenceStub._remove).to.have.been.calledWith(
+        'firebase:persistence-token:api-key:test-app'
+      );
+    });
+
+    it('orders async updates correctly', async () => {
+      const tokens: FirebaseToken[] = Array.from({ length: 5 }, (_, i) => ({
+        token: `token-${i}`,
+        expirationTime: Date.now() + i
+      }));
+
+      persistenceStub._set.callsFake(() => {
+        return new Promise(resolve => {
+          setTimeout(() => resolve(), 1);
+        });
+      });
+
+      await Promise.all(tokens.map(t => auth._updateFirebaseToken(t)));
+
+      for (let i = 0; i < tokens.length; i++) {
+        expect(persistenceStub._set.getCall(i)).to.have.been.calledWith(
+          'firebase:persistence-token:api-key:test-app',
+          {
+            token: tokens[i].token,
+            expirationTime: tokens[i].expirationTime
+          }
+        );
+      }
+    });
+
+    it('throws if persistence._set fails', async () => {
+      persistenceStub._set.rejects(new Error('fail'));
+      await expect(auth._updateFirebaseToken(token)).to.be.rejectedWith('fail');
+    });
+
+    it('throws if persistence._remove fails', async () => {
+      persistenceStub._remove.rejects(new Error('remove fail'));
+      await expect(auth._updateFirebaseToken(null)).to.be.rejectedWith(
+        'remove fail'
+      );
+    });
+  });
+
+  describe('#_initializeWithPersistence', () => {
+    let mockToken: FirebaseToken;
+    let persistenceManager: any;
+    let subscription: any;
+    let authImpl: AuthImpl;
+
+    beforeEach(() => {
+      mockToken = {
+        token: 'test-token',
+        expirationTime: 123456789
+      };
+
+      persistenceManager = {
+        getFirebaseToken: sinon.stub().resolves(mockToken),
+        getCurrentUser: sinon.stub().resolves(null),
+        setCurrentUser: sinon.stub().resolves(),
+        removeCurrentUser: sinon.stub().resolves(),
+        getPersistence: sinon.stub().returns('LOCAL')
+      };
+
+      subscription = {
+        next: sinon.spy()
+      };
+
+      sinon.stub(PersistenceUserManager, 'create').resolves(persistenceManager);
+
+      authImpl = new AuthImpl(
+        FAKE_APP,
+        FAKE_HEARTBEAT_CONTROLLER_PROVIDER,
+        FAKE_APP_CHECK_CONTROLLER_PROVIDER,
+        {
+          apiKey: FAKE_APP.options.apiKey!,
+          apiHost: DefaultConfig.API_HOST,
+          apiScheme: DefaultConfig.API_SCHEME,
+          tokenApiHost: DefaultConfig.TOKEN_API_HOST,
+          clientPlatform: ClientPlatform.BROWSER,
+          sdkClientVersion: 'v'
+        }
+      );
+
+      (authImpl as any).firebaseTokenSubscription = subscription;
+    });
+
+    afterEach(() => {
+      sinon.restore();
+    });
+
+    it('should load the firebaseToken from persistence and set it', async () => {
+      await authImpl._initializeWithPersistence([
+        persistenceStub as PersistenceInternal
+      ]);
+
+      expect(persistenceManager.getFirebaseToken).to.have.been.called;
+      expect((authImpl as any).firebaseToken).to.eql(mockToken);
+      expect(subscription.next).to.have.been.calledWith(mockToken);
+    });
+
+    it('should set firebaseToken to null if getFirebaseToken returns undefined', async () => {
+      persistenceManager.getFirebaseToken.resolves(undefined);
+
+      await authImpl._initializeWithPersistence([
+        persistenceStub as PersistenceInternal
+      ]);
+
+      expect((authImpl as any).firebaseToken).to.be.null;
+      expect(subscription.next).to.have.been.calledWith(null);
+    });
+
+    it('should set firebaseToken to null if getFirebaseToken returns null', async () => {
+      persistenceManager.getFirebaseToken.resolves(null);
+
+      await authImpl._initializeWithPersistence([
+        persistenceStub as PersistenceInternal
+      ]);
+
+      expect((authImpl as any).firebaseToken).to.be.null;
+      expect(subscription.next).to.have.been.calledWith(null);
+    });
+  });
+
   describe('#signOut', () => {
     it('sets currentUser to null, calls remove', async () => {
       await auth._updateCurrentUser(testUser(auth, 'test'));

packages/auth/src/core/auth/auth_impl.ts

@@ -106,6 +106,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
   private persistenceManager?: PersistenceUserManager;
   private redirectPersistenceManager?: PersistenceUserManager;
   private authStateSubscription = new Subscription<User>(this);
+  private firebaseTokenSubscription = new Subscription<FirebaseToken>(this);
   private idTokenSubscription = new Subscription<User>(this);
   private readonly beforeStateQueue = new AuthMiddlewareQueue(this);
   private redirectUser: UserInternal | null = null;
@@ -195,6 +196,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
       }
 
       await this.initializeCurrentUser(popupRedirectResolver);
+      await this.initializeFirebaseToken();
 
       this.lastNotifiedUid = this.currentUser?.uid || null;
 
@@ -403,6 +405,12 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     return this.directlySetCurrentUser(user);
   }
 
+  private async initializeFirebaseToken(): Promise<void> {
+    this.firebaseToken =
+      (await this.persistenceManager?.getFirebaseToken()) ?? null;
+    this.firebaseTokenSubscription.next(this.firebaseToken);
+  }
+
   useDeviceLanguage(): void {
     this.languageCode = _getUserLanguage();
   }
@@ -461,6 +469,12 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     firebaseToken: FirebaseToken | null
   ): Promise<void> {
     this.firebaseToken = firebaseToken;
+    this.firebaseTokenSubscription.next(firebaseToken);
+    if (firebaseToken) {
+      await this.assertedPersistence.setFirebaseToken(firebaseToken);
+    } else {
+      await this.assertedPersistence.removeFirebaseToken();
+    }
   }
 
   async signOut(): Promise<void> {

packages/auth/src/core/persistence/persistence_user_manager.ts

@@ -17,6 +17,7 @@
 
 import { getAccountInfo } from '../../api/account_management/account';
 import { ApiKey, AppName, AuthInternal } from '../../model/auth';
+import { FirebaseToken } from '../../model/public_types';
 import { UserInternal } from '../../model/user';
 import { PersistedBlob, PersistenceInternal } from '../persistence';
 import { UserImpl } from '../user/user_impl';
@@ -27,7 +28,8 @@ export const enum KeyName {
   AUTH_USER = 'authUser',
   AUTH_EVENT = 'authEvent',
   REDIRECT_USER = 'redirectUser',
-  PERSISTENCE_USER = 'persistence'
+  PERSISTENCE_USER = 'persistence',
+  PERSISTENCE_TOKEN = 'persistence-token'
 }
 export const enum Namespace {
   PERSISTENCE = 'firebase'
@@ -44,6 +46,7 @@ export function _persistenceKeyName(
 export class PersistenceUserManager {
   private readonly fullUserKey: string;
   private readonly fullPersistenceKey: string;
+  private readonly firebaseTokenPersistenceKey: string;
   private readonly boundEventHandler: () => void;
 
   private constructor(
@@ -58,6 +61,11 @@ export class PersistenceUserManager {
       config.apiKey,
       name
     );
+    this.firebaseTokenPersistenceKey = _persistenceKeyName(
+      KeyName.PERSISTENCE_TOKEN,
+      config.apiKey,
+      name
+    );
     this.boundEventHandler = auth._onStorageEvent.bind(auth);
     this.persistence._addListener(this.fullUserKey, this.boundEventHandler);
   }
@@ -66,6 +74,32 @@ export class PersistenceUserManager {
     return this.persistence._set(this.fullUserKey, user.toJSON());
   }
 
+  setFirebaseToken(firebaseToken: FirebaseToken): Promise<void> {
+    return this.persistence._set(this.firebaseTokenPersistenceKey, {
+      token: firebaseToken.token,
+      expirationTime: firebaseToken.expirationTime
+    });
+  }
+
+  async getFirebaseToken(): Promise<FirebaseToken | null> {
+    const blob = await this.persistence._get<PersistedBlob>(
+      this.firebaseTokenPersistenceKey
+    );
+    if (!blob) {
+      return null;
+    }
+    const token = blob.token as string;
+    const expirationTime = blob.expirationTime as number;
+    return {
+      token,
+      expirationTime
+    };
+  }
+
+  removeFirebaseToken(): Promise<void> {
+    return this.persistence._remove(this.firebaseTokenPersistenceKey);
+  }
+
   async getCurrentUser(): Promise<UserInternal | null> {
     const blob = await this.persistence._get<PersistedBlob | string>(
       this.fullUserKey