fix: migrate 2fa to custom auth

packages/auth/server/index.ts

@@ -13,6 +13,7 @@ import { oauthRoute } from './routes/oauth';
 import { passkeyRoute } from './routes/passkey';
 import { sessionRoute } from './routes/session';
 import { signOutRoute } from './routes/sign-out';
+import { twoFactorRoute } from './routes/two-factor';
 import type { HonoAuthContext } from './types/context';
 
 // Note: You must chain routes for Hono RPC client to work.
@@ -46,7 +47,8 @@ export const auth = new Hono<HonoAuthContext>()
   .route('/callback', callbackRoute)
   .route('/oauth', oauthRoute)
   .route('/email-password', emailPasswordRoute)
-  .route('/passkey', passkeyRoute);
+  .route('/passkey', passkeyRoute)
+  .route('/two-factor', twoFactorRoute);
 
 /**
  * Handle errors.

packages/auth/server/lib/utils/get-session.ts

@@ -8,7 +8,6 @@ import { validateSessionToken } from '../session/session';
 import { getSessionCookie } from '../session/session-cookies';
 
 export const getSession = async (c: Context | Request): Promise<SessionValidationResult> => {
-  // Todo: Make better
   const sessionId = await getSessionCookie(mapRequestToContextForCookie(c));
 
   if (!sessionId) {
@@ -29,7 +28,6 @@ export const getRequiredSession = async (c: Context | Request) => {
     return { session, user };
   }
 
-  // Todo: Test if throwing errors work
   if (c instanceof Request) {
     throw new Error('Unauthorized');
   }
@@ -37,9 +35,11 @@ export const getRequiredSession = async (c: Context | Request) => {
   throw new AppError(AuthenticationErrorCode.Unauthorized);
 };
 
+/**
+ * Todo: Rethink, this is pretty sketchy.
+ */
 const mapRequestToContextForCookie = (c: Context | Request) => {
   if (c instanceof Request) {
-    // c.req.raw.headers.
     const partialContext = {
       req: {
         raw: c,

packages/auth/server/routes/two-factor.ts

@@ -0,0 +1,151 @@
+import { sValidator } from '@hono/standard-validator';
+import { Hono } from 'hono';
+
+import { AppError } from '@documenso/lib/errors/app-error';
+import { disableTwoFactorAuthentication } from '@documenso/lib/server-only/2fa/disable-2fa';
+import { enableTwoFactorAuthentication } from '@documenso/lib/server-only/2fa/enable-2fa';
+import { setupTwoFactorAuthentication } from '@documenso/lib/server-only/2fa/setup-2fa';
+import { viewBackupCodes } from '@documenso/lib/server-only/2fa/view-backup-codes';
+import { prisma } from '@documenso/prisma';
+
+import { AuthenticationErrorCode } from '../lib/errors/error-codes';
+import { getRequiredSession } from '../lib/utils/get-session';
+import type { HonoAuthContext } from '../types/context';
+import {
+  ZDisableTwoFactorRequestSchema,
+  ZEnableTwoFactorRequestSchema,
+  ZViewTwoFactorRecoveryCodesRequestSchema,
+} from './two-factor.types';
+
+export const twoFactorRoute = new Hono<HonoAuthContext>()
+  /**
+   * Setup two factor authentication.
+   */
+  .post('/setup', async (c) => {
+    const { user } = await getRequiredSession(c);
+
+    const result = await setupTwoFactorAuthentication({
+      user,
+    });
+
+    return c.json({
+      success: true,
+      secret: result.secret,
+      uri: result.uri,
+    });
+  })
+
+  /**
+   * Enable two factor authentication.
+   */
+  .post('/enable', sValidator('json', ZEnableTwoFactorRequestSchema), async (c) => {
+    const requestMetadata = c.get('requestMetadata');
+
+    const { user: sessionUser } = await getRequiredSession(c);
+
+    const user = await prisma.user.findFirst({
+      where: {
+        id: sessionUser.id,
+      },
+      select: {
+        id: true,
+        email: true,
+        twoFactorEnabled: true,
+        twoFactorSecret: true,
+      },
+    });
+
+    if (!user) {
+      throw new AppError(AuthenticationErrorCode.InvalidRequest);
+    }
+
+    const { code } = c.req.valid('json');
+
+    const result = await enableTwoFactorAuthentication({
+      user,
+      code,
+      requestMetadata,
+    });
+
+    return c.json({
+      success: true,
+      recoveryCodes: result.recoveryCodes,
+    });
+  })
+
+  /**
+   * Disable two factor authentication.
+   */
+  .post('/disable', sValidator('json', ZDisableTwoFactorRequestSchema), async (c) => {
+    const requestMetadata = c.get('requestMetadata');
+
+    const { user: sessionUser } = await getRequiredSession(c);
+
+    const user = await prisma.user.findFirst({
+      where: {
+        id: sessionUser.id,
+      },
+      select: {
+        id: true,
+        email: true,
+        twoFactorEnabled: true,
+        twoFactorSecret: true,
+        twoFactorBackupCodes: true,
+      },
+    });
+
+    if (!user) {
+      throw new AppError(AuthenticationErrorCode.InvalidRequest);
+    }
+
+    const { totpCode, backupCode } = c.req.valid('json');
+
+    await disableTwoFactorAuthentication({
+      user,
+      totpCode,
+      backupCode,
+      requestMetadata,
+    });
+
+    return c.text('OK', 201);
+  })
+
+  /**
+   * View backup codes.
+   */
+  .post(
+    '/view-recovery-codes',
+    sValidator('json', ZViewTwoFactorRecoveryCodesRequestSchema),
+    async (c) => {
+      const { user: sessionUser } = await getRequiredSession(c);
+
+      const user = await prisma.user.findFirst({
+        where: {
+          id: sessionUser.id,
+        },
+        select: {
+          id: true,
+          email: true,
+          twoFactorEnabled: true,
+          twoFactorSecret: true,
+          twoFactorBackupCodes: true,
+        },
+      });
+
+      if (!user) {
+        throw new AppError(AuthenticationErrorCode.InvalidRequest);
+      }
+
+      const { token } = c.req.valid('json');
+
+      const backupCodes = await viewBackupCodes({
+        user,
+        token,
+      });
+
+      return c.json({
+        success: true,
+        backupCodes,
+      });
+    },
+  );

packages/auth/server/routes/two-factor.types.ts

@@ -0,0 +1,22 @@
+import { z } from 'zod';
+
+export const ZEnableTwoFactorRequestSchema = z.object({
+  code: z.string().min(6).max(6),
+});
+
+export type TEnableTwoFactorRequestSchema = z.infer<typeof ZEnableTwoFactorRequestSchema>;
+
+export const ZDisableTwoFactorRequestSchema = z.object({
+  totpCode: z.string().trim().optional(),
+  backupCode: z.string().trim().optional(),
+});
+
+export type TDisableTwoFactorRequestSchema = z.infer<typeof ZDisableTwoFactorRequestSchema>;
+
+export const ZViewTwoFactorRecoveryCodesRequestSchema = z.object({
+  token: z.string().trim().min(1),
+});
+
+export type TViewTwoFactorRecoveryCodesRequestSchema = z.infer<
+  typeof ZViewTwoFactorRecoveryCodesRequestSchema
+>;