llm-chat/backend/middleware/rateLimit.js

const rateLimit = require('express-rate-limit');
const { Logger } = require('./logger');

/**
 * 通用速率限制配置
 */
const createRateLimit = (options = {}) => {
  const defaultOptions = {
    windowMs: parseInt(process.env.RATE_LIMIT_WINDOW_MS) || 60000, // 1分钟
    max: parseInt(process.env.RATE_LIMIT_MAX_REQUESTS) || 100, // 每分钟最多100个请求
    message: {
      success: false,
      error: '请求过于频繁，请稍后再试',
      code: 'RATE_LIMIT_EXCEEDED',
      retryAfter: Math.ceil(options.windowMs / 1000) || 60
    },
    standardHeaders: true, // 返回标准的 `RateLimit-*` 头部
    legacyHeaders: false, // 禁用旧的 `X-RateLimit-*` 头部
    skip: (req) => {
      // 跳过健康检查和API信息接口
      return req.path === '/api/health' || req.path === '/api';
    },
    handler: (req, res, next, options) => {
      Logger.warn('触发速率限制', {
        ip: req.ip,
        path: req.path,
        method: req.method,
        userAgent: req.get('User-Agent')
      });
      
      res.status(options.statusCode).json(options.message);
    },
    keyGenerator: (req) => {
      // 基于IP地址生成限制键
      return req.ip;
    }
  };

  return rateLimit({
    ...defaultOptions,
    ...options
  });
};

/**
 * 全局速率限制 - 适用于所有API端点
 */
const globalRateLimit = createRateLimit({
  windowMs: 60000, // 1分钟
  max: 100, // 每分钟最多100个请求
  message: {
    success: false,
    error: '请求过于频繁，请稍后再试',
    code: 'RATE_LIMIT_EXCEEDED',
    retryAfter: 60
  }
});

/**
 * 聊天API专用速率限制 - 更严格的限制
 */
const chatRateLimit = createRateLimit({
  windowMs: 60000, // 1分钟
  max: 20, // 每分钟最多20次聊天请求
  message: {
    success: false,
    error: '聊天请求过于频繁，请稍后再试',
    code: 'CHAT_RATE_LIMIT_EXCEEDED',
    retryAfter: 60
  },
  skip: (req) => {
    // 只对聊天相关的POST请求进行限制
    return !(req.path.startsWith('/api/chat') && req.method === 'POST');
  }
});

/**
 * 创建对话速率限制
 */
const conversationCreateRateLimit = createRateLimit({
  windowMs: 300000, // 5分钟
  max: 10, // 每5分钟最多创建10个对话
  message: {
    success: false,
    error: '创建对话过于频繁，请稍后再试',
    code: 'CONVERSATION_CREATE_RATE_LIMIT_EXCEEDED',
    retryAfter: 300
  },
  skip: (req) => {
    return !(req.path === '/api/conversations' && req.method === 'POST');
  }
});

/**
 * 流式聊天专用速率限制
 */
const streamChatRateLimit = createRateLimit({
  windowMs: 60000, // 1分钟
  max: 10, // 每分钟最多10次流式请求
  message: {
    success: false,
    error: '流式聊天请求过于频繁，请稍后再试',
    code: 'STREAM_RATE_LIMIT_EXCEEDED',
    retryAfter: 60
  },
  skip: (req) => {
    return !(req.path === '/api/chat/stream' && req.method === 'POST');
  }
});

/**
 * 搜索API速率限制
 */
const searchRateLimit = createRateLimit({
  windowMs: 60000, // 1分钟
  max: 30, // 每分钟最多30次搜索
  message: {
    success: false,
    error: '搜索请求过于频繁，请稍后再试',
    code: 'SEARCH_RATE_LIMIT_EXCEEDED',
    retryAfter: 60
  },
  skip: (req) => {
    return !req.path.includes('/search');
  }
});

/**
 * 基于用户会话的速率限制
 */
