feat(auth): add two-factor authentication support

6 files changed, 276 insertions, 11 deletions

diff --git a/src/app/settings/two-factor-form.tsx b/src/app/settings/two-factor-form.tsx
new file mode 100644
index 0000000..5dccada
--- /dev/null
+++ b/src/app/settings/two-factor-form.tsx
@@ -0,0 +1,202 @@
+"use client"
+
+import { useState } from "react"
+import { useRouter } from "next/navigation"
+import { Shield, Loader2, Copy } from "lucide-react"
+import { toast } from "sonner"
+import Image from "next/image"
+
+import { Button } from "@/components/ui/button"
+import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"
+import { Input } from "@/components/ui/input"
+import { Label } from "@/components/ui/label"
+import {
+  Dialog,
+  DialogContent,
+  DialogDescription,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog"
+
+interface TwoFactorFormProps {
+  twoFactorEnabled: boolean
+}
+
+export function TwoFactorForm({ twoFactorEnabled: initialTwoFactorEnabled }: TwoFactorFormProps) {
+  const router = useRouter()
+  const [twoFactorEnabled, setTwoFactorEnabled] = useState(initialTwoFactorEnabled)
+  const [is2FALoading, setIs2FALoading] = useState(false)
+  const [setupData, setSetupData] = useState<{ secret: string; qrCode: string } | null>(null)
+  const [verificationCode, setVerificationCode] = useState("")
+  const [isDialogOpen, setIsDialogOpen] = useState(false)
+
+  const start2FASetup = async () => {
+    setIs2FALoading(true)
+    try {
+      const res = await fetch("/api/user/2fa", { method: "PUT" })
+      const data = await res.json()
+      if (data.error) throw new Error(data.error)
+      setSetupData(data)
+      setIsDialogOpen(true)
+    } catch {
+      toast.error("Failed to start 2FA setup")
+    } finally {
+      setIs2FALoading(false)
+    }
+  }
+
+  const verifyAndEnable2FA = async () => {
+    if (verificationCode.length !== 6) {
+      toast.error("Please enter a 6-digit code")
+      return
+    }
+
+    setIs2FALoading(true)
+    try {
+      const res = await fetch("/api/user/2fa", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          code: verificationCode,
+          secret: setupData?.secret,
+        }),
+      })
+      const data = await res.json()
+      
+      if (data.error) throw new Error(data.error)
+      
+      setTwoFactorEnabled(true)
+      setIsDialogOpen(false)
+      toast.success("Two-factor authentication enabled")
+      router.refresh()
+    } catch {
+      toast.error("Invalid verification code")
+    } finally {
+      setIs2FALoading(false)
+      setVerificationCode("")
+    }
+  }
+
+  const disable2FA = async () => {
+    if (!confirm("Are you sure you want to disable 2FA? This will make your account less secure.")) return
+
+    setIs2FALoading(true)
+    try {
+      const res = await fetch("/api/user/2fa", { method: "DELETE" })
+      const data = await res.json()
+      
+      if (data.error) throw new Error(data.error)
+      
+      setTwoFactorEnabled(false)
+      toast.success("Two-factor authentication disabled")
+      router.refresh()
+    } catch {
+      toast.error("Failed to disable 2FA")
+    } finally {
+      setIs2FALoading(false)
+    }
+  }
+
+  const copyToClipboard = () => {
+    if (setupData?.secret) {
+      navigator.clipboard.writeText(setupData.secret)
+      toast.success("Secret copied to clipboard")
+    }
+  }
+
+  return (
+    <Card>
+      <CardHeader>
+        <CardTitle className="flex items-center">
+          <Shield className="mr-2 h-5 w-5" />
+          Two-Factor Authentication
+        </CardTitle>
+        <CardDescription>
+          Add an extra layer of security to your account
+        </CardDescription>
+      </CardHeader>
+      <CardContent className="space-y-4">
+        <div className="flex items-center justify-between">
+          <div className="space-y-1">
+            <p className="font-medium">Status: {twoFactorEnabled ? "Enabled" : "Disabled"}</p>
+            <p className="text-sm text-muted-foreground">
+              {twoFactorEnabled 
+                ? "Your account is secured with 2FA." 
+                : "Protect your account by enabling 2FA."}
+            </p>
+          </div>
+          {twoFactorEnabled ? (
+            <Button 
+              variant="destructive" 
+              onClick={disable2FA} 
+              disabled={is2FALoading}
+            >
+              {is2FALoading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+              Disable 2FA
+            </Button>
+          ) : (
+            <Dialog open={isDialogOpen} onOpenChange={setIsDialogOpen}>
+              <DialogTrigger asChild>
+                <Button onClick={start2FASetup} disabled={is2FALoading}>
+                  {is2FALoading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+                  Enable 2FA
+                </Button>
+              </DialogTrigger>
+              <DialogContent className="sm:max-w-md">
+                <DialogHeader>
+                  <DialogTitle>Set up Two-Factor Authentication</DialogTitle>
+                  <DialogDescription>
+                    Scan the QR code with your authenticator app (like Google Authenticator or Authy).
+                  </DialogDescription>
+                </DialogHeader>
+                
+                {setupData && (
+                  <div className="flex flex-col items-center space-y-4 py-4">
+                    <div className="relative w-48 h-48">
+                      <Image 
+                        src={setupData.qrCode} 
+                        alt="2FA QR Code" 
+                        fill
+                        style={{ objectFit: "contain" }}
+                      />
+                    </div>
+                    
+                    <div className="flex items-center space-x-2">
+                      <code className="bg-muted px-2 py-1 rounded text-sm">
+                        {setupData.secret}
+                      </code>
+                      <Button size="icon" variant="ghost" onClick={copyToClipboard}>
+                        <Copy className="h-4 w-4" />
+                      </Button>
+                    </div>
+
+                    <div className="w-full space-y-2">
+                      <Label htmlFor="code">Verification Code</Label>
+                      <Input
+                        id="code"
+                        placeholder="Enter 6-digit code"
+                        value={verificationCode}
+                        onChange={(e) => setVerificationCode(e.target.value.slice(0, 6))}
+                        maxLength={6}
+                      />
+                    </div>
+
+                    <Button 
+                      className="w-full" 
+                      onClick={verifyAndEnable2FA}
+                      disabled={verificationCode.length !== 6 || is2FALoading}
+                    >
+                      {is2FALoading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+                      Verify & Enable
+                    </Button>
+                  </div>
+                )}
+              </DialogContent>
+            </Dialog>
+          )}
+        </div>
+      </CardContent>
+    </Card>
+  )
+}
diff --git a/src/lib/auth-helpers.ts b/src/lib/auth-helpers.ts
new file mode 100644
index 0000000..b2d7488
--- /dev/null
+++ b/src/lib/auth-helpers.ts
@@ -0,0 +1,27 @@
+import { authenticator } from "otplib"
+
+interface TwoFactorCheck {
+  twoFactorEnabled?: boolean
+  twoFactorSecret?: string
+}
+
+export function verifyTwoFactor(
+  user: TwoFactorCheck, 
+  code?: string
+): void {
+  if (user.twoFactorEnabled) {
+    // If the user signed up but hasn't set up 2FA yet (secret is missing),
+    // we can either skip 2FA or treat it as disabled.
+    // Here we treat it as disabled if no secret is present.
+    if (user.twoFactorSecret) {
+      if (!code) {
+        throw new Error("2FA_REQUIRED")
+      }
+
+      const isValid = authenticator.check(code, user.twoFactorSecret)
+      if (!isValid) {
+        throw new Error("Invalid 2FA Code")
+      }
+    }
+  }
+}
diff --git a/src/lib/auth.ts b/src/lib/auth.ts
index ad47d5f..91cf0cb 100644
--- a/src/lib/auth.ts
+++ b/src/lib/auth.ts
@@ -1,10 +1,10 @@
 import { type NextAuthOptions } from "next-auth"
 import CredentialsProvider from "next-auth/providers/credentials"
 import bcrypt from "bcryptjs"
-import { authenticator } from "otplib"
 import dbConnect from "./mongodb"
 import User from "@/model/User"
 import { loginSchema } from "./validation"
+import { verifyTwoFactor } from "./auth-helpers"
 
 export const authOptions: NextAuthOptions = {
   providers: [
@@ -32,16 +32,10 @@ export const authOptions: NextAuthOptions = {
             const isPasswordValid = await bcrypt.compare(password, user.password)
             if (!isPasswordValid) return null
 
-            if (user.twoFactorEnabled) {
-              if (!twoFactorCode) {
-                throw new Error("2FA_REQUIRED")
-              }
-
-              const isValid = authenticator.check(twoFactorCode, user.twoFactorSecret)
-              if (!isValid) {
-                throw new Error("Invalid 2FA Code")
-              }
-            }
+            verifyTwoFactor({
+              twoFactorEnabled: user.twoFactorEnabled,
+              twoFactorSecret: user.twoFactorSecret
+            }, twoFactorCode)
 
             return {
                 id: user._id.toString(),
diff --git a/src/lib/validation.ts b/src/lib/validation.ts
index bc5a440..b580cad 100644
--- a/src/lib/validation.ts
+++ b/src/lib/validation.ts
@@ -35,6 +35,7 @@ export const loginSchema = z.object({
     .string()
     .length(6, 'Code must be 6 digits')
     .optional()
+    .or(z.literal(''))
 })
 
 // Profile update schema (reusing name and email from registerSchema)
diff --git a/src/proxy.ts b/src/proxy.ts
new file mode 100644
index 0000000..796acfa
--- /dev/null
+++ b/src/proxy.ts
@@ -0,0 +1,15 @@
+import { withAuth } from "next-auth/middleware"
+
+export default withAuth({
+  callbacks: {
+    authorized: ({ token }) => !!token,
+  },
+})
+
+export const config = {
+  matcher: [
+    "/settings/:path*",
+    "/dashboard/:path*",
+    "/api/user/:path*"
+  ]
+}
diff --git a/src/services/user.service.ts b/src/services/user.service.ts
new file mode 100644
index 0000000..17ad380
--- /dev/null
+++ b/src/services/user.service.ts
@@ -0,0 +1,26 @@
+import dbConnect from "@/lib/mongodb"
+import User from "@/model/User"
+
+interface UserDocument {
+  name: string
+  email: string
+  profileImage?: { url: string }
+  twoFactorEnabled?: boolean
+}
+
+export async function getUserProfile(email: string) {
+  await dbConnect()
+  
+  const user = await User.findOne({ email }).lean() as unknown as UserDocument
+  
+  if (!user) {
+    return null
+  }
+
+  return {
+    name: user.name,
+    email: user.email,
+    profileImage: user.profileImage,
+    twoFactorEnabled: user.twoFactorEnabled,
+  }
+}