From 878b37a73d10feb15ff53f59045fca4b0c0dd5ce Mon Sep 17 00:00:00 2001
From: Ryder <ryder777777@users.noreply.github.com>
Date: Wed, 21 Oct 2026 14:32:50 +0000
Subject: [PATCH] feat(security): add ApiKeyRateLimitGuard with Retry-After
 headers (#413)

---
 .../guards/api-key-rate-limit.guard.spec.ts   | 48 +++++++++++
 src/common/guards/api-key-rate-limit.guard.ts | 82 +++++++++++++++++++
 2 files changed, 130 insertions(+)
 create mode 100644 src/common/guards/api-key-rate-limit.guard.spec.ts
 create mode 100644 src/common/guards/api-key-rate-limit.guard.ts

diff --git a/src/common/guards/api-key-rate-limit.guard.spec.ts b/src/common/guards/api-key-rate-limit.guard.spec.ts
new file mode 100644
index 00000000..d8607f58
--- /dev/null
+++ b/src/common/guards/api-key-rate-limit.guard.spec.ts
@@ -0,0 +1,48 @@
+import { describe, it, expect, vi } from 'vitest';
+import { ExecutionContext, HttpException, HttpStatus } from '@nestjs/common';
+import { ApiKeyRateLimitGuard, ApiKeyRateLimitStore } from './api-key-rate-limit.guard';
+
+function makeContext(ip: string, apiKey?: string, setHeader = vi.fn()): ExecutionContext {
+  return {
+    switchToHttp: () => ({
+      getRequest: () => ({
+        ip,
+        headers: apiKey ? { 'x-api-key': apiKey } : {},
+      }),
+      getResponse: () => ({ setHeader }),
+    }),
+  } as unknown as ExecutionContext;
+}
+
+describe('ApiKeyRateLimitGuard (#413)', () => {
+  it('allows requests within configured IP and API key thresholds', async () => {
+    const store: ApiKeyRateLimitStore = {
+      increment: vi.fn().mockResolvedValue({ count: 2, ttlSeconds: 55 }),
+    };
+    const guard = new ApiKeyRateLimitGuard(store, {
+      windowSeconds: 60,
+      maxRequestsPerIp: 10,
+      maxRequestsPerApiKey: 20,
+    });
+
+    expect(await guard.canActivate(makeContext('10.0.0.1', 'ak_live_1234567890'))).toBe(true);
+    expect(store.increment).toHaveBeenCalledTimes(2);
+  });
+
+  it('throws HTTP 429 and sets Retry-After header when IP limit is exceeded', async () => {
+    const setHeader = vi.fn();
+    const store: ApiKeyRateLimitStore = {
+      increment: vi.fn().mockResolvedValue({ count: 11, ttlSeconds: 42 }),
+    };
+    const guard = new ApiKeyRateLimitGuard(store, {
+      windowSeconds: 60,
+      maxRequestsPerIp: 10,
+    });
+
+    await expect(guard.canActivate(makeContext('10.0.0.1', undefined, setHeader))).rejects.toSatisfy(
+      (err: unknown) =>
+        err instanceof HttpException && err.getStatus() === HttpStatus.TOO_MANY_REQUESTS,
+    );
+    expect(setHeader).toHaveBeenCalledWith('Retry-After', '42');
+  });
+});
diff --git a/src/common/guards/api-key-rate-limit.guard.ts b/src/common/guards/api-key-rate-limit.guard.ts
new file mode 100644
index 00000000..c57802fd
--- /dev/null
+++ b/src/common/guards/api-key-rate-limit.guard.ts
@@ -0,0 +1,82 @@
+import {
+  CanActivate,
+  ExecutionContext,
+  HttpException,
+  HttpStatus,
+  Injectable,
+} from '@nestjs/common';
+
+export interface ApiKeyRateLimitStore {
+  increment(key: string, windowSeconds: number): Promise<{ count: number; ttlSeconds: number }>;
+}
+
+export interface ApiKeyRateLimitOptions {
+  windowSeconds?: number;
+  maxRequestsPerIp?: number;
+  maxRequestsPerApiKey?: number;
+}
+
+@Injectable()
+export class ApiKeyRateLimitGuard implements CanActivate {
+  private readonly windowSeconds: number;
+  private readonly maxPerIp: number;
+  private readonly maxPerKey: number;
+
+  constructor(
+    private readonly store: ApiKeyRateLimitStore,
+    options: ApiKeyRateLimitOptions = {},
+  ) {
+    this.windowSeconds = options.windowSeconds ?? 60;
+    this.maxPerIp = options.maxRequestsPerIp ?? 60;
+    this.maxPerKey = options.maxRequestsPerApiKey ?? 120;
+  }
+
+  async canActivate(context: ExecutionContext): Promise<boolean> {
+    const http = context.switchToHttp();
+    const req = http.getRequest<{
+      ip?: string;
+      headers?: Record<string, string | string[] | undefined>;
+    }>();
+    const res = http.getResponse<{ setHeader?: (name: string, value: string) => void }>();
+
+    const ip = (req.ip || 'unknown').trim();
+    const rawHeader = req.headers?.['x-api-key'];
+    const apiKeyPrefix =
+      typeof rawHeader === 'string' && rawHeader.trim().length > 0
+        ? rawHeader.trim().slice(0, 12)
+        : undefined;
+
+    const ipBucket = await this.store.increment(`rl:apikey:ip:${ip}`, this.windowSeconds);
+    if (ipBucket.count > this.maxPerIp) {
+      res?.setHeader?.('Retry-After', String(Math.max(1, ipBucket.ttlSeconds)));
+      throw new HttpException(
+        {
+          statusCode: HttpStatus.TOO_MANY_REQUESTS,
+          message: 'Too many API key authentication requests from this IP address',
+          retryAfter: Math.max(1, ipBucket.ttlSeconds),
+        },
+        HttpStatus.TOO_MANY_REQUESTS,
+      );
+    }
+
+    if (apiKeyPrefix) {
+      const keyBucket = await this.store.increment(
+        `rl:apikey:id:${apiKeyPrefix}`,
+        this.windowSeconds,
+      );
+      if (keyBucket.count > this.maxPerKey) {
+        res?.setHeader?.('Retry-After', String(Math.max(1, keyBucket.ttlSeconds)));
+        throw new HttpException(
+          {
+            statusCode: HttpStatus.TOO_MANY_REQUESTS,
+            message: 'API key rate limit exceeded',
+            retryAfter: Math.max(1, keyBucket.ttlSeconds),
+          },
+          HttpStatus.TOO_MANY_REQUESTS,
+        );
+      }
+    }
+
+    return true;
+  }
+}