Partition pre-authentication rate limits by API-key public ID

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<string, TokenBucket>();
 
   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", {

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);
+  });
 });

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);
   });
 });