feat: implement TOTP-based two-factor authentication flow (#801)

* feat: initial code

* fix: mfa verification flow and data refresh after completion

* refactor: improve ui for security tab & mfa enable modal

* refactor: ui for mfa verify & redirect non-users

* fix: improve QR code readability in dark mode

* fix: improve QR code readability in dark mode

* feat: improvements on 2fa

* feat: lint fix

* fix: enhance MFA verification error handling with Supabase error codes

---------

Co-authored-by: Foad Kesheh <[email protected]>

app/login/page.tsx

@@ -138,6 +138,19 @@ export default async function Login({
       return redirect(`/login?message=4`)
     }
 
+    const { data: aalData, error: aalError } =
+      await supabase.auth.mfa.getAuthenticatorAssuranceLevel()
+    if (aalError) {
+      return redirect(`/login?message=auth`)
+    }
+
+    if (
+      aalData.nextLevel === "aal2" &&
+      aalData.nextLevel !== aalData.currentLevel
+    ) {
+      return redirect(`/login/verify`)
+    }
+
     const { data: homeWorkspace, error: homeWorkspaceError } = await supabase
       .from("workspaces")
       .select("*")

app/login/verify/mfa-verification.tsx

@@ -0,0 +1,165 @@
+"use client"
+
+import { useContext, useState } from "react"
+import { Brand } from "@/components/ui/brand"
+import { Button } from "@/components/ui/button"
+import { Input } from "@/components/ui/input"
+import { IconAlertCircle } from "@tabler/icons-react"
+import { supabase } from "@/lib/supabase/browser-client"
+import { useRouter } from "next/navigation"
+import { PentestGPTContext } from "@/context/context"
+import { getHomeWorkspaceByUserId } from "@/db/workspaces"
+
+interface MFAVerificationProps {
+  onVerify: (code: string) => Promise<{ success: boolean } | void>
+}
+
+export function MFAVerification({ onVerify }: MFAVerificationProps) {
+  const router = useRouter()
+  const [verifyCode, setVerifyCode] = useState("")
+  const [error, setError] = useState("")
+  const [isVerifying, setIsVerifying] = useState(false)
+  const { user } = useContext(PentestGPTContext)
+  const { fetchStartingData } = useContext(PentestGPTContext)
+
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault()
+    setError("")
+    setIsVerifying(true)
+
+    if (!user) {
+      router.push("/login")
+      return
+    }
+
+    try {
+      const result = await onVerify(verifyCode)
+      if (result?.success) {
+        await fetchStartingData()
+        const homeWorkspaceId = await getHomeWorkspaceByUserId(user.id)
+        router.push(`/${homeWorkspaceId}/chat`)
+      }
+    } catch (err) {
+      setIsVerifying(false)
+
+      // Handle Supabase Auth specific errors
+      if (err instanceof Error) {
+        const error = err as any // for accessing .code property
+
+        switch (error.code) {
+          case "mfa_verification_failed":
+            setError("Invalid verification code. Please try again.")
+            break
+          case "mfa_challenge_expired":
+            setError("Verification code has expired. Please request a new one.")
+            break
+          case "over_request_rate_limit":
+            setError(
+              "Too many attempts. Please wait a few minutes before trying again."
+            )
+            break
+          case "mfa_totp_verify_not_enabled":
+            setError(
+              "MFA verification is currently disabled. Please contact support."
+            )
+            break
+          case "mfa_verification_rejected":
+            setError(
+              "Verification was rejected. Please try again or contact support."
+            )
+            break
+          default:
+            // Check for specific error messages as fallback
+            if (error.message?.includes("Invalid one-time password")) {
+              setError("Invalid verification code. Please try again.")
+            } else if (error.message?.includes("rate limit")) {
+              setError("Too many attempts. Please wait a moment and try again.")
+            } else {
+              console.error("MFA Verification Error:", error)
+              setError(
+                "Unable to verify code. Please try again or contact support."
+              )
+            }
+        }
+      } else {
+        console.error("Unknown MFA Error:", err)
+        setError("An unexpected error occurred. Please try again.")
+      }
+    }
+  }
+
+  const handleSignOut = async () => {
+    await supabase.auth.signOut({ scope: "local" })
+    router.push("/login")
+    router.refresh()
+  }
+
+  return (
+    <div className="mx-4 w-full max-w-md">
+      <form
+        onSubmit={handleSubmit}
+        className="animate-in flex w-full flex-1 flex-col justify-center gap-3"
+      >
+        <div className="mb-6 flex justify-center">
+          <Brand />
+        </div>
+
+        <h2 className="text-center text-2xl font-semibold">
+          Two-Factor Authentication
+        </h2>
+        <p className="text-muted-foreground mb-4 text-center">
+          Enter the 6-digit code from your authenticator app to continue
+        </p>
+
+        <div className="flex justify-center">
+          <div className="flex w-full max-w-[280px] gap-2">
+            <Input
+              type="text"
+              inputMode="numeric"
+              pattern="[0-9]*"
+              maxLength={6}
+              className="text-center text-lg tracking-widest"
+              value={verifyCode}
+              onChange={e => setVerifyCode(e.target.value.replace(/\D/g, ""))}
+              placeholder="000000"
+              disabled={isVerifying}
+              required
+            />
+          </div>
+        </div>
+
+        {error && (
+          <div className="flex items-center justify-center gap-2 text-sm text-red-500">
+            <IconAlertCircle size={16} />
+            <span>{error}</span>
+          </div>
+        )}
+
+        <div className="mt-4 flex flex-col gap-2">
+          <Button
+            type="submit"
+            className="w-full"
+            disabled={verifyCode.length !== 6 || isVerifying}
+          >
+            {isVerifying ? "Verifying..." : "Verify"}
+          </Button>
+          <button
+            type="button"
+            onClick={handleSignOut}
+            className="text-muted-foreground hover:text-foreground text-sm transition-colors"
+          >
+            Cancel and sign out
+          </button>
+          <a
+            href="https://help.hackerai.co"
+            target="_blank"
+            rel="noopener noreferrer"
+            className="text-muted-foreground hover:text-foreground text-center text-sm transition-colors"
+          >
+            Need help? Visit our Help Center
+          </a>
+        </div>
+      </form>
+    </div>
+  )
+}

app/login/verify/page.tsx

@@ -0,0 +1,67 @@
+import { createClient } from "@/lib/supabase/server"
+import { redirect } from "next/navigation"
+import { MFAVerification } from "./mfa-verification"
+
+export default async function VerifyMFA() {
+  const supabase = await createClient()
+
+  const [
+    { data: mfaCheck, error: mfaError },
+    {
+      data: { user }
+    }
+  ] = await Promise.all([supabase.rpc("check_mfa"), supabase.auth.getUser()])
+
+  // Handle MFA check error
+  if (mfaError) {
+    console.error("MFA check failed:", mfaError)
+    throw mfaError
+  }
+
+  // Redirect if user doesn't need MFA or isn't authenticated
+  if (mfaCheck) {
+    return redirect("/")
+  } else if (!user) {
+    return redirect("/login")
+  }
+
+  const verifyMFA = async (code: string): Promise<{ success: boolean }> => {
+    "use server"
+
+    const supabase = await createClient()
+
+    try {
+      const { data: factors, error: factorsError } =
+        await supabase.auth.mfa.listFactors()
+      if (factorsError) throw factorsError
+
+      const totpFactor = factors.totp[0]
+      if (!totpFactor) {
+        throw new Error("No TOTP factors found!")
+      }
+
+      const { data: challenge, error: challengeError } =
+        await supabase.auth.mfa.challenge({
+          factorId: totpFactor.id
+        })
+      if (challengeError) throw challengeError
+
+      const { error: verifyError } = await supabase.auth.mfa.verify({
+        factorId: totpFactor.id,
+        challengeId: challenge.id,
+        code
+      })
+      if (verifyError) throw verifyError
+
+      return { success: true }
+    } catch (error) {
+      throw error
+    }
+  }
+
+  return (
+    <div className="flex min-h-screen items-center justify-center">
+      <MFAVerification onVerify={verifyMFA} />
+    </div>
+  )
+}

app/setup/page.tsx

@@ -7,11 +7,9 @@ import { supabase } from "@/lib/supabase/browser-client"
 import { TablesUpdate } from "@/supabase/types"
 import { useRouter } from "next/navigation"
 import { useContext, useEffect } from "react"
-import { fetchHostedModels } from "@/lib/models/fetch-models"
 
 export default function SetupPage() {
-  const { setProfile, setAvailableHostedModels, refreshTeamMembers } =
-    useContext(PentestGPTContext)
+  const { setProfile, fetchStartingData } = useContext(PentestGPTContext)
   const router = useRouter()
 
   useEffect(() => {
@@ -28,6 +26,10 @@ export default function SetupPage() {
       const profile = await getProfileByUserId(user.id)
       setProfile(profile)
 
+      if (!profile) {
+        throw new Error("Profile not found")
+      }
+
       if (!profile.has_onboarded) {
         const updateProfilePayload: TablesUpdate<"profiles"> = {
           ...profile,
@@ -36,17 +38,12 @@ export default function SetupPage() {
         await updateProfile(profile.id, updateProfilePayload)
       }
 
-      const data = await fetchHostedModels()
-      if (data) {
-        setAvailableHostedModels(data.hostedModels)
-      }
-
-      refreshTeamMembers()
+      await fetchStartingData()
 
       const homeWorkspaceId = await getHomeWorkspaceByUserId(user.id)
       router.push(`/${homeWorkspaceId}/chat`)
     })()
-  }, [router, setProfile, setAvailableHostedModels, refreshTeamMembers])
+  }, [])
 
   return null
 }

