feat: add two factor auth (#643)

Add two factor authentication for users who wish to enhance the security of their accounts.

packages/lib/server-only/2fa/disable-2fa.ts

@@ -0,0 +1,48 @@
+import { compare } from 'bcrypt';
+
+import { prisma } from '@documenso/prisma';
+import { User } from '@documenso/prisma/client';
+
+import { ErrorCode } from '../../next-auth/error-codes';
+import { validateTwoFactorAuthentication } from './validate-2fa';
+
+type DisableTwoFactorAuthenticationOptions = {
+  user: User;
+  backupCode: string;
+  password: string;
+};
+
+export const disableTwoFactorAuthentication = async ({
+  backupCode,
+  user,
+  password,
+}: DisableTwoFactorAuthenticationOptions) => {
+  if (!user.password) {
+    throw new Error(ErrorCode.USER_MISSING_PASSWORD);
+  }
+
+  const isCorrectPassword = await compare(password, user.password);
+
+  if (!isCorrectPassword) {
+    throw new Error(ErrorCode.INCORRECT_PASSWORD);
+  }
+
+  const isValid = await validateTwoFactorAuthentication({ backupCode, user });
+
+  if (!isValid) {
+    throw new Error(ErrorCode.INCORRECT_TWO_FACTOR_BACKUP_CODE);
+  }
+
+  await prisma.user.update({
+    where: {
+      id: user.id,
+    },
+    data: {
+      twoFactorEnabled: false,
+      twoFactorBackupCodes: null,
+      twoFactorSecret: null,
+    },
+  });
+
+  return true;
+};

packages/lib/server-only/2fa/enable-2fa.ts

@@ -0,0 +1,47 @@
+import { ErrorCode } from '@documenso/lib/next-auth/error-codes';
+import { prisma } from '@documenso/prisma';
+import { User } from '@documenso/prisma/client';
+
+import { getBackupCodes } from './get-backup-code';
+import { verifyTwoFactorAuthenticationToken } from './verify-2fa-token';
+
+type EnableTwoFactorAuthenticationOptions = {
+  user: User;
+  code: string;
+};
+
+export const enableTwoFactorAuthentication = async ({
+  user,
+  code,
+}: EnableTwoFactorAuthenticationOptions) => {
+  if (user.identityProvider !== 'DOCUMENSO') {
+    throw new Error(ErrorCode.INCORRECT_IDENTITY_PROVIDER);
+  }
+
+  if (user.twoFactorEnabled) {
+    throw new Error(ErrorCode.TWO_FACTOR_ALREADY_ENABLED);
+  }
+
+  if (!user.twoFactorSecret) {
+    throw new Error(ErrorCode.TWO_FACTOR_SETUP_REQUIRED);
+  }
+
+  const isValidToken = await verifyTwoFactorAuthenticationToken({ user, totpCode: code });
+
+  if (!isValidToken) {
+    throw new Error(ErrorCode.INCORRECT_TWO_FACTOR_CODE);
+  }
+
+  const updatedUser = await prisma.user.update({
+    where: {
+      id: user.id,
+    },
+    data: {
+      twoFactorEnabled: true,
+    },
+  });
+
+  const recoveryCodes = getBackupCodes({ user: updatedUser });
+
+  return { recoveryCodes };
+};

packages/lib/server-only/2fa/get-backup-code.ts

@@ -0,0 +1,38 @@
+import { z } from 'zod';
+
+import { User } from '@documenso/prisma/client';
+
+import { DOCUMENSO_ENCRYPTION_KEY } from '../../constants/crypto';
+import { symmetricDecrypt } from '../../universal/crypto';
+
+interface GetBackupCodesOptions {
+  user: User;
+}
+
+const ZBackupCodeSchema = z.array(z.string());
+
+export const getBackupCodes = ({ user }: GetBackupCodesOptions) => {
+  const key = DOCUMENSO_ENCRYPTION_KEY;
+
+  if (!user.twoFactorEnabled) {
+    throw new Error('User has not enabled 2FA');
+  }
+
+  if (!user.twoFactorBackupCodes) {
+    throw new Error('User has no backup codes');
+  }
+
+  const secret = Buffer.from(symmetricDecrypt({ key, data: user.twoFactorBackupCodes })).toString(
+    'utf-8',
+  );
+
+  const data = JSON.parse(secret);
+
+  const result = ZBackupCodeSchema.safeParse(data);
+
+  if (result.success) {
+    return result.data;
+  }
+
+  return null;
+};

packages/lib/server-only/2fa/is-2fa-availble.ts

@@ -0,0 +1,17 @@
+import { User } from '@documenso/prisma/client';
+
+import { DOCUMENSO_ENCRYPTION_KEY } from '../../constants/crypto';
+
+type IsTwoFactorAuthenticationEnabledOptions = {
+  user: User;
+};
+
+export const isTwoFactorAuthenticationEnabled = ({
+  user,
+}: IsTwoFactorAuthenticationEnabledOptions) => {
+  return (
+    user.twoFactorEnabled &&
+    user.identityProvider === 'DOCUMENSO' &&
+    typeof DOCUMENSO_ENCRYPTION_KEY === 'string'
+  );
+};

packages/lib/server-only/2fa/setup-2fa.ts

@@ -0,0 +1,76 @@
+import { base32 } from '@scure/base';
+import { compare } from 'bcrypt';
+import crypto from 'crypto';
+import { createTOTPKeyURI } from 'oslo/otp';
+
+import { ErrorCode } from '@documenso/lib/next-auth/error-codes';
+import { prisma } from '@documenso/prisma';
+import { User } from '@documenso/prisma/client';
+
+import { DOCUMENSO_ENCRYPTION_KEY } from '../../constants/crypto';
+import { symmetricEncrypt } from '../../universal/crypto';
+
+type SetupTwoFactorAuthenticationOptions = {
+  user: User;
+  password: string;
+};
+
+const ISSUER = 'Documenso';
+
+export const setupTwoFactorAuthentication = async ({
+  user,
+  password,
+}: SetupTwoFactorAuthenticationOptions) => {
+  const key = DOCUMENSO_ENCRYPTION_KEY;
+
+  if (!key) {
+    throw new Error(ErrorCode.MISSING_ENCRYPTION_KEY);
+  }
+
+  if (user.identityProvider !== 'DOCUMENSO') {
+    throw new Error(ErrorCode.INCORRECT_IDENTITY_PROVIDER);
+  }
+
+  if (!user.password) {
+    throw new Error(ErrorCode.USER_MISSING_PASSWORD);
+  }
+
+  const isCorrectPassword = await compare(password, user.password);
+
+  if (!isCorrectPassword) {
+    throw new Error(ErrorCode.INCORRECT_PASSWORD);
+  }
+
+  const secret = crypto.randomBytes(10);
+
+  const backupCodes = new Array(10)
+    .fill(null)
+    .map(() => crypto.randomBytes(5).toString('hex'))
+    .map((code) => `${code.slice(0, 5)}-${code.slice(5)}`.toUpperCase());
+
+  const accountName = user.email;
+  const uri = createTOTPKeyURI(ISSUER, accountName, secret);
+  const encodedSecret = base32.encode(secret);
+
+  await prisma.user.update({
+    where: {
+      id: user.id,
+    },
+    data: {
+      twoFactorEnabled: false,
+      twoFactorBackupCodes: symmetricEncrypt({
+        data: JSON.stringify(backupCodes),
+        key: key,
+      }),
+      twoFactorSecret: symmetricEncrypt({
+        data: encodedSecret,
+        key: key,
+      }),
+    },
+  });
+
+  return {
+    secret: encodedSecret,
+    uri,
+  };
+};

packages/lib/server-only/2fa/validate-2fa.ts

@@ -0,0 +1,35 @@
+import { User } from '@documenso/prisma/client';
+
+import { ErrorCode } from '../../next-auth/error-codes';
+import { verifyTwoFactorAuthenticationToken } from './verify-2fa-token';
+import { verifyBackupCode } from './verify-backup-code';
+
+type ValidateTwoFactorAuthenticationOptions = {
+  totpCode?: string;
+  backupCode?: string;
+  user: User;
+};
+
+export const validateTwoFactorAuthentication = async ({
+  backupCode,
+  totpCode,
+  user,
+}: ValidateTwoFactorAuthenticationOptions) => {
+  if (!user.twoFactorEnabled) {
+    throw new Error(ErrorCode.TWO_FACTOR_SETUP_REQUIRED);
+  }
+
+  if (!user.twoFactorSecret) {
+    throw new Error(ErrorCode.TWO_FACTOR_MISSING_SECRET);
+  }
+
+  if (totpCode) {
+    return await verifyTwoFactorAuthenticationToken({ user, totpCode });
+  }
+
+  if (backupCode) {
+    return await verifyBackupCode({ user, backupCode });
+  }
+
+  throw new Error(ErrorCode.TWO_FACTOR_MISSING_CREDENTIALS);
+};

packages/lib/server-only/2fa/verify-2fa-token.ts

@@ -0,0 +1,33 @@
+import { base32 } from '@scure/base';
+import { TOTPController } from 'oslo/otp';
+
+import { User } from '@documenso/prisma/client';
+
+import { DOCUMENSO_ENCRYPTION_KEY } from '../../constants/crypto';
+import { symmetricDecrypt } from '../../universal/crypto';
+
+const totp = new TOTPController();
+
+type VerifyTwoFactorAuthenticationTokenOptions = {
+  user: User;
+  totpCode: string;
+};
+
+export const verifyTwoFactorAuthenticationToken = async ({
+  user,
+  totpCode,
+}: VerifyTwoFactorAuthenticationTokenOptions) => {
+  const key = DOCUMENSO_ENCRYPTION_KEY;
+
+  if (!user.twoFactorSecret) {
+    throw new Error('user missing 2fa secret');
+  }
+
+  const secret = Buffer.from(symmetricDecrypt({ key, data: user.twoFactorSecret })).toString(
+    'utf-8',
+  );
+
+  const isValidToken = await totp.verify(totpCode, base32.decode(secret));
+
+  return isValidToken;
+};

packages/lib/server-only/2fa/verify-backup-code.ts

@@ -0,0 +1,18 @@
+import { User } from '@documenso/prisma/client';
+
+import { getBackupCodes } from './get-backup-code';
+
+type VerifyBackupCodeParams = {
+  user: User;
+  backupCode: string;
+};
+
+export const verifyBackupCode = async ({ user, backupCode }: VerifyBackupCodeParams) => {
+  const userBackupCodes = await getBackupCodes({ user });
+
+  if (!userBackupCodes) {
+    throw new Error('User has no backup codes');
+  }
+
+  return userBackupCodes.includes(backupCode);
+};