Make useMfa and AuthnDialog more reusable and error proof.

gravitational · Dec 16, 2024 · 5497280 · 5497280
1 parent a8c8dac
commit 5497280
Show file tree

Hide file tree

Showing 8 changed files with 248 additions and 183 deletions.
diff --git a/web/packages/teleport/src/Console/DocumentKubeExec/DocumentKubeExec.tsx b/web/packages/teleport/src/Console/DocumentKubeExec/DocumentKubeExec.tsx
@@ -22,7 +22,7 @@ import { Box, Indicator } from 'design';
 
 import * as stores from 'teleport/Console/stores/types';
 import { Terminal, TerminalRef } from 'teleport/Console/DocumentSsh/Terminal';
-import { useMfa } from 'teleport/lib/useMfa';
+import { useMfaTty } from 'teleport/lib/useMfa';
 import useKubeExecSession from 'teleport/Console/DocumentKubeExec/useKubeExecSession';
 
 import Document from 'teleport/Console/Document';
@@ -39,11 +39,11 @@ export default function DocumentKubeExec({ doc, visible }: Props) {
   const terminalRef = useRef<TerminalRef>();
   const { tty, status, closeDocument, sendKubeExecData } =
     useKubeExecSession(doc);
-  const mfa = useMfa(tty);
+  const mfa = useMfaTty(tty);
   useEffect(() => {
     // when switching tabs or closing tabs, focus on visible terminal
     terminalRef.current?.focus();
-  }, [visible, mfa.mfaChallenge]);
+  }, [visible, mfa.challenge]);
   const theme = useTheme();
 
   const terminal = (
@@ -63,7 +63,13 @@ export default function DocumentKubeExec({ doc, visible }: Props) {
           <Indicator />
         </Box>
       )}
