2fa POC (EE)

apps/server/package.json

@@ -71,6 +71,7 @@
     "nestjs-kysely": "^1.2.0",
     "nodemailer": "^7.0.3",
     "openid-client": "^5.7.1",
+    "otpauth": "9.4.0",
     "passport-google-oauth20": "^2.0.0",
     "passport-jwt": "^4.0.1",
     "pg": "^8.16.0",

apps/server/src/core/auth/auth.controller.ts

@@ -22,25 +22,43 @@ import { PasswordResetDto } from './dto/password-reset.dto';
 import { VerifyUserTokenDto } from './dto/verify-user-token.dto';
 import { FastifyReply } from 'fastify';
 import { validateSsoEnforcement } from './auth.util';
+import { TwoFAVerifyDto } from '../../ee/2fa/dto/2fa.dto';
 
 @Controller('auth')
 export class AuthController {
   constructor(
-    private authService: AuthService,
-    private environmentService: EnvironmentService,
+    private readonly authService: AuthService,
+    private readonly environmentService: EnvironmentService,
   ) {}
 
   @HttpCode(HttpStatus.OK)
   @Post('login')
   async login(
-    @AuthWorkspace() workspace: Workspace,
     @Res({ passthrough: true }) res: FastifyReply,
-    @Body() loginInput: LoginDto,
+    @Body() loginDto: LoginDto,
   ) {
-    validateSsoEnforcement(workspace);
+    const { token, workspace } = await this.authService.login(
+      loginDto,
+      this.environmentService.getWorkspaceId(),
+    );
 
-    const authToken = await this.authService.login(loginInput, workspace.id);
-    this.setAuthCookie(res, authToken);
+    this.setAuthCookie(res, token);
+    return workspace;
+  }
+
+  @HttpCode(HttpStatus.OK)
+  @Post('register')
+  async register(
+    @Res({ passthrough: true }) res: FastifyReply,
+    @Body() createUserDto: CreateUserDto,
+  ) {
+    const { token, workspace } = await this.authService.register(
+      createUserDto,
+      this.environmentService.getWorkspaceId(),
+    );
+
+    this.setAuthCookie(res, token);
+    return workspace;
   }
 
   @UseGuards(SetupGuard)
@@ -71,34 +89,36 @@ export class AuthController {
   @HttpCode(HttpStatus.OK)
   @Post('forgot-password')
   async forgotPassword(
-    @Body() forgotPasswordDto: ForgotPasswordDto,
+    @Body() dto: ForgotPasswordDto,
     @AuthWorkspace() workspace: Workspace,
   ) {
-    validateSsoEnforcement(workspace);
-    return this.authService.forgotPassword(forgotPasswordDto, workspace);
+    return this.authService.forgotPassword(dto, workspace);
   }
 
   @HttpCode(HttpStatus.OK)
   @Post('password-reset')
   async passwordReset(
     @Res({ passthrough: true }) res: FastifyReply,
-    @Body() passwordResetDto: PasswordResetDto,
-    @AuthWorkspace() workspace: Workspace,
+    @Body() dto: PasswordResetDto,
   ) {
-    const authToken = await this.authService.passwordReset(
-      passwordResetDto,
-      workspace.id,
+    const { token, workspace } = await this.authService.passwordReset(
+      dto,
+      this.environmentService.getWorkspaceId(),
     );
-    this.setAuthCookie(res, authToken);
+
+    this.setAuthCookie(res, token);
+    return workspace;
   }
 
   @HttpCode(HttpStatus.OK)
-  @Post('verify-token')
-  async verifyResetToken(
-    @Body() verifyUserTokenDto: VerifyUserTokenDto,
-    @AuthWorkspace() workspace: Workspace,
+  @Post('verify-user-token')
+  async verifyUserToken(
+    @Body() dto: VerifyUserTokenDto,
   ) {
-    return this.authService.verifyUserToken(verifyUserTokenDto, workspace.id);
+    return this.authService.verifyUserToken(
+      dto,
+      this.environmentService.getWorkspaceId(),
+    );
   }
 
   @UseGuards(JwtAuthGuard)
@@ -111,11 +131,20 @@ export class AuthController {
     return this.authService.getCollabToken(user.id, workspace.id);
   }
 
-  @UseGuards(JwtAuthGuard)
   @HttpCode(HttpStatus.OK)
-  @Post('logout')
-  async logout(@Res({ passthrough: true }) res: FastifyReply) {
-    res.clearCookie('authToken');
+  @Post('login-with-2fa')
+  async loginWith2FA(
+    @Res({ passthrough: true }) res: FastifyReply,
+    @Body() dto: { userId: string; twoFAToken: string },
+  ) {
+    const { token, workspace } = await this.authService.loginWith2FA(
+      dto.userId,
+      this.environmentService.getWorkspaceId(),
+      dto.twoFAToken,
+    );
+
+    this.setAuthCookie(res, token);
+    return workspace;
   }
 
   setAuthCookie(res: FastifyReply, token: string) {

apps/server/src/core/auth/auth.module.ts

@@ -1,9 +1,9 @@
 import { Module } from '@nestjs/common';
 import { AuthController } from './auth.controller';
 import { AuthService } from './services/auth.service';
+import { SignupService } from './services/signup.service';
 import { JwtStrategy } from './strategies/jwt.strategy';
 import { WorkspaceModule } from '../workspace/workspace.module';
-import { SignupService } from './services/signup.service';
 import { TokenModule } from './token.module';
 
 @Module({

apps/server/src/core/auth/dto/update-user.dto.ts

@@ -0,0 +1,13 @@
+@IsOptional()
+is2faEnabled?: boolean;
+
+@IsOptional()
+@IsString()
+twofaSecret?: string;
+
+@IsOptional()
+@IsString()
+twofaMethod?: string;
+
+@IsOptional()
+twofaBackupCodes?: any; 

apps/server/src/core/auth/services/auth.service.ts

@@ -29,6 +29,7 @@ import { InjectKysely } from 'nestjs-kysely';
 import { executeTx } from '@docmost/db/utils';
 import { VerifyUserTokenDto } from '../dto/verify-user-token.dto';
 import { DomainService } from '../../../integrations/environment/domain.service';
+import { EnvironmentService } from '../../../integrations/environment/environment.service';
 
 @Injectable()
 export class AuthService {
@@ -39,6 +40,7 @@ export class AuthService {
     private userTokenRepo: UserTokenRepo,
     private mailService: MailService,
     private domainService: DomainService,
+    private environmentService: EnvironmentService,
     @InjectKysely() private readonly db: KyselyDB,
   ) {}
 
@@ -61,6 +63,42 @@ export class AuthService {
       throw new UnauthorizedException(errorMessage);
     }
 
+    // Check if 2FA is enabled for this user
+    if (user.is2faEnabled) {
+      // Return a special response indicating 2FA is required
+      // The frontend should then call the 2FA verify endpoint
+      return {
+        requires2FA: true,
+        userId: user.id,
+        message: '2FA code required',
+      };
+    }
+
+    user.lastLoginAt = new Date();
+    await this.userRepo.updateLastLogin(user.id, workspaceId);
+
+    return this.tokenService.generateAccessToken(user);
+  }
+
+  async loginWith2FA(userId: string, workspaceId: string, twoFAToken: string) {
+    const user = await this.userRepo.findById(userId, workspaceId, {
+      includePassword: true,
+    });
+
+    if (!user || user?.deletedAt) {
+      throw new UnauthorizedException('User not found');
+    }
+
+    if (!user.is2faEnabled) {
+      throw new UnauthorizedException('2FA not enabled for this user');
+    }
+
+    // Verify 2FA token
+    const isValid2FA = await this.verify2FAToken(userId, workspaceId, twoFAToken);
+    if (!isValid2FA) {
+      throw new UnauthorizedException('Invalid 2FA code');
+    }
+
     user.lastLoginAt = new Date();
     await this.userRepo.updateLastLogin(user.id, workspaceId);
 
@@ -229,4 +267,55 @@ export class AuthService {
     );
     return { token };
   }
+
+  async enable2FA(userId: string, workspaceId: string, secret: string) {
+    await this.userRepo.updateUser({
+      is2faEnabled: true,
+      twofaSecret: secret,
+      twofaMethod: 'totp',
+    }, userId, workspaceId);
+  }
+
+  async disable2FA(userId: string, workspaceId: string) {
+    await this.userRepo.updateUser({
+      is2faEnabled: false,
+      twofaSecret: null,
+      twofaMethod: null,
+    }, userId, workspaceId);
+  }
+
+  async is2FAAvailable(workspace: Workspace): Promise<boolean> {
+    if (this.environmentService.isCloud()) {
+      return true; // Available in cloud
+    }
+    
+    // For self-hosted, check if it's a valid EE license
+    try {
+      const { LicenseService } = await import('@docmost/ee/licence/license.service');
+      const licenseService = new LicenseService();
+      return licenseService.isValidEELicense(workspace.licenseKey);
+    } catch {
+      return false; // EE module not available
+    }
+  }
+
+  async verify2FAToken(userId: string, workspaceId: string, token: string): Promise<boolean> {
+    try {
+      const { TwoFactorAuthService } = await import('@docmost/ee/2fa/2fa.service');
+      const twoFactorAuthService = new TwoFactorAuthService();
+      
+      const user = await this.userRepo.findById(userId, workspaceId);
+      if (!user?.twofaSecret) {
+        return false;
+      }
+      
+      return twoFactorAuthService.verifyToken(user.twofaSecret, token);
+    } catch {
+      return false; // 2FA not available
+    }
+  }
+
+  async getUserById(userId: string, workspaceId: string) {
+    return this.userRepo.findById(userId, workspaceId);
+  }
 }

apps/server/src/core/user/user.service.ts

@@ -88,6 +88,19 @@ export class UserService {
       user.locale = updateUserDto.locale;
     }
 
+    if (typeof updateUserDto.is2faEnabled !== 'undefined') {
+      user.is2faEnabled = updateUserDto.is2faEnabled;
+    }
+    if (typeof updateUserDto.twofaSecret !== 'undefined') {
+      user.twofaSecret = updateUserDto.twofaSecret;
+    }
+    if (typeof updateUserDto.twofaMethod !== 'undefined') {
+      user.twofaMethod = updateUserDto.twofaMethod;
+    }
+    if (typeof updateUserDto.twofaBackupCodes !== 'undefined') {
+      user.twofaBackupCodes = updateUserDto.twofaBackupCodes;
+    }
+
     delete updateUserDto.confirmPassword;
 
     await this.userRepo.updateUser(updateUserDto, userId, workspace.id);

apps/server/src/database/migrations/20250624T000000-add-2fa-to-users.ts

@@ -0,0 +1,21 @@
+import { Kysely } from 'kysely';
+
+export async function up(db: Kysely<any>): Promise<void> {
+  await db.schema
+    .alterTable('users')
+    .addColumn('is_2fa_enabled', 'boolean', (col) => col.notNull().defaultTo(false))
+    .addColumn('twofa_secret', 'varchar')
+    .addColumn('twofa_method', 'varchar')
+    .addColumn('twofa_backup_codes', 'jsonb')
+    .execute();
+}
+
+export async function down(db: Kysely<any>): Promise<void> {
+  await db.schema
+    .alterTable('users')
+    .dropColumn('is_2fa_enabled')
+    .dropColumn('twofa_secret')
+    .dropColumn('twofa_method')
+    .dropColumn('twofa_backup_codes')
+    .execute();
+} 

apps/server/src/database/repos/user/user.repo.ts

@@ -33,6 +33,10 @@ export class UserRepo {
     'createdAt',
     'updatedAt',
     'deletedAt',
+    'is2faEnabled',
+    'twofaSecret',
+    'twofaMethod',
+    'twofaBackupCodes',
   ];
 
   async findById(

apps/server/src/database/types/db.d.ts

@@ -266,6 +266,10 @@ export interface Users {
   timezone: string | null;
   updatedAt: Generated<Timestamp>;
   workspaceId: string | null;
+  is2faEnabled: boolean;
+  twofaSecret: string | null;
+  twofaMethod: string | null;
+  twofaBackupCodes: Json | null;
 }
 
 export interface UserTokens {

apps/server/src/database/types/entity.types.ts

@@ -33,7 +33,12 @@ export type UpdatableWorkspaceInvitation = Updateable<
 >;
 
 // User
-export type User = Selectable<Users>;
+export type User = Selectable<Users> & {
+  is2faEnabled: boolean;
+  twofaSecret: string | null;
+  twofaMethod: string | null;
+  twofaBackupCodes: any | null;
+};
 export type InsertableUser = Insertable<Users>;
 export type UpdatableUser = Updateable<Omit<Users, 'id'>>;
 

pnpm-lock.yaml

@@ -531,6 +531,9 @@ importers:
       openid-client:
         specifier: ^5.7.1
         version: 5.7.1
+      otpauth:
+        specifier: 9.4.0
+        version: 9.4.0
       passport-google-oauth20:
         specifier: ^2.0.0
         version: 2.0.0
@@ -2840,6 +2843,10 @@ packages:
     cpu: [x64]
     os: [win32]
 
+  '@noble/hashes@1.7.1':
+    resolution: {integrity: sha512-B8XBPsn4vT/KJAGqDzbwztd+6Yte3P4V7iafm24bxgDe/mlRuK6xmWPuCNrKt2vDafZ8MfJLlchDG/vYafQEjQ==}
+    engines: {node: ^14.21.3 || >=16}
+
   '@node-saml/node-saml@5.0.1':
     resolution: {integrity: sha512-YQzFPEC+CnsfO9AFYnwfYZKIzOLx3kITaC1HrjHVLTo6hxcQhc+LgHODOMvW4VCV95Gwrz1MshRUWCPzkDqmnA==}
     engines: {node: '>= 18'}
@@ -7620,6 +7627,9 @@ packages:
     resolution: {integrity: sha512-D2FR03Vir7FIu45XBY20mTb+/ZSWB00sjU9jdQXt83gDrI4Ztz5Fs7/yy74g2N5SVQY4xY1qDr4rNddwYRVX0g==}
     engines: {node: '>=0.10.0'}
 
+  otpauth@9.4.0:
+    resolution: {integrity: sha512-fHIfzIG5RqCkK9cmV8WU+dPQr9/ebR5QOwGZn2JAr1RQF+lmAuLL2YdtdqvmBjNmgJlYk3KZ4a0XokaEhg1Jsw==}
+
   p-limit@2.3.0:
     resolution: {integrity: sha512-//88mFWSJx8lxCzwdAABTJL2MyWB12+eIY7MDL2SqLmAkeKU9qxRvWuSyTjm3FUmpBEMuFfckAIqEaVGUDxb6w==}
     engines: {node: '>=6'}
@@ -12502,6 +12512,8 @@ snapshots:
   '@next/swc-win32-x64-msvc@14.2.10':
     optional: true
 
+  '@noble/hashes@1.7.1': {}
+
   '@node-saml/node-saml@5.0.1':
     dependencies:
       '@types/debug': 4.1.12
@@ -18155,6 +18167,10 @@ snapshots:
 
   os-tmpdir@1.0.2: {}
 
+  otpauth@9.4.0:
+    dependencies:
+      '@noble/hashes': 1.7.1
+
   p-limit@2.3.0:
     dependencies:
       p-try: 2.2.0