i/server/test/middleware.ipRateLimit.test.js

// 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);
    });
});