goodgo-platform/apps/api/src/modules/auth/presentation/controllers/mfa.controller.ts

import {
  Body,
  Controller,
  Delete,
  Get,
  Post,
  Res,
  UseGuards,
} from '@nestjs/common';
import { type CommandBus, type QueryBus } from '@nestjs/cqrs';
import { ApiTags, ApiOperation, ApiResponse, ApiBearerAuth } from '@nestjs/swagger';
import { Throttle } from '@nestjs/throttler';
import { type Response } from 'express';
import { EndpointRateLimit, EndpointRateLimitGuard } from '@modules/shared';
import { DisableMfaCommand } from '../../application/commands/disable-mfa/disable-mfa.command';
import { SetupMfaCommand } from '../../application/commands/setup-mfa/setup-mfa.command';
import { type SetupMfaResultDto } from '../../application/commands/setup-mfa/setup-mfa.handler';
import { UseBackupCodeCommand } from '../../application/commands/use-backup-code/use-backup-code.command';
import { VerifyMfaChallengeCommand } from '../../application/commands/verify-mfa-challenge/verify-mfa-challenge.command';
import { VerifyMfaSetupCommand } from '../../application/commands/verify-mfa-setup/verify-mfa-setup.command';
import { type VerifyMfaSetupResultDto } from '../../application/commands/verify-mfa-setup/verify-mfa-setup.handler';
import { type MfaStatusDto } from '../../application/queries/get-mfa-status/get-mfa-status.handler';
import { GetMfaStatusQuery } from '../../application/queries/get-mfa-status/get-mfa-status.query';
import { type TokenService, type JwtPayload, type TokenPair } from '../../infrastructure/services/token.service';
import { CurrentUser } from '../decorators/current-user.decorator';
import {
  type VerifyMfaSetupDto,
  type VerifyMfaChallengeDto,
  type UseBackupCodeDto,
  type DisableMfaDto,
} from '../dto/mfa.dto';
import { JwtAuthGuard } from '../guards/jwt-auth.guard';

const IS_TEST = process.env['NODE_ENV'] === 'test';
const IS_PRODUCTION = process.env['NODE_ENV'] === 'production';
const MFA_RATE_LIMIT = IS_TEST ? 10_000 : 5;
const ACCESS_TOKEN_MAX_AGE = 15 * 60 * 1000;
const REFRESH_TOKEN_MAX_AGE = 30 * 24 * 60 * 60 * 1000;
const AUTH_COOKIE_MAX_AGE = 30 * 24 * 60 * 60 * 1000;

function setAuthCookies(res: Response, tokens: TokenPair): void {
  res.cookie('access_token', tokens.accessToken, {
    httpOnly: true,
    secure: IS_PRODUCTION,
    sameSite: 'strict',
    path: '/',
    maxAge: ACCESS_TOKEN_MAX_AGE,
  });
  res.cookie('refresh_token', tokens.refreshToken, {
    httpOnly: true,
    secure: IS_PRODUCTION,
    sameSite: 'strict',
    path: '/auth',
    maxAge: REFRESH_TOKEN_MAX_AGE,
  });
  res.cookie('goodgo_authenticated', '1', {
    httpOnly: false,
    secure: IS_PRODUCTION,
    sameSite: 'lax',
    path: '/',
    maxAge: AUTH_COOKIE_MAX_AGE,
  });
}

@ApiTags('auth')
@Controller('auth/mfa')
export class MfaController {
  constructor(
    private readonly commandBus: CommandBus,
    private readonly queryBus: QueryBus,
    private readonly tokenService: TokenService,
  ) {}

