whale-town-end/src/business/security/guards/throttle.guard.ts

/**
 * 频率限制守卫
 * 
 * 功能描述：
 * - 实现API接口的频率限制功能
 * - 基于IP地址进行限制
 * - 支持自定义限制规则
 * 
 * 使用场景：
 * - 防止API滥用
 * - 登录暴力破解防护
 * - 验证码发送频率控制
 * 
 * @author kiro-ai
 * @version 1.0.0
 * @since 2025-12-24
 */

import {
  Injectable,
  CanActivate,
  ExecutionContext,
  HttpException,
  HttpStatus,
  Logger
} from '@nestjs/common';
import { Reflector } from '@nestjs/core';
import { Request } from 'express';
import { THROTTLE_KEY, ThrottleConfig } from '../decorators/throttle.decorator';

/**
 * 频率限制记录接口
 */
interface ThrottleRecord {
  /** 请求次数 */
  count: number;
  /** 窗口开始时间 */
  windowStart: number;
  /** 最后请求时间 */
  lastRequest: number;
}

/**
 * 频率限制响应接口
 */
interface ThrottleResponse {
  /** 请求是否成功 */
  success: boolean;
  /** 响应消息 */
  message: string;
  /** 错误代码 */
  error_code: string;
  /** 限制信息 */
  throttle_info: {
    /** 限制次数 */
    limit: number;
    /** 时间窗口（秒） */
    window_seconds: number;
    /** 当前请求次数 */
    current_requests: number;
    /** 重置时间 */
    reset_time: string;
  };
}

@Injectable()
export class ThrottleGuard implements CanActivate {
  private readonly logger = new Logger(ThrottleGuard.name);
  
  /**
   * 存储频率限制记录
   * Key: IP地址 + 路径
   * Value: 限制记录
   */
  private readonly records = new Map<string, ThrottleRecord>();

  /**
   * 清理过期记录的间隔（毫秒）
   */
  private readonly cleanupInterval = 60000; // 1分钟

  constructor(private readonly reflector: Reflector) {
    // 启动定期清理任务
    this.startCleanupTask();
  }

  /**
   * 守卫检查函数
   * 
   * @param context 执行上下文
   * @returns 是否允许通过
   */
  async canActivate(context: ExecutionContext): Promise<boolean> {
    // 1. 获取频率限制配置
    const throttleConfig = this.getThrottleConfig(context);
    
    if (!throttleConfig) {
      // 没有配置频率限制，直接通过
      return true;
    }

    // 2. 获取请求信息
    const request = context.switchToHttp().getRequest<Request>();
    const key = this.generateKey(request, throttleConfig);
    
    // 3. 检查频率限制
    const isAllowed = this.checkThrottle(key, throttleConfig);
    
    if (!isAllowed) {
      // 4. 记录被限制的请求
      this.logger.warn('请求被频率限制', {
        operation: 'throttle_limit',
        method: request.method,
        url: request.url,
        ip: request.ip,
        userAgent: request.get('User-Agent'),
        limit: throttleConfig.limit,
        ttl: throttleConfig.ttl,
        timestamp: new Date().toISOString()
      });

      // 5. 抛出频率限制异常
      const record = this.records.get(key);
      const resetTime = new Date(record!.windowStart + throttleConfig.ttl * 1000);
      
      const response: ThrottleResponse = {
        success: false,
        message: throttleConfig.message || '请求过于频繁，请稍后再试',
        error_code: 'TOO_MANY_REQUESTS',
        throttle_info: {
          limit: throttleConfig.limit,
          window_seconds: throttleConfig.ttl,
          current_requests: record!.count,
          reset_time: resetTime.toISOString()
        }
      };

      throw new HttpException(response, HttpStatus.TOO_MANY_REQUESTS);
    }

    return true;
  }

