@geekmidas/rate-limit 0.1.0

package/src/__tests__/rate-limit.spec.ts

@@ -0,0 +1,321 @@
+import { InMemoryCache } from '@geekmidas/cache/memory';
+import { describe, expect, it, vi } from 'vitest';
+import {
+  type RateLimitConfig,
+  type RateLimitContext,
+  type RateLimitData,
+  TooManyRequestsError,
+  checkRateLimit,
+  defaultKeyGenerator,
+  getRateLimitHeaders,
+} from '../index';
+
+describe('Rate Limiting', () => {
+  const createContext = (
+    overrides?: Partial<RateLimitContext>,
+  ): RateLimitContext => ({
+    header: vi.fn((key: string) => {
+      if (key === 'x-forwarded-for') return '192.168.1.1';
+      return undefined;
+    }),
+    services: [],
+    logger: {
+      debug: vi.fn(),
+      info: vi.fn(),
+      warn: vi.fn(),
+      error: vi.fn(),
+      child: vi.fn().mockReturnThis(),
+    } as any,
+    session: {},
+    path: '/api/test',
+    method: 'GET',
+    ...overrides,
+  });
+
+  describe('defaultKeyGenerator', () => {
+    it('should generate key with IP from x-forwarded-for', () => {
+      const ctx = createContext();
+      const key = defaultKeyGenerator(ctx);
+      expect(key).toBe('rate-limit:GET:/api/test:192.168.1.1');
+    });
+
+    it('should use x-real-ip when x-forwarded-for is not available', () => {
+      const ctx = createContext({
+        header: vi.fn((key: string) => {
+          if (key === 'x-real-ip') return '10.0.0.1';
+          return undefined;
+        }),
+      });
+      const key = defaultKeyGenerator(ctx);
+      expect(key).toBe('rate-limit:GET:/api/test:10.0.0.1');
+    });
+
+    it('should use unknown when no IP header is available', () => {
+      const ctx = createContext({
+        header: vi.fn(() => undefined),
+      });
+      const key = defaultKeyGenerator(ctx);
+      expect(key).toBe('rate-limit:GET:/api/test:unknown');
+    });
+
+    it('should handle comma-separated IPs in x-forwarded-for', () => {
+      const ctx = createContext({
+        header: vi.fn((key: string) => {
+          if (key === 'x-forwarded-for')
+            return '192.168.1.1, 10.0.0.1, 172.16.0.1';
+          return undefined;
+        }),
+      });
+      const key = defaultKeyGenerator(ctx);
+      expect(key).toBe('rate-limit:GET:/api/test:192.168.1.1');
+    });
+  });
+
+  describe('checkRateLimit', () => {
+    it('should allow requests within the limit', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 5,
+        windowMs: 60000, // 1 minute
+        cache,
+      };
+      const ctx = createContext();
+
+      const info = await checkRateLimit(config, ctx);
+      expect(info.count).toBe(1);
+      expect(info.remaining).toBe(4);
+      expect(info.limit).toBe(5);
+    });
+
+    it('should increment count on subsequent requests', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 5,
+        windowMs: 60000,
+        cache,
+      };
+      const ctx = createContext();
+
+      // First request
+      await checkRateLimit(config, ctx);
+
+      // Second request
+      const info = await checkRateLimit(config, ctx);
+      expect(info.count).toBe(2);
+      expect(info.remaining).toBe(3);
+    });
+
+    it('should throw TooManyRequestsError when limit is exceeded', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 2,
+        windowMs: 60000,
+        cache,
+      };
+      const ctx = createContext();
+
+      // Use up the limit
+      await checkRateLimit(config, ctx);
+      await checkRateLimit(config, ctx);
+
+      // Exceed the limit
+      await expect(checkRateLimit(config, ctx)).rejects.toThrow(
+        TooManyRequestsError,
+      );
+    });
+
+    it('should use custom message when rate limit is exceeded', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 1,
+        windowMs: 60000,
+        cache,
+        message: 'Custom rate limit message',
+      };
+      const ctx = createContext();
+
+      await checkRateLimit(config, ctx);
+
+      try {
+        await checkRateLimit(config, ctx);
+      } catch (error) {
+        expect(error).toBeInstanceOf(TooManyRequestsError);
+        expect((error as TooManyRequestsError).message).toBe(
+          'Custom rate limit message',
+        );
+      }
+    });
+
+    it('should skip rate limiting when skip function returns true', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 1,
+        windowMs: 60000,
+        cache,
+        skip: vi.fn().mockResolvedValue(true),
+      };
+      const ctx = createContext();
+
+      // Should not throw even though limit would be exceeded
+      await checkRateLimit(config, ctx);
+      const info = await checkRateLimit(config, ctx);
+
+      expect(info.count).toBe(0);
+      expect(info.remaining).toBe(1);
+      expect(config.skip).toHaveBeenCalledWith(ctx);
+    });
+
+    it('should call custom handler when rate limit is exceeded', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const handler = vi.fn();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 1,
+        windowMs: 60000,
+        cache,
+        handler,
+      };
+      const ctx = createContext();
+
+      await checkRateLimit(config, ctx);
+
+      try {
+        await checkRateLimit(config, ctx);
+      } catch {
+        // Expected to throw
+      }
+
+      expect(handler).toHaveBeenCalledWith(
+        ctx,
+        expect.objectContaining({
+          count: 2,
+          limit: 1,
+          remaining: 0,
+        }),
+      );
+    });
+
+    it('should use custom key generator', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const keyGenerator = vi.fn().mockReturnValue('custom-key');
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 5,
+        windowMs: 60000,
+        cache,
+        keyGenerator,
+      };
+      const ctx = createContext();
+
+      await checkRateLimit(config, ctx);
+
+      expect(keyGenerator).toHaveBeenCalledWith(ctx);
+    });
+
+    it('should reset count after window expires', async () => {
+      const cache = new InMemoryCache<RateLimitData>();
+      const config: RateLimitConfig<RateLimitData> = {
+        limit: 2,
+        windowMs: 100, // 100ms window
+        cache,
+      };
+      const ctx = createContext();
+
+      // Use up the limit
+      await checkRateLimit(config, ctx);
+      await checkRateLimit(config, ctx);
+
+      // Wait for window to expire
+      await new Promise((resolve) => setTimeout(resolve, 150));
+
+      // Should be allowed again
+      const info = await checkRateLimit(config, ctx);
+      expect(info.count).toBe(1);
+      expect(info.remaining).toBe(1);
+    });
+  });
+
+  describe('getRateLimitHeaders', () => {
+    it('should generate standard headers by default', () => {
+      const info = {
+        count: 3,
+        limit: 10,
+        remaining: 7,
+        resetTime: Date.now() + 60000,
+        retryAfter: 60000,
+      };
+      const config: RateLimitConfig = {
+        limit: 10,
+        windowMs: 60000,
+        cache: new InMemoryCache(),
+      };
+
+      const headers = getRateLimitHeaders(info, config);
+
+      expect(headers['X-RateLimit-Limit']).toBe('10');
+      expect(headers['X-RateLimit-Remaining']).toBe('7');
+      expect(headers['X-RateLimit-Reset']).toMatch(/^\d{4}-\d{2}-\d{2}T/);
+      expect(headers['Retry-After']).toBeUndefined();
+    });
+
+    it('should include Retry-After when limit is exceeded', () => {
+      const info = {
+        count: 11,
+        limit: 10,
+        remaining: 0,
+        resetTime: Date.now() + 60000,
+        retryAfter: 60000,
+      };
+      const config: RateLimitConfig = {
+        limit: 10,
+        windowMs: 60000,
+        cache: new InMemoryCache(),
+      };
+
+      const headers = getRateLimitHeaders(info, config);
+
+      expect(headers['Retry-After']).toBe('60');
+    });
+
+    it('should not include standard headers when disabled', () => {
+      const info = {
+        count: 3,
+        limit: 10,
+        remaining: 7,
+        resetTime: Date.now() + 60000,
+        retryAfter: 60000,
+      };
+      const config: RateLimitConfig = {
+        limit: 10,
+        windowMs: 60000,
+        cache: new InMemoryCache(),
+        standardHeaders: false,
+      };
+
+      const headers = getRateLimitHeaders(info, config);
+
+      expect(headers['X-RateLimit-Limit']).toBeUndefined();
+      expect(headers['X-RateLimit-Remaining']).toBeUndefined();
+      expect(headers['X-RateLimit-Reset']).toBeUndefined();
+    });
+
+    it('should include legacy headers when enabled', () => {
+      const info = {
+        count: 3,
+        limit: 10,
+        remaining: 7,
+        resetTime: Date.now() + 60000,
+        retryAfter: 60000,
+      };
+      const config: RateLimitConfig = {
+        limit: 10,
+        windowMs: 60000,
+        cache: new InMemoryCache(),
+        legacyHeaders: true,
+      };
+
+      const headers = getRateLimitHeaders(info, config);
+
+      expect(headers['X-RateLimit-Retry-After']).toBe('60000');
+      expect(headers['X-RateLimit-Reset-After']).toBe('60');
+    });
+  });
+});

