Fix up the redirect persistence storage in Auth Compat (#4537)

* Fix bug where immediate redirects in compat were broken; also save redirect persistence before link and reauth as well

* Formatting

Author: sam-gc

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

@@ -43,12 +43,22 @@ describe('auth compat', () => {
       sinon.restore;
     });
 
-    it('saves the persistence into session storage if available', () => {
-      const authCompat = new Auth(app, underlyingAuth);
+    it('saves the persistence into session storage if available', async () => {
       if (typeof self !== 'undefined') {
+        underlyingAuth._initializationPromise = Promise.resolve();
         sinon.stub(underlyingAuth, '_getPersistence').returns('TEST');
+        sinon
+          .stub(underlyingAuth, '_initializationPromise')
+          .value(Promise.resolve());
+        sinon.stub(
+          exp._getInstance<exp.PopupRedirectResolverInternal>(
+            CompatPopupRedirectResolver
+          ),
+          '_openRedirect'
+        );
+        const authCompat = new Auth(app, underlyingAuth);
         // eslint-disable-next-line @typescript-eslint/no-floating-promises
-        authCompat.signInWithRedirect(new exp.GoogleAuthProvider());
+        await authCompat.signInWithRedirect(new exp.GoogleAuthProvider());
         expect(
           sessionStorage.getItem('firebase:persistence:api-key:undefined')
         ).to.eq('TEST');

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

@@ -25,7 +25,12 @@ import {
   Unsubscribe
 } from '@firebase/util';
 
-import { _validatePersistenceArgument, Persistence } from './persistence';
+import {
+  _validatePersistenceArgument,
+  Persistence,
+  _getPersistenceFromRedirect,
+  _savePersistenceForRedirect
+} from './persistence';
 import { _isPopupRedirectSupported } from './platform';
 import { CompatPopupRedirectResolver } from './popup_redirect';
 import { User } from './user';
@@ -35,7 +40,6 @@ import {
 } from './user_credential';
 import { unwrap, Wrapper } from './wrap';
 
-const PERSISTENCE_KEY = 'persistence';
 const _assert: typeof exp._assert = exp._assert;
 
 export class Auth
@@ -51,7 +55,7 @@ export class Auth
     // Note this is slightly different behavior: in this case, the stored
     // persistence is checked *first* rather than last. This is because we want
     // the fallback (if no user is found) to be the stored persistence type
-    const storedPersistence = this.getPersistenceFromRedirect();
+    const storedPersistence = _getPersistenceFromRedirect(this.auth);
     const persistences = storedPersistence ? [storedPersistence] : [];
     persistences.push(exp.indexedDBLocalPersistence);
 
@@ -294,7 +298,8 @@ export class Auth
       this.auth,
       exp.AuthErrorCode.OPERATION_NOT_SUPPORTED
     );
-    this.savePersistenceForRedirect();
+
+    await _savePersistenceForRedirect(this.auth);
     return exp.signInWithRedirect(
       this.auth,
       provider as exp.AuthProvider,
@@ -313,49 +318,6 @@ export class Auth
   _delete(): Promise<void> {
     return this.auth._delete();
   }
-
-  private savePersistenceForRedirect(): void {
-    const win = getSelfWindow();
-    const key = exp._persistenceKeyName(
-      PERSISTENCE_KEY,
-      this.auth.config.apiKey,
-      this.auth.name
-    );
-    if (win?.sessionStorage) {
-      win.sessionStorage.setItem(key, this.auth._getPersistence());
-    }
-  }
-
-  private getPersistenceFromRedirect(): exp.Persistence | null {
-    const win = getSelfWindow();
-    if (!win?.sessionStorage) {
-      return null;
-    }
-
-    const key = exp._persistenceKeyName(
-      PERSISTENCE_KEY,
-      this.auth.config.apiKey,
-      this.auth.name
-    );
-    const persistence = win.sessionStorage.getItem(key);
-
-    switch (persistence) {
-      case exp.inMemoryPersistence.type:
-        return exp.inMemoryPersistence;
-      case exp.indexedDBLocalPersistence.type:
-        return exp.indexedDBLocalPersistence;
-      case exp.browserSessionPersistence.type:
-        return exp.browserSessionPersistence;
-      case exp.browserLocalPersistence.type:
-        return exp.browserLocalPersistence;
-      default:
-        return null;
-    }
-  }
-}
-
-function getSelfWindow(): Window | null {
-  return typeof window !== 'undefined' ? window : null;
 }
 
 function wrapObservers(

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

@@ -15,7 +15,8 @@
  * limitations under the License.
  */
 
-import { _assert, AuthErrorCode, Auth } from '@firebase/auth-exp/internal';
+import { AuthInternal } from '@firebase/auth-exp/dist/esm5/src/model/auth';
+import * as exp from '@firebase/auth-exp/internal';
 import { isIndexedDBAvailable, isNode, isReactNative } from '@firebase/util';
 import { _isWebStorageSupported, _isWorker } from './platform';
 
@@ -25,26 +26,30 @@ export const Persistence = {
   SESSION: 'SESSION'
 };
 
+const _assert: typeof exp._assert = exp._assert;
+
+const PERSISTENCE_KEY = 'persistence';
+
 /**
  * Validates that an argument is a valid persistence value. If an invalid type
  * is specified, an error is thrown synchronously.
  */
 export function _validatePersistenceArgument(
-  auth: Auth,
+  auth: exp.Auth,
   persistence: string
 ): void {
   _assert(
     Object.values(Persistence).includes(persistence),
     auth,
-    AuthErrorCode.INVALID_PERSISTENCE
+    exp.AuthErrorCode.INVALID_PERSISTENCE
   );
   // Validate if the specified type is supported in the current environment.
   if (isReactNative()) {
     // This is only supported in a browser.
     _assert(
       persistence !== Persistence.SESSION,
       auth,
-      AuthErrorCode.UNSUPPORTED_PERSISTENCE
+      exp.AuthErrorCode.UNSUPPORTED_PERSISTENCE
     );
     return;
   }
@@ -53,7 +58,7 @@ export function _validatePersistenceArgument(
     _assert(
       persistence === Persistence.NONE,
       auth,
-      AuthErrorCode.UNSUPPORTED_PERSISTENCE
+      exp.AuthErrorCode.UNSUPPORTED_PERSISTENCE
     );
     return;
   }
@@ -64,14 +69,63 @@ export function _validatePersistenceArgument(
       persistence === Persistence.NONE ||
         (persistence === Persistence.LOCAL && isIndexedDBAvailable()),
       auth,
-      AuthErrorCode.UNSUPPORTED_PERSISTENCE
+      exp.AuthErrorCode.UNSUPPORTED_PERSISTENCE
     );
     return;
   }
   // This is restricted by what the browser supports.
   _assert(
     persistence === Persistence.NONE || _isWebStorageSupported(),
     auth,
-    AuthErrorCode.UNSUPPORTED_PERSISTENCE
+    exp.AuthErrorCode.UNSUPPORTED_PERSISTENCE
+  );
+}
+
+export async function _savePersistenceForRedirect(
+  auth: AuthInternal
+): Promise<void> {
+  await auth._initializationPromise;
+
+  const win = getSelfWindow();
+  const key = exp._persistenceKeyName(
+    PERSISTENCE_KEY,
+    auth.config.apiKey,
+    auth.name
+  );
+  if (win?.sessionStorage) {
+    win.sessionStorage.setItem(key, auth._getPersistence());
+  }
+}
+
+export function _getPersistenceFromRedirect(
+  auth: AuthInternal
+): exp.Persistence | null {
+  const win = getSelfWindow();
+  if (!win?.sessionStorage) {
+    return null;
+  }
+
+  const key = exp._persistenceKeyName(
+    PERSISTENCE_KEY,
+    auth.config.apiKey,
+    auth.name
   );
+  const persistence = win.sessionStorage.getItem(key);
+
+  switch (persistence) {
+    case exp.inMemoryPersistence.type:
+      return exp.inMemoryPersistence;
+    case exp.indexedDBLocalPersistence.type:
+      return exp.indexedDBLocalPersistence;
+    case exp.browserSessionPersistence.type:
+      return exp.browserSessionPersistence;
+    case exp.browserLocalPersistence.type:
+      return exp.browserLocalPersistence;
+    default:
+      return null;
+  }
+}
+
+function getSelfWindow(): Window | null {
+  return typeof window !== 'undefined' ? window : null;
 }

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

@@ -17,6 +17,7 @@
 
 import * as exp from '@firebase/auth-exp/internal';
 import * as compat from '@firebase/auth-types';
+import { _savePersistenceForRedirect } from './persistence';
 import { CompatPopupRedirectResolver } from './popup_redirect';
 import {
   convertConfirmationResult,
@@ -96,7 +97,8 @@ export class User implements compat.User, Wrapper<exp.User> {
       )
     );
   }
-  linkWithRedirect(provider: compat.AuthProvider): Promise<void> {
+  async linkWithRedirect(provider: compat.AuthProvider): Promise<void> {
+    await _savePersistenceForRedirect(exp._castAuth(this.auth));
     return exp.linkWithRedirect(
       this.user,
       provider as exp.AuthProvider,
@@ -144,7 +146,10 @@ export class User implements compat.User, Wrapper<exp.User> {
       )
     );
   }
-  reauthenticateWithRedirect(provider: compat.AuthProvider): Promise<void> {
+  async reauthenticateWithRedirect(
+    provider: compat.AuthProvider
+  ): Promise<void> {
+    await _savePersistenceForRedirect(exp._castAuth(this.auth));
     return exp.reauthenticateWithRedirect(
       this.user,
       provider as exp.AuthProvider,

packages-exp/auth-exp/internal/index.ts

@@ -33,7 +33,7 @@ export {
 } from '../src/model/popup_redirect';
 export { UserCredentialInternal, UserParameters } from '../src/model/user';
 export { registerAuth } from '../src/core/auth/register';
-export { DefaultConfig, AuthImpl } from '../src/core/auth/auth_impl';
+export { DefaultConfig, AuthImpl, _castAuth } from '../src/core/auth/auth_impl';
 
 export { ClientPlatform, _getClientVersion } from '../src/core/util/version';