From 6d6cc9cf0b208e9869df2fb1b5121a66f806fcbf Mon Sep 17 00:00:00 2001 From: Rostyslav Manko <144740362+RostyslavManko@users.noreply.github.com> Date: Thu, 9 Jan 2025 17:33:51 -0500 Subject: [PATCH] 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 --- app/login/page.tsx | 13 ++ app/login/verify/mfa-verification.tsx | 165 ++++++++++++++++++ app/login/verify/page.tsx | 67 +++++++ app/setup/page.tsx | 17 +- components/utility/global-state.tsx | 24 +-- components/utility/mfa/mfa-disable-modal.tsx | 44 +++++ components/utility/mfa/mfa-enable-modal.tsx | 99 +++++++++++ components/utility/mfa/use-mfa.ts | 150 ++++++++++++++++ .../profile-tabs/data-controls-tab.tsx | 2 +- .../utility/profile-tabs/security-tab.tsx | 164 ++++++++++++++++- .../utility/profile-tabs/subscription-tab.tsx | 4 +- context/context.tsx | 4 +- db/profile.ts | 4 +- db/workspaces.ts | 4 +- middleware.ts | 46 +++-- .../20250106015643_add_mfa_policies.sql | 104 +++++++++++ 16 files changed, 864 insertions(+), 47 deletions(-) create mode 100644 app/login/verify/mfa-verification.tsx create mode 100644 app/login/verify/page.tsx create mode 100644 components/utility/mfa/mfa-disable-modal.tsx create mode 100644 components/utility/mfa/mfa-enable-modal.tsx create mode 100644 components/utility/mfa/use-mfa.ts create mode 100644 supabase/migrations/20250106015643_add_mfa_policies.sql diff --git a/app/login/page.tsx b/app/login/page.tsx index b62f9521..d0ca79a0 100644 --- a/app/login/page.tsx +++ b/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("*") diff --git a/app/login/verify/mfa-verification.tsx b/app/login/verify/mfa-verification.tsx new file mode 100644 index 00000000..ae20a973 --- /dev/null +++ b/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 ( +
+
+
+ +
+ +

+ Two-Factor Authentication +

+

+ Enter the 6-digit code from your authenticator app to continue +

+ +
+
+ setVerifyCode(e.target.value.replace(/\D/g, ""))} + placeholder="000000" + disabled={isVerifying} + required + /> +
+
+ + {error && ( +
+ + {error} +
+ )} + +
+ + + + Need help? Visit our Help Center + +
+
+
+ ) +} diff --git a/app/login/verify/page.tsx b/app/login/verify/page.tsx new file mode 100644 index 00000000..e16301c0 --- /dev/null +++ b/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 ( +
+ +
+ ) +} diff --git a/app/setup/page.tsx b/app/setup/page.tsx index a2cf5e94..d7dc3186 100644 --- a/app/setup/page.tsx +++ b/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 } diff --git a/components/utility/global-state.tsx b/components/utility/global-state.tsx index 5840bc58..a0d5d00b 100644 --- a/components/utility/global-state.tsx +++ b/components/utility/global-state.tsx @@ -181,17 +181,7 @@ export const GlobalState: FC = ({ 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 = ({ 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 = ({ 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 = ({ children }) => { // PROFILE STORE profile, setProfile, + fetchStartingData, // CONTENT TYPE STORE contentType, diff --git a/components/utility/mfa/mfa-disable-modal.tsx b/components/utility/mfa/mfa-disable-modal.tsx new file mode 100644 index 00000000..c3444f71 --- /dev/null +++ b/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 = ({ + isOpen, + onClose, + onDisable +}) => { + return ( + + + + Confirm Disable 2FA + +
+

+ Your identity has been verified. Are you sure you want to disable + two-factor authentication? This will make your account less secure. +

+
+ + +
+
+ + + ) +} diff --git a/components/utility/mfa/mfa-enable-modal.tsx b/components/utility/mfa/mfa-enable-modal.tsx new file mode 100644 index 00000000..e8842fcd --- /dev/null +++ b/components/utility/mfa/mfa-enable-modal.tsx @@ -0,0 +1,99 @@ +import { FC } from "react" +import { DialogPanel, DialogTitle, Dialog } from "@headlessui/react" +import { Button } from "@/components/ui/button" +import { Label } from "@/components/ui/label" + +interface MFAEnableModalProps { + isOpen: boolean + onClose: () => void + onVerify: (code: string) => Promise + qrCode: string + verifyCode: string + setVerifyCode: (code: string) => void + error: string + secret: string + showSecret: boolean + setShowSecret: (show: boolean) => void +} + +export const MFAEnableModal: FC = ({ + isOpen, + onClose, + onVerify, + qrCode, + verifyCode, + setVerifyCode, + error, + secret, + showSecret, + setShowSecret +}) => { + return ( + + + ) +} diff --git a/components/utility/mfa/use-mfa.ts b/components/utility/mfa/use-mfa.ts new file mode 100644 index 00000000..9218582d --- /dev/null +++ b/components/utility/mfa/use-mfa.ts @@ -0,0 +1,150 @@ +import { useState, useEffect } from "react" +import { supabase } from "@/lib/supabase/browser-client" +import { toast } from "sonner" + +export const useMFA = () => { + const [isLoading, setIsLoading] = useState(true) + const [factors, setFactors] = useState([]) + const [factorId, setFactorId] = useState("") + const [qrCode, setQrCode] = useState("") + const [secret, setSecret] = useState("") + const [error, setError] = useState("") + const [isEnrolling, setIsEnrolling] = useState(false) + + const fetchFactors = async () => { + try { + const { data, error } = await supabase.auth.mfa.listFactors() + if (error) throw error + setFactors([...data.totp, ...data.phone]) + } catch (error) { + console.error("Error fetching MFA factors:", error) + toast.error("Failed to load MFA status") + } finally { + setIsLoading(false) + } + } + + const startEnrollment = async () => { + try { + setIsEnrolling(true) + + // First check for and delete any existing unverified TOTP factors + const { data: existingFactors, error: listError } = + await supabase.auth.mfa.listFactors() + if (listError) throw listError + + const unverifiedFactor = existingFactors.all.find( + f => f.factor_type === "totp" && f.status === "unverified" + ) + + if (unverifiedFactor) { + // Delete the unverified factor + const { error: unenrollError } = await supabase.auth.mfa.unenroll({ + factorId: unverifiedFactor.id + }) + if (unenrollError) throw unenrollError + } + + // Create new factor + const { data, error } = await supabase.auth.mfa.enroll({ + factorType: "totp" + }) + if (error) throw error + + setFactorId(data.id) + setQrCode(data.totp.qr_code) + setSecret(data.totp.secret) + } catch (error) { + console.error("Error enrolling MFA:", error) + toast.error("Failed to start MFA enrollment") + throw error + } finally { + setIsEnrolling(false) + } + } + + const verifyMFA = async (code: string) => { + setError("") + try { + const { data: challengeData, error: challengeError } = + await supabase.auth.mfa.challenge({ factorId }) + if (challengeError) throw challengeError + + const { error: verifyError } = await supabase.auth.mfa.verify({ + factorId, + challengeId: challengeData.id, + code + }) + if (verifyError) throw verifyError + + await supabase.auth.refreshSession() + toast.success("MFA enabled successfully") + await fetchFactors() + } catch (error: any) { + setError(error.message) + throw error + } + } + + const verifyBeforeUnenroll = async (code: string) => { + setError("") + try { + if (!factors[0]?.id) throw new Error("No MFA factor found") + + const { data: challengeData, error: challengeError } = + await supabase.auth.mfa.challenge({ factorId: factors[0].id }) + if (challengeError) throw challengeError + + const { error: verifyError } = await supabase.auth.mfa.verify({ + factorId: factors[0].id, + challengeId: challengeData.id, + code + }) + if (verifyError) throw verifyError + + // Only verify, don't unenroll yet + return true + } catch (error: any) { + setError(error.message) + toast.error(error.message) + throw error + } + } + + const unenrollMFA = async () => { + try { + if (!factors[0]?.id) throw new Error("No MFA factor found") + + const { error: unenrollError } = await supabase.auth.mfa.unenroll({ + factorId: factors[0].id + }) + if (unenrollError) throw unenrollError + + await supabase.auth.refreshSession() + toast.success("2FA disabled successfully") + await fetchFactors() + } catch (error: any) { + setError(error.message) + toast.error(error.message) + throw error + } + } + + useEffect(() => { + fetchFactors() + }, []) + + return { + isLoading, + factors, + factorId, + qrCode, + secret, + error, + isEnrolling, + startEnrollment, + verifyMFA, + verifyBeforeUnenroll, + unenrollMFA + } +} diff --git a/components/utility/profile-tabs/data-controls-tab.tsx b/components/utility/profile-tabs/data-controls-tab.tsx index d35df744..ea66d346 100644 --- a/components/utility/profile-tabs/data-controls-tab.tsx +++ b/components/utility/profile-tabs/data-controls-tab.tsx @@ -48,7 +48,7 @@ export const DataControlsTab: FC = ({ onClose={() => setIsSharedChatsPopupOpen(false)} className="relative z-50" > -