  /**
   * 获取频率限制配置
   * 
   * @param context 执行上下文
   * @returns 频率限制配置或null
   */
  private getThrottleConfig(context: ExecutionContext): ThrottleConfig | null {
    // 从方法装饰器获取配置
    const methodConfig = this.reflector.get<ThrottleConfig>(
      THROTTLE_KEY,
      context.getHandler()
    );

    if (methodConfig) {
      return methodConfig;
    }

    // 从类装饰器获取配置
    const classConfig = this.reflector.get<ThrottleConfig>(
      THROTTLE_KEY,
      context.getClass()
    );

    return classConfig || null;
  }

  /**
   * 生成限制键
   * 
   * @param request 请求对象
   * @param config 频率限制配置
   * @returns 限制键
   */
  private generateKey(request: Request, config: ThrottleConfig): string {
    const ip = request.ip || 'unknown';
    const path = request.route?.path || request.url;
    const method = request.method;
    
    // 根据限制类型生成不同的键
    if (config.type === 'user') {
      // 基于用户的限制（需要从JWT中获取用户ID）
      const userId = this.extractUserId(request);
      return `user:${userId}:${method}:${path}`;
    } else {
      // 基于IP的限制（默认）
      return `ip:${ip}:${method}:${path}`;
    }
  }

  /**
   * 检查频率限制
   * 
   * @param key 限制键
   * @param config 频率限制配置
   * @returns 是否允许通过
   */
  private checkThrottle(key: string, config: ThrottleConfig): boolean {
    const now = Date.now();
    const windowMs = config.ttl * 1000;
    
    let record = this.records.get(key);
    
    if (!record) {
      // 第一次请求
      this.records.set(key, {
        count: 1,
        windowStart: now,
        lastRequest: now
      });
      return true;
    }

    // 检查是否需要重置窗口
    if (now - record.windowStart >= windowMs) {
      // 重置窗口
      record.count = 1;
      record.windowStart = now;
      record.lastRequest = now;
      return true;
    }

    // 在当前窗口内
    if (record.count >= config.limit) {
      // 超过限制
      return false;
    }

    // 增加计数
    record.count++;
    record.lastRequest = now;
    return true;
  }

  /**
   * 从请求中提取用户ID
   * 
   * @param request 请求对象
   * @returns 用户ID
   */
  private extractUserId(request: Request): string {
    // 这里应该从JWT token中提取用户ID
    // 简化实现，使用IP作为fallback
    const authHeader = request.get('Authorization');
    if (authHeader && authHeader.startsWith('Bearer ')) {
      try {
        // 这里应该解析JWT token获取用户ID
        // 简化实现，返回token的hash
        const token = authHeader.substring(7);
        return Buffer.from(token).toString('base64').substring(0, 10);
      } catch (error) {
        // JWT解析失败，使用IP
        return request.ip || 'unknown';
      }
    }
    
    return request.ip || 'unknown';
  }

  /**
   * 启动清理任务
   */
  private startCleanupTask(): void {
    setInterval(() => {
      this.cleanupExpiredRecords();
    }, this.cleanupInterval);
  }

  /**
   * 清理过期记录
   */
  private cleanupExpiredRecords(): void {
    const now = Date.now();
    const maxAge = 3600000; // 1小时
    
    for (const [key, record] of this.records.entries()) {
      if (now - record.lastRequest > maxAge) {
        this.records.delete(key);
      }
    }
  }

  /**
   * 获取当前记录统计
   * 
   * @returns 记录统计信息
   */
  getStats() {
    return {
      totalRecords: this.records.size,
      records: Array.from(this.records.entries()).map(([key, record]) => ({
        key,
        count: record.count,
        windowStart: new Date(record.windowStart).toISOString(),
        lastRequest: new Date(record.lastRequest).toISOString()
      }))
    };
  }

  /**
   * 清除所有记录
   */
  clearAllRecords(): void {
    this.records.clear();
  }

  /**
   * 清除指定键的记录
   * 
   * @param key 限制键
   */
  clearRecord(key: string): void {
    this.records.delete(key);
  }
}