added two factor to login endpoint

server/auth/index.ts

@@ -10,6 +10,8 @@ export const lucia = new Lucia(adapter, {
     getUserAttributes: (attributes) => {
         return {
             email: attributes.email,
+            twoFactorEnabled: attributes.twoFactorEnabled,
+            twoFactorSecret: attributes.twoFactorSecret,
         };
     },
     // getSessionAttributes: (attributes) => {
@@ -42,6 +44,8 @@ declare module "lucia" {
 interface DatabaseUserAttributes {
     email: string;
     passwordHash: string;
+    twoFactorEnabled: boolean;
+    twoFactorSecret: string | null;
 }
 
 interface DatabaseSessionAttributes {

server/db/schema.ts

@@ -71,6 +71,10 @@ export const users = sqliteTable("user", {
     id: text("id").primaryKey(), // has to be id not userId for lucia
     email: text("email").notNull().unique(),
     passwordHash: text("passwordHash").notNull(),
+    twoFactorEnabled: integer("twoFactorEnabled", { mode: "boolean" })
+        .notNull()
+        .default(false),
+    twoFactorSecret: text("twoFactorSecret"),
 });
 
 // Sessions table

server/routers/auth/login.ts

@@ -2,7 +2,6 @@ import { verify } from "@node-rs/argon2";
 import lucia from "@server/auth";
 import db from "@server/db";
 import { users } from "@server/db/schema";
-import logger from "@server/logger";
 import HttpCode from "@server/types/HttpCode";
 import response from "@server/utils/response";
 import { eq } from "drizzle-orm";
@@ -10,12 +9,21 @@ import { NextFunction, Request, Response } from "express";
 import createHttpError from "http-errors";
 import { z } from "zod";
 import { fromError } from "zod-validation-error";
+import { decodeHex } from "oslo/encoding";
+import { TOTPController } from "oslo/otp";
 
 export const loginBodySchema = z.object({
     email: z.string().email(),
     password: z.string(),
+    code: z.string().optional(),
 });
 
+export type LoginBody = z.infer<typeof loginBodySchema>;
+
+export type LoginResponse = {
+    codeRequested?: boolean;
+};
+
 export async function login(
     req: Request,
     res: Response,
@@ -32,7 +40,7 @@ export async function login(
         );
     }
 
-    const { email, password } = parsedBody.data;
+    const { email, password, code } = parsedBody.data;
 
     const sessionId = req.cookies[lucia.sessionCookieName];
     const { session: existingSession } = await lucia.validateSession(sessionId);
@@ -70,7 +78,7 @@ export async function login(
         parallelism: 1,
     });
     if (!validPassword) {
-        await new Promise((resolve) => setTimeout(resolve, 500)); // delay to prevent brute force attacks
+        await new Promise((resolve) => setTimeout(resolve, 250)); // delay to prevent brute force attacks
         return next(
             createHttpError(
                 HttpCode.BAD_REQUEST,
@@ -79,6 +87,44 @@ export async function login(
         );
     }
 
+    if (existingUser.twoFactorEnabled) {
+        if (!code) {
+            return res.status(HttpCode.ACCEPTED).send(
+                response<{ codeRequested: boolean }>({
+                    data: { codeRequested: true },
+                    success: true,
+                    error: false,
+                    message: "Two-factor authentication required",
+                    status: HttpCode.ACCEPTED,
+                }),
+            );
+        }
+
+        if (!existingUser.twoFactorSecret) {
+            return next(
+                createHttpError(
+                    HttpCode.INTERNAL_SERVER_ERROR,
+                    "Failed to authenticate user",
+                ),
+            );
+        }
+
+        const validOTP = await new TOTPController().verify(
+            code,
+            decodeHex(existingUser.twoFactorSecret),
+        );
+
+        if (!validOTP) {
+            await new Promise((resolve) => setTimeout(resolve, 250)); // delay to prevent brute force attacks
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    "The two-factor code you entered is incorrect",
+                ),
+            );
+        }
+    }
+
     const session = await lucia.createSession(existingUser.id, {});
     res.appendHeader(
         "Set-Cookie",