mm_os/middleware/rate_limit/index.js

这是超级美眉服务端框架，用于快速构建应用程序。

/**
 * API速率限制中间件
 * 用于防止DoS攻击，限制客户端在一定时间内的请求频率
 * @param {Object} server 服务实例
 * @param {Object} config 配置参数
 */
module.exports = function(server, config) {
    // 初始化速率限制配置
    const cg = {
        // 默认配置
        windowMs: 60,      // 时间窗口，调整为1分钟（更精细的限制）
        maxRequests: 5000,         // 每个时间窗口内的最大请求数，从1000提高到5000
        message: '请求过于频繁，请稍后再试', // 超过限制时的提示信息
        statusCode: 429,          // 超过限制时的HTTP状态码
        // 白名单路径，这些路径不受速率限制
        whitelistPaths: [],
        // 合并用户配置
        ...(config && config.rate_limit ? config.rate_limit : {})
    };
    
    // 生成基于IP的唯一标识符
    function getClientKey(ctx) {
        // 优先使用X-Forwarded-For头部（如果有代理的话）
        const forwardedFor = ctx.headers['x-forwarded-for'];
        if (forwardedFor) {
            // 通常格式为：X-Forwarded-For: client, proxy1, proxy2
            return forwardedFor.split(',')[0].trim();
        }
        // 直接使用IP地址
        return ctx.ip;
    }
    
    // 使用Redis存储请求计数
    async function incrementRequestWithRedis(clientKey) {
        try {
            // 检查Redis是否可用
            if ($.cache && typeof $.cache.addInt === 'function') {
                const key = `rate_limit:${clientKey}`;
                
                // 使用addInt方法增加计数（mm_redis包提供的方法）
                const count = await $.cache.addInt(key, 1);
                
                // 设置过期时间
                await $.cache.ttl(key, Math.ceil(cg.windowMs));
                
                return count || 0;
            }
        } catch (error) {
            $.log.error('Redis速率限制失败:', error);
        }
        
        // Redis不可用时返回null，将使用内存存储
        return null;
    }

    // 中间件主逻辑
    server.use(async (ctx, next) => {
        try {
            // 跳过静态文件和favicon请求
            if (ctx.path === '/favicon.ico' || ctx.path.startsWith('/static/')) {
                await next();
                return;
            }
            
            // 检查是否是白名单路径
            if (cg.whitelistPaths && cg.whitelistPaths.some(path => 
                ctx.path === path || ctx.path.startsWith(path + '/'))) {
                await next();
                return;
            }
            
            // 获取客户端唯一标识符
            const clientKey = getClientKey(ctx);
            
            let requestCount = await incrementRequestWithRedis(clientKey);
            
            // 设置响应头部，告知客户端当前的限制状态
            ctx.set('X-RateLimit-Limit', cg.maxRequests);
            ctx.set('X-RateLimit-Remaining', Math.max(0, cg.maxRequests - requestCount));
            
            // 检查是否超过限制
            if (requestCount > cg.maxRequests) {
                $.log.warn(`API速率限制触发: ${clientKey} 请求 ${ctx.path} 次数过多`);
                
                ctx.status = cg.statusCode;
                ctx.body = {
                    code: cg.statusCode,
                    msg: cg.message
                };
                
                // 记录超过限制的请求
                if ($.log && $.log.warn) {
                    $.log.warn(`速率限制触发: IP=${clientKey}, Path=${ctx.path}, Method=${ctx.method}`);
                }
                
                return;
            }
            
            // 继续处理请求
            await next();
        } catch (error) {
            $.log.error('速率限制中间件错误:', error);
            // 确保请求可以继续处理
            await next();
        }
    });
    
    // 记录中间件初始化信息
    $.log.info(`速率限制中间件已加载: ${cg.maxRequests}请求/${cg.windowMs}秒`);
    
    return server;
};