Expose a callback function that is triggered for token refresh

Author: mansisampat

packages/auth/demo/src/index.js

@@ -149,6 +149,14 @@ async function getActiveUserBlocking() {
   }
 }
 
+class TokenRefreshHandlerImpl {
+  refreshToken() {
+    log("inside here");
+    console.log("inside handler");
+    return;
+  }
+}
+
 /**
  * Refreshes the current user data in the UI, displaying a user info box if
  * a user is signed in, or removing it.
@@ -2092,6 +2100,8 @@ function initApp() {
     popupRedirectResolver: browserPopupRedirectResolver,
     tenantConfig: tenantConfig
   });
+  const tokenRefreshHandler = new TokenRefreshHandlerImpl(); 
+  regionalAuth.setTokenRefreshHandler(tokenRefreshHandler);
 
   const firebaseTokenStatus = document.getElementById('firebase-token-status');
   setTimeout(() => {

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

@@ -49,7 +49,7 @@ import {
   Subscribe
 } from '@firebase/util';
 
-import { AuthInternal, ConfigInternal } from '../../model/auth';
+import { AuthInternal, ConfigInternal, TokenRefreshHandler } from '../../model/auth';
 import { PopupRedirectResolverInternal } from '../../model/popup_redirect';
 import { UserInternal } from '../../model/user';
 import {
@@ -102,6 +102,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
   currentUser: User | null = null;
   emulatorConfig: EmulatorConfig | null = null;
   firebaseToken: FirebaseToken | null = null;
+  tokenRefreshHandler?: TokenRefreshHandler;
   private operations = Promise.resolve();
   private persistenceManager?: PersistenceUserManager;
   private redirectPersistenceManager?: PersistenceUserManager;
@@ -112,6 +113,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
@@ -159,6 +161,10 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     this.tenantConfig = tenantConfig;
   }
 
+  setTokenRefreshHandler(tokenRefreshHandler: TokenRefreshHandler): void {
+    this.tokenRefreshHandler = tokenRefreshHandler;
+  }
+
   _initializeWithPersistence(
     persistenceHierarchy: PersistenceInternal[],
     popupRedirectResolver?: PopupRedirectResolver
@@ -196,7 +202,7 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
       }
 
       await this.initializeCurrentUser(popupRedirectResolver);
-      await this.initializeFirebaseToken();
+      this.firebaseToken = await this.getFirebaseAccessToken(false);
 
       this.lastNotifiedUid = this.currentUser?.uid || null;
 
@@ -210,6 +216,33 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     return this._initializationPromise;
   }
 
+  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)
+        await this.tokenRefreshHandler.refreshToken();
+        return this.getFirebaseAccessToken(false);
+      }
+
+      return null;
+  }
+
   /**
    * If the persistence is changed in another window, the user manager will let us know
    */
@@ -240,6 +273,20 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     await this._updateCurrentUser(user, /* skipBeforeStateCallbacks */ true);
   }
 
+  private isFirebaseAccessTokenValid(
+    firebaseToken: FirebaseToken | null
+  ): boolean {
+    if(
+      firebaseToken &&
+      firebaseToken.expirationTime &&
+      (Date.now() >
+        firebaseToken.expirationTime - this.TOKEN_EXPIRATION_BUFFER)
+    ) {
+      return false;
+    }
+    return true;
+  }
+
   private async initializeCurrentUserFromIdToken(
     idToken: string
   ): Promise<void> {

packages/auth/src/core/auth/firebase_internal.ts

@@ -28,10 +28,17 @@ interface TokenListener {
   (tok: string | null): unknown;
 }
 
+// interface FirebaseAccessTokenRefreshListener {
+//   (tok: string | null): unknown;
+// }
+
 export class AuthInterop implements FirebaseAuthInternal {
   private readonly TOKEN_EXPIRATION_BUFFER = 30_000;
   private readonly internalListeners: Map<TokenListener, Unsubscribe> =
     new Map();
+  // private readonly firebaseAccessTokenRefreshListeners:
+  //   Map<FirebaseAccessTokenRefreshListener, Unsubscribe> = 
+  //   new Map();
 
   constructor(private readonly auth: AuthInternal) {}
 
@@ -51,7 +58,11 @@ export class AuthInterop implements FirebaseAuthInternal {
           'Refresh token is not a valid operation for Regional Auth instance initialized.'
         );
       }
-      return this.getTokenForRegionalAuth();
+      const firebaseAccessToken = (await this.auth.getFirebaseAccessToken())?.token
+      if (!firebaseAccessToken) {
+        return null;
+      }
+      return { accessToken: firebaseAccessToken ?? null};
     }
     if (!this.auth.currentUser) {
       return null;
@@ -76,6 +87,23 @@ export class AuthInterop implements FirebaseAuthInternal {
     this.updateProactiveRefresh();
   }
 
+  // addFirebaseAccessTokenRefreshListener(
+  //   listener: FirebaseAccessTokenRefreshListener): void {
+  //     this.assertAuthConfigured();
+  //     this.assertRegionalAuthConfigured();
+  //     if (this.firebaseAccessTokenRefreshListeners.has(listener)) {
+  //       return;
+  //     }
+
+  //     const unsubscribe = this.auth.onFirebaseAccessTokenRefreshTriggered(
+  //       firebaseAccessToken => {
+  //         listener(
+  //           (firebaseAccessToken as FirebaseToken)?.token || null
+  //         );
+  //       });
+  //     this.firebaseAccessTokenRefreshListeners.set(listener, unsubscribe);   
+  // }
+
   removeAuthTokenListener(listener: TokenListener): void {
     this.assertAuthConfigured();
     const unsubscribe = this.internalListeners.get(listener);
@@ -95,6 +123,13 @@ export class AuthInterop implements FirebaseAuthInternal {
     );
   }
 
+  // private assertRegionalAuthConfigured(): void {
+  //   _assert(
+  //     this.auth.tenantConfig,
+  //     AuthErrorCode.DEPENDENT_SDK_INIT_BEFORE_AUTH)
+  //   ;
+  // }
+
   private updateProactiveRefresh(): void {
     if (this.internalListeners.size > 0) {
       this.auth._startProactiveRefresh();

packages/auth/src/model/auth.ts

@@ -58,6 +58,10 @@ export interface ConfigInternal extends Config {
   clientPlatform: ClientPlatform;
 }
 
+export interface TokenRefreshHandler {
+  refreshToken(): void | Promise<void>
+}
+
 /**
  * UserInternal and AuthInternal reference each other, so both of them are included in the public typings.
  * In order to exclude them, we mark them as internal explicitly.
@@ -67,6 +71,7 @@ export interface ConfigInternal extends Config {
 export interface AuthInternal extends Auth {
   currentUser: User | null;
   emulatorConfig: EmulatorConfig | null;
+  tokenRefreshHandler?: TokenRefreshHandler;
   _agentRecaptchaConfig: RecaptchaConfig | null;
   _tenantRecaptchaConfigs: Record<string, RecaptchaConfig>;
   _projectPasswordPolicy: PasswordPolicy | null;
@@ -75,6 +80,7 @@ export interface AuthInternal extends Auth {
   _isInitialized: boolean;
   _initializationPromise: Promise<void> | null;
   _persistenceManagerAvailable: Promise<void>;
+  getFirebaseAccessToken(forceRefresh?: boolean): Promise<FirebaseToken | null>;
   _updateCurrentUser(user: UserInternal | null): Promise<void>;
   _updateFirebaseToken(firebaseToken: FirebaseToken | null): Promise<void>;