feat: created reusable mfa flow

frontend/src/pages/login/select-organization.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useEffect } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { useTranslation } from "react-i18next";
 import Head from "next/head";
 import Image from "next/image";
@@ -12,15 +12,18 @@ import jwt_decode from "jwt-decode";
 
 import { createNotification } from "@app/components/notifications";
 import { IsCliLoginSuccessful } from "@app/components/utilities/attemptCliLogin";
+import SecurityClient from "@app/components/utilities/SecurityClient";
 import { Button, Spinner } from "@app/components/v2";
 import { SessionStorageKeys } from "@app/const";
 import { useUser } from "@app/context";
+import { useToggle } from "@app/hooks";
 import { useGetOrganizations, useLogoutUser, useSelectOrganization } from "@app/hooks/api";
 import { UserAgentType } from "@app/hooks/api/auth/types";
 import { Organization } from "@app/hooks/api/types";
 import { AuthMethod } from "@app/hooks/api/users/types";
 import { getAuthToken, isLoggedIn } from "@app/reactQuery";
 import { navigateUserToOrg } from "@app/views/Login/Login.utils";
+import { Mfa } from "@app/views/Login/Mfa";
 
 const LoadingScreen = () => {
   return (
@@ -37,7 +40,10 @@ export default function LoginPage() {
 
   const organizations = useGetOrganizations();
   const selectOrg = useSelectOrganization();
+
   const { user, isLoading: userLoading } = useUser();
+  const [shouldShowMfa, toggleShowMfa] = useToggle(false);
+  const [mfaSuccessCallback, setMfaSuccessCallback] = useState<() => void>(() => {});
 
   const queryParams = new URLSearchParams(window.location.search);
   const callbackPort = queryParams.get("callback_port");
@@ -77,11 +83,19 @@ export default function LoginPage() {
         return;
       }
 
-      const { token } = await selectOrg.mutateAsync({
+      const { token, isMfaEnabled } = await selectOrg.mutateAsync({
         organizationId: organization.id,
         userAgent: callbackPort ? UserAgentType.CLI : undefined
       });
 
+      if (isMfaEnabled) {
+        SecurityClient.setMfaToken(token);
+        toggleShowMfa.on();
+
+        setMfaSuccessCallback(() => () => handleSelectOrganization(organization));
+        return;
+      }
+
       if (callbackPort) {
         const privateKey = localStorage.getItem("PRIVATE_KEY");
 
@@ -178,56 +192,57 @@ export default function LoginPage() {
         <meta property="og:title" content={t("login.og-title") ?? ""} />
         <meta name="og:description" content={t("login.og-description") ?? ""} />
       </Head>
-      <div className="mx-auto mt-20 w-fit rounded-lg border-2 border-mineshaft-500 p-10 shadow-lg">
-        <Link href="/">
-          <div className="mb-4 flex justify-center">
-            <Image src="/images/gradientLogo.svg" height={90} width={120} alt="Infisical logo" />
-          </div>
-        </Link>
-        <form
-          onSubmit={() => console.log("submit")}
-          className="mx-auto flex w-full flex-col items-center justify-center"
-        >
-          <div className="mb-8 space-y-2">
-            <h1 className="bg-gradient-to-b from-white to-bunker-200 bg-clip-text text-center text-2xl font-medium text-transparent">
-              Choose your organization
-            </h1>
-
-            <div className="space-y-1">
-              <p className="text-md text-center text-gray-500">
-                You&lsquo;re currently logged in as <strong>{user.username}</strong>
-              </p>
-              <p className="text-md text-center text-gray-500">
-                Not you?{" "}
-                <Button variant="link" onClick={handleLogout} className="font-semibold">
-                  Change account
-                </Button>
-              </p>
+      {shouldShowMfa ? (
+        <Mfa successCallback={mfaSuccessCallback} closeMfa={() => toggleShowMfa.off()} />
+      ) : (
+        <div className="mx-auto mt-20 w-fit rounded-lg border-2 border-mineshaft-500 p-10 shadow-lg">
+          <Link href="/">
+            <div className="mb-4 flex justify-center">
+              <Image src="/images/gradientLogo.svg" height={90} width={120} alt="Infisical logo" />
             </div>
-          </div>
-          <div className="mt-2 w-1/4 min-w-[21.2rem] space-y-4 rounded-md text-center md:min-w-[25.1rem] lg:w-1/4">
-            {organizations.isLoading ? (
-              <Spinner />
-            ) : (
-              organizations.data?.map((org) => (
-                // eslint-disable-next-line jsx-a11y/click-events-have-key-events, jsx-a11y/no-static-element-interactions
-                <div
-                  onClick={() => handleSelectOrganization(org)}
-                  key={org.id}
-                  className="group flex cursor-pointer items-center justify-between rounded-md bg-mineshaft-700 px-4 py-3 capitalize text-gray-200 shadow-md transition-colors hover:bg-mineshaft-600"
-                >
-                  <p className="truncate transition-colors">{org.name}</p>
+          </Link>
+          <form className="mx-auto flex w-full flex-col items-center justify-center">
+            <div className="mb-8 space-y-2">
+              <h1 className="bg-gradient-to-b from-white to-bunker-200 bg-clip-text text-center text-2xl font-medium text-transparent">
+                Choose your organization
+              </h1>
 
-                  <FontAwesomeIcon
-                    icon={faArrowRight}
-                    className="text-gray-400 transition-all group-hover:translate-x-2 group-hover:text-primary-500"
-                  />
-                </div>
-              ))
-            )}
-          </div>
-        </form>
-      </div>
+              <div className="space-y-1">
+                <p className="text-md text-center text-gray-500">
+                  You&lsquo;re currently logged in as <strong>{user.username}</strong>
+                </p>
+                <p className="text-md text-center text-gray-500">
+                  Not you?{" "}
+                  <Button variant="link" onClick={handleLogout} className="font-semibold">
+                    Change account
+                  </Button>
+                </p>
+              </div>
+            </div>
+            <div className="mt-2 w-1/4 min-w-[21.2rem] space-y-4 rounded-md text-center md:min-w-[25.1rem] lg:w-1/4">
+              {organizations.isLoading ? (
+                <Spinner />
+              ) : (
+                organizations.data?.map((org) => (
+                  // eslint-disable-next-line jsx-a11y/click-events-have-key-events, jsx-a11y/no-static-element-interactions
+                  <div
+                    onClick={() => handleSelectOrganization(org)}
+                    key={org.id}
+                    className="group flex cursor-pointer items-center justify-between rounded-md bg-mineshaft-700 px-4 py-3 capitalize text-gray-200 shadow-md transition-colors hover:bg-mineshaft-600"
+                  >
+                    <p className="truncate transition-colors">{org.name}</p>
+
+                    <FontAwesomeIcon
+                      icon={faArrowRight}
+                      className="text-gray-400 transition-all group-hover:translate-x-2 group-hover:text-primary-500"
+                    />
+                  </div>
+                ))
+              )}
+            </div>
+          </form>
+        </div>
+      )}
 
       <div className="pb-28" />
     </div>

frontend/src/views/Login/Mfa.tsx

@@ -0,0 +1,156 @@
+import { useState } from "react";
+import ReactCodeInput from "react-code-input";
+import Image from "next/image";
+import Link from "next/link";
+import { useRouter } from "next/router";
+import { t } from "i18next";
+
+import Error from "@app/components/basic/Error";
+import SecurityClient from "@app/components/utilities/SecurityClient";
+import { Button } from "@app/components/v2";
+import { useUser } from "@app/context";
+import { useSendMfaToken } from "@app/hooks/api";
+import { verifyMfaToken } from "@app/hooks/api/auth/queries";
+
+// The style for the verification code input
+const codeInputProps = {
+  inputStyle: {
+    fontFamily: "monospace",
+    margin: "4px",
+    MozAppearance: "textfield",
+    width: "48px",
+    borderRadius: "5px",
+    fontSize: "24px",
+    height: "48px",
+    paddingLeft: "7",
+    backgroundColor: "#0d1117",
+    color: "white",
+    border: "1px solid #2d2f33",
+    textAlign: "center",
+    outlineColor: "#8ca542",
+    borderColor: "#2d2f33"
+  }
+} as const;
+
+type Props = {
+  successCallback: () => void;
+  closeMfa: () => void;
+};
+
+export const Mfa = ({ successCallback, closeMfa }: Props) => {
+  const [mfaCode, setMfaCode] = useState("");
+  const router = useRouter();
+  const [isLoading, setIsLoading] = useState(false);
+  const [isLoadingResend, setIsLoadingResend] = useState(false);
+  const [triesLeft, setTriesLeft] = useState<number | undefined>(undefined);
+  const { user } = useUser();
+
+  const sendMfaToken = useSendMfaToken();
+
+  const verifyMfa = async () => {
+    if (!user.email) {
+      return;
+    }
+
+    setIsLoading(true);
+    try {
+      const { token } = await verifyMfaToken({
+        email: user.email,
+        mfaCode
+      });
+
+      SecurityClient.setMfaToken("");
+      SecurityClient.setToken(token);
+
+      successCallback();
+      closeMfa();
+    } catch (error) {
+      if (triesLeft) {
+        setTriesLeft((left) => {
+          if (triesLeft === 1) {
+            router.push("/");
+
+            SecurityClient.setMfaToken("");
+            SecurityClient.setToken("");
+          }
+          return (left as number) - 1;
+        });
+      } else {
+        setTriesLeft(2);
+      }
+    } finally {
+      setIsLoading(false);
+    }
+  };
+
+  const handleResendMfaCode = async () => {
+    if (!user?.email) {
+      return;
+    }
+
+    try {
+      setIsLoadingResend(true);
+      await sendMfaToken.mutateAsync({ email: user.email });
+      setIsLoadingResend(false);
+    } catch (err) {
+      console.error(err);
+      setIsLoadingResend(false);
+    }
+  };
+
+  return (
+    <div className="mx-auto w-max pb-4 pt-4 md:mb-16 md:px-8">
+      <Link href="/">
+        <div className="mb-4 flex justify-center">
+          <Image src="/images/gradientLogo.svg" height={90} width={120} alt="Infisical logo" />
+        </div>
+      </Link>
+      <p className="text-l flex justify-center text-bunker-300">{t("mfa.step2-message")}</p>
+      <p className="text-l my-1 flex justify-center font-semibold text-bunker-300">{user.email}</p>
+      <div className="mx-auto hidden w-max min-w-[20rem] md:block">
+        <ReactCodeInput
+          name=""
+          inputMode="tel"
+          type="text"
+          fields={6}
+          onChange={setMfaCode}
+          className="mt-6 mb-2"
+          {...codeInputProps}
+        />
+      </div>
+      {typeof triesLeft === "number" && (
+        <Error text={`Invalid code. You have ${triesLeft} attempt(s) remaining.`} />
+      )}
+      <div className="mx-auto mt-2 flex w-1/4 min-w-[20rem] max-w-xs flex-col items-center justify-center text-center text-sm md:max-w-md md:text-left lg:w-[19%]">
+        <div className="text-l w-full py-1 text-lg">
+          <Button
+            onClick={() => verifyMfa()}
+            size="sm"
+            isFullWidth
+            className="h-14"
+            colorSchema="primary"
+            variant="outline_bg"
+            isLoading={isLoading}
+          >
+            {String(t("mfa.verify"))}
+          </Button>
+        </div>
+      </div>
+      <div className="mx-auto flex max-h-24 w-full max-w-md flex-col items-center justify-center pt-2">
+        <div className="flex flex-row items-baseline gap-1 text-sm">
+          <span className="text-bunker-400">{t("signup.step2-resend-alert")}</span>
+          <div className="text-md mt-2 flex flex-row text-bunker-400">
+            <button disabled={isLoadingResend} onClick={handleResendMfaCode} type="button">
+              <span className="cursor-pointer duration-200 hover:text-bunker-200 hover:underline hover:decoration-primary-700 hover:underline-offset-4">
+                {isLoadingResend
+                  ? t("signup.step2-resend-progress")
+                  : t("signup.step2-resend-submit")}
+              </span>
+            </button>
+          </div>
+        </div>
+        <p className="pb-2 text-sm text-bunker-400">{t("signup.step2-spam-alert")}</p>
+      </div>
+    </div>
+  );
+};