Rate limit dynamic client registration per source address
Replace the process-wide registration counter used without Redis with a shared per-key fixed-window limiter, so one requester can no longer exhaust the registration budget for every other client on the process. The in-memory fallback is bounded (expired windows are swept, then the oldest keys are shed) so memory stays constant under key churn. Also fixes checkMcpRateLimit inverting redis.rateLimit's result: it resolves true once the counter reaches its threshold, so with Redis the MCP endpoint was denying requests below the limit and allowing those above it. Both call sites now share the same limiter and semantics.
This commit is contained in:
1 parent
67449d76ef
commit
d2b4aa174e
5 files changed
+229
-60
No files matched your search
@@ -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);
|
||||
});
|
||||
@@ -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,
|
||||
});
|
||||
|
||||
@@ -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<string, { count: number; resetAt: number }>();
|
||||
|
||||
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<boolean> {
|
||||
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();
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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<string, { count: number; resetAt: number }>();
|
||||
|
||||
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<boolean> {
|
||||
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();
|
||||
}
|
||||
Reference in new issue
Block a user