Remove buggy useEffect logic; Move mfaEmitter into useMfa for straigh…

…tforward state sharing for file transfers.
gravitational · Feb 1, 2025 · dc34081 · dc34081
1 parent 2004893
commit dc34081
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 55 deletions.
diff --git a/web/packages/teleport/src/Console/DocumentDb/DocumentDb.tsx b/web/packages/teleport/src/Console/DocumentDb/DocumentDb.tsx
@@ -24,7 +24,7 @@ import AuthnDialog from 'teleport/components/AuthnDialog';
 import Document from 'teleport/Console/Document';
 import { Terminal, TerminalRef } from 'teleport/Console/DocumentSsh/Terminal';
 import * as stores from 'teleport/Console/stores/types';
-import { useMfaEmitter } from 'teleport/lib/useMfa';
+import { useMfa } from 'teleport/lib/useMfa';
 
 import { ConnectDialog } from './ConnectDialog';
 import { useDbSession } from './useDbSession';
@@ -37,7 +37,8 @@ type Props = {
 export function DocumentDb({ doc, visible }: Props) {
   const terminalRef = useRef<TerminalRef>();
   const { tty, status, closeDocument, sendDbConnectData } = useDbSession(doc);
-  const mfa = useMfaEmitter(tty);
+  const mfa = useMfa({ emitterSender: tty });
+
   useEffect(() => {
     // when switching tabs or closing tabs, focus on visible terminal
     terminalRef.current?.focus();

diff --git a/web/packages/teleport/src/Console/DocumentKubeExec/DocumentKubeExec.tsx b/web/packages/teleport/src/Console/DocumentKubeExec/DocumentKubeExec.tsx
@@ -25,7 +25,7 @@ import Document from 'teleport/Console/Document';
 import useKubeExecSession from 'teleport/Console/DocumentKubeExec/useKubeExecSession';
 import { Terminal, TerminalRef } from 'teleport/Console/DocumentSsh/Terminal';
 import * as stores from 'teleport/Console/stores/types';
-import { useMfaEmitter } from 'teleport/lib/useMfa';
+import { useMfa } from 'teleport/lib/useMfa';
 
 import KubeExecData from './KubeExecDataDialog';
 
@@ -38,7 +38,7 @@ export default function DocumentKubeExec({ doc, visible }: Props) {
   const terminalRef = useRef<TerminalRef>();
   const { tty, status, closeDocument, sendKubeExecData } =
     useKubeExecSession(doc);
-  const mfa = useMfaEmitter(tty);
+  const mfa = useMfa({ emitterSender: tty });
   useEffect(() => {
     // when switching tabs or closing tabs, focus on visible terminal
     terminalRef.current?.focus();

diff --git a/web/packages/teleport/src/Console/DocumentSsh/DocumentSsh.tsx b/web/packages/teleport/src/Console/DocumentSsh/DocumentSsh.tsx
@@ -30,7 +30,7 @@ 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 { useMfa } from 'teleport/lib/useMfa';
 import { MfaChallengeScope } from 'teleport/services/auth/auth';
 
 import { useConsoleContext } from '../consoleContextProvider';
@@ -54,14 +54,13 @@ 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 = useMfa({
+    emitterSender: tty,
     req: {
       scope: MfaChallengeScope.USER_SESSION,
     },
   });
-  const ft = useFileTransfer(tty, session, doc, ftMfa);
+  const ft = useFileTransfer(tty, session, doc, mfa);
   const theme = useTheme();
 
   function handleCloseFileTransfer() {
@@ -74,13 +73,10 @@ 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 (mfa.attempt.status === 'processing') {
       terminalRef.current?.focus();
     }
-  }, [visible, ttyMfa.attempt.status, ftMfa.attempt.status]);
+  }, [visible, mfa.attempt.status]);
 
   const onSearchClose = useCallback(() => {
     setShowSearch(false);
@@ -141,8 +137,8 @@ function DocumentSsh({ doc, visible }: PropTypes) {
           <Indicator />
         </Box>
       )}
-      <AuthnDialog mfaState={ttyMfa} onClose={closeDocument} />
-      <AuthnDialog mfaState={ftMfa} />
+      {/* TODO: don't close doc when doing file transfer */}
+      <AuthnDialog mfaState={mfa} onClose={closeDocument} />
       {status === 'initialized' && terminal}
     </Document>
   );

diff --git a/web/packages/teleport/src/DesktopSession/useDesktopSession.tsx b/web/packages/teleport/src/DesktopSession/useDesktopSession.tsx
@@ -24,7 +24,7 @@ import useAttempt from 'shared/hooks/useAttemptNext';
 
 import type { UrlDesktopParams } from 'teleport/config';
 import { ButtonState } from 'teleport/lib/tdp';
-import { useMfaEmitter } from 'teleport/lib/useMfa';
+import { useMfa } from 'teleport/lib/useMfa';
 import desktopService from 'teleport/services/desktops';
 import userService from 'teleport/services/user';
 
@@ -129,7 +129,7 @@ export default function useDesktopSession() {
   });
   const tdpClient = clientCanvasProps.tdpClient;
 
-  const mfa = useMfaEmitter(tdpClient);
+  const mfa = useMfa({ emitterSender: tdpClient });
 
   const onShareDirectory = () => {
     try {

diff --git a/web/packages/teleport/src/lib/useMfa.ts b/web/packages/teleport/src/lib/useMfa.ts
@@ -37,7 +37,9 @@ import {
 
 export type MfaProps = {
   req?: CreateAuthenticateChallengeRequest;
-  isMfaRequired?: boolean | null;
+  // Watch for MFA events from the emitter sender and handle them
+  // through the same MFA state.
+  emitterSender?: EventEmitterMfaSender;
 };
 
 type mfaResponsePromiseWithResolvers = {
@@ -52,19 +54,11 @@ type mfaResponsePromiseWithResolvers = {
  * be used to display options to the user and prompt for them to complete
  * the MFA check.
  */
-export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
+export function useMfa({ req, emitterSender }: MfaProps): MfaState {
   const [mfaRequired, setMfaRequired] = useState<boolean>();
   const [options, setMfaOptions] = useState<MfaOption[]>();
   const [challenge, setMfaChallenge] = useState<MfaAuthenticateChallenge>();
 
-  useEffect(() => {
-    setMfaRequired(isMfaRequired);
-  }, [isMfaRequired]);
-
-  useEffect(() => {
-    setMfaRequired(null);
-  }, [req?.isMfaRequiredRequest]);
-
   // getResponse is used to initiate MFA authentication.
   //   1. Check if MFA is required by getting a new MFA challenge
   //   2. If MFA is required, set the challenge in the MFA state and wait for it to
@@ -157,8 +151,29 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
     [challenge, setMfaAttempt]
   );
 
+  // Handle MFA challenges from the emitterSender, if provided.
+  useEffect(() => {
+    if (emitterSender) {
+      const challengeHandler = async (challengeJson: string) => {
+        // set Mfa required for other uses of this MfaState (e.g. file transfers)
+        setMfaRequired(true);
+
+        const challenge = parseMfaChallengeJson(JSON.parse(challengeJson));
+        const resp = await getChallengeResponse(challenge);
+        emitterSender.sendChallengeResponse(resp);
+      };
+
+      emitterSender?.on(TermEvent.MFA_CHALLENGE, challengeHandler);
+      return () => {
+        emitterSender?.removeListener(
+          TermEvent.MFA_CHALLENGE,
+          challengeHandler
+        );
+      };
+    }
+  }, [getChallengeResponse, emitterSender]);
+
   return {
-    required: mfaRequired,
     options,
     challenge,
     getChallengeResponse,
@@ -168,32 +183,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
   };
 }
 
-export function useMfaEmitter(emitterSender: EventEmitterMfaSender): MfaState {
-  const [mfaRequired, setMfaRequired] = useState(false);
-
-  const mfa = useMfa({ isMfaRequired: mfaRequired });
-
-  useEffect(() => {
-    const challengeHandler = async (challengeJson: string) => {
-      // set Mfa required for other uses of this MfaState (e.g. file transfers)
-      setMfaRequired(true);
-
-      const challenge = parseMfaChallengeJson(JSON.parse(challengeJson));
-      const resp = await mfa.getChallengeResponse(challenge);
-      emitterSender.sendChallengeResponse(resp);
-    };
-
-    emitterSender?.on(TermEvent.MFA_CHALLENGE, challengeHandler);
-    return () => {
-      emitterSender?.removeListener(TermEvent.MFA_CHALLENGE, challengeHandler);
-    };
-  }, [mfa, emitterSender]);
-
-  return mfa;
-}
-
 export type MfaState = {
-  required: boolean;
   options: MfaOption[];
   challenge: MfaAuthenticateChallenge;
   // Generally you wouldn't pass in a challenge, unless you already
@@ -209,7 +199,6 @@ export type MfaState = {
 // used for testing
 export function makeDefaultMfaState(): MfaState {
   return {
-    required: true,
     options: null,
     challenge: null,
     getChallengeResponse: async () => null,