package/src/index.ts

@@ -0,0 +1,290 @@
+import type { Cache } from '@geekmidas/cache';
+import type { Logger } from '@geekmidas/logger';
+import type { Service, ServiceRecord } from '@geekmidas/services';
+
+/**
+ * Error thrown when rate limit is exceeded
+ */
+export class TooManyRequestsError extends Error {
+  public readonly statusCode = 429;
+  public readonly retryAfter?: number;
+
+  constructor(message?: string, retryAfter?: number) {
+    super(message || 'Too many requests, please try again later.');
+    this.name = 'TooManyRequestsError';
+    this.retryAfter = retryAfter;
+  }
+}
+
+/**
+ * Rate limit configuration for an endpoint
+ */
+export interface RateLimitConfig<T = RateLimitData> {
+  /**
+   * Maximum number of requests allowed in the window
+   */
+  limit: number;
+
+  /**
+   * Time window in milliseconds
+   */
+  windowMs: number;
+
+  /**
+   * Cache instance to store rate limit data
+   */
+  cache: Cache<T>;
+
+  /**
+   * Key generator function to identify clients
+   * Defaults to using IP address
+   */
+  keyGenerator?: RateLimitKeyGenerator;
+
+  /**
+   * Skip rate limiting for certain requests
+   */
+  skip?: RateLimitSkipFn;
+
+  /**
+   * Optional message to return when rate limit is exceeded
+   */
+  message?: string;
+
+  /**
+   * Optional custom handler when rate limit is exceeded
+   */
+  handler?: RateLimitExceededHandler;
+
+  /**
+   * Whether to include rate limit headers in response
+   * @default true
+   */
+  standardHeaders?: boolean;
+
+  /**
+   * Whether to include legacy rate limit headers
+   * @default false
+   */
+  legacyHeaders?: boolean;
+}
+
+/**
+ * Context for rate limiting decisions
+ */
+export interface RateLimitContext<
+  TServices extends Service[] = [],
+  TLogger extends Logger = Logger,
+  TSession = unknown,
+> {
+  header: (key: string) => string | undefined;
+  services: ServiceRecord<TServices>;
+  logger: TLogger;
+  session: TSession;
+  path: string;
+  method: string;
+}
+
+/**
+ * Function to generate a unique key for rate limiting
+ */
+export type RateLimitKeyGenerator<
+  TServices extends Service[] = [],
+  TLogger extends Logger = Logger,
+  TSession = unknown,
+> = (
+  ctx: RateLimitContext<TServices, TLogger, TSession>,
+) => string | Promise<string>;
+
+/**
+ * Function to determine if rate limiting should be skipped
+ */
+export type RateLimitSkipFn<
+  TServices extends Service[] = [],
+  TLogger extends Logger = Logger,
+  TSession = unknown,
+> = (
+  ctx: RateLimitContext<TServices, TLogger, TSession>,
+) => boolean | Promise<boolean>;
+
+/**
+ * Handler for when rate limit is exceeded
+ */
+export type RateLimitExceededHandler<
+  TServices extends Service[] = [],
+  TLogger extends Logger = Logger,
+  TSession = unknown,
+> = (
+  ctx: RateLimitContext<TServices, TLogger, TSession>,
+  info: RateLimitInfo,
+) => void | Promise<void>;
+
+/**
+ * Information about current rate limit status
+ */
+export interface RateLimitInfo {
+  /**
+   * Current request count in the window
+   */
+  count: number;
+
+  /**
+   * Maximum allowed requests
+   */
+  limit: number;
+
+  /**
+   * Remaining requests allowed
+   */
+  remaining: number;
+
+  /**
+   * Time when the window resets (Unix timestamp)
+   */
+  resetTime: number;
+
+  /**
+   * Time until reset in milliseconds
+   */
+  retryAfter: number;
+}
+
+/**
+ * Headers to be set on responses
+ */
+export interface RateLimitHeaders {
+  'X-RateLimit-Limit'?: string;
+  'X-RateLimit-Remaining'?: string;
+  'X-RateLimit-Reset'?: string;
+  'Retry-After'?: string;
+  'X-RateLimit-Retry-After'?: string;
+  'X-RateLimit-Reset-After'?: string;
+}
+
+/**
+ * Data stored in cache for rate limiting
+ */
+export interface RateLimitData {
+  count: number;
+  resetTime: number;
+}
+
+/**
+ * Default key generator using IP address
+ */
+export const defaultKeyGenerator: RateLimitKeyGenerator = (ctx) => {
+  // Try various headers for IP address
+  const ip =
+    ctx.header('x-forwarded-for')?.split(',')[0]?.trim() ||
+    ctx.header('x-real-ip') ||
+    ctx.header('x-client-ip') ||
+    ctx.header('cf-connecting-ip') ||
+    'unknown';
+
+  return `rate-limit:${ctx.method}:${ctx.path}:${ip}`;
+};
+
+/**
+ * Check rate limit and throw error if exceeded
+ */
+export async function checkRateLimit<
+  TServices extends Service[] = [],
+  TLogger extends Logger = Logger,
+  TSession = unknown,
+>(
+  config: RateLimitConfig<RateLimitData>,
+  ctx: RateLimitContext<TServices, TLogger, TSession>,
+): Promise<RateLimitInfo> {
+  // Check if we should skip rate limiting
+  if (config.skip && (await config.skip(ctx))) {
+    return {
+      count: 0,
+      limit: config.limit,
+      remaining: config.limit,
+      resetTime: Date.now() + config.windowMs,
+      retryAfter: config.windowMs,
+    };
+  }
+
+  // Generate key for this request
+  const keyGenerator = config.keyGenerator || defaultKeyGenerator;
+  const key = await keyGenerator(ctx);
+
+  // Get current data from cache
+  const now = Date.now();
+  let data = await config.cache.get(key);
+
+  // If no data or window expired, create new entry
+  if (!data || data.resetTime <= now) {
+    const resetTime = now + config.windowMs;
+    data = { count: 1, resetTime };
+
+    // Store with TTL matching the window
+    const ttlSeconds = Math.ceil(config.windowMs / 1000);
+    await config.cache.set(key, data, ttlSeconds);
+  } else {
+    // Increment count
+    data.count++;
+
+    // Calculate remaining TTL
+    const remainingMs = data.resetTime - now;
+    const ttlSeconds = Math.ceil(remainingMs / 1000);
+    await config.cache.set(key, data, ttlSeconds);
+  }
+
+  // Calculate rate limit info
+  const info: RateLimitInfo = {
+    count: data.count,
+    limit: config.limit,
+    remaining: Math.max(0, config.limit - data.count),
+    resetTime: data.resetTime,
+    retryAfter: data.resetTime - now,
+  };
+
+  // Check if limit exceeded
+  if (data.count > config.limit) {
+    // Call custom handler if provided
+    if (config.handler) {
+      await config.handler(ctx, info);
+    }
+
+    // Throw rate limit error
+    const retryAfterSeconds = Math.ceil(info.retryAfter / 1000);
+    throw new TooManyRequestsError(
+      config.message || 'Too many requests, please try again later.',
+      retryAfterSeconds,
+    );
+  }
+
+  return info;
+}
+
+/**
+ * Generate rate limit headers
+ */
+export function getRateLimitHeaders(
+  info: RateLimitInfo,
+  config: RateLimitConfig,
+): RateLimitHeaders {
+  const headers: RateLimitHeaders = {};
+
+  if (config.standardHeaders !== false) {
+    headers['X-RateLimit-Limit'] = info.limit.toString();
+    headers['X-RateLimit-Remaining'] = info.remaining.toString();
+    headers['X-RateLimit-Reset'] = new Date(info.resetTime).toISOString();
+  }
+
+  if (config.legacyHeaders) {
+    headers['X-RateLimit-Retry-After'] = info.retryAfter.toString();
+    headers['X-RateLimit-Reset-After'] = Math.ceil(
+      info.retryAfter / 1000,
+    ).toString();
+  }
+
+  // Always set Retry-After when limit is exceeded
+  if (info.remaining === 0) {
+    headers['Retry-After'] = Math.ceil(info.retryAfter / 1000).toString();
+  }
+
+  return headers;
+}