sentryagent-idp/tests/unit/middleware/rateLimit.test.ts

/**
 * Unit tests for src/middleware/rateLimit.ts
 *
 * Covers:
 *   - Redis path: RateLimiterRedis honours limit, sets headers, calls next()
 *   - Fallback path: RateLimiterMemory used when Redis is disabled
 *   - 429 response shape: Retry-After header, RateLimitError passed to next()
 *   - Prometheus counter incremented on rejection
 */

import { Request, Response, NextFunction } from 'express';
import { RateLimitError } from '../../../src/utils/errors';

// ── Mocks ─────────────────────────────────────────────────────────────────────

/** Controls whether the mocked ioredis client is returned (Redis enabled path). */
let mockRedisEnabled = true;

const mockIoredisClient = { status: 'ready' };

jest.mock('../../../src/infrastructure/redisClient', () => ({
  getRateLimitRedisClient: jest.fn(() => (mockRedisEnabled ? mockIoredisClient : null)),
}));

/** Tracks the last RateLimiterRedis / RateLimiterMemory consume call. */
const mockConsume = jest.fn();

/** Factory stubs — return the same mock consume regardless of class. */
jest.mock('rate-limiter-flexible', () => {
  class MockRateLimiterRedis {
    readonly points: number = 100;
    readonly consume = mockConsume;
  }

  class MockRateLimiterMemory {
    readonly points: number = 100;
    readonly consume = mockConsume;
  }

  class MockRateLimiterRes extends Error {
    remainingPoints: number;
    msBeforeNext: number;
    consumedPoints: number;
    isFirstInDuration: boolean;
    constructor(opts?: Partial<{ remainingPoints: number; msBeforeNext: number }>) {
      super('Too Many Requests');
      this.remainingPoints = opts?.remainingPoints ?? 0;
      this.msBeforeNext = opts?.msBeforeNext ?? 30000;
      this.consumedPoints = 101;
      this.isFirstInDuration = false;
    }
  }

  return {
    RateLimiterRedis: MockRateLimiterRedis,
    RateLimiterMemory: MockRateLimiterMemory,
    RateLimiterRes: MockRateLimiterRes,
    RateLimiterAbstract: class {},
  };
});

/** Stub for the Prometheus counter so we can assert increments. */
const mockCounterInc = jest.fn();
jest.mock('../../../src/metrics/registry', () => ({
  rateLimitHitsTotal: { inc: (...args: unknown[]) => mockCounterInc(...args) },
}));

// ── Import after mocks are in place ──────────────────────────────────────────

import { rateLimitMiddleware, _resetRateLimiterForTests } from '../../../src/middleware/rateLimit';

// ── Helpers ───────────────────────────────────────────────────────────────────

function buildMocks(path = '/api/v1/agents'): {
  req: Partial<Request>;
  res: Partial<Response>;
  next: jest.Mock;
} {
  return {
    req: {
      user: { client_id: 'agent-123', sub: 'agent-123', scope: '', jti: 'jti', iat: 0, exp: 0 },
      ip: '127.0.0.1',
      path,
    },
    res: {
      setHeader: jest.fn(),
    },
    next: jest.fn(),
  };
}

/** Builds a successful RateLimiterRes result (request allowed). */
function makeAllowedResult(remaining = 99, msBeforeNext = 30000): Record<string, unknown> {
  return {
    remainingPoints: remaining,
    msBeforeNext,
    consumedPoints: 100 - remaining,
    isFirstInDuration: false,
  };
}

/** Returns the MockRateLimiterRes class from the mock module. */
function getMockRateLimiterRes(): new (opts?: { msBeforeNext?: number; remainingPoints?: number }) => Error {
  return (jest.requireMock('rate-limiter-flexible') as {
    RateLimiterRes: new (opts?: { msBeforeNext?: number; remainingPoints?: number }) => Error;
  }).RateLimiterRes;
}

// ── Tests ─────────────────────────────────────────────────────────────────────

