diff --git a/src/authentication/middleware.ts b/src/authentication/middleware.ts index fae3d64..41e9ebc 100644 --- a/src/authentication/middleware.ts +++ b/src/authentication/middleware.ts @@ -27,7 +27,7 @@ export interface RateLimitPolicy { export class AuthenticationRateLimiter { private readonly requestsPerMinute: number; private readonly burst: number; - private bucket: TokenBucket; + private readonly buckets = new Map(); public constructor( policy: RateLimitPolicy, @@ -44,26 +44,30 @@ export class AuthenticationRateLimiter { ) { throw new Error("Authentication rate-limit policy is invalid"); } - this.bucket = { tokens: this.burst, updatedAt: this.now() }; } - public consume(): RateLimitDecision { + public consume(publicId: string): RateLimitDecision { const now = this.now(); - const elapsed = Math.max(0, now - this.bucket.updatedAt); - this.bucket.tokens = Math.min( + const bucket = this.buckets.get(publicId) ?? { + tokens: this.burst, + updatedAt: now, + }; + const elapsed = Math.max(0, now - bucket.updatedAt); + bucket.tokens = Math.min( this.burst, - this.bucket.tokens + (elapsed * this.requestsPerMinute) / 60_000, + bucket.tokens + (elapsed * this.requestsPerMinute) / 60_000, ); - this.bucket.updatedAt = now; - if (this.bucket.tokens >= 1) { - this.bucket.tokens -= 1; + bucket.updatedAt = now; + this.buckets.set(publicId, bucket); + if (bucket.tokens >= 1) { + bucket.tokens -= 1; return { allowed: true, retryAfterSeconds: 0 }; } return { allowed: false, retryAfterSeconds: Math.max( 1, - Math.ceil(((1 - this.bucket.tokens) * 60) / this.requestsPerMinute), + Math.ceil(((1 - bucket.tokens) * 60) / this.requestsPerMinute), ), }; } @@ -202,8 +206,10 @@ export function rateLimitAuthentication( return async (context, next) => { const authorization = context.req.header("authorization"); const match = /^Bearer ([^\s]+)$/.exec(authorization ?? ""); - if (match !== null && parseApiKeyToken(match[1] ?? "") !== undefined) { - const decision = limiter.consume(); + const parsed = + match === null ? undefined : parseApiKeyToken(match[1] ?? ""); + if (parsed !== undefined) { + const decision = limiter.consume(parsed.publicId); if (!decision.allowed) { const requestId = context.get("requestId"); logger.emit("request_rate_limited", { diff --git a/tests/security/authentication/bearer-authentication.test.ts b/tests/security/authentication/bearer-authentication.test.ts index be655dd..c84b8fa 100644 --- a/tests/security/authentication/bearer-authentication.test.ts +++ b/tests/security/authentication/bearer-authentication.test.ts @@ -133,11 +133,12 @@ describe("pre-authentication rate limiting", () => { authenticationBurst: 1, }, }).app; + const token = createApiKeyToken().token; const request = () => app.request("/mcp", { method: "POST", headers: { - authorization: `Bearer ${createApiKeyToken().token}`, + authorization: `Bearer ${token}`, "content-type": "application/json", host: "127.0.0.1", }, @@ -150,4 +151,37 @@ describe("pre-authentication rate limiting", () => { expect(limited.headers.get("retry-after")).toBe("60"); expect(calls).toBe(1); }); + + it("does not let unknown public ids exhaust another key's limit", async () => { + let calls = 0; + const app = createTestApplication({ + authenticator: { + authenticate: () => { + calls += 1; + return Promise.resolve(undefined); + }, + }, + rateLimit: { + accountRequestsPerMinute: 60_000, + apiKeyRequestsPerMinute: 60_000, + burst: 1000, + authenticationRequestsPerMinute: 1, + authenticationBurst: 1, + }, + }).app; + const request = (token: string) => + app.request("/mcp", { + method: "POST", + headers: { + authorization: `Bearer ${token}`, + "content-type": "application/json", + host: "127.0.0.1", + }, + body: initializeBody, + }); + + expect((await request(createApiKeyToken().token)).status).toBe(401); + expect((await request(createApiKeyToken().token)).status).toBe(401); + expect(calls).toBe(2); + }); }); diff --git a/tests/unit/authentication/api-key.test.ts b/tests/unit/authentication/api-key.test.ts index beddce4..7ec61cc 100644 --- a/tests/unit/authentication/api-key.test.ts +++ b/tests/unit/authentication/api-key.test.ts @@ -96,7 +96,7 @@ describe("account and API-key rate policy", () => { }); describe("authentication rate policy", () => { - it("caps unauthenticated database-bound attempts and refills globally", () => { + it("caps and refills database-bound attempts per public id", () => { let now = 0; const limiter = new AuthenticationRateLimiter( { @@ -109,13 +109,14 @@ describe("authentication rate policy", () => { () => now, ); - expect(limiter.consume().allowed).toBe(true); - expect(limiter.consume().allowed).toBe(true); - expect(limiter.consume()).toEqual({ + expect(limiter.consume("public-a").allowed).toBe(true); + expect(limiter.consume("public-a").allowed).toBe(true); + expect(limiter.consume("public-a")).toEqual({ allowed: false, retryAfterSeconds: 1, }); + expect(limiter.consume("public-b").allowed).toBe(true); now = 1000; - expect(limiter.consume().allowed).toBe(true); + expect(limiter.consume("public-a").allowed).toBe(true); }); });