Fix useMfa error handling

web/packages/teleport/src/components/AuthnDialog/AuthnDialog.tsx

@@ -23,21 +23,26 @@ import { Cross, FingerprintSimple } from 'design/Icon';
 import { guessProviderType } from 'shared/components/ButtonSso';
 import { SSOIcon } from 'shared/components/ButtonSso/ButtonSso';
 
-import { MfaState } from 'teleport/lib/useMfa';
+import { MfaCanceledError, MfaState } from 'teleport/lib/useMfa';
 import { MFA_OPTION_TOTP } from 'teleport/services/mfa';
 
 export type Props = {
   mfaState: MfaState;
   replaceErrorText?: string;
+  // onClose is an optional function to perform additional operations
+  // upon closing the dialog. e.g. close a shell session
   onClose?: () => void;
 };
 
 export default function AuthnDialog({
-  mfaState: { options, challenge, submit, attempt, resetAttempt },
+  mfaState: { options, challenge, submit, attempt, cancelAttempt },
   replaceErrorText,
-  onClose,
+  onClose = () => {},
 }: Props) {
-  if (!challenge && attempt.status !== 'error') return;
+  const showError =
+    attempt.status === 'error' && !(attempt.error instanceof MfaCanceledError);
+
+  if (!challenge && !showError) return;
 
   // TODO(Joerger): TOTP should be pretty easy to support here with a small button -> form flow.
   const onlyTotpAvailable =
@@ -50,7 +55,7 @@ export default function AuthnDialog({
         <ButtonIcon
           data-testid="close-dialog"
           onClick={() => {
-            resetAttempt();
+            cancelAttempt();
             onClose();
           }}
         >

web/packages/teleport/src/lib/useMfa.test.tsx

@@ -27,7 +27,7 @@ import {
   MfaChallengeResponse,
 } from 'teleport/services/mfa';
 
-import { useMfa } from './useMfa';
+import { MfaCanceledError, useMfa } from './useMfa';
 
 const mockChallenge: MfaAuthenticateChallenge = {
   webauthnPublicKey: {} as PublicKeyCredentialRequestOptions,
@@ -232,14 +232,15 @@ describe('useMfa', () => {
       expect(auth.getMfaChallenge).toHaveBeenCalled();
     });
 
-    mfa.current.resetAttempt();
+    mfa.current.cancelAttempt();
 
-    await expect(respPromise).rejects.toThrow(
-      new Error('MFA attempt cancelled by user')
-    );
+    await expect(respPromise).rejects.toThrow(new MfaCanceledError());
 
     await waitFor(() => {
       expect(mfa.current.attempt.status).toEqual('error');
     });
+    expect(
+      mfa.current.attempt.status === 'error' && mfa.current.attempt.error
+    ).toEqual(new MfaCanceledError());
   });
 });

web/packages/teleport/src/lib/useMfa.ts

@@ -43,7 +43,7 @@ export type MfaProps = {
 type mfaResponsePromiseWithResolvers = {
   promise: Promise<MfaChallengeResponse>;
   resolve: (v: MfaChallengeResponse) => void;
-  reject: (v?: any) => void;
+  reject: (err: Error) => void;
 };
 
 /**
@@ -57,9 +57,6 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
   const [options, setMfaOptions] = useState<MfaOption[]>();
   const [challenge, setMfaChallenge] = useState<MfaAuthenticateChallenge>();
 
-  const mfaResponsePromiseWithResolvers =
-    useRef<mfaResponsePromiseWithResolvers>();
-
   useEffect(() => {
     setMfaRequired(isMfaRequired);
   }, [isMfaRequired]);
@@ -77,7 +74,6 @@ 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) => {
@@ -90,63 +86,63 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
           return;
         }
 
-        // Set mfa requirement and options after we get a challenge for the first time.
-        if (!mfaRequired) setMfaRequired(true);
-        if (!options) setMfaOptions(getMfaChallengeOptions(challenge));
-
         // Prepare a new promise to collect the mfa response retrieved
         // through the submit function.
-        let resolve, reject;
+        let resolve: (value: MfaChallengeResponse) => void;
+        let reject: (err: Error) => void;
         const promise = new Promise<MfaChallengeResponse>((res, rej) => {
           resolve = res;
           reject = rej;
         });
 
-        mfaResponsePromiseWithResolvers.current = {
+        mfaResponseRef.current = {
           promise,
           resolve,
           reject,
         };
 
+        // Set mfa requirement and options after we get a challenge for the first time.
+        setMfaRequired(true);
+        setMfaOptions(getMfaChallengeOptions(challenge));
+
         setMfaChallenge(challenge);
         try {
           return await promise;
         } finally {
-          mfaResponsePromiseWithResolvers.current = null;
           setMfaChallenge(null);
         }
       },
-      [req, mfaResponsePromiseWithResolvers, options, mfaRequired]
+      [req, mfaRequired]
     )
   );
 
-  const resetAttempt = () => {
-    if (mfaResponsePromiseWithResolvers.current)
-      mfaResponsePromiseWithResolvers.current.reject(
-        new Error('MFA attempt cancelled by user')
-      );
-    mfaResponsePromiseWithResolvers.current = null;
-    setMfaChallenge(null);
-    setMfaAttempt(makeEmptyAttempt());
+  const mfaResponseRef = useRef<mfaResponsePromiseWithResolvers>();
+
+  const cancelAttempt = () => {
+    if (mfaResponseRef.current) {
+      mfaResponseRef.current.reject(new MfaCanceledError());
+    }
   };
 
   const getChallengeResponse = useCallback(
     async (challenge?: MfaAuthenticateChallenge) => {
       const [resp, err] = await getResponse(challenge);
+
       if (err) throw err;
+
       return resp;
     },
     [getResponse]
   );
 
   const submit = useCallback(
     async (mfaType?: DeviceType, totpCode?: string) => {
-      if (!mfaResponsePromiseWithResolvers.current) {
+      if (!mfaResponseRef.current) {
         throw new Error('submit called without an in flight MFA attempt');
       }
 
       try {
-        await mfaResponsePromiseWithResolvers.current.resolve(
+        await mfaResponseRef.current.resolve(
           await auth.getMfaChallengeResponse(challenge, mfaType, totpCode)
         );
       } catch (err) {
@@ -158,7 +154,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
         });
       }
     },
-    [challenge, mfaResponsePromiseWithResolvers, setMfaAttempt]
+    [challenge, setMfaAttempt]
   );
 
   return {
@@ -168,7 +164,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
     getChallengeResponse,
     submit,
     attempt,
-    resetAttempt,
+    cancelAttempt,
   };
 }
 
@@ -207,7 +203,7 @@ export type MfaState = {
   ) => Promise<MfaChallengeResponse>;
   submit: (mfaType?: DeviceType, totpCode?: string) => Promise<void>;
   attempt: Attempt<any>;
-  resetAttempt: () => void;
+  cancelAttempt: () => void;
 };
 
 // used for testing
@@ -219,6 +215,12 @@ export function makeDefaultMfaState(): MfaState {
     getChallengeResponse: async () => null,
     submit: () => null,
     attempt: makeEmptyAttempt(),
-    resetAttempt: () => null,
+    cancelAttempt: () => null,
   };
 }
+
+export class MfaCanceledError extends Error {
+  constructor() {
+    super('User canceled MFA attempt');
+  }
+}