From 9c37036a39b0179ee85a6ae0955e0dbc824f34ec Mon Sep 17 00:00:00 2001
From: Milo Schwartz <mschwartz10612@gmail.com>
Date: Sun, 22 Dec 2024 14:38:17 -0500
Subject: [PATCH] add 2fa form to login

---
 src/components/LoginForm.tsx | 218 ++++++++++++++++++++++++++---------
 1 file changed, 162 insertions(+), 56 deletions(-)

diff --git a/src/components/LoginForm.tsx b/src/components/LoginForm.tsx
index 9da6f15..8d8e2e6 100644
--- a/src/components/LoginForm.tsx
+++ b/src/components/LoginForm.tsx
@@ -12,14 +12,14 @@ import {
     FormField,
     FormItem,
     FormLabel,
-    FormMessage,
+    FormMessage
 } from "@/components/ui/form";
 import {
     Card,
     CardContent,
     CardDescription,
     CardHeader,
-    CardTitle,
+    CardTitle
 } from "@/components/ui/card";
 import { Alert, AlertDescription } from "@/components/ui/alert";
 import { LoginResponse } from "@server/routers/auth";
@@ -29,6 +29,12 @@ import { formatAxiosError } from "@app/lib/utils";
 import { LockIcon } from "lucide-react";
 import { createApiClient } from "@app/api";
 import { useEnvContext } from "@app/hooks/useEnvContext";
+import {
+    InputOTP,
+    InputOTPGroup,
+    InputOTPSeparator,
+    InputOTPSlot
+} from "./ui/input-otp";
 
 type LoginFormProps = {
     redirect?: string;
@@ -39,7 +45,11 @@ const formSchema = z.object({
     email: z.string().email({ message: "Invalid email address" }),
     password: z
         .string()
-        .min(8, { message: "Password must be at least 8 characters" }),
+        .min(8, { message: "Password must be at least 8 characters" })
+});
+
+const mfaSchema = z.object({
+    code: z.string().length(6, { message: "Invalid code" })
 });
 
 export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
@@ -50,17 +60,26 @@ export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
     const [error, setError] = useState<string | null>(null);
     const [loading, setLoading] = useState(false);
 
+    const [mfaRequested, setMfaRequested] = useState(false);
+
     const form = useForm<z.infer<typeof formSchema>>({
         resolver: zodResolver(formSchema),
         defaultValues: {
             email: "",
-            password: "",
-        },
+            password: ""
+        }
     });
 
