From 5ee3d03525ce9b89d384d1ab84a7d234d98d2f5b Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Thu, 19 Dec 2024 12:25:04 -0800
Subject: [PATCH] Replace Promise.withResolvers for compatiblity with older
 browers; Fix bug where MFA couldn't be retried after a failed attempt; Add
 extra tests.

---
 web/packages/teleport/src/lib/useMfa.test.tsx | 47 +++++++++++++++++--
 web/packages/teleport/src/lib/useMfa.ts       | 47 ++++++++++++++-----
 2 files changed, 80 insertions(+), 14 deletions(-)

diff --git a/web/packages/teleport/src/lib/useMfa.test.tsx b/web/packages/teleport/src/lib/useMfa.test.tsx
index 9d2516808f107..886fbff25a662 100644
--- a/web/packages/teleport/src/lib/useMfa.test.tsx
+++ b/web/packages/teleport/src/lib/useMfa.test.tsx
@@ -62,7 +62,13 @@ const mockChallengeReq: CreateAuthenticateChallengeRequest = {
 };
 
 describe('useMfa', () => {
-  beforeEach(() => jest.spyOn(console, 'error').mockImplementation());
+  beforeEach(() => {
+    jest.spyOn(console, 'error').mockImplementation();
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
 
   test('mfa required', async () => {
     jest.spyOn(auth, 'getMfaChallenge').mockResolvedValueOnce(mockChallenge);
@@ -184,7 +190,11 @@ describe('useMfa', () => {
       throw err;
     });
 
-    const { result: mfa } = renderHook(() => useMfa({}));
+    const { result: mfa } = renderHook(() =>
+      useMfa({
+        req: mockChallengeReq,
+      })
+    );
 
     const respPromise = mfa.current.getChallengeResponse();
     await waitFor(() => {
@@ -192,7 +202,6 @@ describe('useMfa', () => {
     });
     await mfa.current.submit('webauthn');
 
-    await expect(respPromise).rejects.toThrow(err);
     await waitFor(() => {
       expect(mfa.current.attempt).toEqual({
         status: 'error',
@@ -201,5 +210,37 @@ describe('useMfa', () => {
         data: null,
       });
     });
+
+    // After an error, the mfa response promise remains in an unresolved state,
+    // allowing for retries.
+    jest
+      .spyOn(auth, 'getMfaChallengeResponse')
+      .mockResolvedValueOnce(mockResponse);
+    await mfa.current.submit('webauthn');
+    expect(await respPromise).toEqual(mockResponse);
+  });
+
+  test('reset mfa attempt', async () => {
+    jest.spyOn(auth, 'getMfaChallenge').mockResolvedValue(mockChallenge);
+    const { result: mfa } = renderHook(() =>
+      useMfa({
+        req: mockChallengeReq,
+      })
+    );
+
+    const respPromise = mfa.current.getChallengeResponse();
+    await waitFor(() => {
+      expect(auth.getMfaChallenge).toHaveBeenCalled();
+    });
+
+    mfa.current.resetAttempt();
+
+    await expect(respPromise).rejects.toThrow(
+      new Error('MFA attempt cancelled by user')
+    );
+
+    await waitFor(() => {
+      expect(mfa.current.attempt.status).toEqual('error');
+    });
   });
 });
diff --git a/web/packages/teleport/src/lib/useMfa.ts b/web/packages/teleport/src/lib/useMfa.ts
index 262b9eb9ac10a..50f852c3768e3 100644
--- a/web/packages/teleport/src/lib/useMfa.ts
+++ b/web/packages/teleport/src/lib/useMfa.ts
@@ -39,6 +39,12 @@ export type MfaProps = {
   isMfaRequired?: boolean | null;
 };
 
+type mfaResponsePromiseWithResolvers = {
+  promise: Promise<MfaChallengeResponse>;
+  resolve: (v: MfaChallengeResponse) => void;
+  reject: (v?: any) => void;
+};
+
 /**
  * Use the returned object to request MFA checks with a shared state.
  * When MFA authentication is in progress, the object's properties can
@@ -48,10 +54,10 @@ export type MfaProps = {
 export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
   const [mfaRequired, setMfaRequired] = useState<boolean>();
   const [options, setMfaOptions] = useState<MfaOption[]>();
-
   const [challenge, setMfaChallenge] = useState<MfaAuthenticateChallenge>();
-  const mfaResponsePromise =
-    useRef<PromiseWithResolvers<MfaChallengeResponse>>();
+
+  const mfaResponsePromiseWithResolvers =
+    useRef<mfaResponsePromiseWithResolvers>();
 
   useEffect(() => {
     setMfaRequired(isMfaRequired);
@@ -70,6 +76,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
   //   4. Receive the mfa response through the mfaResponsePromise ref and return it.
   //
   // The caller should also display errors seen in attempt.
+
   const [attempt, getResponse, setMfaAttempt] = useAsync(
     useCallback(
       async (challenge?: MfaAuthenticateChallenge) => {
@@ -88,22 +95,36 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
 
         // Prepare a new promise to collect the mfa response retrieved
         // through the submit function.
-        mfaResponsePromise.current = Promise.withResolvers();
+        let resolve, reject;
+        const promise = new Promise<MfaChallengeResponse>((res, rej) => {
+          resolve = res;
+          reject = rej;
+        });
+
+        mfaResponsePromiseWithResolvers.current = {
+          promise,
+          resolve,
+          reject,
+        };
+
         setMfaChallenge(challenge);
         try {
-          return await mfaResponsePromise.current.promise;
+          return await promise;
         } finally {
-          mfaResponsePromise.current = null;
+          mfaResponsePromiseWithResolvers.current = null;
           setMfaChallenge(null);
         }
       },
-      [req, mfaResponsePromise, options, mfaRequired]
+      [req, mfaResponsePromiseWithResolvers, options, mfaRequired]
     )
   );
 
   const resetAttempt = () => {
-    if (mfaResponsePromise.current) mfaResponsePromise.current.reject();
-    mfaResponsePromise.current = null;
+    if (mfaResponsePromiseWithResolvers.current)
+      mfaResponsePromiseWithResolvers.current.reject(
+        new Error('MFA attempt cancelled by user')
+      );
+    mfaResponsePromiseWithResolvers.current = null;
     setMfaChallenge(null);
     setMfaAttempt(makeEmptyAttempt());
   };
@@ -119,8 +140,12 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
 
   const submit = useCallback(
     async (mfaType?: DeviceType, totpCode?: string) => {
+      if (!mfaResponsePromiseWithResolvers.current) {
+        throw new Error('submit called without an in flight MFA attempt');
+      }
+
       try {
-        await mfaResponsePromise.current.resolve(
+        await mfaResponsePromiseWithResolvers.current.resolve(
           await auth.getMfaChallengeResponse(challenge, mfaType, totpCode)
         );
       } catch (err) {
@@ -132,7 +157,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
         });
       }
     },
-    [challenge, mfaResponsePromise, setMfaAttempt]
+    [challenge, mfaResponsePromiseWithResolvers, setMfaAttempt]
   );
 
   return {