Implement token refresh mechanism for BYO-CIAM

Author: mansisampat

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

@@ -38,7 +38,9 @@ import {
   Unsubscribe,
   PasswordValidationStatus,
   TenantConfig,
-  FirebaseToken
+  FirebaseToken,
+  RefreshIdpTokenResult,
+  TokenRefreshHandler
 } from '../../model/public_types';
 import {
   createSubscribe,
@@ -85,6 +87,7 @@ import { PasswordPolicyInternal } from '../../model/password_policy';
 import { PasswordPolicyImpl } from './password_policy_impl';
 import { getAccountInfo } from '../../api/account_management/account';
 import { UserImpl } from '../user/user_impl';
+import { exchangeToken } from '../strategies/exhange_token';
 
 interface AsyncAction {
   (): Promise<void>;
@@ -102,6 +105,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
   currentUser: User | null = null;
   emulatorConfig: EmulatorConfig | null = null;
   firebaseToken: FirebaseToken | null = null;
+  private tokenRefreshHandler?: TokenRefreshHandler;
   private operations = Promise.resolve();
   private persistenceManager?: PersistenceUserManager;
   private redirectPersistenceManager?: PersistenceUserManager;
@@ -112,6 +116,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
   private redirectUser: UserInternal | null = null;
   private isProactiveRefreshEnabled = false;
   private readonly EXPECTED_PASSWORD_POLICY_SCHEMA_VERSION: number = 1;
+  private readonly TOKEN_EXPIRATION_BUFFER = 30_000;
 
   // Any network calls will set this to true and prevent subsequent emulator
   // initialization
@@ -210,6 +215,10 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     return this._initializationPromise;
   }
 
+  setTokenRefreshHandler(tokenRefreshHandler: TokenRefreshHandler): void {
+      this.tokenRefreshHandler = tokenRefreshHandler;
+  }
+
   /**
    * If the persistence is changed in another window, the user manager will let us know
    */
@@ -240,6 +249,35 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     await this._updateCurrentUser(user, /* skipBeforeStateCallbacks */ true);
   }
 
+  async getFirebaseAccessToken(forceRefresh?: boolean): 
+    Promise<FirebaseToken | null> {
+      const firebaseAccessToken =
+        (await this.persistenceManager?.getFirebaseToken()) ?? null;
+
+      if (
+        firebaseAccessToken &&
+        this.isFirebaseAccessTokenValid(firebaseAccessToken) &&
+        !forceRefresh
+      ) {
+        this.firebaseToken = firebaseAccessToken;
+        this.firebaseTokenSubscription.next(this.firebaseToken);
+        return firebaseAccessToken;
+      }
+
+      if (firebaseAccessToken && this.tokenRefreshHandler) {
+        // Resets the Firebase Access Token to null i.e. logs out the user.
+        await this._updateFirebaseToken(null);
+        // Awaits for the callback method to execute. The callback method
+        // is responsible for performing the exchangeToken(auth, valid3pIdpToken)
+        const result: RefreshIdpTokenResult = await this.tokenRefreshHandler.refreshIdpToken();
+        _assert(result.idToken && result.idpConfigId, AuthErrorCode.INVALID_CREDENTIAL);
+        await exchangeToken(this, result.idpConfigId, result.idToken);
+        return this.getFirebaseAccessToken(false);
+      }
+
+      return null;
+  }
+
   private async initializeCurrentUserFromIdToken(
     idToken: string
   ): Promise<void> {
@@ -405,6 +443,20 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     return this.directlySetCurrentUser(user);
   }
 
+  private isFirebaseAccessTokenValid(
+    firebaseToken: FirebaseToken | null
+  ): boolean {
+    if(
+      firebaseToken &&
+      firebaseToken.expirationTime &&
+      (Date.now() >
+        firebaseToken.expirationTime - this.TOKEN_EXPIRATION_BUFFER)
+    ) {
+      return false;
+    }
+    return true;
+  }
+
   private async initializeFirebaseToken(): Promise<void> {
     this.firebaseToken =
       (await this.persistenceManager?.getFirebaseToken()) ?? null;

packages/auth/src/model/auth.ts

@@ -67,6 +67,7 @@ export interface ConfigInternal extends Config {
 export interface AuthInternal extends Auth {
   currentUser: User | null;
   emulatorConfig: EmulatorConfig | null;
+  getFirebaseAccessToken(forceRefresh?: boolean): Promise<FirebaseToken | null>;
   _agentRecaptchaConfig: RecaptchaConfig | null;
   _tenantRecaptchaConfigs: Record<string, RecaptchaConfig>;
   _projectPasswordPolicy: PasswordPolicy | null;

packages/auth/src/model/public_types.ts

@@ -274,7 +274,7 @@ export interface Auth {
    * @param tokenRefreshHandler - An object that implements the `TokenRefreshHandler`
    *   interface, providing the logic to refresh the IDP token.
    */
-   setTokenRefreshHandler(tokenRefreshHandler: TokenRefreshHandler): void;
+  setTokenRefreshHandler(tokenRefreshHandler: TokenRefreshHandler): void;
 
   /**
    * The {@link Auth} instance's language code.