-    async function onSubmit(values: z.infer<typeof formSchema>) {
-        const { email, password } = values;
+    const mfaForm = useForm<z.infer<typeof mfaSchema>>({
+        resolver: zodResolver(mfaSchema),
+        defaultValues: {
+            code: ""
+        }
+    });
 
+    async function onSubmit(values: any) {
+        const { email, password } = form.getValues();
+        const { code } = mfaForm.getValues()
 
         setLoading(true);
 
@@ -68,18 +87,30 @@ export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
             .post<AxiosResponse<LoginResponse>>("/auth/login", {
                 email,
                 password,
+                code
             })
             .catch((e) => {
                 console.error(e);
                 setError(
-                    formatAxiosError(e, "An error occurred while logging in"),
+                    formatAxiosError(e, "An error occurred while logging in")
                 );
             });
 
-        if (res && res.status === 200) {
+        if (res) {
             setError(null);
 
-            if (res.data?.data?.emailVerificationRequired) {
+            const data = res.data.data;
+
+            console.log(data);
+
+            if (data?.codeRequested) {
+                setMfaRequested(true);
+                setLoading(false);
+                mfaForm.reset();
+                return;
+            }
+
+            if (data?.emailVerificationRequired) {
                 if (redirect) {
                     router.push(`/auth/verify-email?redirect=${redirect}`);
                 } else {
@@ -97,51 +128,126 @@ export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
     }
 
     return (
-        <Form {...form}>
-            <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-8">
-                <FormField
-                    control={form.control}
-                    name="email"
-                    render={({ field }) => (
-                        <FormItem>
-                            <FormLabel>Email</FormLabel>
-                            <FormControl>
-                                <Input
-                                    placeholder="Enter your email"
-                                    {...field}
-                                />
-                            </FormControl>
-                            <FormMessage />
-                        </FormItem>
-                    )}
-                />
-                <FormField
-                    control={form.control}
-                    name="password"
-                    render={({ field }) => (
-                        <FormItem>
-                            <FormLabel>Password</FormLabel>
-                            <FormControl>
-                                <Input
-                                    type="password"
-                                    placeholder="Enter your password"
-                                    {...field}
-                                />
-                            </FormControl>
-                            <FormMessage />
-                        </FormItem>
-                    )}
-                />
-                {error && (
-                    <Alert variant="destructive">
-                        <AlertDescription>{error}</AlertDescription>
-                    </Alert>
-                )}
-                <Button type="submit" className="w-full" loading={loading}>
-                    <LockIcon className="w-4 h-4 mr-2" />
-                    Login
-                </Button>
-            </form>
-        </Form>
+        <div>
+            {!mfaRequested && (
+                <Form {...form}>
+                    <form
+                        onSubmit={form.handleSubmit(onSubmit)}
+                        className="space-y-8"
+                    >
+                        <FormField
+                            control={form.control}
+                            name="email"
+                            render={({ field }) => (
+                                <FormItem>
+                                    <FormLabel>Email</FormLabel>
+                                    <FormControl>
+                                        <Input
+                                            placeholder="Enter your email"
+                                            {...field}
+                                        />
+                                    </FormControl>
+                                    <FormMessage />
+                                </FormItem>
+                            )}
+                        />
+                        <FormField
+                            control={form.control}
+                            name="password"
+                            render={({ field }) => (
+                                <FormItem>
+                                    <FormLabel>Password</FormLabel>
+                                    <FormControl>
+                                        <Input
+                                            type="password"
+                                            placeholder="Enter your password"
+                                            {...field}
+                                        />
+                                    </FormControl>
+                                    <FormMessage />
+                                </FormItem>
+                            )}
+                        />
+                        {error && (
+                            <Alert variant="destructive">
+                                <AlertDescription>{error}</AlertDescription>
+                            </Alert>
+                        )}
+                        <Button
+                            type="submit"
+                            className="w-full"
+                            loading={loading}
+                        >
+                            <LockIcon className="w-4 h-4 mr-2" />
+                            Login
+                        </Button>
+                    </form>
+                </Form>
+            )}
+
+            {mfaRequested && (
+                <Form {...mfaForm}>
+                    <form
+                        onSubmit={mfaForm.handleSubmit(onSubmit)}
+                        className="space-y-8"
+                    >
+                        <FormField
+                            control={mfaForm.control}
+                            name="code"
+                            render={({ field }) => (
+                                <FormItem>
+                                    <FormLabel>Authenticator Code</FormLabel>
+                                    <FormControl>
+                                        <div className="flex justify-center">
+                                            <InputOTP maxLength={6} {...field}>
+                                                <InputOTPGroup>
+                                                    <InputOTPSlot index={0} />
+                                                    <InputOTPSlot index={1} />
+                                                    <InputOTPSlot index={2} />
+                                                </InputOTPGroup>
+                                                <InputOTPSeparator />
+                                                <InputOTPGroup>
+                                                    <InputOTPSlot index={3} />
+                                                    <InputOTPSlot index={4} />
+                                                    <InputOTPSlot index={5} />
+                                                </InputOTPGroup>
+                                            </InputOTP>
+                                        </div>
+                                    </FormControl>
+                                    <FormMessage />
+                                </FormItem>
+                            )}
+                        />
+                        {error && (
+                            <Alert variant="destructive">
+                                <AlertDescription>{error}</AlertDescription>
+                            </Alert>
+                        )}
+
+                        <div className="space-y-4">
+                            <Button
+                                type="submit"
+                                className="w-full"
+                                loading={loading}
+                            >
+                                <LockIcon className="w-4 h-4 mr-2" />
+                                Submit Code
+                            </Button>
+                            <Button
+                                type="button"
+                                className="w-full"
+                                variant="outline"
+                                onClick={() => {
+                                    setMfaRequested(false);
+                                    mfaForm.reset();
+                                }}
+                            >
+                                Back to Login
+                            </Button>
+                        </div>
+                    </form>
+                </Form>
+            )}
+        </div>
     );
 }