const createSessionRateLimit = (options = {}) => {
  const sessions = new Map();
  
  return (req, res, next) => {
    const sessionId = req.headers['x-session-id'] || req.ip;
    const now = Date.now();
    const windowMs = options.windowMs || 60000;
    const maxRequests = options.max || 50;
    
    // 清理过期的会话记录
    for (const [id, data] of sessions.entries()) {
      if (now - data.windowStart > windowMs) {
        sessions.delete(id);
      }
    }
    
    // 获取或创建会话记录
    let sessionData = sessions.get(sessionId);
    if (!sessionData || now - sessionData.windowStart > windowMs) {
      sessionData = {
        windowStart: now,
        requestCount: 0
      };
      sessions.set(sessionId, sessionData);
    }
    
    // 检查是否超过限制
    if (sessionData.requestCount >= maxRequests) {
      Logger.warn('会话速率限制触发', {
        sessionId,
        requestCount: sessionData.requestCount,
        maxRequests,
        path: req.path
      });
      
      return res.status(429).json({
        success: false,
        error: '会话请求过于频繁，请稍后再试',
        code: 'SESSION_RATE_LIMIT_EXCEEDED',
        retryAfter: Math.ceil((windowMs - (now - sessionData.windowStart)) / 1000)
      });
    }
    
    // 增加请求计数
    sessionData.requestCount++;
    
    // 设置响应头
    res.set({
      'X-RateLimit-Limit': maxRequests,
      'X-RateLimit-Remaining': Math.max(0, maxRequests - sessionData.requestCount),
      'X-RateLimit-Reset': new Date(sessionData.windowStart + windowMs).toISOString()
    });
    
    next();
  };
};

/**
 * 自适应速率限制 - 根据服务器负载动态调整
 */
const createAdaptiveRateLimit = (baseOptions = {}) => {
  let currentLoad = 0;
  
  // 定期检查服务器负载
  setInterval(() => {
    const usage = process.cpuUsage();
    const memUsage = process.memoryUsage();
    
    // 简单的负载计算（可以根据需要改进）
    const cpuLoad = (usage.user + usage.system) / 1000000; // 转换为秒
    const memLoad = memUsage.heapUsed / memUsage.heapTotal;
    
    currentLoad = Math.max(cpuLoad / 100, memLoad); // 标准化到 0-1
  }, 5000); // 每5秒检查一次
  
  return createRateLimit({
    ...baseOptions,
    max: (req) => {
      const baseMax = baseOptions.max || 100;
      
      // 根据负载调整限制
      if (currentLoad > 0.8) {
        return Math.floor(baseMax * 0.5); // 高负载时减少50%
      } else if (currentLoad > 0.6) {
        return Math.floor(baseMax * 0.7); // 中等负载时减少30%
      } else {
        return baseMax; // 低负载时保持原限制
      }
    }
  });
};

/**
 * 获取客户端真实IP地址
 */
const getClientIP = (req) => {
  return req.headers['x-forwarded-for'] || 
         req.headers['x-real-ip'] || 
         req.connection.remoteAddress || 
         req.socket.remoteAddress ||
         (req.connection.socket ? req.connection.socket.remoteAddress : null) ||
         req.ip;
};

/**
 * IP白名单中间件
 */
const createIPWhitelist = (whitelist = []) => {
  return (req, res, next) => {
    const clientIP = getClientIP(req);
    
    // 本地开发环境跳过IP检查
    if (process.env.NODE_ENV === 'development') {
      return next();
    }
    
    if (whitelist.length > 0 && !whitelist.includes(clientIP)) {
      Logger.warn('IP不在白名单中', { clientIP, path: req.path });
      
      return res.status(403).json({
        success: false,
        error: '访问被拒绝',
        code: 'IP_NOT_WHITELISTED'
      });
    }
    
    next();
  };
};

module.exports = {
  globalRateLimit,
  chatRateLimit,
  conversationCreateRateLimit,
  streamChatRateLimit,
  searchRateLimit,
  createRateLimit,
  createSessionRateLimit,
  createAdaptiveRateLimit,
  createIPWhitelist,
  getClientIP
};