Feat: 2fa (#1783)

* preplan otp, better qrcode library

* add 2fa as feature

* add totp generation

* working totp lifecycle

* don't allow disabled user to log in

not a security issue as permission handler would fail anyway

* require 2fa on login

if enabled

* update packages

* fix typo

* remove console.logs

src/server/database/migrations/0000_short_skin.sql

@@ -80,6 +80,8 @@ CREATE TABLE `users_table` (
 	`email` text,
 	`name` text NOT NULL,
 	`role` integer NOT NULL,
+	`totp_key` text,
+	`totp_verified` integer NOT NULL,
 	`enabled` integer NOT NULL,
 	`created_at` text DEFAULT (CURRENT_TIMESTAMP) NOT NULL,
 	`updated_at` text DEFAULT (CURRENT_TIMESTAMP) NOT NULL

src/server/database/migrations/meta/0000_snapshot.json

@@ -1,7 +1,7 @@
 {
   "version": "6",
   "dialect": "sqlite",
-  "id": "8c2af02b-c4bd-4880-a9ad-b38805636208",
+  "id": "91f8ccee-7842-4fd3-bb84-f43e00466b20",
   "prevId": "00000000-0000-0000-0000-000000000000",
   "tables": {
     "clients_table": {
@@ -558,6 +558,20 @@
           "notNull": true,
           "autoincrement": false
         },
+        "totp_key": {
+          "name": "totp_key",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "totp_verified": {
+          "name": "totp_verified",
+          "type": "integer",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
         "enabled": {
           "name": "enabled",
           "type": "integer",

src/server/database/migrations/meta/0001_snapshot.json

@@ -1,6 +1,6 @@
 {
-  "id": "a61263b1-9af1-4d2e-99e9-80d08127b545",
-  "prevId": "8c2af02b-c4bd-4880-a9ad-b38805636208",
+  "id": "0224c6a5-3456-402d-a40d-0821637015da",
+  "prevId": "91f8ccee-7842-4fd3-bb84-f43e00466b20",
   "version": "6",
   "dialect": "sqlite",
   "tables": {
@@ -558,6 +558,20 @@
           "notNull": true,
           "autoincrement": false
         },
+        "totp_key": {
+          "name": "totp_key",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false,
+          "autoincrement": false
+        },
+        "totp_verified": {
+          "name": "totp_verified",
+          "type": "integer",
+          "primaryKey": false,
+          "notNull": true,
+          "autoincrement": false
+        },
         "enabled": {
           "name": "enabled",
           "type": "integer",

src/server/database/migrations/meta/_journal.json

@@ -5,14 +5,14 @@
     {
       "idx": 0,
       "version": "6",
-      "when": 1741355094140,
+      "when": 1743490907551,
       "tag": "0000_short_skin",
       "breakpoints": true
     },
     {
       "idx": 1,
       "version": "6",
-      "when": 1741355098159,
+      "when": 1743490912488,
       "tag": "0001_classy_the_stranger",
       "breakpoints": true
     }

src/server/database/repositories/user/schema.ts

@@ -10,6 +10,8 @@ export const user = sqliteTable('users_table', {
   email: text(),
   name: text().notNull(),
   role: int().$type<Role>().notNull(),
+  totpKey: text('totp_key'),
+  totpVerified: int('totp_verified', { mode: 'boolean' }).notNull(),
   enabled: int({ mode: 'boolean' }).notNull(),
   createdAt: text('created_at')
     .notNull()

src/server/database/repositories/user/service.ts

@@ -1,7 +1,24 @@
 import { eq, sql } from 'drizzle-orm';
+import { TOTP } from 'otpauth';
 import { user } from './schema';
+import type { UserType } from './types';
 import type { DBType } from '#db/sqlite';
 
+type LoginResult =
+  | {
+      success: true;
+      user: UserType;
+    }
+  | {
+      success: false;
+      error:
+        | 'INCORRECT_CREDENTIALS'
+        | 'TOTP_REQUIRED'
+        | 'USER_DISABLED'
+        | 'INVALID_TOTP_CODE'
+        | 'UNEXPECTED_ERROR';
+    };
+
 function createPreparedStatement(db: DBType) {
   return {
     findAll: db.query.user.findMany().prepare(),
@@ -21,6 +38,14 @@ function createPreparedStatement(db: DBType) {
       })
       .where(eq(user.id, sql.placeholder('id')))
       .prepare(),
+    updateKey: db
+      .update(user)
+      .set({
+        totpKey: sql.placeholder('key') as never as string,
+        totpVerified: false,
+      })
+      .where(eq(user.id, sql.placeholder('id')))
+      .prepare(),
   };
 }
 
@@ -67,6 +92,7 @@ export class UserService {
         email: null,
         name: 'Administrator',
         role: userCount === 0 ? roles.ADMIN : roles.CLIENT,
+        totpVerified: false,
         enabled: true,
       });
     });
@@ -105,4 +131,121 @@ export class UserService {
         .execute();
     });
   }
+
+  updateTotpKey(id: ID, key: string | null) {
+    return this.#statements.updateKey.execute({ id, key });
+  }
+
+  login(username: string, password: string, code: string | undefined) {
+    return this.#db.transaction(async (tx): Promise<LoginResult> => {
+      const txUser = await tx.query.user
+        .findFirst({ where: eq(user.username, username) })
+        .execute();
+
+      if (!txUser) {
+        return { success: false, error: 'INCORRECT_CREDENTIALS' };
+      }
+
+      const passwordValid = await isPasswordValid(password, txUser.password);
+
+      if (!passwordValid) {
+        return { success: false, error: 'INCORRECT_CREDENTIALS' };
+      }
+
+      if (txUser.totpVerified) {
+        if (!code) {
+          return { success: false, error: 'TOTP_REQUIRED' };
+        } else {
+          if (!txUser.totpKey) {
+            return { success: false, error: 'UNEXPECTED_ERROR' };
+          }
+
+          const totp = new TOTP({
+            issuer: 'wg-easy',
+            label: txUser.username,
+            algorithm: 'SHA1',
+            digits: 6,
+            period: 30,
+            secret: txUser.totpKey,
+          });
+
+          const valid = totp.validate({ token: code, window: 1 });
+
+          if (valid === null) {
+            return { success: false, error: 'INVALID_TOTP_CODE' };
+          }
+        }
+      }
+
+      if (!txUser.enabled) {
+        return { success: false, error: 'USER_DISABLED' };
+      }
+
+      return { success: true, user: txUser };
+    });
+  }
+
+  verifyTotp(id: ID, code: string) {
+    return this.#db.transaction(async (tx) => {
+      const txUser = await tx.query.user
+        .findFirst({ where: eq(user.id, id) })
+        .execute();
+
+      if (!txUser) {
+        throw new Error('User not found');
+      }
+
+      if (!txUser.totpKey) {
+        throw new Error('TOTP key is not set');
+      }
+
+      const totp = new TOTP({
+        issuer: 'wg-easy',
+        label: txUser.username,
+        algorithm: 'SHA1',
+        digits: 6,
+        period: 30,
+        secret: txUser.totpKey,
+      });
+
+      const valid = totp.validate({ token: code, window: 1 });
+
+      if (valid === null) {
+        throw new Error('Invalid TOTP code');
+      }
+
+      await tx
+        .update(user)
+        .set({ totpVerified: true })
+        .where(eq(user.id, id))
+        .execute();
+    });
+  }
+
+  deleteTotpKey(id: ID, currentPassword: string) {
+    return this.#db.transaction(async (tx) => {
+      const txUser = await tx.query.user
+        .findFirst({ where: eq(user.id, id) })
+        .execute();
+
+      if (!txUser) {
+        throw new Error('User not found');
+      }
+
+      const passwordValid = await isPasswordValid(
+        currentPassword,
+        txUser.password
+      );
+
+      if (!passwordValid) {
+        throw new Error('Invalid password');
+      }
+
+      await tx
+        .update(user)
+        .set({ totpKey: null, totpVerified: false })
+        .where(eq(user.id, id))
+        .execute();
+    });
+  }
 }

src/server/database/repositories/user/types.ts

@@ -16,10 +16,16 @@ const password = z
 
 const remember = z.boolean({ message: t('zod.user.remember') });
 
+const totpCode = z
+  .string({ message: t('zod.user.totpCode') })
+  .min(6, t('zod.user.totpCode'))
+  .pipe(safeStringRefine);
+
 export const UserLoginSchema = z.object({
   username: username,
   password: password,
   remember: remember,
+  totpCode: totpCode.optional(),
 });
 
 export const UserSetupSchema = z
@@ -58,3 +64,17 @@ export const UserUpdatePasswordSchema = z
   .refine((val) => val.newPassword === val.confirmPassword, {
     message: t('zod.user.passwordMatch'),
   });
+
+export const UserUpdateTotpSchema = z.union([
+  z.object({
+    type: z.literal('setup'),
+  }),
+  z.object({
+    type: z.literal('create'),
+    code: totpCode,
+  }),
+  z.object({
+    type: z.literal('delete'),
+    currentPassword: password,
+  }),
+]);