diff --git a/src/common/guards/api-key-throttler.guard.spec.ts b/src/common/guards/api-key-throttler.guard.spec.ts new file mode 100644 index 0000000..7768912 --- /dev/null +++ b/src/common/guards/api-key-throttler.guard.spec.ts @@ -0,0 +1,515 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { + ApiKeyThrottlerGuard, + API_KEY_THROTTLE_LIMIT_KEY, + API_KEY_THROTTLE_TTL_KEY, +} from './api-key-throttler.guard'; +import { ExecutionContext } from '@nestjs/common'; +import { ThrottlerException } from '@nestjs/throttler'; +import { API_KEY_HEADER } from '../constants/headers'; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** Minimal mock of ApiKeyService.verify() */ +function mockApiKeyService(overrides: { verify?: ReturnType } = {}) { + return { + verify: overrides.verify ?? vi.fn(), + } as any; +} + +/** Build a mock ThrottlerStorage with an in-memory counter map. */ +function mockStorage() { + const store = new Map(); + + return { + store, + increment: vi.fn( + async ( + key: string, + _ttl: number, + limit: number, + _blockDuration: number, + _throttlerName: string, + ) => { + const existing = store.get(key); + const now = Date.now(); + if (!existing || now - existing.createdAt > _ttl) { + store.set(key, { hits: 1, createdAt: now }); + return { totalHits: 1, timeToExpire: _ttl, isBlocked: false, timeToBlockExpire: 0 }; + } + existing.hits += 1; + const timeToExpire = _ttl - (now - existing.createdAt); + const isBlocked = existing.hits > limit; + return { + totalHits: existing.hits, + timeToExpire: Math.max(0, timeToExpire), + isBlocked, + timeToBlockExpire: isBlocked ? timeToExpire : 0, + }; + }, + ), + } as any; +} + +function mockReflector(overrides: Record = {}) { + return { + getAllAndOverride: vi.fn((key: string) => overrides[key] ?? undefined), + } as any; +} + +/** Fake Express request */ +function fakeReq( + headers: Record = {}, + ip = '127.0.0.1', +) { + return { + headers, + ip, + socket: { remoteAddress: '127.0.0.1' }, + } as any; +} + +function fakeRes() { + const headers: Record = {}; + return { + setHeader: vi.fn((k: string, v: string) => { headers[k] = v; }), + get headers() { return headers; }, + header: vi.fn((k: string, v: string) => { headers[k] = String(v); }), + status: vi.fn().mockReturnThis(), + json: vi.fn(), + } as any; +} + +/** Minimal ExecutionContext for HTTP requests */ +function fakeContext(req: any, res: any): ExecutionContext { + return { + getHandler: () => () => {}, + getClass: () => class {}, + switchToHttp: () => ({ + getRequest: () => req, + getResponse: () => res, + }), + getArgByIndex: vi.fn(), + getArgs: vi.fn(), + getType: () => 'http', + switchToRpc: vi.fn(), + switchToWs: vi.fn(), + } as unknown as ExecutionContext; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe('ApiKeyThrottlerGuard', () => { + let guard: ApiKeyThrottlerGuard; + let storage: ReturnType; + let reflector: ReturnType; + let apiKeyService: ReturnType; + let req: any; + let res: any; + + beforeEach(() => { + vi.clearAllMocks(); + storage = mockStorage(); + reflector = mockReflector(); + apiKeyService = mockApiKeyService(); + + guard = new ApiKeyThrottlerGuard( + { throttlers: [{ name: 'default', ttl: 60_000, limit: 100 }] } as any, + storage, + reflector, + apiKeyService, + ); + + req = fakeReq(); + res = fakeRes(); + }); + + // ----------------------------------------------------------------------- + // getTracker — API key extraction + // ----------------------------------------------------------------------- + describe('getTracker()', () => { + it('returns org: when a valid x-api-key header is present', async () => { + const orgId = 'org-abc-123'; + apiKeyService.verify.mockResolvedValue({ organizationId: orgId }); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_test123' }); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe(`org:${orgId}`); + expect(apiKeyService.verify).toHaveBeenCalledWith('ak_live_test123'); + }); + + it('extracts Bearer token from Authorization header when x-api-key is absent', async () => { + const orgId = 'org-bearer-456'; + apiKeyService.verify.mockResolvedValue({ organizationId: orgId }); + req = fakeReq({ authorization: 'Bearer ak_live_bearer_token' }); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe(`org:${orgId}`); + expect(apiKeyService.verify).toHaveBeenCalledWith('ak_live_bearer_token'); + }); + + it('prefers x-api-key header over Authorization Bearer token', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-primary' }); + req = fakeReq({ + [API_KEY_HEADER]: 'ak_live_primary', + authorization: 'Bearer ak_live_secondary', + }); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('org:org-primary'); + expect(apiKeyService.verify).toHaveBeenCalledTimes(1); + expect(apiKeyService.verify).toHaveBeenCalledWith('ak_live_primary'); + }); + + it('falls back to IP tracking when no API key is present', async () => { + req = fakeReq({}, '192.168.1.100'); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:192.168.1.100'); + expect(apiKeyService.verify).not.toHaveBeenCalled(); + }); + + it('falls back to IP tracking when API key verification returns null', async () => { + apiKeyService.verify.mockResolvedValue(null); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_invalid' }, '10.0.0.1'); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:10.0.0.1'); + }); + + it('falls back to IP when verification throws', async () => { + apiKeyService.verify.mockRejectedValue(new Error('db timeout')); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_error' }, '172.16.0.1'); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:172.16.0.1'); + }); + + it('uses x-forwarded-for header for IP fallback when available', async () => { + req = fakeReq({ 'x-forwarded-for': '203.0.113.50, 70.41.3.18' }, '127.0.0.1'); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:203.0.113.50'); + }); + + it('returns "anonymous" when no IP information is available', async () => { + req = { headers: {}, ip: undefined, socket: {} }; + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:anonymous'); + }); + + it('ignores Authorization header when scheme is not Bearer', async () => { + req = fakeReq({ authorization: 'Basic dXNlcjpwYXNz' }); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:127.0.0.1'); + expect(apiKeyService.verify).not.toHaveBeenCalled(); + }); + + it('ignores Authorization header when token part is empty', async () => { + req = fakeReq({ authorization: 'Bearer ' }); + + const tracker = await guard['getTracker'](req); + + expect(tracker).toBe('ip:127.0.0.1'); + }); + }); + + // ----------------------------------------------------------------------- + // handleRequest — rate-limit enforcement + // ----------------------------------------------------------------------- + describe('handleRequest()', () => { + it('passes when under the limit and sets rate-limit headers', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-1' }); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_test' }); + const ctx = fakeContext(req, res); + + const result = await (guard as any).handleRequest({ + context: ctx, + limit: 100, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 100 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + + expect(result).toBe(true); + expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Limit', '100'); + expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Remaining', '99'); + expect(res.setHeader).toHaveBeenCalledWith( + 'X-RateLimit-Reset', + expect.any(String), + ); + }); + + it('throws ThrottlerException when limit is exceeded', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-2' }); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_burst' }); + const ctx = fakeContext(req, res); + + // Exhaust the limit + for (let i = 0; i < 100; i++) { + await (guard as any).handleRequest({ + context: ctx, + limit: 100, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 100 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + } + + // 101st request should throw + await expect( + (guard as any).handleRequest({ + context: ctx, + limit: 100, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 100 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }), + ).rejects.toThrow(ThrottlerException); + }); + + it('skips enforcement for non-default throttler names', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-3' }); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_skip' }); + const ctx = fakeContext(req, res); + + const result = await (guard as any).handleRequest({ + context: ctx, + limit: 10, + ttl: 60_000, + throttler: { name: 'auth', ttl: 60_000, limit: 10 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + + expect(result).toBe(true); + // Storage should NOT have been called + expect(storage.increment).not.toHaveBeenCalled(); + }); + + it('applies per-route limit override from metadata', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-4' }); + req = fakeReq({ [API_KEY_HEADER]: 'ak_live_override' }); + reflector = mockReflector({ + [API_KEY_THROTTLE_LIMIT_KEY]: 500, + }); + // Re-create guard with the new reflector + guard = new ApiKeyThrottlerGuard( + { throttlers: [{ name: 'default', ttl: 60_000, limit: 100 }] } as any, + storage, + reflector, + apiKeyService, + ); + + const ctx = fakeContext(req, res); + + await (guard as any).handleRequest({ + context: ctx, + limit: 100, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 100 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + + expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Limit', '500'); + expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Remaining', '499'); + }); + }); + + // ----------------------------------------------------------------------- + // Independent org tracking + // ----------------------------------------------------------------------- + describe('independent org tracking', () => { + it('tracks different organizations independently', async () => { + apiKeyService.verify.mockImplementation(async (key: string) => { + if (key === 'ak_live_orgA') return { organizationId: 'org-A' }; + if (key === 'ak_live_orgB') return { organizationId: 'org-B' }; + return null; + }); + + // Exhaust org A + const reqA = fakeReq({ [API_KEY_HEADER]: 'ak_live_orgA' }); + const resA = fakeRes(); + const ctxA = fakeContext(reqA, resA); + + for (let i = 0; i < 5; i++) { + await (guard as any).handleRequest({ + context: ctxA, + limit: 5, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 5 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + } + + // Org A should now be blocked on the 6th request + await expect( + (guard as any).handleRequest({ + context: ctxA, + limit: 5, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 5 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }), + ).rejects.toThrow(ThrottlerException); + + // Org B should still have a clean slate + const reqB = fakeReq({ [API_KEY_HEADER]: 'ak_live_orgB' }); + const resB = fakeRes(); + const ctxB = fakeContext(reqB, resB); + + const result = await (guard as any).handleRequest({ + context: ctxB, + limit: 5, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 5 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + + expect(result).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // API key vs IP fallback tracking independence + // ----------------------------------------------------------------------- + describe('api key vs ip fallback', () => { + it('tracks API key requests separately from IP requests', async () => { + apiKeyService.verify.mockResolvedValue({ organizationId: 'org-key' }); + + // Exhaust the IP-based counter (no API key) + const reqIp = fakeReq({}, '10.0.0.50'); + const resIp = fakeRes(); + const ctxIp = fakeContext(reqIp, resIp); + + for (let i = 0; i < 3; i++) { + await (guard as any).handleRequest({ + context: ctxIp, + limit: 3, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 3 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + } + + // IP tracker is now blocked + await expect( + (guard as any).handleRequest({ + context: ctxIp, + limit: 3, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 3 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }), + ).rejects.toThrow(ThrottlerException); + + // API key from same IP should still pass (different tracker key) + const reqKey = fakeReq({ [API_KEY_HEADER]: 'ak_live_orgkey' }, '10.0.0.50'); + const resKey = fakeRes(); + const ctxKey = fakeContext(reqKey, resKey); + + const result = await (guard as any).handleRequest({ + context: ctxKey, + limit: 3, + ttl: 60_000, + throttler: { name: 'default', ttl: 60_000, limit: 3 }, + blockDuration: 0, + getTracker: guard['getTracker'].bind(guard), + generateKey: guard['generateKey'].bind(guard), + }); + + expect(result).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Metadata keys export + // ----------------------------------------------------------------------- + describe('metadata keys', () => { + it('exports API_KEY_THROTTLE_LIMIT_KEY', () => { + expect(API_KEY_THROTTLE_LIMIT_KEY).toBe('astroid:apiKeyThrottleLimit'); + }); + + it('exports API_KEY_THROTTLE_TTL_KEY', () => { + expect(API_KEY_THROTTLE_TTL_KEY).toBe('astroid:apiKeyThrottleTtl'); + }); + }); + + // ----------------------------------------------------------------------- + // Error response format + // ----------------------------------------------------------------------- + describe('throwThrottlingException()', () => { + it('throws a ThrottlerException with structured message', async () => { + await expect( + (guard as any).throwThrottlingException({}), + ).rejects.toThrow(ThrottlerException); + }); + }); + + // ----------------------------------------------------------------------- + // extractBearerToken (private) + // ----------------------------------------------------------------------- + describe('extractBearerToken()', () => { + it('extracts token from Bearer scheme', () => { + const result = (guard as any).extractBearerToken({ + headers: { authorization: 'Bearer abc123' }, + }); + expect(result).toBe('abc123'); + }); + + it('returns undefined for non-Bearer schemes', () => { + const result = (guard as any).extractBearerToken({ + headers: { authorization: 'Basic dXNlcjpwYXNz' }, + }); + expect(result).toBeUndefined(); + }); + + it('returns undefined when authorization header is missing', () => { + const result = (guard as any).extractBearerToken({ + headers: {}, + }); + expect(result).toBeUndefined(); + }); + + it('returns undefined when Bearer token is empty', () => { + const result = (guard as any).extractBearerToken({ + headers: { authorization: 'Bearer ' }, + }); + expect(result).toBeUndefined(); + }); + }); +}); diff --git a/src/common/guards/api-key-throttler.guard.ts b/src/common/guards/api-key-throttler.guard.ts new file mode 100644 index 0000000..e1f2890 --- /dev/null +++ b/src/common/guards/api-key-throttler.guard.ts @@ -0,0 +1,202 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { ThrottlerException, ThrottlerGuard, ThrottlerRequest } from '@nestjs/throttler'; +import { Reflector } from '@nestjs/core'; +import { Request } from 'express'; +import { InjectThrottlerOptions, InjectThrottlerStorage } from '@nestjs/throttler'; +import { ThrottlerModuleOptions, ThrottlerStorage } from '@nestjs/throttler'; +import { API_KEY_HEADER } from '../constants/headers'; +import { ApiKeyService } from '../../modules/developer/api-key.service'; + +/** + * Metadata key for per-route limit overrides. + * Set via `@SetMetadata(API_KEY_THROTTLE_LIMIT_KEY, 200)` to grant an org + * a custom quota above the default 100 req/min. + */ +export const API_KEY_THROTTLE_LIMIT_KEY = 'astroid:apiKeyThrottleLimit'; + +/** + * Metadata key for per-route TTL overrides (seconds). + */ +export const API_KEY_THROTTLE_TTL_KEY = 'astroid:apiKeyThrottleTtl'; + +/** + * Custom Throttler Guard that enforces per-organization rate limits keyed to + * verified API keys rather than client IP addresses. + * + * ## How it works + * + * 1. Extracts the raw API key from the `x-api-key` header (primary) or the + * `Authorization` header (secondary, expects `Bearer `). + * 2. Calls `ApiKeyService.verify()` to validate the key — checking revocation, + * expiry, and updating `lastUsedAt`. + * 3. On success, uses `org:` as the rate-limit tracker so all + * requests from the same organisation share a single counter. + * 4. Falls back to `ip:` when no valid API key is present. + * 5. Supports dynamic per-route overrides via `@SetMetadata` for subscription + * tiers that grant higher quotas. + * + * ## Usage + * + * ```ts + * @UseGuards(ApiKeyThrottlerGuard) + * @Controller('agents') + * export class AgentController { ... } + * ``` + * + * Or per-route: + * ```ts + * @UseGuards(ApiKeyThrottlerGuard) + * @SetMetadata(API_KEY_THROTTLE_LIMIT_KEY, 500) + * @Get('premium') + * async premiumEndpoint() { ... } + * ``` + * + * ## 429 Response + * + * Returns the standard Astroid error envelope with `RATE_LIMITED` code and the + * standard `Retry-After`, `X-RateLimit-*` headers. + */ +@Injectable() +export class ApiKeyThrottlerGuard extends ThrottlerGuard { + private readonly logger = new Logger(ApiKeyThrottlerGuard.name); + + constructor( + @InjectThrottlerOptions() + options: ThrottlerModuleOptions, + @InjectThrottlerStorage() + storageService: ThrottlerStorage, + reflector: Reflector, + private readonly apiKeyService: ApiKeyService, + ) { + super(options, storageService, reflector); + } + + /** + * Resolve the rate-limiting key for this request. + * + * Priority: + * 1. Verified API key → `org:` + * 2. IP address fallback → `ip:` + */ + protected override async getTracker(req: Record): Promise { + const request = req as unknown as Request; + + // --- 1. Try API key from x-api-key header --- + const rawApiKey = + (request.headers?.[API_KEY_HEADER] as string | undefined) ?? + this.extractBearerToken(request); + + if (rawApiKey) { + try { + const verified = await this.apiKeyService.verify(rawApiKey); + if (verified) { + return `org:${verified.organizationId}`; + } + } catch { + // Verification failed — fall through to IP fallback. + this.logger.debug('API key verification failed; falling back to IP tracking'); + } + } + + // --- 2. IP fallback --- + const forwarded = request.headers?.['x-forwarded-for']; + const rawIp = + (Array.isArray(forwarded) + ? forwarded[0] + : typeof forwarded === 'string' + ? forwarded.split(',')[0]?.trim() + : undefined) ?? + request.ip ?? + request.socket?.remoteAddress ?? + 'anonymous'; + return `ip:${rawIp}`; + } + + /** + * Override handleRequest to: + * - Apply per-route limit/TTL overrides from metadata + * - Enforce only the first throttler (our single-key limiter) + * - Set standard rate-limit headers + * - Throw a structured 429 with Retry-After + */ + protected override async handleRequest(requestProps: ThrottlerRequest): Promise { + const { context, throttler } = requestProps; + + // Only enforce the first (or a named 'api-key') throttler entry to avoid + // double-counting against tier-based throttlers. + if (throttler.name !== 'default' && throttler.name !== 'api-key') { + return true; + } + + // Read per-route overrides from metadata. + const routeLimit = this.reflector.getAllAndOverride( + API_KEY_THROTTLE_LIMIT_KEY, + [context.getHandler(), context.getClass()], + ); + const routeTtl = this.reflector.getAllAndOverride( + API_KEY_THROTTLE_TTL_KEY, + [context.getHandler(), context.getClass()], + ); + + const limit = routeLimit ?? (await this.resolveThrottlerValue(requestProps.limit, context)); + const ttl = routeTtl ?? (await this.resolveThrottlerValue(requestProps.ttl, context)); + + const { req, res } = this.getRequestResponse(context); + const tracker = await requestProps.getTracker(req, context); + const key = requestProps.generateKey(context, tracker, throttler.name); + + const { totalHits, timeToExpire, isBlocked, timeToBlockExpire } = + await this.storageService.increment(key, ttl, limit, 0, throttler.name); + + if (isBlocked) { + res.setHeader('Retry-After', String(timeToBlockExpire)); + throw new ThrottlerException( + `Rate limit exceeded. Retry after ${Math.ceil(timeToBlockExpire / 1000)}s.`, + ); + } + + // Set standard rate-limit response headers. + res.setHeader('X-RateLimit-Limit', String(limit)); + res.setHeader('X-RateLimit-Remaining', String(Math.max(0, limit - totalHits))); + res.setHeader('X-RateLimit-Reset', String(timeToExpire)); + + return true; + } + + /** + * Override the default error message to deliver a structured envelope + * compatible with AllExceptionsFilter / ErrorCode.RATE_LIMITED. + */ + protected override async throwThrottlingException( + _context: any, + _detail?: any, + ): Promise { + throw new ThrottlerException( + 'Rate limit exceeded. Please slow down your requests.', + ); + } + + /** + * Extract a Bearer token from the Authorization header. + * Returns `undefined` if the header is missing or not Bearer-scheme. + */ + private extractBearerToken(request: Request): string | undefined { + const auth = request.headers?.authorization; + if (!auth || typeof auth !== 'string') return undefined; + + const [scheme, token] = auth.split(' ', 2); + if (scheme?.toLowerCase() !== 'bearer' || !token) return undefined; + return token; + } + + /** + * Resolve a throttler value that may be a function or a plain number. + * (Mirrors the private `resolveValue` in the base ThrottlerGuard.) + */ + private async resolveThrottlerValue( + value: number | ((context: unknown) => number | Promise), + context: unknown, + ): Promise { + return typeof value === 'function' ? (value as (ctx: unknown) => number | Promise)(context) : value; + } +}