3 files changed, 33 insertions, 11 deletions
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)