From e77d07442c4f51edd22bce74671ca6cce1072d8e Mon Sep 17 00:00:00 2001
From: joerger <bjoerger@goteleport.com>
Date: Thu, 2 Jan 2025 13:32:12 -0800
Subject: [PATCH] Address comments.

---
 .../teleport/src/JoinTokens/JoinTokens.tsx    |  4 +-
 .../teleport/src/MFAContext/MFAContext.tsx    | 44 +++++++++++--------
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/web/packages/teleport/src/JoinTokens/JoinTokens.tsx b/web/packages/teleport/src/JoinTokens/JoinTokens.tsx
index 8af5200a07681..43cac62115b21 100644
--- a/web/packages/teleport/src/JoinTokens/JoinTokens.tsx
+++ b/web/packages/teleport/src/JoinTokens/JoinTokens.tsx
@@ -75,9 +75,7 @@ export const JoinTokens = () => {
   const [editingToken, setEditingToken] = useState<JoinToken | null>(null);
   const [tokenToDelete, setTokenToDelete] = useState<JoinToken | null>(null);
   const [joinTokensAttempt, runJoinTokensAttempt, setJoinTokensAttempt] =
-    useAsync(async () => {
-      return ctx.joinTokenService.fetchJoinTokens(null);
-    });
+    useAsync(() => ctx.joinTokenService.fetchJoinTokens(null));
 
   const resources = useResources(
     joinTokensAttempt.data?.items.map(makeTokenResource) || [],
diff --git a/web/packages/teleport/src/MFAContext/MFAContext.tsx b/web/packages/teleport/src/MFAContext/MFAContext.tsx
index bba93cfd177d8..03d0754566a9c 100644
--- a/web/packages/teleport/src/MFAContext/MFAContext.tsx
+++ b/web/packages/teleport/src/MFAContext/MFAContext.tsx
@@ -1,4 +1,5 @@
-import { PropsWithChildren, createContext, useCallback, useRef } from 'react';
+import { createContext, PropsWithChildren, useCallback, useState } from 'react';
+
 import AuthnDialog from 'teleport/components/AuthnDialog';
 import { useMfa } from 'teleport/lib/useMfa';
 import auth, { MfaChallengeScope } from 'teleport/services/auth/auth';
@@ -11,30 +12,35 @@ export interface MfaContextValue {
 export const MfaContext = createContext<MfaContextValue>(null);
 
 export const MfaContextProvider = ({ children }: PropsWithChildren) => {
-  const allowReuse = useRef(false);
-  const adminMfa = useMfa({
-    req: {
-      scope: MfaChallengeScope.ADMIN_ACTION,
-      allowReuse: allowReuse.current,
-      isMfaRequiredRequest: {
-        admin_action: {},
-      },
-    },
-  });
+  const adminMfa = useMfa({});
 
   const getAdminActionMfaResponse = useCallback(
     async (reusable: boolean = false) => {
-      allowReuse.current = reusable;
-      return (await adminMfa.getChallengeResponse()) || {}; // return an empty challenge to prevent mfa retry.
+      const chal = await auth.getMfaChallenge({
+        scope: MfaChallengeScope.ADMIN_ACTION,
+        allowReuse: reusable,
+        isMfaRequiredRequest: {
+          admin_action: {},
+        },
+      });
+
+      const res = await adminMfa.getChallengeResponse(chal);
+      if (!res) {
+        return {}; // return an empty challenge to prevent mfa retry.
+      }
+
+      return res;
     },
-    [adminMfa, allowReuse]
+    [adminMfa]
   );
 
-  const mfaCtx = {
-    getAdminActionMfaResponse,
-    useMfa,
-  };
-  auth.setMfaContext(mfaCtx);
+  const [mfaCtx, setMfaCtx] = useState<MfaContextValue>();
+
+  if (!mfaCtx) {
+    const mfaCtx = { getAdminActionMfaResponse };
+    setMfaCtx(mfaCtx);
+    auth.setMfaContext(mfaCtx);
+  }
 
   return (
     <MfaContext.Provider value={mfaCtx}>