components/utility/global-state.tsx

@@ -181,17 +181,7 @@ export const GlobalState: FC<GlobalStateProps> = ({ children }) => {
   }, [])
 
   useEffect(() => {
-    ;(async () => {
-      const profile = await fetchStartingData()
-
-      if (profile) {
-        const hostedModelRes = await fetchHostedModels()
-        if (!hostedModelRes) return
-
-        setEnvKeyMap(hostedModelRes.envKeyMap)
-        setAvailableHostedModels(hostedModelRes.hostedModels)
-      }
-    })()
+    fetchStartingData()
   }, [])
 
   const updateSubscription = useCallback(
@@ -221,6 +211,10 @@ export const GlobalState: FC<GlobalStateProps> = ({ children }) => {
       setUserEmail(userFromAuth.email || "Not available")
 
       const profile = await getProfileByUserId(userFromAuth.id)
+      if (!profile) {
+        return
+      }
+
       setProfile(profile)
 
       if (!profile.has_onboarded) {
@@ -264,7 +258,12 @@ export const GlobalState: FC<GlobalStateProps> = ({ children }) => {
         updateSubscription(subscription)
       }
 
-      return profile
+      const hostedModelRes = await fetchHostedModels()
+
+      if (hostedModelRes) {
+        setEnvKeyMap(hostedModelRes.envKeyMap)
+        setAvailableHostedModels(hostedModelRes.hostedModels)
+      }
     }
   }
 
@@ -448,6 +447,7 @@ export const GlobalState: FC<GlobalStateProps> = ({ children }) => {
         // PROFILE STORE
         profile,
         setProfile,
+        fetchStartingData,
 
         // CONTENT TYPE STORE
         contentType,

components/utility/mfa/mfa-disable-modal.tsx

@@ -0,0 +1,44 @@
+import { FC } from "react"
+import {
+  Dialog,
+  DialogContent,
+  DialogHeader,
+  DialogTitle
+} from "@/components/ui/dialog"
+import { Button } from "@/components/ui/button"
+
+interface MFADisableModalProps {
+  isOpen: boolean
+  onClose: () => void
+  onDisable: () => void
+}
+
+export const MFADisableModal: FC<MFADisableModalProps> = ({
+  isOpen,
+  onClose,
+  onDisable
+}) => {
+  return (
+    <Dialog open={isOpen} onOpenChange={onClose}>
+      <DialogContent>
+        <DialogHeader>
+          <DialogTitle>Confirm Disable 2FA</DialogTitle>
+        </DialogHeader>
+        <div className="space-y-4">
+          <p className="text-muted-foreground">
+            Your identity has been verified. Are you sure you want to disable
+            two-factor authentication? This will make your account less secure.
+          </p>
+          <div className="flex justify-end gap-2">
+            <Button variant="outline" onClick={onClose}>
+              Cancel
+            </Button>
+            <Button variant="destructive" onClick={onDisable}>
+              Disable 2FA
+            </Button>
+          </div>
+        </div>
+      </DialogContent>
+    </Dialog>
+  )
+}