beforeEach(() => {
  jest.clearAllMocks();
  _resetRateLimiterForTests();
  mockRedisEnabled = true;
  process.env['RATE_LIMIT_WINDOW_MS'] = '60000';
  process.env['RATE_LIMIT_MAX_REQUESTS'] = '100';
});

describe('rateLimitMiddleware — Redis path (REDIS_RATE_LIMIT_ENABLED=true)', () => {
  it('calls next() without error when request is under the limit', async () => {
    mockConsume.mockResolvedValue(makeAllowedResult(99));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(next).toHaveBeenCalledWith();
    expect(next).not.toHaveBeenCalledWith(expect.any(Error));
  });

  it('sets X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset headers on success', async () => {
    mockConsume.mockResolvedValue(makeAllowedResult(50));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Limit', expect.any(Number));
    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Remaining', 50);
    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Reset', expect.any(Number));
  });

  it('keys the limiter by client_id from req.user', async () => {
    mockConsume.mockResolvedValue(makeAllowedResult(99));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(mockConsume).toHaveBeenCalledWith('agent-123');
  });

  it('falls back to req.ip when req.user is not set', async () => {
    mockConsume.mockResolvedValue(makeAllowedResult(99));
    const { req, res, next } = buildMocks();
    req.user = undefined;

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(mockConsume).toHaveBeenCalledWith('127.0.0.1');
  });
});

describe('rateLimitMiddleware — 429 response shape', () => {
  it('calls next(RateLimitError) when limit is exceeded', async () => {
    const MockRateLimiterRes = getMockRateLimiterRes();
    mockConsume.mockRejectedValue(new MockRateLimiterRes({ msBeforeNext: 45000 }));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(next).toHaveBeenCalledWith(expect.any(RateLimitError));
  });

  it('sets Retry-After header on rejection', async () => {
    const MockRateLimiterRes = getMockRateLimiterRes();
    mockConsume.mockRejectedValue(new MockRateLimiterRes({ msBeforeNext: 30000 }));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(res.setHeader).toHaveBeenCalledWith('Retry-After', 30);
  });

  it('sets X-RateLimit-Remaining to 0 on rejection', async () => {
    const MockRateLimiterRes = getMockRateLimiterRes();
    mockConsume.mockRejectedValue(new MockRateLimiterRes({ msBeforeNext: 30000 }));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Remaining', 0);
  });

  it('increments agentidp_rate_limit_hits_total with endpoint label on rejection', async () => {
    const MockRateLimiterRes = getMockRateLimiterRes();
    mockConsume.mockRejectedValue(new MockRateLimiterRes({ msBeforeNext: 10000 }));
    const { req, res, next } = buildMocks('/api/v1/agents');

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(mockCounterInc).toHaveBeenCalledWith({ endpoint: '/api/v1/agents' });
  });
});

describe('rateLimitMiddleware — fallback path (Redis disabled)', () => {
  beforeEach(() => {
    mockRedisEnabled = false;
    _resetRateLimiterForTests();
  });

  it('calls next() without error when request is under the limit (memory limiter)', async () => {
    mockConsume.mockResolvedValue(makeAllowedResult(99));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(next).toHaveBeenCalledWith();
    expect(next).not.toHaveBeenCalledWith(expect.any(Error));
  });

  it('calls next(RateLimitError) when memory limiter is exceeded', async () => {
    const MockRateLimiterRes = getMockRateLimiterRes();
    mockConsume.mockRejectedValue(new MockRateLimiterRes({ msBeforeNext: 60000 }));
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(next).toHaveBeenCalledWith(expect.any(RateLimitError));
  });
});

describe('rateLimitMiddleware — unexpected errors', () => {
  it('passes non-RateLimiterRes errors to next() as-is', async () => {
    const unexpectedError = new Error('Redis network failure');
    mockConsume.mockRejectedValue(unexpectedError);
    const { req, res, next } = buildMocks();

    await rateLimitMiddleware(req as Request, res as Response, next as NextFunction);

    expect(next).toHaveBeenCalledWith(unexpectedError);
    expect(next).not.toHaveBeenCalledWith(expect.any(RateLimitError));
  });
});