feat(EE): MFA implementation (#1381)

* feat(EE): MFA implementation for enterprise edition
- Add TOTP-based two-factor authentication
- Add backup codes support
- Add MFA enforcement at workspace level
- Add MFA setup and challenge UI pages
- Support MFA for login and password reset flows
- Add MFA validation for secure pages
* fix types
* remove unused object
* sync
* remove unused type
* sync
* refactor: rename MFA enabled field to is_enabled
* sync

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

@@ -6,6 +6,7 @@ import {
   Post,
   Res,
   UseGuards,
+  Logger,
 } from '@nestjs/common';
 import { LoginDto } from './dto/login.dto';
 import { AuthService } from './services/auth.service';
@@ -22,12 +23,16 @@ 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 { ModuleRef } from '@nestjs/core';
 
 @Controller('auth')
 export class AuthController {
+  private readonly logger = new Logger(AuthController.name);
+
   constructor(
     private authService: AuthService,
     private environmentService: EnvironmentService,
+    private moduleRef: ModuleRef,
   ) {}
 
   @HttpCode(HttpStatus.OK)
@@ -39,6 +44,45 @@ export class AuthController {
   ) {
     validateSsoEnforcement(workspace);
 
+    let MfaModule: any;
+    let isMfaModuleReady = false;
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      MfaModule = require('./../../ee/mfa/services/mfa.service');
+      isMfaModuleReady = true;
+    } catch (err) {
+      this.logger.debug(
+        'MFA module requested but EE module not bundled in this build',
+      );
+      isMfaModuleReady = false;
+    }
+    if (isMfaModuleReady) {
+      const mfaService = this.moduleRef.get(MfaModule.MfaService, {
+        strict: false,
+      });
+
+      const mfaResult = await mfaService.checkMfaRequirements(
+        loginInput,
+        workspace,
+        res,
+      );
+
+      if (mfaResult) {
+        // If user has MFA enabled OR workspace enforces MFA, require MFA verification
+        if (mfaResult.userHasMfa || mfaResult.requiresMfaSetup) {
+          return {
+            userHasMfa: mfaResult.userHasMfa,
+            requiresMfaSetup: mfaResult.requiresMfaSetup,
+            isMfaEnforced: mfaResult.isMfaEnforced,
+          };
+        } else if (mfaResult.authToken) {
+          // User doesn't have MFA and workspace doesn't require it
+          this.setAuthCookie(res, mfaResult.authToken);
+          return;
+        }
+      }
+    }
+
     const authToken = await this.authService.login(loginInput, workspace.id);
     this.setAuthCookie(res, authToken);
   }
@@ -85,11 +129,22 @@ export class AuthController {
     @Body() passwordResetDto: PasswordResetDto,
     @AuthWorkspace() workspace: Workspace,
   ) {
-    const authToken = await this.authService.passwordReset(
+    const result = await this.authService.passwordReset(
       passwordResetDto,
-      workspace.id,
+      workspace,
     );
-    this.setAuthCookie(res, authToken);
+
+    if (result.requiresLogin) {
+      return {
+        requiresLogin: true,
+      };
+    }
+
+    // Set auth cookie if no MFA is required
+    this.setAuthCookie(res, result.authToken);
+    return {
+      requiresLogin: false,
+    };
   }
 
   @HttpCode(HttpStatus.OK)

apps/server/src/core/auth/dto/jwt-payload.ts

@@ -3,6 +3,7 @@ export enum JwtType {
   COLLAB = 'collab',
   EXCHANGE = 'exchange',
   ATTACHMENT = 'attachment',
+  MFA_TOKEN = 'mfa_token',
 }
 export type JwtPayload = {
   sub: string;
@@ -30,3 +31,8 @@ export type JwtAttachmentPayload = {
   type: 'attachment';
 };
 
+export interface JwtMfaTokenPayload {
+  sub: string;
+  workspaceId: string;
+  type: 'mfa_token';
+}

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

@@ -47,7 +47,7 @@ export class AuthService {
       includePassword: true,
     });
 
-    const errorMessage = 'email or password does not match';
+    const errorMessage = 'Email or password does not match';
     if (!user || user?.deletedAt) {
       throw new UnauthorizedException(errorMessage);
     }
@@ -156,10 +156,13 @@ export class AuthService {
     });
   }
 
