Fix MFA always required for file transfers

web/packages/teleport/src/Console/DocumentSsh/DocumentSsh.tsx

@@ -25,12 +25,13 @@ import {
   FileTransferActionBar,
   FileTransferContextProvider,
   FileTransferRequests,
+  useFileTransferContext,
 } from 'shared/components/FileTransfer';
 import { TerminalSearch } from 'shared/components/TerminalSearch';
 
 import AuthnDialog from 'teleport/components/AuthnDialog';
 import * as stores from 'teleport/Console/stores';
-import { useMfa, useMfaEmitter } from 'teleport/lib/useMfa';
+import { useMfaEmitter } from 'teleport/lib/useMfa';
 import { MfaChallengeScope } from 'teleport/services/auth/auth';
 
 import { useConsoleContext } from '../consoleContextProvider';
@@ -54,14 +55,17 @@ function DocumentSsh({ doc, visible }: PropTypes) {
   const { tty, status, closeDocument, session } = useSshSession(doc);
   const [showSearch, setShowSearch] = useState(false);
 
-  const ttyMfa = useMfaEmitter(tty);
-  const ftMfa = useMfa({
-    isMfaRequired: ttyMfa.required,
+  const mfa = useMfaEmitter(tty, {
+    // The MFA requirement will be determined by whether we do/don't get
+    // an mfa challenge over the event emitter at session start.
+    isMfaRequired: false,
     req: {
       scope: MfaChallengeScope.USER_SESSION,
     },
   });
-  const ft = useFileTransfer(tty, session, doc, ftMfa);
+  const ft = useFileTransfer(tty, session, doc, mfa);
+  const { openedDialog: ftOpenedDialog } = useFileTransferContext();
+
   const theme = useTheme();
 
   function handleCloseFileTransfer() {
@@ -73,14 +77,12 @@ function DocumentSsh({ doc, visible }: PropTypes) {
   }
 
   useEffect(() => {
-    // when switching tabs or closing tabs, focus on visible terminal
-    if (
-      ttyMfa.attempt.status === 'processing' ||
-      ftMfa.attempt.status === 'processing'
-    ) {
+    // If an MFA attempt starts while switching tabs or closing tabs,
+    // automatically focus on visible terminal.
+    if (mfa.challenge) {
       terminalRef.current?.focus();
     }
-  }, [visible, ttyMfa.attempt.status, ftMfa.attempt.status]);
+  }, [visible, mfa.challenge]);
 
   const onSearchClose = useCallback(() => {
     setShowSearch(false);
@@ -141,8 +143,15 @@ function DocumentSsh({ doc, visible }: PropTypes) {
           <Indicator />
         </Box>
       )}
-      <AuthnDialog mfaState={ttyMfa} onClose={closeDocument} />
-      <AuthnDialog mfaState={ftMfa} />
+      <AuthnDialog
+        mfaState={mfa}
+        onClose={() => {
+          // Don't close the ssh doc if this is just a file transfer request.
+          if (!ftOpenedDialog) {
+            closeDocument();
+          }
+        }}
+      />
       {status === 'initialized' && terminal}
     </Document>
   );

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

@@ -17,7 +17,7 @@
  */
 
 import { renderHook, waitFor } from '@testing-library/react';
-import { useState } from 'react';
+import { act } from 'react';
 
 import { CreateAuthenticateChallengeRequest } from 'teleport/services/auth';
 import auth, { MfaChallengeScope } from 'teleport/services/auth/auth';
@@ -61,10 +61,6 @@ const mockChallengeReq: CreateAuthenticateChallengeRequest = {
 };
 
 describe('useMfa', () => {
-  beforeEach(() => {
-    jest.spyOn(console, 'error').mockImplementation();
-  });
-
   afterEach(() => {
     jest.clearAllMocks();
   });
@@ -80,27 +76,25 @@ describe('useMfa', () => {
       })
     );
 
-    const respPromise = mfa.current.getChallengeResponse();
-    await waitFor(() => {
-      expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
+    let resp;
+    await act(async () => {
+      resp = mfa.current.getChallengeResponse();
     });
 
+    expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
     expect(mfa.current.options).toEqual([MFA_OPTION_WEBAUTHN]);
-    expect(mfa.current.required).toEqual(true);
+    expect(mfa.current.mfaRequired).toEqual(true);
     expect(mfa.current.challenge).toEqual(mockChallenge);
     expect(mfa.current.attempt.status).toEqual('processing');
 
-    await mfa.current.submit('webauthn');
-    await waitFor(() => {
-      expect(auth.getMfaChallengeResponse).toHaveBeenCalledWith(
-        mockChallenge,
-        'webauthn',
-        undefined
-      );
-    });
+    await act(async () => mfa.current.submit('webauthn'));
+    expect(auth.getMfaChallengeResponse).toHaveBeenCalledWith(
+      mockChallenge,
+      'webauthn',
+      undefined
+    );
 
-    const resp = await respPromise;
-    expect(resp).toEqual(mockResponse);
+    expect(await resp).toEqual(mockResponse);
     expect(mfa.current.challenge).toEqual(null);
     expect(mfa.current.attempt.status).toEqual('success');
   });
@@ -116,51 +110,63 @@ describe('useMfa', () => {
 
     // If a challenge is not returned, an empty mfa response should be returned
     // early and the requirement changed to false for future calls.
-    const resp = await mfa.current.getChallengeResponse();
+    const resp = await act(async () => {
+      return mfa.current.getChallengeResponse();
+    });
     expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
     expect(resp).toEqual(undefined);
-    await waitFor(() => expect(mfa.current.required).toEqual(false));
+    await waitFor(() => expect(mfa.current.mfaRequired).toEqual(false));
   });
 
-  test('adaptable mfa requirement state', async () => {
-    jest.spyOn(auth, 'getMfaChallenge').mockResolvedValue(null);
-
-    let isMfaRequired: boolean;
-    let setMfaRequired: (b: boolean) => void;
+  test('adaptable mfa requirement', async () => {
+    jest.spyOn(auth, 'getMfaChallenge').mockResolvedValueOnce(mockChallenge);
+    jest
+      .spyOn(auth, 'getMfaChallengeResponse')
+      .mockResolvedValueOnce(mockResponse);
 
-    let req: CreateAuthenticateChallengeRequest;
-    let setReq: (r: CreateAuthenticateChallengeRequest) => void;
+    const { result: mfa } = renderHook(() =>
+      useMfa({
+        isMfaRequired: false,
+        req: mockChallengeReq,
+      })
+    );
 
-    const { result: mfa } = renderHook(() => {
-      [isMfaRequired, setMfaRequired] = useState<boolean>(null);
-      [req, setReq] =
-        useState<CreateAuthenticateChallengeRequest>(mockChallengeReq);
+    // The mfa requirement can be initialized to a non-undefined value to skip
+    // the mfa check when it isn't needed, e.g. the requirement was predetermined.
+    expect(mfa.current.mfaRequired).toEqual(false);
 
-      return useMfa({
-        req: req,
-        isMfaRequired: isMfaRequired,
-      });
+    // The mfa required state can be changed directly, in case the requirement
+    // need to be updated by the caller.
+    await act(async () => {
+      mfa.current.setMfaRequired(true);
+    });
+    await waitFor(() => {
+      expect(mfa.current.mfaRequired).toEqual(true);
     });
 
-    // mfaRequired should change when the isMfaRequired arg changes, allowing
-    // callers to propagate mfa required late (e.g. per-session MFA for file transfers)
-    setMfaRequired(false);
-    await waitFor(() => expect(mfa.current.required).toEqual(false));
-
-    setMfaRequired(true);
-    await waitFor(() => expect(mfa.current.required).toEqual(true));
+    // The mfa required state can be changed back to undefined (default) or null to force
+    // the next mfa attempt to re-check mfa required / attempt to get a challenge.
+    await act(async () => {
+      mfa.current.setMfaRequired(null);
+    });
+    await waitFor(() => {
+      expect(mfa.current.mfaRequired).toEqual(null);
+    });
 
-    setMfaRequired(null);
-    await waitFor(() => expect(mfa.current.required).toEqual(null));
+    // mfa required will be updated to true as usual once the server returns an mfa challenge.
+    await act(async () => {
+      mfa.current.getChallengeResponse();
+    });
+    await waitFor(() => {
+      expect(mfa.current.challenge).toBeDefined();
+    });
+    await act(async () => {
+      return mfa.current.submit();
+    });
 
-    // If isMfaRequiredRequest changes, the mfaRequired value should be reset.
-    setReq({
-      ...mockChallengeReq,
-      isMfaRequiredRequest: {
-        admin_action: {},
-      },
+    await waitFor(() => {
+      expect(mfa.current.mfaRequired).toEqual(true);
     });
-    await waitFor(() => expect(mfa.current.required).toEqual(null));
   });
 
   test('mfa challenge error', async () => {
@@ -171,14 +177,15 @@ describe('useMfa', () => {
 
     const { result: mfa } = renderHook(() => useMfa({}));
 
-    await expect(mfa.current.getChallengeResponse).rejects.toThrow(err);
-    await waitFor(() => {
-      expect(mfa.current.attempt).toEqual({
-        status: 'error',
-        statusText: err.message,
-        error: err,
-        data: null,
-      });
+    await act(async () => {
+      await expect(mfa.current.getChallengeResponse()).rejects.toThrow(err);
+    });
+
+    expect(mfa.current.attempt).toEqual({
+      status: 'error',
+      statusText: err.message,
+      error: err,
+      data: null,
     });
   });
 
@@ -195,28 +202,29 @@ describe('useMfa', () => {
       })
     );
 
-    const respPromise = mfa.current.getChallengeResponse();
-    await waitFor(() => {
-      expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
+    let resp;
+    await act(async () => {
+      resp = mfa.current.getChallengeResponse();
     });
-    await mfa.current.submit('webauthn');
 
-    await waitFor(() => {
-      expect(mfa.current.attempt).toEqual({
-        status: 'error',
-        statusText: err.message,
-        error: err,
-        data: null,
-      });
+    expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
+
+    await act(async () => mfa.current.submit('webauthn'));
+    expect(mfa.current.attempt).toEqual({
+      status: 'error',
+      statusText: err.message,
+      error: err,
+      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);
+
+    await act(async () => mfa.current.submit('webauthn'));
+    expect(await resp).toEqual(mockResponse);
   });
 
   test('reset mfa attempt', async () => {
@@ -227,18 +235,22 @@ describe('useMfa', () => {
       })
     );
 
-    const respPromise = mfa.current.getChallengeResponse();
-    await waitFor(() => {
-      expect(auth.getMfaChallenge).toHaveBeenCalled();
+    let resp: Promise<MfaChallengeResponse>;
+    await act(async () => {
+      resp = mfa.current.getChallengeResponse();
     });
 
-    mfa.current.cancelAttempt();
+    // Before calling mfa.current.cancelAttempt(), we need to write code that handles rejection of
+    // resp. Otherwise the test is going to fail because of unhandled promise rejection.
+    // eslint-disable-next-line jest/valid-expect
+    const expectedRespRejection = expect(resp).rejects.toEqual(
+      new MfaCanceledError()
+    );
 
-    await expect(respPromise).rejects.toThrow(new MfaCanceledError());
+    await act(async () => mfa.current.cancelAttempt());
 
-    await waitFor(() => {
-      expect(mfa.current.attempt.status).toEqual('error');
-    });
+    await expectedRespRejection;
+    expect(mfa.current.attempt.status).toEqual('error');
     expect(
       mfa.current.attempt.status === 'error' && mfa.current.attempt.error
     ).toEqual(new MfaCanceledError());