add 2fa form to login

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 mfaForm = useForm<z.infer<typeof mfaSchema>>({
-        const { email, password } = values;
+        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,8 +128,13 @@ export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
     }
 
     return (
+        <div>
+            {!mfaRequested && (
                 <Form {...form}>
-            <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-8">
+                    <form
+                        onSubmit={form.handleSubmit(onSubmit)}
+                        className="space-y-8"
+                    >
                         <FormField
                             control={form.control}
                             name="email"
@@ -137,11 +173,81 @@ export default function LoginForm({ redirect, onLogin }: LoginFormProps) {
                                 <AlertDescription>{error}</AlertDescription>
                             </Alert>
                         )}
-                <Button type="submit" className="w-full" loading={loading}>
+                        <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>
     );
 }