-  async passwordReset(passwordResetDto: PasswordResetDto, workspaceId: string) {
+  async passwordReset(
+    passwordResetDto: PasswordResetDto,
+    workspace: Workspace,
+  ) {
     const userToken = await this.userTokenRepo.findById(
       passwordResetDto.token,
-      workspaceId,
+      workspace.id,
     );
 
     if (
@@ -170,7 +173,9 @@ export class AuthService {
       throw new BadRequestException('Invalid or expired token');
     }
 
-    const user = await this.userRepo.findById(userToken.userId, workspaceId);
+    const user = await this.userRepo.findById(userToken.userId, workspace.id, {
+      includeUserMfa: true,
+    });
     if (!user || user.deletedAt) {
       throw new NotFoundException('User not found');
     }
@@ -183,7 +188,7 @@ export class AuthService {
           password: newPasswordHash,
         },
         user.id,
-        workspaceId,
+        workspace.id,
         trx,
       );
 
@@ -201,7 +206,18 @@ export class AuthService {
       template: emailTemplate,
     });
 
-    return this.tokenService.generateAccessToken(user);
+    // Check if user has MFA enabled or workspace enforces MFA
+    const userHasMfa = user?.['mfa']?.isEnabled || false;
+    const workspaceEnforcesMfa = workspace.enforceMfa || false;
+
+    if (userHasMfa || workspaceEnforcesMfa) {
+      return {
+        requiresLogin: true,
+      };
+    }
+
+    const authToken = await this.tokenService.generateAccessToken(user);
+    return { authToken };
   }
 
   async verifyUserToken(

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

@@ -9,6 +9,7 @@ import {
   JwtAttachmentPayload,
   JwtCollabPayload,
   JwtExchangePayload,
+  JwtMfaTokenPayload,
   JwtPayload,
   JwtType,
 } from '../dto/jwt-payload';
@@ -76,6 +77,22 @@ export class TokenService {
     return this.jwtService.sign(payload, { expiresIn: '1h' });
   }
 
+  async generateMfaToken(
+    user: User,
+    workspaceId: string,
+  ): Promise<string> {
+    if (user.deactivatedAt || user.deletedAt) {
+      throw new ForbiddenException();
+    }
+
+    const payload: JwtMfaTokenPayload = {
+      sub: user.id,
+      workspaceId,
+      type: JwtType.MFA_TOKEN,
+    };
+    return this.jwtService.sign(payload, { expiresIn: '5m' });
+  }
+
   async verifyJwt(token: string, tokenType: string) {
     const payload = await this.jwtService.verifyAsync(token, {
       secret: this.environmentService.getAppSecret(),

apps/server/src/core/auth/strategies/jwt.strategy.ts

@@ -6,6 +6,7 @@ import { JwtPayload, JwtType } from '../dto/jwt-payload';
 import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
 import { UserRepo } from '@docmost/db/repos/user/user.repo';
 import { FastifyRequest } from 'fastify';
+import { extractBearerTokenFromHeader } from '../../../common/helpers';
 
 @Injectable()
 export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
@@ -18,7 +19,7 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
   ) {
     super({
       jwtFromRequest: (req: FastifyRequest) => {
-        return req.cookies?.authToken || this.extractTokenFromHeader(req);
+        return req.cookies?.authToken || extractBearerTokenFromHeader(req);
       },
       ignoreExpiration: false,
       secretOrKey: environmentService.getAppSecret(),
@@ -48,9 +49,4 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
 
     return { user, workspace };
   }
-
-  private extractTokenFromHeader(request: FastifyRequest): string | undefined {
-    const [type, token] = request.headers.authorization?.split(' ') ?? [];
-    return type === 'Bearer' ? token : undefined;
-  }
 }

apps/server/src/core/workspace/controllers/workspace.controller.ts

@@ -29,7 +29,8 @@ import WorkspaceAbilityFactory from '../../casl/abilities/workspace-ability.fact
 import {
   WorkspaceCaslAction,
   WorkspaceCaslSubject,
-} from '../../casl/interfaces/workspace-ability.type';import { FastifyReply } from 'fastify';
+} from '../../casl/interfaces/workspace-ability.type';
+import { FastifyReply } from 'fastify';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
 import { CheckHostnameDto } from '../dto/check-hostname.dto';
 import { RemoveWorkspaceUserDto } from '../dto/remove-workspace-user.dto';
@@ -257,17 +258,27 @@ export class WorkspaceController {
     @AuthWorkspace() workspace: Workspace,
     @Res({ passthrough: true }) res: FastifyReply,
   ) {
-    const authToken = await this.workspaceInvitationService.acceptInvitation(
+    const result = await this.workspaceInvitationService.acceptInvitation(
       acceptInviteDto,
       workspace,
     );
 
-    res.setCookie('authToken', authToken, {
+    if (result.requiresLogin) {
+      return {
+        requiresLogin: true,
+      };
+    }
+
+    res.setCookie('authToken', result.authToken, {
       httpOnly: true,
       path: '/',
       expires: this.environmentService.getCookieExpiresIn(),
       secure: this.environmentService.isHttps(),
     });
+
+    return {
+      requiresLogin: false,
+    };
   }
 
   @Public()

apps/server/src/core/workspace/dto/update-workspace.dto.ts

@@ -14,4 +14,8 @@ export class UpdateWorkspaceDto extends PartialType(CreateWorkspaceDto) {
   @IsOptional()
   @IsBoolean()
   enforceSso: boolean;
+
+  @IsOptional()
+  @IsBoolean()
+  enforceMfa: boolean;
 }

apps/server/src/core/workspace/services/workspace-invitation.service.ts

@@ -177,7 +177,14 @@ export class WorkspaceInvitationService {
     }
   }
 
-  async acceptInvitation(dto: AcceptInviteDto, workspace: Workspace) {
+  async acceptInvitation(
+    dto: AcceptInviteDto,
+    workspace: Workspace,
+  ): Promise<{
+    authToken?: string;
+    requiresLogin?: boolean;
+    message?: string;
+  }> {
     const invitation = await this.db
       .selectFrom('workspaceInvitations')
       .selectAll()
@@ -289,7 +296,14 @@ export class WorkspaceInvitationService {
       });
     }
 
-    return this.tokenService.generateAccessToken(newUser);
+    if (workspace.enforceMfa) {
+      return {
+        requiresLogin: true,
+      };
+    }
+
+    const authToken = await this.tokenService.generateAccessToken(newUser);
+    return { authToken };
   }
 
   async resendInvitation(