[Auth] Fix persistence issue when manually setting/using session storage (#5239)

* Add session storage to the list of persistence in getAuth

* test

* Wrap up fixes to persistence layer

* Formatting

* Add comments

* Fix compatability layer

* Formatting

* PR feedback; reinstate readonly persistence layer checks

* Formatting

* PR feedback

* Formatting

Author: sam-gc

packages-exp/auth-compat-exp/src/auth.test.ts

@@ -65,7 +65,7 @@ describe('auth compat', () => {
         providerStub.getImmediate.returns(underlyingAuth);
         const authCompat = new Auth(
           app,
-          (providerStub as unknown) as Provider<'auth-exp'>
+          providerStub as unknown as Provider<'auth-exp'>
         );
         // eslint-disable-next-line @typescript-eslint/no-floating-promises
         await authCompat.signInWithRedirect(new exp.GoogleAuthProvider());
@@ -83,15 +83,16 @@ describe('auth compat', () => {
         );
         providerStub.isInitialized.returns(false);
         providerStub.initialize.returns(underlyingAuth);
-        new Auth(app, (providerStub as unknown) as Provider<'auth-exp'>);
+        new Auth(app, providerStub as unknown as Provider<'auth-exp'>);
         // eslint-disable-next-line @typescript-eslint/no-floating-promises
         expect(providerStub.initialize).to.have.been.calledWith({
           options: {
             popupRedirectResolver: CompatPopupRedirectResolver,
             persistence: [
               exp.inMemoryPersistence,
               exp.indexedDBLocalPersistence,
-              exp.browserLocalPersistence
+              exp.browserLocalPersistence,
+              exp.browserSessionPersistence
             ]
           }
         });

packages-exp/auth-compat-exp/src/auth.ts

@@ -67,7 +67,8 @@ export class Auth
 
       for (const persistence of [
         exp.indexedDBLocalPersistence,
-        exp.browserLocalPersistence
+        exp.browserLocalPersistence,
+        exp.browserSessionPersistence
       ]) {
         if (!persistences.includes(persistence)) {
           persistences.push(persistence);

packages-exp/auth-exp/index.ts

@@ -27,6 +27,7 @@ import { initializeAuth } from './src';
 import { registerAuth } from './src/core/auth/register';
 import { ClientPlatform } from './src/core/util/version';
 import { browserLocalPersistence } from './src/platform_browser/persistence/local_storage';
+import { browserSessionPersistence } from './src/platform_browser/persistence/session_storage';
 import { indexedDBLocalPersistence } from './src/platform_browser/persistence/indexed_db';
 import { browserPopupRedirectResolver } from './src/platform_browser/popup_redirect';
 import { Auth } from './src/model/public_types';
@@ -136,7 +137,11 @@ export function getAuth(app: FirebaseApp = getApp()): Auth {
 
   return initializeAuth(app, {
     popupRedirectResolver: browserPopupRedirectResolver,
-    persistence: [indexedDBLocalPersistence, browserLocalPersistence]
+    persistence: [
+      indexedDBLocalPersistence,
+      browserLocalPersistence,
+      browserSessionPersistence
+    ]
   });
 }
 

packages-exp/auth-exp/src/core/persistence/index.ts

@@ -44,4 +44,6 @@ export interface PersistenceInternal extends Persistence {
   _remove(key: string): Promise<void>;
   _addListener(key: string, listener: StorageEventListener): void;
   _removeListener(key: string, listener: StorageEventListener): void;
+  // Should this persistence allow migration up the chosen hierarchy?
+  _shouldAllowMigration?: boolean;
 }

packages-exp/auth-exp/src/core/persistence/persistence_user_manager.test.ts

@@ -35,7 +35,8 @@ import { KeyName, PersistenceUserManager } from './persistence_user_manager';
 chai.use(sinonChai);
 
 function makePersistence(
-  type = PersistenceType.NONE
+  type = PersistenceType.NONE,
+  shouldAllowMigration = false
 ): {
   persistence: PersistenceInternal;
   stub: sinon.SinonStubbedInstance<PersistenceInternal>;
@@ -49,7 +50,8 @@ function makePersistence(
     },
     _remove: async () => {},
     _addListener(_key: string, _listener: StorageEventListener) {},
-    _removeListener(_key: string, _listener: StorageEventListener) {}
+    _removeListener(_key: string, _listener: StorageEventListener) {},
+    _shouldAllowMigration: shouldAllowMigration
   };
 
   const stub = sinon.stub(persistence);
@@ -69,7 +71,7 @@ describe('core/persistence/persistence_user_manager', () => {
       expect(manager.persistence).to.eq(_getInstance(inMemoryPersistence));
     });
 
-    it('chooses the first one available', async () => {
+    it('chooses the first one with a user', async () => {
       const a = makePersistence();
       const b = makePersistence();
       const c = makePersistence();
@@ -78,11 +80,31 @@ describe('core/persistence/persistence_user_manager', () => {
       a.stub._isAvailable.resolves(false);
       a.stub._get.onFirstCall().resolves(testUser(auth, 'uid').toJSON());
       b.stub._isAvailable.resolves(true);
+      a.stub._get.onFirstCall().resolves(testUser(auth, 'uid-b').toJSON());
 
       const out = await PersistenceUserManager.create(auth, search);
       expect(a.stub._isAvailable).to.have.been.calledOnce;
       expect(b.stub._isAvailable).to.have.been.calledOnce;
-      expect(c.stub._isAvailable).to.not.have.been.called;
+      expect(c.stub._isAvailable).to.have.been.calledOnce;
+
+      // a should not be chosen since it is not available (despite having a user).
+      expect(out.persistence).to.eq(a.persistence);
+    });
+
+    it('defaults to first available persistence if no user', async () => {
+      const a = makePersistence();
+      const b = makePersistence();
+      const c = makePersistence();
+      const search = [a.persistence, b.persistence, c.persistence];
+      const auth = await testAuth();
+      a.stub._isAvailable.resolves(false);
+      b.stub._isAvailable.resolves(true);
+      c.stub._isAvailable.resolves(true);
+
+      const out = await PersistenceUserManager.create(auth, search);
+      expect(a.stub._isAvailable).to.have.been.calledOnce;
+      expect(b.stub._isAvailable).to.have.been.calledOnce;
+      expect(c.stub._isAvailable).to.have.been.calledOnce;
 
       // a should not be chosen since it is not available (despite having a user).
       expect(out.persistence).to.eq(b.persistence);
@@ -108,14 +130,16 @@ describe('core/persistence/persistence_user_manager', () => {
       expect((await out.getCurrentUser())!.uid).to.eq(user.uid);
     });
 
-    it('migrate found user to the selected persistence and clear others', async () => {
-      const a = makePersistence();
-      const b = makePersistence();
-      const c = makePersistence();
+    it('migrate found user to higher order persistence, if applicable', async () => {
+      const a = makePersistence(PersistenceType.NONE, true);
+      const b = makePersistence(PersistenceType.NONE, true);
+      const c = makePersistence(PersistenceType.NONE, true);
       const search = [a.persistence, b.persistence, c.persistence];
       const auth = await testAuth();
       const user = testUser(auth, 'uid');
       a.stub._isAvailable.resolves(true);
+      b.stub._isAvailable.resolves(true);
+      c.stub._isAvailable.resolves(true);
       b.stub._get.resolves(user.toJSON());
       c.stub._get.resolves(testUser(auth, 'wrong-uid').toJSON());
 
@@ -143,8 +167,46 @@ describe('core/persistence/persistence_user_manager', () => {
       expect((await out.getCurrentUser())!.uid).to.eq(user.uid);
     });
 
+    it('migrate found user to available persistence, if applicable', async () => {
+      const a = makePersistence(PersistenceType.NONE, true);
+      const b = makePersistence(PersistenceType.NONE, true);
+      const c = makePersistence(PersistenceType.NONE, true);
+      const search = [a.persistence, b.persistence, c.persistence];
+      const auth = await testAuth();
+      const user = testUser(auth, 'uid');
+      a.stub._isAvailable.resolves(false); // Important
+      b.stub._isAvailable.resolves(true);
+      c.stub._isAvailable.resolves(true);
+      a.stub._get.resolves(user.toJSON());
+      c.stub._get.resolves(testUser(auth, 'wrong-uid').toJSON());
+
+      let persistedUserInB: PersistenceValue | null = null;
+      b.stub._set.callsFake(async (_, value) => {
+        persistedUserInB = value;
+      });
+      b.stub._get.callsFake(async () => persistedUserInB);
+
+      const out = await PersistenceUserManager.create(auth, search);
+      expect(b.stub._set).to.have.been.calledOnceWith(
+        'firebase:authUser:test-api-key:test-app',
+        user.toJSON()
+      );
+      expect(a.stub._set).to.not.have.been.called;
+      expect(c.stub._set).to.not.have.been.called;
+      expect(a.stub._remove).to.have.been.calledOnceWith(
+        'firebase:authUser:test-api-key:test-app'
+      );
+      expect(c.stub._remove).to.have.been.calledOnceWith(
+        'firebase:authUser:test-api-key:test-app'
+      );
+
+      expect(out.persistence).to.eq(b.persistence);
+      expect((await out.getCurrentUser())!.uid).to.eq(user.uid);
+    });
+
     it('uses default user key if none provided', async () => {
       const { stub, persistence } = makePersistence();
+      stub._isAvailable.resolves(true);
       await PersistenceUserManager.create(auth, [persistence]);
       expect(stub._get).to.have.been.calledWith(
         'firebase:authUser:test-api-key:test-app'
@@ -153,6 +215,7 @@ describe('core/persistence/persistence_user_manager', () => {
 
     it('uses user key if provided', async () => {
       const { stub, persistence } = makePersistence();
+      stub._isAvailable.resolves(true);
       await PersistenceUserManager.create(
         auth,
         [persistence],
@@ -176,7 +239,7 @@ describe('core/persistence/persistence_user_manager', () => {
       expect(out.persistence).to.eq(_getInstance(inMemoryPersistence));
       expect(a.stub._get).to.have.been.calledOnce;
       expect(b.stub._get).to.have.been.calledOnce;
-      expect(c.stub._get).to.have.been.called;
+      expect(c.stub._get).to.have.been.calledOnce;
     });
   });
 
@@ -235,10 +298,8 @@ describe('core/persistence/persistence_user_manager', () => {
       });
 
       it('removes current user & sets it in the new persistene', async () => {
-        const {
-          persistence: nextPersistence,
-          stub: nextStub
-        } = makePersistence();
+        const { persistence: nextPersistence, stub: nextStub } =
+          makePersistence();
         const auth = await testAuth();
         const user = testUser(auth, 'uid');
         persistenceStub._get.returns(Promise.resolve(user.toJSON()));
@@ -259,10 +320,8 @@ describe('core/persistence/persistence_user_manager', () => {
         const user = testUser(auth, 'uid');
         stub._get.returns(Promise.resolve(user.toJSON()));
 
-        const {
-          persistence: nextPersistence,
-          stub: nextStub
-        } = makePersistence(PersistenceType.LOCAL);
+        const { persistence: nextPersistence, stub: nextStub } =
+          makePersistence(PersistenceType.LOCAL);
 
         // This should migrate the user even if both has type LOCAL. For example, developer may want
         // to switch from localStorage to indexedDB (both type LOCAL) and we should honor that.

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

@@ -113,52 +113,77 @@ export class PersistenceUserManager {
       );
     }
 
-    // Use the first persistence that supports a full read-write roundtrip (or fallback to memory).
-    let chosenPersistence = _getInstance<PersistenceInternal>(
-      inMemoryPersistence
-    );
-    for (const persistence of persistenceHierarchy) {
-      if (await persistence._isAvailable()) {
-        chosenPersistence = persistence;
-        break;
-      }
-    }
+    // Eliminate any persistences that are not available
+    const availablePersistences = (
+      await Promise.all(
+        persistenceHierarchy.map(async persistence => {
+          if (await persistence._isAvailable()) {
+            return persistence;
+          }
+          return undefined;
+        })
+      )
+    ).filter(persistence => persistence) as PersistenceInternal[];
+
+    // Fall back to the first persistence listed, or in memory if none available
+    let selectedPersistence =
+      availablePersistences[0] ||
+      _getInstance<PersistenceInternal>(inMemoryPersistence);
 
-    // However, attempt to migrate users stored in other persistences (in the hierarchy order).
-    let userToMigrate: UserInternal | null = null;
     const key = _persistenceKeyName(userKey, auth.config.apiKey, auth.name);
+
+    // Pull out the existing user, setting the chosen persistence to that
+    // persistence if the user exists.
+    let userToMigrate: UserInternal | null = null;
+    // Note, here we check for a user in _all_ persistences, not just the
+    // ones deemed available. If we can migrate a user out of a broken
+    // persistence, we will (but only if that persistence supports migration).
     for (const persistence of persistenceHierarchy) {
-      // We attempt to call _get without checking _isAvailable since here we don't care if the full
-      // round-trip (read+write) is supported. We'll take the first one that we can read or give up.
       try {
-        const blob = await persistence._get<PersistedBlob>(key); // throws if unsupported
+        const blob = await persistence._get<PersistedBlob>(key);
         if (blob) {
           const user = UserImpl._fromJSON(auth, blob); // throws for unparsable blob (wrong format)
-          if (persistence !== chosenPersistence) {
+          if (persistence !== selectedPersistence) {
             userToMigrate = user;
           }
+          selectedPersistence = persistence;
           break;
         }
       } catch {}
     }
 
+    // If we find the user in a persistence that does support migration, use
+    // that migration path (of only persistences that support migration)
+    const migrationHierarchy = availablePersistences.filter(
+      p => p._shouldAllowMigration
+    );
+
+    // If the persistence does _not_ allow migration, just finish off here
+    if (
+      !selectedPersistence._shouldAllowMigration ||
+      !migrationHierarchy.length
+    ) {
+      return new PersistenceUserManager(selectedPersistence, auth, userKey);
+    }
+
+    selectedPersistence = migrationHierarchy[0];
     if (userToMigrate) {
       // This normally shouldn't throw since chosenPersistence.isAvailable() is true, but if it does
       // we'll just let it bubble to surface the error.
-      await chosenPersistence._set(key, userToMigrate.toJSON());
+      await selectedPersistence._set(key, userToMigrate.toJSON());
     }
 
     // Attempt to clear the key in other persistences but ignore errors. This helps prevent issues
     // such as users getting stuck with a previous account after signing out and refreshing the tab.
     await Promise.all(
       persistenceHierarchy.map(async persistence => {
-        if (persistence !== chosenPersistence) {
+        if (persistence !== selectedPersistence) {
           try {
             await persistence._remove(key);
           } catch {}
         }
       })
     );
-    return new PersistenceUserManager(chosenPersistence, auth, userKey);
+    return new PersistenceUserManager(selectedPersistence, auth, userKey);
   }
 }

packages-exp/auth-exp/src/platform_browser/auth.test.ts

@@ -324,6 +324,7 @@ describe('core/auth/initializeAuth', () => {
         const stub = sinon.stub(
           _getInstance<PersistenceInternal>(browserSessionPersistence)
         );
+        stub._isAvailable.returns(Promise.resolve(true));
         stub._remove.returns(Promise.resolve());
         completeRedirectFnStub.returns(Promise.reject(new Error('no')));
 

packages-exp/auth-exp/src/platform_browser/persistence/indexed_db.ts

@@ -158,6 +158,7 @@ class IndexedDBLocalPersistence implements InternalPersistence {
 
   type = PersistenceType.LOCAL;
   db?: IDBDatabase;
+  readonly _shouldAllowMigration = true;
 
   private readonly listeners: Record<string, Set<StorageEventListener>> = {};
   private readonly localCache: Record<string, PersistenceValue | null> = {};

packages-exp/auth-exp/src/platform_browser/persistence/local_storage.ts

@@ -46,7 +46,8 @@ const IE10_LOCAL_STORAGE_SYNC_DELAY = 10;
 
 class BrowserLocalPersistence
   extends BrowserPersistenceClass
-  implements InternalPersistence {
+  implements InternalPersistence
+{
   static type: 'LOCAL' = 'LOCAL';
 
   constructor() {
@@ -69,6 +70,7 @@ class BrowserLocalPersistence
     _iframeCannotSyncWebStorage() && _isIframe();
   // Whether to use polling instead of depending on window events
   private readonly fallbackToPolling = _isMobileBrowser();
+  readonly _shouldAllowMigration = true;
 
   private forAllChangedKeys(
     cb: (key: string, oldValue: string | null, newValue: string | null) => void