-      {mfa.mfaChallenge && <AuthnDialog mfa={mfa} onCancel={closeDocument} />}
+      <AuthnDialog
+        {...mfa}
+        resetAttempt={() => {
+          mfa.resetAttempt();
+          closeDocument();
+        }}
+      />
 
       {status === 'waiting-for-exec-data' && (
         <KubeExecData onExec={sendKubeExecData} onClose={closeDocument} />

diff --git a/web/packages/teleport/src/Console/DocumentSsh/DocumentSsh.tsx b/web/packages/teleport/src/Console/DocumentSsh/DocumentSsh.tsx
@@ -32,7 +32,8 @@ import { TerminalSearch } from 'shared/components/TerminalSearch';
 import * as stores from 'teleport/Console/stores';
 
 import AuthnDialog from 'teleport/components/AuthnDialog';
-import { useMfa } from 'teleport/lib/useMfa';
+import { useMfa, useMfaTty } from 'teleport/lib/useMfa';
+import { MfaChallengeScope } from 'teleport/services/auth/auth';
 
 import Document from '../Document';
 
@@ -56,8 +57,15 @@ function DocumentSsh({ doc, visible }: PropTypes) {
   const terminalRef = useRef<TerminalRef>();
   const { tty, status, closeDocument, session } = useSshSession(doc);
   const [showSearch, setShowSearch] = useState(false);
-  const mfa = useMfa(tty);
-  const ft = useFileTransfer(tty, session, doc, mfa.mfaRequired);
+
+  const ttyMfa = useMfaTty(tty);
+  const ftMfa = useMfa({
+    isMfaRequired: ttyMfa.required,
+    req: {
+      scope: MfaChallengeScope.USER_SESSION,
+    },
+  });
+  const ft = useFileTransfer(tty, session, doc, ftMfa);
   const theme = useTheme();
 
   function handleCloseFileTransfer() {
@@ -71,7 +79,7 @@ function DocumentSsh({ doc, visible }: PropTypes) {
   useEffect(() => {
     // when switching tabs or closing tabs, focus on visible terminal
     terminalRef.current?.focus();
-  }, [visible, mfa.mfaChallenge]);
+  }, [visible, ftMfa.challenge]);
 
   const onSearchClose = useCallback(() => {
     setShowSearch(false);
@@ -111,13 +119,9 @@ function DocumentSsh({ doc, visible }: PropTypes) {
             beforeClose={() =>
               window.confirm('Are you sure you want to cancel file transfers?')
             }
-            errorText={
-              ft.mfaAttempt.status === 'error' ? ft.mfaAttempt.statusText : null
-            }
             afterClose={handleCloseFileTransfer}
             transferHandlers={{
-              getDownloader: ft.getDownloader,
-              getUploader: ft.getUploader,
+              ...ft,
             }}
           />
         </>
@@ -136,15 +140,14 @@ function DocumentSsh({ doc, visible }: PropTypes) {
           <Indicator />
         </Box>
       )}
-      {mfa.mfaChallenge && <AuthnDialog {...mfa} onCancel={closeDocument} />}
-      {ft.mfaChallenge && (
-        <AuthnDialog
-          mfaChallenge={ft.mfaChallenge}
-          submitMfa={ft.submitMfa}
-          submitAttempt={ft.submitMfaAttempt}
-          onCancel={ft.clearMfaChallenge}
-        />
-      )}
+      <AuthnDialog
+        {...ttyMfa}
+        resetAttempt={() => {
+          ttyMfa.resetAttempt();
+          closeDocument();
+        }}
+      />
+      <AuthnDialog {...ftMfa} />
       {status === 'initialized' && terminal}
     </Document>
   );

diff --git a/web/packages/teleport/src/Console/DocumentSsh/useFileTransfer.ts b/web/packages/teleport/src/Console/DocumentSsh/useFileTransfer.ts
@@ -16,22 +16,17 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import { useCallback, useEffect, useRef, useState } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 import { useFileTransferContext } from 'shared/components/FileTransfer';
 
 import { DocumentSsh } from 'teleport/Console/stores';
 import { EventType } from 'teleport/lib/term/enums';
 import Tty from 'teleport/lib/term/tty';
 import { Session } from 'teleport/services/session';
 
-import { useAsync } from 'shared/hooks/useAsync';
 import cfg from 'teleport/config';
-import auth, { MfaChallengeScope } from 'teleport/services/auth/auth';
-import {
-  DeviceType,
-  MfaAuthenticateChallenge,
-  MfaChallengeResponse,
-} from 'teleport/services/mfa';
+
+import { MfaState } from 'teleport/lib/useMfa';
 
 import { useConsoleContext } from '../consoleContextProvider';
 
@@ -59,7 +54,7 @@ export const useFileTransfer = (
   tty: Tty,
   session: Session,
   currentDoc: DocumentSsh,
-  addMfaToScpUrls: boolean
+  mfa: MfaState
 ) => {
   const { filesStore } = useFileTransferContext();
   const startTransfer = filesStore.start;
@@ -70,60 +65,13 @@ export const useFileTransfer = (
   >([]);
   const { clusterId, serverId, login } = currentDoc;
 
-  const [mfaChallenge, setMfaChallenge] = useState<MfaAuthenticateChallenge>();
-  const mfaResponsePromise =
-    useRef<PromiseWithResolvers<MfaChallengeResponse>>();
-
-  const clearMfaChallenge = () => {
-    setMfaChallenge(null);
-  };
-
-  const [mfaAttempt, getMfaResponse] = useAsync(
-    useCallback(async () => {
-      const challenge = await auth.getMfaChallenge({
-        scope: MfaChallengeScope.USER_SESSION,
-        isMfaRequiredRequest: {
-          node: {
-            node_name: serverId,
-            login: login,
-          },
-        },
-      });
-      if (!challenge) return;
-
-      setMfaChallenge(challenge);
-      mfaResponsePromise.current = Promise.withResolvers();
-      const mfaResponse = await mfaResponsePromise.current.promise;
-      return mfaResponse;
-    }, [setMfaChallenge, clearMfaChallenge, mfaResponsePromise])
-  );
-
-  const [submitMfaAttempt, submitMfa] = useAsync(
-    useCallback(
-      async (mfaType?: DeviceType) => {
-        try {
-          const resp = auth.getMfaChallengeResponse(mfaChallenge, mfaType);
-          mfaResponsePromise.current.resolve(resp);
-        } catch (err) {
-          mfaResponsePromise.current.reject(err);
-          throw err;
-        }
-      },
-      [mfaChallenge, mfaResponsePromise]
-    )
-  );
-
   const download = useCallback(
     async (
       location: string,
       abortController: AbortController,
       moderatedSessionParams?: ModeratedSessionParams
     ) => {
-      let mfaResponse: MfaChallengeResponse;
-      if (addMfaToScpUrls) {
-        [mfaResponse] = await getMfaResponse();
-      }
-
+      const mfaResponse = await mfa.getChallengeResponse();
       const url = cfg.getScpUrl({
         location,
         clusterId,
@@ -144,7 +92,7 @@ export const useFileTransfer = (
       }
       return getHttpFileTransferHandlers().download(url, abortController);
     },
-    [clusterId, login, serverId, addMfaToScpUrls]
+    [clusterId, login, serverId, mfa]
   );
 
   const upload = useCallback(
@@ -154,10 +102,7 @@ export const useFileTransfer = (
       abortController: AbortController,
       moderatedSessionParams?: ModeratedSessionParams
     ) => {
-      let mfaResponse: MfaChallengeResponse;
-      if (addMfaToScpUrls) {
-        [mfaResponse] = await getMfaResponse();
-      }
+      const mfaResponse = await mfa.getChallengeResponse();
 
       const url = cfg.getScpUrl({
         location,
@@ -178,7 +123,7 @@ export const useFileTransfer = (
       }
       return getHttpFileTransferHandlers().upload(url, file, abortController);
     },
-    [clusterId, serverId, login, addMfaToScpUrls]
+    [clusterId, serverId, login, mfa]
   );
 
   /*
@@ -317,11 +262,6 @@ export const useFileTransfer = (
   }
 
   return {
-    mfaAttempt,
-    mfaChallenge,
-    clearMfaChallenge,
-    submitMfa,
-    submitMfaAttempt,
     fileTransferRequests,
     getUploader,
     getDownloader,

diff --git a/web/packages/teleport/src/DesktopSession/DesktopSession.tsx b/web/packages/teleport/src/DesktopSession/DesktopSession.tsx
@@ -184,11 +184,12 @@ export function DesktopSession(props: State) {
 const MfaDialog = ({ mfa }: { mfa: MfaState }) => {
   return (
     <AuthnDialog
-      mfa={mfa}
-      onCancel={() => {
-        mfa.submitAttempt.status = 'failed';
-        mfa.submitAttempt.statusText =
-          'This session requires multi factor authentication to continue. Please hit "Retry" and follow the prompts given by your browser to complete authentication.';
+      {...mfa}
+      replaceErrorText={
+        'This session requires multi factor authentication to continue. Please hit try again and follow the prompts given by your browser to complete authentication.'
+      }
+      resetAttempt={() => {
+        mfa.resetAttempt();
       }}
     />
   );
@@ -294,7 +295,7 @@ const nextScreenState = (
 
   // Otherwise, calculate a new screen state.
   const showAnotherSessionActive = showAnotherSessionActiveDialog;
-  const showMfa = webauthn.mfaChallenge;
+  const showMfa = webauthn.challenge;
   const showAlert =
     fetchAttempt.status === 'failed' || // Fetch attempt failed
     tdpConnection.status === 'failed' || // TDP connection failed

diff --git a/web/packages/teleport/src/DesktopSession/useDesktopSession.tsx b/web/packages/teleport/src/DesktopSession/useDesktopSession.tsx
@@ -22,7 +22,7 @@ import { useParams } from 'react-router';
 import useAttempt from 'shared/hooks/useAttemptNext';
 
 import { ButtonState } from 'teleport/lib/tdp';
-import { useMfa } from 'teleport/lib/useMfa';
+import { useMfaTty } from 'teleport/lib/useMfa';
 import desktopService from 'teleport/services/desktops';
 import userService from 'teleport/services/user';
 
@@ -130,7 +130,7 @@ export default function useDesktopSession() {
   });
   const tdpClient = clientCanvasProps.tdpClient;
 
-  const mfa = useMfa(tdpClient);
+  const mfa = useMfaTty(tdpClient);
 
   const onShareDirectory = () => {
     try {

diff --git a/web/packages/teleport/src/components/AuthnDialog/AuthnDialog.test.tsx b/web/packages/teleport/src/components/AuthnDialog/AuthnDialog.test.tsx
@@ -52,7 +52,7 @@ describe('AuthnDialog', () => {
 
   test('renders single option dialog', () => {
     const mfa = makeMockState({
-      mfaChallenge: {
+      challenge: {
         ssoChallenge: mockSsoChallenge,
       },
     });
@@ -68,7 +68,7 @@ describe('AuthnDialog', () => {
 
   test('renders multi option dialog', () => {
     const mfa = makeMockState({
-      mfaChallenge: {
+      challenge: {
         ssoChallenge: mockSsoChallenge,
         webauthnPublicKey: {
           challenge: new ArrayBuffer(1),
@@ -90,7 +90,7 @@ describe('AuthnDialog', () => {
   test('displays error text when provided', () => {
     const errorText = 'Authentication failed';
     const mfa = makeMockState({
-      mfaChallenge: {},
+      challenge: {},
       submitAttempt: {
         status: 'error',
         statusText: errorText,
@@ -106,27 +106,27 @@ describe('AuthnDialog', () => {
 
   test('sso button renders with callback', async () => {
     const mfa = makeMockState({
-      mfaChallenge: {
+      challenge: {
         ssoChallenge: mockSsoChallenge,
       },
-      submitMfa: jest.fn(),
+      submit: jest.fn(),
     });
     render(<AuthnDialog mfa={mfa} onCancel={mockOnCancel} />);
     const ssoButton = screen.getByText('Okta');
     fireEvent.click(ssoButton);
-    expect(mfa.submitMfa).toHaveBeenCalledTimes(1);
+    expect(mfa.submit).toHaveBeenCalledTimes(1);
   });
 
   test('webauthn button renders with callback', async () => {
     const mfa = makeMockState({
-      mfaChallenge: {
+      challenge: {
         webauthnPublicKey: { challenge: new ArrayBuffer(0) },
       },
-      submitMfa: jest.fn(),
+      submit: jest.fn(),
     });
     render(<AuthnDialog mfa={mfa} onCancel={mockOnCancel} />);
     const webauthn = screen.getByText('Passkey or MFA Device');
     fireEvent.click(webauthn);
-    expect(mfa.submitMfa).toHaveBeenCalledTimes(1);
+    expect(mfa.submit).toHaveBeenCalledTimes(1);
   });
 });