diff --git a/src/app/api/oauth/register/route.test.ts b/src/app/api/oauth/register/route.test.ts new file mode 100644 index 000000000..ff26e8185 --- /dev/null +++ b/src/app/api/oauth/register/route.test.ts @@ -0,0 +1,70 @@ +import { beforeEach, expect, test, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + parseRequest: vi.fn(), + createOauthClient: vi.fn(), +})); + +vi.mock('@/lib/request', () => ({ + parseRequest: mocks.parseRequest, +})); + +vi.mock('@/queries/prisma/oauth', () => ({ + createOauthClient: mocks.createOauthClient, +})); + +vi.mock('@/lib/crypto', () => ({ + uuid: () => 'client-id', +})); + +vi.mock('@/lib/redis', () => ({ + default: { enabled: false, client: {} }, +})); + +import { resetRateLimits } from '@/lib/rate-limit'; +import { POST } from './route'; + +const body = { + client_name: 'Test client', + redirect_uris: ['http://127.0.0.1:3000/callback'], +}; + +function register(ip: string) { + return POST( + new Request('https://umami.example/api/oauth/register', { + method: 'POST', + headers: { 'x-forwarded-for': ip, 'content-type': 'application/json' }, + body: JSON.stringify(body), + }), + ); +} + +beforeEach(() => { + resetRateLimits(); + delete process.env.OAUTH_DISABLED; + delete process.env.OAUTH_DISABLE_DCR; + delete process.env.DISABLE_LOGIN; + + mocks.parseRequest.mockReset(); + mocks.createOauthClient.mockReset(); + + mocks.parseRequest.mockResolvedValue({ body }); + mocks.createOauthClient.mockImplementation(async input => ({ + ...input, + createdAt: new Date(), + })); +}); + +test('limits registrations per source address without affecting other sources', async () => { + for (let i = 0; i < 20; i++) { + expect((await register('203.0.113.1')).status).toBe(201); + } + + const limited = await register('203.0.113.1'); + + expect(limited.status).toBe(429); + await expect(limited.json()).resolves.toMatchObject({ error: 'temporarily_unavailable' }); + + expect((await register('203.0.113.2')).status).toBe(201); + expect(mocks.createOauthClient).toHaveBeenCalledTimes(21); +}); diff --git a/src/app/api/oauth/register/route.ts b/src/app/api/oauth/register/route.ts index e14297397..f0e678c48 100644 --- a/src/app/api/oauth/register/route.ts +++ b/src/app/api/oauth/register/route.ts @@ -4,31 +4,14 @@ import { isDynamicRegistrationEnabled } from '@/lib/oauth/config'; import { OAuthError, oauthErrorResponse, oauthJson } from '@/lib/oauth/errors'; import { corsPreflight } from '@/lib/oauth/metadata'; import { isAcceptableRedirectUri } from '@/lib/oauth/redirect'; -import redis from '@/lib/redis'; +import { checkRateLimit } from '@/lib/rate-limit'; import { parseRequest } from '@/lib/request'; import { notFound } from '@/lib/response'; import { createOauthClient } from '@/queries/prisma/oauth'; import { clientRegistrationRequestSchema } from '../schema'; const REGISTRATIONS_PER_HOUR = 20; -const REGISTRATION_WINDOW_MS = 60 * 60 * 1000; -let registrationWindow = { count: 0, expiresAt: 0 }; - -function isLocalRegistrationLimited() { - const now = Date.now(); - - if (now >= registrationWindow.expiresAt) { - registrationWindow = { count: 0, expiresAt: now + REGISTRATION_WINDOW_MS }; - } - - // A process-wide cap keeps both registrations and limiter memory bounded without Redis. - if (registrationWindow.count >= REGISTRATIONS_PER_HOUR) { - return true; - } - - registrationWindow.count++; - return false; -} +const REGISTRATION_WINDOW_SECONDS = 60 * 60; export function OPTIONS() { return corsPreflight(); @@ -55,15 +38,15 @@ export async function POST(request: Request) { } try { - const limited = redis.enabled - ? await redis.client.rateLimit( - `oauth:register:${getIpAddress(request.headers) ?? 'unknown'}`, - REGISTRATIONS_PER_HOUR + 1, - REGISTRATION_WINDOW_MS / 1000, - ) - : isLocalRegistrationLimited(); + // Budget is per source address so one requester cannot exhaust registration for every other + // client. Enforced via Redis when available, otherwise a bounded per-key in-process window. + const allowed = await checkRateLimit( + `oauth:register:${getIpAddress(request.headers) ?? 'unknown'}`, + REGISTRATIONS_PER_HOUR, + REGISTRATION_WINDOW_SECONDS, + ); - if (limited) { + if (!allowed) { throw new OAuthError('temporarily_unavailable', 'Too many registrations. Try again later.', { status: 429, }); diff --git a/src/lib/mcp/rate-limit.ts b/src/lib/mcp/rate-limit.ts index 2895b43fe..24eaeabe2 100644 --- a/src/lib/mcp/rate-limit.ts +++ b/src/lib/mcp/rate-limit.ts @@ -1,9 +1,7 @@ -import redis from '@/lib/redis'; +import { checkRateLimit, resetRateLimits } from '@/lib/rate-limit'; export const MCP_RATE_LIMIT_WINDOW_SECONDS = 60; -const memory = new Map(); - export function getMcpRateLimit() { const configured = Number(process.env.MCP_RATE_LIMIT); @@ -12,43 +10,17 @@ export function getMcpRateLimit() { /** * Per-user, per-client limit for MCP requests. Keyed by identity rather than IP because remote - * MCP platforms share egress addresses. Uses Redis when available, otherwise a process-local - * fixed window (sufficient for single-instance self-hosted deployments). + * MCP platforms share egress addresses. Resolves `true` while the request is within the limit. */ -export async function checkMcpRateLimit( +export function checkMcpRateLimit( userId: string, clientId: string, limit = getMcpRateLimit(), windowSeconds = MCP_RATE_LIMIT_WINDOW_SECONDS, ): Promise { - const key = `mcp:rate:${userId}:${clientId}`; - - if (redis.enabled) { - return redis.client.rateLimit(key, limit, windowSeconds); - } - - const now = Date.now(); - const entry = memory.get(key); - - if (!entry || entry.resetAt <= now) { - memory.set(key, { count: 1, resetAt: now + windowSeconds * 1000 }); - - if (memory.size > 10_000) { - for (const [k, v] of memory) { - if (v.resetAt <= now) { - memory.delete(k); - } - } - } - - return true; - } - - entry.count += 1; - - return entry.count <= limit; + return checkRateLimit(`mcp:rate:${userId}:${clientId}`, limit, windowSeconds); } export function resetMcpRateLimits() { - memory.clear(); + resetRateLimits(); } diff --git a/src/lib/rate-limit.test.ts b/src/lib/rate-limit.test.ts new file mode 100644 index 000000000..36258bdfc --- /dev/null +++ b/src/lib/rate-limit.test.ts @@ -0,0 +1,73 @@ +import { beforeEach, describe, expect, test, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + redis: { enabled: false, client: { rateLimit: vi.fn() } }, +})); + +vi.mock('@/lib/redis', () => ({ + default: mocks.redis, +})); + +import { checkRateLimit, RATE_LIMIT_MAX_KEYS, resetRateLimits } from './rate-limit'; + +beforeEach(() => { + resetRateLimits(); + mocks.redis.enabled = false; + mocks.redis.client.rateLimit.mockReset(); +}); + +describe('checkRateLimit (in-memory)', () => { + test('gives every key its own budget', async () => { + for (let i = 0; i < 3; i++) { + await expect(checkRateLimit('a', 3, 60)).resolves.toBe(true); + } + + await expect(checkRateLimit('a', 3, 60)).resolves.toBe(false); + await expect(checkRateLimit('b', 3, 60)).resolves.toBe(true); + }); + + test('resets after the window', async () => { + vi.useFakeTimers(); + + try { + await expect(checkRateLimit('a', 1, 1)).resolves.toBe(true); + await expect(checkRateLimit('a', 1, 1)).resolves.toBe(false); + + vi.advanceTimersByTime(1100); + + await expect(checkRateLimit('a', 1, 1)).resolves.toBe(true); + } finally { + vi.useRealTimers(); + } + }); + + test('keeps tracking new keys once the key cap is reached', async () => { + for (let i = 0; i < RATE_LIMIT_MAX_KEYS + 50; i++) { + await expect(checkRateLimit(`key-${i}`, 1, 60)).resolves.toBe(true); + } + + // The newest key is still tracked and limited independently. + await expect(checkRateLimit(`key-${RATE_LIMIT_MAX_KEYS + 49}`, 1, 60)).resolves.toBe(false); + await expect(checkRateLimit('fresh', 1, 60)).resolves.toBe(true); + }); +}); + +describe('checkRateLimit (redis)', () => { + test('allows exactly `limit` hits per window', async () => { + mocks.redis.enabled = true; + + let count = 0; + mocks.redis.client.rateLimit.mockImplementation(async (_key: string, threshold: number) => { + count += 1; + + return count >= threshold; + }); + + for (let i = 0; i < 3; i++) { + await expect(checkRateLimit('a', 3, 60)).resolves.toBe(true); + } + + await expect(checkRateLimit('a', 3, 60)).resolves.toBe(false); + expect(mocks.redis.client.rateLimit).toHaveBeenCalledWith('a', 4, 60); + }); +}); diff --git a/src/lib/rate-limit.ts b/src/lib/rate-limit.ts new file mode 100644 index 000000000..a6e3300fd --- /dev/null +++ b/src/lib/rate-limit.ts @@ -0,0 +1,71 @@ +import redis from '@/lib/redis'; + +/** Upper bound on tracked keys when Redis is unavailable, so limiter memory stays bounded. */ +export const RATE_LIMIT_MAX_KEYS = 10_000; + +const memory = new Map(); + +function evictMemoryEntries(now: number) { + for (const [key, entry] of memory) { + if (entry.resetAt <= now) { + memory.delete(key); + } + } + + // Still full after dropping expired windows: shed the oldest windows. This fails open for the + // evicted keys (their budget restarts) rather than denying anyone, which keeps a flood of + // distinct keys from turning the limiter itself into a denial-of-service vector. + while (memory.size >= RATE_LIMIT_MAX_KEYS) { + const oldest = memory.keys().next().value; + + if (oldest === undefined) { + break; + } + + memory.delete(oldest); + } +} + +/** + * Fixed-window rate limiter keyed by a caller-supplied identity (IP address, user/client, ...). + * Resolves `true` while the key is within `limit` hits for the current window. + * + * Uses Redis when configured so the window is shared across instances. Otherwise falls back to a + * bounded process-local map in which every key still gets its own independent budget, so one + * source can never exhaust the limit for another. + */ +export async function checkRateLimit( + key: string, + limit: number, + windowSeconds: number, +): Promise { + if (redis.enabled) { + // `rateLimit` resolves true once the counter reaches its threshold, so pass `limit + 1` to + // allow exactly `limit` hits per window. + return !(await redis.client.rateLimit(key, limit + 1, windowSeconds)); + } + + const now = Date.now(); + const entry = memory.get(key); + + if (entry && entry.resetAt > now) { + entry.count += 1; + + return entry.count <= limit; + } + + // Delete before re-inserting so the key moves to the newest position for eviction ordering. + memory.delete(key); + + if (memory.size >= RATE_LIMIT_MAX_KEYS) { + evictMemoryEntries(now); + } + + memory.set(key, { count: 1, resetAt: now + windowSeconds * 1000 }); + + return limit >= 1; +} + +export function resetRateLimits() { + memory.clear(); +}