  @UseGuards(JwtAuthGuard)
  @Post('setup')
  @ApiBearerAuth('JWT')
  @ApiOperation({ summary: 'Generate TOTP secret and QR code for MFA setup' })
  @ApiResponse({ status: 201, description: 'TOTP secret and QR code generated' })
  @ApiResponse({ status: 400, description: 'MFA already enabled' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  async setup(@CurrentUser() user: JwtPayload): Promise<SetupMfaResultDto> {
    return this.commandBus.execute(new SetupMfaCommand(user.sub));
  }

  @UseGuards(JwtAuthGuard)
  @Post('verify-setup')
  @ApiBearerAuth('JWT')
  @ApiOperation({ summary: 'Verify TOTP setup with first code and enable MFA' })
  @ApiResponse({ status: 201, description: 'MFA enabled, backup codes returned' })
  @ApiResponse({ status: 400, description: 'Invalid TOTP code or MFA not set up' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  async verifySetup(
    @CurrentUser() user: JwtPayload,
    @Body() dto: VerifyMfaSetupDto,
  ): Promise<VerifyMfaSetupResultDto> {
    return this.commandBus.execute(
      new VerifyMfaSetupCommand(user.sub, dto.totpCode),
    );
  }

  @Throttle({ default: { ttl: 3_600_000, limit: MFA_RATE_LIMIT } })
  @EndpointRateLimit({ limit: IS_TEST ? 10_000 : 5, windowSeconds: 60, keyStrategy: 'ip' })
  @UseGuards(EndpointRateLimitGuard)
  @Post('challenge')
  @ApiOperation({ summary: 'Verify TOTP code during login MFA challenge' })
  @ApiResponse({ status: 201, description: 'MFA verified, auth tokens returned' })
  @ApiResponse({ status: 401, description: 'Invalid TOTP code or expired challenge' })
  async verifyChallenge(
    @Body() dto: VerifyMfaChallengeDto,
    @Res({ passthrough: true }) res: Response,
  ): Promise<{ message: string; accessToken: string; refreshToken: string }> {
    const tokens: TokenPair = await this.commandBus.execute(
      new VerifyMfaChallengeCommand(dto.challengeId, dto.totpCode),
    );
    setAuthCookies(res, tokens);
    return {
      message: 'Xác thực MFA thành công',
      accessToken: tokens.accessToken,
      refreshToken: tokens.refreshToken,
    };
  }

  @Throttle({ default: { ttl: 3_600_000, limit: MFA_RATE_LIMIT } })
  @EndpointRateLimit({ limit: IS_TEST ? 10_000 : 5, windowSeconds: 60, keyStrategy: 'ip' })
  @UseGuards(EndpointRateLimitGuard)
  @Post('backup-codes')
  @ApiOperation({ summary: 'Use a backup code during MFA challenge' })
  @ApiResponse({ status: 201, description: 'Backup code accepted, auth tokens returned' })
  @ApiResponse({ status: 401, description: 'Invalid backup code or expired challenge' })
  async useBackupCode(
    @Body() dto: UseBackupCodeDto,
    @Res({ passthrough: true }) res: Response,
  ): Promise<{ message: string; accessToken: string; refreshToken: string; remainingBackupCodes: number }> {
    const result = await this.commandBus.execute(
      new UseBackupCodeCommand(dto.challengeId, dto.backupCode),
    );
    setAuthCookies(res, result);
    return {
      message: 'Xác thực bằng mã backup thành công',
      accessToken: result.accessToken,
      refreshToken: result.refreshToken,
      remainingBackupCodes: result.remainingBackupCodes,
    };
  }

  @UseGuards(JwtAuthGuard)
  @Delete()
  @ApiBearerAuth('JWT')
  @ApiOperation({ summary: 'Disable MFA (requires current TOTP code)' })
  @ApiResponse({ status: 200, description: 'MFA disabled' })
  @ApiResponse({ status: 400, description: 'MFA not enabled' })
  @ApiResponse({ status: 401, description: 'Invalid TOTP code' })
  async disable(
    @CurrentUser() user: JwtPayload,
    @Body() dto: DisableMfaDto,
  ): Promise<{ message: string }> {
    return this.commandBus.execute(
      new DisableMfaCommand(user.sub, dto.totpCode),
    );
  }

  @UseGuards(JwtAuthGuard)
  @Get('status')
  @ApiBearerAuth('JWT')
  @ApiOperation({ summary: 'Get MFA status for current user' })
  @ApiResponse({ status: 200, description: 'MFA status returned' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  async getStatus(@CurrentUser() user: JwtPayload): Promise<MfaStatusDto> {
    return this.queryBus.execute(new GetMfaStatusQuery(user.sub));
  }
}