// server/test/middleware.ipRateLimit.test.js — IP 限流中间件测试 // Trae v2.7 加的内存限流器:每 IP 每窗口 max 次 import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; import { ipRateLimit, _clearBuckets } from '../src/middleware/ipRateLimit.js'; function mockReq(headers = {}) { return { headers, socket: { remoteAddress: '127.0.0.1' }, ip: '127.0.0.1', }; } function mockRes() { const headers = {}; const res = { headers, statusCode: 200, set(k, v) { if (typeof k === 'object') Object.assign(headers, k); else headers[k] = v; return this; }, status(code) { this.statusCode = code; return this; }, json(body) { this.body = body; return this; }, }; return res; } describe('ipRateLimit middleware', () => { beforeEach(() => _clearBuckets()); afterEach(() => _clearBuckets()); it('第一次调用正常通过', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 3, name: 't1' }); const req = mockReq(); const res = mockRes(); let nextCalled = false; mw(req, res, () => { nextCalled = true; }); expect(nextCalled).toBe(true); // 第一次调用走"新窗口"分支(line 24-27),不设 headers 直接 next // 第二次调用起才会返回 X-RateLimit-* headers }); it('第二次调用设置 X-RateLimit headers', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 3, name: 't1b' }); mw(mockReq(), mockRes(), () => {}); const res = mockRes(); let nextCalled = false; mw(mockReq(), res, () => { nextCalled = true; }); expect(nextCalled).toBe(true); expect(res.headers['X-RateLimit-Limit']).toBe('3'); expect(res.headers['X-RateLimit-Remaining']).toBe('1'); // b.count=2, max=3, 3-2=1 }); it('max 次内都通过', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 3, name: 't2' }); for (let i = 0; i < 3; i++) { const req = mockReq(); const res = mockRes(); let nextCalled = false; mw(req, res, () => { nextCalled = true; }); expect(nextCalled).toBe(true); } }); it('超过 max 返 429 + Retry-After + RATE_LIMITED', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 2, name: 't3' }); // 前两次通过 for (let i = 0; i < 2; i++) { mw(mockReq(), mockRes(), () => {}); } // 第三次触发 const res = mockRes(); let nextCalled = false; mw(mockReq(), res, () => { nextCalled = true; }); expect(nextCalled).toBe(false); expect(res.statusCode).toBe(429); expect(res.body.error.code).toBe('RATE_LIMITED'); expect(res.headers['Retry-After']).toBeDefined(); expect(res.headers['X-RateLimit-Remaining']).toBe('0'); }); it('不同 IP 互不影响', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 1, name: 't4' }); // IP A 用完配额 mw(mockReq({ 'x-forwarded-for': '1.1.1.1' }), mockRes(), () => {}); const res = mockRes(); let nextCalled = false; mw(mockReq({ 'x-forwarded-for': '2.2.2.2' }), res, () => { nextCalled = true; }); expect(nextCalled).toBe(true); }); it('X-Forwarded-For 优先于 socket 地址', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 1, name: 't5' }); mw(mockReq({ 'x-forwarded-for': '1.1.1.1, 10.0.0.1' }), mockRes(), () => {}); const res = mockRes(); let nextCalled = false; mw(mockReq({ 'x-forwarded-for': '1.1.1.1' }), res, () => { nextCalled = true; }); expect(nextCalled).toBe(false); // 同 IP }); it('窗口过期后重置', () => { vi.useFakeTimers(); const mw = ipRateLimit({ windowMs: 1000, max: 1, name: 't6' }); mw(mockReq(), mockRes(), () => {}); // 立刻再次 → 429 const res1 = mockRes(); mw(mockReq(), res1, () => {}); expect(res1.statusCode).toBe(429); // 时间快进 1.1 秒 → 窗口过期 → 重新允许 vi.advanceTimersByTime(1100); const res2 = mockRes(); let nextCalled = false; mw(mockReq(), res2, () => { nextCalled = true; }); expect(nextCalled).toBe(true); vi.useRealTimers(); }); it('_clearBuckets 测试钩子能清空所有计数', () => { const mw = ipRateLimit({ windowMs: 60_000, max: 1, name: 't7' }); mw(mockReq(), mockRes(), () => {}); const res1 = mockRes(); mw(mockReq(), res1, () => {}); expect(res1.statusCode).toBe(429); _clearBuckets(); const res2 = mockRes(); let nextCalled = false; mw(mockReq(), res2, () => { nextCalled = true; }); expect(nextCalled).toBe(true); }); });