[Auth] Fix redirect middleware (#6165)

* Fix redirect middleware

* Fix compat tests

* Formatting

* PR feedback

Author: sam-gc

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

@@ -175,6 +175,7 @@ describe('popup_redirect/CompatPopupRedirectResolver', () => {
 
 class FakeResolver implements exp.PopupRedirectResolverInternal {
   _completeRedirectFn = async (): Promise<null> => null;
+  _overrideRedirectResult = (): void => {};
   _redirectPersistence = exp.inMemoryPersistence;
   _shouldInitProactively = true;
 

packages/auth-compat/src/popup_redirect.ts

@@ -38,6 +38,7 @@ export class CompatPopupRedirectResolver
     resolver: exp.PopupRedirectResolver,
     bypassAuthState: boolean
   ) => Promise<exp.UserCredential | null> = exp._getRedirectResult;
+  _overrideRedirectResult = exp._overrideRedirectResult;
 
   async _initialize(auth: exp.AuthImpl): Promise<exp.EventManager> {
     await this.selectUnderlyingResolver();

packages/auth/internal/index.ts

@@ -45,6 +45,7 @@ export { TaggedWithTokenResponse } from '../src/model/id_token';
 export { _fail, _assert } from '../src/core/util/assert';
 export { AuthPopup } from '../src/platform_browser/util/popup';
 export { _getRedirectResult } from '../src/platform_browser/strategies/redirect';
+export { _overrideRedirectResult } from '../src/core/strategies/redirect';
 export { cordovaPopupRedirectResolver } from '../src/platform_cordova/popup_redirect/popup_redirect';
 export { FetchProvider } from '../src/core/util/fetch_provider';
 export { SAMLAuthCredential } from '../src/core/credentials/saml';

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

@@ -192,12 +192,14 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
     popupRedirectResolver?: PopupRedirectResolver
   ): Promise<void> {
     // First check to see if we have a pending redirect event.
-    let storedUser =
+    const previouslyStoredUser =
       (await this.assertedPersistence.getCurrentUser()) as UserInternal | null;
+    let futureCurrentUser = previouslyStoredUser;
+    let needsTocheckMiddleware = false;
     if (popupRedirectResolver && this.config.authDomain) {
       await this.getOrInitRedirectPersistenceManager();
       const redirectUserEventId = this.redirectUser?._redirectEventId;
-      const storedUserEventId = storedUser?._redirectEventId;
+      const storedUserEventId = futureCurrentUser?._redirectEventId;
       const result = await this.tryRedirectSignIn(popupRedirectResolver);
 
       // If the stored user (i.e. the old "currentUser") has a redirectId that
@@ -208,44 +210,51 @@ export class AuthImpl implements AuthInternal, _FirebaseService {
         (!redirectUserEventId || redirectUserEventId === storedUserEventId) &&
         result?.user
       ) {
-        storedUser = result.user as UserInternal;
+        futureCurrentUser = result.user as UserInternal;
+        needsTocheckMiddleware = true;
       }
     }
 
     // If no user in persistence, there is no current user. Set to null.
-    if (!storedUser) {
+    if (!futureCurrentUser) {
       return this.directlySetCurrentUser(null);
     }
 
-    if (!storedUser._redirectEventId) {
-      // This isn't a redirect user, we can reload and bail
-      // This will also catch the redirected user, if available, as that method
-      // strips the _redirectEventId
-      return this.reloadAndSetCurrentUserOrClear(storedUser);
+    if (!futureCurrentUser._redirectEventId) {
+      // This isn't a redirect link operation, we can reload and bail.
+      // First though, ensure that we check the middleware is happy.
+      if (needsTocheckMiddleware) {
+        try {
+          await this._runBeforeStateCallbacks(futureCurrentUser);
+        } catch(e) {
+          futureCurrentUser = previouslyStoredUser;
+          // We know this is available since the bit is only set when the
+          // resolver is available
+          this._popupRedirectResolver!._overrideRedirectResult(this, () => Promise.reject(e));
+        }
+      }
+  
+      if (futureCurrentUser) {
+        return this.reloadAndSetCurrentUserOrClear(futureCurrentUser);
+      } else {
+        return this.directlySetCurrentUser(null);
+      }
     }
 
     _assert(this._popupRedirectResolver, this, AuthErrorCode.ARGUMENT_ERROR);
     await this.getOrInitRedirectPersistenceManager();
 
-    // At this point in the flow, this is a redirect user. Run blocking
-    // middleware callbacks before setting the user.
-    try {
-      await this._runBeforeStateCallbacks(storedUser);
-    } catch(e) {
-      return;
-    }
-
     // If the redirect user's event ID matches the current user's event ID,
     // DO NOT reload the current user, otherwise they'll be cleared from storage.
     // This is important for the reauthenticateWithRedirect() flow.
     if (
       this.redirectUser &&
-      this.redirectUser._redirectEventId === storedUser._redirectEventId
+      this.redirectUser._redirectEventId === futureCurrentUser._redirectEventId
     ) {
-      return this.directlySetCurrentUser(storedUser);
+      return this.directlySetCurrentUser(futureCurrentUser);
     }
 
-    return this.reloadAndSetCurrentUserOrClear(storedUser);
+    return this.reloadAndSetCurrentUserOrClear(futureCurrentUser);
   }
 
   private async tryRedirectSignIn(

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

@@ -126,6 +126,8 @@ describe('core/auth/initialize', () => {
     ): Promise<UserCredential | null> {
       return null;
     }
+    async _overrideRedirectResult(): Promise<void> {
+    }
   }
 
   const fakePopupRedirectResolver: PopupRedirectResolver =

packages/auth/src/core/strategies/redirect.ts

@@ -139,6 +139,10 @@ export function _clearRedirectOutcomes(): void {
   redirectOutcomeMap.clear();
 }
 
+export function _overrideRedirectResult(auth: AuthInternal, result: () => Promise<UserCredentialInternal | null>): void {
+  redirectOutcomeMap.set(auth._key(), result);
+}
+
 function resolverPersistence(
   resolver: PopupRedirectResolverInternal
 ): PersistenceInternal {

packages/auth/src/model/popup_redirect.ts

@@ -26,6 +26,7 @@ import { FirebaseError } from '@firebase/util';
 
 import { AuthPopup } from '../platform_browser/util/popup';
 import { AuthInternal } from './auth';
+import { UserCredentialInternal } from './user';
 
 export const enum EventFilter {
   POPUP,
@@ -121,10 +122,11 @@ export interface PopupRedirectResolverInternal extends PopupRedirectResolver {
   _redirectPersistence: Persistence;
   _originValidation(auth: Auth): Promise<void>;
 
-  // This is needed so that auth does not have a hard dependency on redirect
+  // These are needed so that auth does not have a hard dependency on redirect
   _completeRedirectFn: (
     auth: Auth,
     resolver: PopupRedirectResolver,
     bypassAuthState: boolean
   ) => Promise<UserCredential | null>;
+  _overrideRedirectResult: (auth: AuthInternal, resultGetter: () => Promise<UserCredentialInternal|null>) => void;
 }

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

@@ -130,7 +130,8 @@ describe('core/auth/initializeAuth', () => {
     async function initAndWait(
       persistence: Persistence | Persistence[],
       popupRedirectResolver?: PopupRedirectResolver,
-      authDomain = FAKE_APP.options.authDomain
+      authDomain = FAKE_APP.options.authDomain,
+      blockMiddleware = false,
     ): Promise<Auth> {
       const auth = new AuthImpl(FAKE_APP, FAKE_HEARTBEAT_CONTROLLER_PROVIDER, {
         apiKey: FAKE_APP.options.apiKey!,
@@ -146,6 +147,12 @@ describe('core/auth/initializeAuth', () => {
         persistence,
         popupRedirectResolver
       });
+
+      if (blockMiddleware) {
+        auth.beforeAuthStateChanged(() => {
+          throw new Error('blocked');
+        });
+      }
       // Auth initializes async. We can make sure the initialization is
       // flushed by awaiting a method on the queue.
       await auth.setPersistence(inMemoryPersistence);
@@ -408,6 +415,84 @@ describe('core/auth/initializeAuth', () => {
         expect(user).not.to.be.null;
         expect(auth.currentUser).to.eq(user);
       });
+
+      it('does not halt old user load if middleware throws', async () => {
+        const stub = sinon.stub(
+          _getInstance<PersistenceInternal>(inMemoryPersistence)
+        );
+        const oldUser = testUser(oldAuth, 'old-uid');
+        stub._get.returns(Promise.resolve(oldUser.toJSON()));
+        const overrideSpy = sinon.spy(_getInstance<PopupRedirectResolverInternal>(browserPopupRedirectResolver), '_overrideRedirectResult');
+        const auth = await initAndWait(
+          [inMemoryPersistence],
+          browserPopupRedirectResolver,
+          FAKE_APP.options.authDomain,
+          /* blockMiddleware */ true
+        );
+
+        expect(auth.currentUser!.uid).to.eq(oldUser.uid);
+        expect(reload._reloadWithoutSaving).to.have.been.called;
+        expect(overrideSpy).not.to.have.been.called;
+      });
+
+      it('Reloads and uses old user if middleware throws', async () => {
+        const stub = sinon.stub(
+          _getInstance<PersistenceInternal>(inMemoryPersistence)
+        );
+        const oldUser = testUser(oldAuth, 'old-uid');
+        stub._get.returns(Promise.resolve(oldUser.toJSON()));
+        const overrideSpy = sinon.spy(_getInstance<PopupRedirectResolverInternal>(browserPopupRedirectResolver), '_overrideRedirectResult');
+
+        let user: UserInternal | null = null;
+        completeRedirectFnStub.callsFake((auth: AuthInternal) => {
+          user = testUser(auth, 'uid', '[email protected]');
+          return Promise.resolve(
+            new UserCredentialImpl({
+              operationType: OperationType.SIGN_IN,
+              user,
+              providerId: null
+            })
+          );
+        });
+
+        const auth = await initAndWait(
+          [inMemoryPersistence],
+          browserPopupRedirectResolver,
+          FAKE_APP.options.authDomain,
+          /* blockMiddleware */ true
+        );
+        expect(user).not.to.be.null;
+        expect(auth.currentUser!.uid).to.eq(oldUser.uid);
+        expect(reload._reloadWithoutSaving).to.have.been.called;
+        expect(overrideSpy).to.have.been.called;
+      });
+
+      it('Nulls current user if redirect blocked by middleware', async () => {
+        const stub = sinon.stub(
+          _getInstance<PersistenceInternal>(inMemoryPersistence)
+        );
+        stub._get.returns(Promise.resolve(null));
+        completeRedirectFnStub.callsFake((auth: AuthInternal) => {
+          const user = testUser(auth, 'uid', '[email protected]');
+          return Promise.resolve(
+            new UserCredentialImpl({
+              operationType: OperationType.SIGN_IN,
+              user,
+              providerId: null
+            })
+          );
+        });
+
+        const auth = await initAndWait(
+          [inMemoryPersistence],
+          browserPopupRedirectResolver,
+          FAKE_APP.options.authDomain,
+          /* blockMiddleware */ true
+        );
+        expect(completeRedirectFnStub).to.have.been.called;
+        expect(auth.currentUser).to.be.null;
+        expect(reload._reloadWithoutSaving).not.to.have.been.called;
+      });
     });
   });
 });

packages/auth/src/platform_browser/popup_redirect.ts

@@ -38,6 +38,7 @@ import { _open, AuthPopup } from './util/popup';
 import { _getRedirectResult } from './strategies/redirect';
 import { _getRedirectUrl } from '../core/util/handler';
 import { _isIOS, _isMobileBrowser, _isSafari } from '../core/util/browser';
+import { _overrideRedirectResult } from '../core/strategies/redirect';
 
 /**
  * The special web storage event
@@ -176,6 +177,8 @@ class BrowserPopupRedirectResolver implements PopupRedirectResolverInternal {
   }
 
   _completeRedirectFn = _getRedirectResult;
+
+  _overrideRedirectResult = _overrideRedirectResult;
 }
 
 /**

packages/auth/src/platform_cordova/popup_redirect/popup_redirect.ts

@@ -42,7 +42,7 @@ import {
 } from './events';
 import { AuthEventManager } from '../../core/auth/auth_event_manager';
 import { _getRedirectResult } from '../../platform_browser/strategies/redirect';
-import { _clearRedirectOutcomes } from '../../core/strategies/redirect';
+import { _clearRedirectOutcomes, _overrideRedirectResult } from '../../core/strategies/redirect';
 import { _cordovaWindow } from '../plugins';
 
 /**
@@ -58,6 +58,7 @@ class CordovaPopupRedirectResolver implements PopupRedirectResolverInternal {
   private readonly originValidationPromises: Record<string, Promise<void>> = {};
 
   _completeRedirectFn = _getRedirectResult;
+  _overrideRedirectResult = _overrideRedirectResult;
 
   async _initialize(auth: AuthInternal): Promise<CordovaAuthEventManager> {
     const key = auth._key();