fix(webapp): bucket restricted additional API keys by environment too

Author: carderne

.server-changes/additional-api-key-rate-limit-bucket.md

@@ -3,4 +3,4 @@ area: webapp
 type: fix
 ---
 
-Additional environment API keys now share a single rate limit bucket per environment instead of each key getting its own, so minting more keys no longer multiplies an environment's effective API rate limit. Root and legacy keys are unaffected. The bucket is keyed on the stable environment id resolved from the already-authenticated request, so no extra lookup is added.
+Additional environment API keys now share a single rate limit bucket per environment instead of each key getting its own, so minting more keys — including scope-restricted keys — no longer multiplies an environment's effective API rate limit. Root and legacy keys are unaffected. The bucket is keyed on the stable environment id (resolved scope-agnostically for bucketing only, never as an auth decision), and the result is cached per key so no extra per-request lookup is added.

apps/webapp/app/models/runtimeEnvironment.server.ts

@@ -301,6 +301,60 @@ export async function findEnvironmentByApiKeyWithResolution(
   return resolveEnvironmentByApiKey(apiKey, branchName, tx, additionalApiKeyLookupEnabled);
 }
 
+export type AdditionalApiKeyRateLimitScope = {
+  environmentId: string;
+  // Organization rate limiter override (bucket size), if configured.
+  apiRateLimiterConfig: unknown;
+};
+
+/**
+ * Resolve ONLY the environment id (and its organization's rate limiter config)
+ * for an additional API key, for RATE-LIMIT BUCKETING.
+ *
+ * Deliberately scope-agnostic: unlike `findEnvironmentByApiKey`, a
+ * scope-restricted additional key still resolves here, so every additional key
+ * for an environment shares that environment's rate limit bucket. This is NOT
+ * an authentication or authorization decision and must never be used as one —
+ * request auth still goes through the RBAC bearer controller, which enforces
+ * scopes. Revoked and expired keys are excluded so they cannot keep a bucket
+ * warm.
+ */
+export async function resolveAdditionalApiKeyRateLimitScope(
+  apiKey: string,
+  tx: PrismaClientOrTransaction = $replica
+): Promise<AdditionalApiKeyRateLimitScope | null> {
+  if (!isAdditionalApiKey(apiKey)) {
+    return null;
+  }
+
+  const now = new Date();
+
+  const match = await tx.apiKey.findFirst({
+    where: {
+      keyHash: hashApiKey(apiKey),
+      revokedAt: null,
+      OR: [{ expiresAt: null }, { expiresAt: { gt: now } }],
+    },
+    select: {
+      runtimeEnvironment: {
+        select: {
+          id: true,
+          organization: { select: { apiRateLimiterConfig: true } },
+        },
+      },
+    },
+  });
+
+  if (!match?.runtimeEnvironment) {
+    return null;
+  }
+
+  return {
+    environmentId: match.runtimeEnvironment.id,
+    apiRateLimiterConfig: match.runtimeEnvironment.organization.apiRateLimiterConfig,
+  };
+}
+
 /**
  * @deprecated We don't use public API keys (`pk_*` tokens) anymore — public
  * access goes through public JWTs (see `isPublicJWT` / `validatePublicJwtKey`).

apps/webapp/app/services/apiRateLimit.server.ts

@@ -1,6 +1,7 @@
 import { tryCatch } from "@trigger.dev/core/v3";
 import { isAdditionalApiKey } from "@trigger.dev/core/v3/apiKeys";
 import { env } from "~/env.server";
+import { resolveAdditionalApiKeyRateLimitScope } from "~/models/runtimeEnvironment.server";
 import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server";
 import { authenticateAuthorizationHeader } from "./apiAuth.server";
 import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
@@ -30,6 +31,34 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     maxItems: 1000,
   },
   limiterConfigOverride: async (authorizationValue) => {
+    const rawApiKey = authorizationValue.replace(/^Bearer /, "");
+
+    // Additional API keys (`tr_*_sk_*`) share their environment's rate limit
+    // bucket rather than each getting their own. Keying on the stable
+    // environment id (not the secret key, which can rotate) keeps a single
+    // bucket per environment no matter how many additional keys exist.
+    //
+    // Resolve scope-agnostically for bucketing: a restricted additional key
+    // authenticates at the route level via the RBAC controller (and fails
+    // closed in the legacy header auth below), but for rate limiting it must
+    // still land on its environment's shared bucket — otherwise minting many
+    // restricted keys would multiply the effective limit. This is NOT an auth
+    // decision. The whole override result is cached per key by the
+    // middleware's SWR cache, so no separate lookup or Redis mapping is needed.
+    if (isAdditionalApiKey(rawApiKey)) {
+      const scope = await resolveAdditionalApiKeyRateLimitScope(rawApiKey);
+
+      // Unknown/revoked/expired key: fall back to the default per-key bucket.
+      if (!scope) {
+        return;
+      }
+
+      return {
+        config: scope.apiRateLimiterConfig,
+        identifier: scope.environmentId,
+      };
+    }
+
     const authenticatedEnv = await authenticateAuthorizationHeader(authorizationValue, {
       allowPublicKey: true,
       allowJWT: true,
@@ -49,15 +78,10 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
       };
     }
 
-    // Additional API keys (`tr_*_sk_*`) share their environment's rate limit
-    // bucket rather than each getting their own.
-    const apiKey = authenticatedEnv.apiKey;
-    const identifier =
-      apiKey && isAdditionalApiKey(apiKey) ? authenticatedEnv.environment.id : undefined;
-
+    // Root/legacy keys keep per-key (hashed header) bucketing: they already map
+    // 1:1 to an environment.
     return {
       config: authenticatedEnv.environment.organization.apiRateLimiterConfig,
-      identifier,
     };
   },
   pathMatchers: [/^\/api/],

apps/webapp/test/authorizationRateLimitMiddleware.test.ts

@@ -222,43 +222,38 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware",
     }
   );
 
-  redisTest(
-    "should key per token when no identifier is supplied",
-    async ({ redisOptions }) => {
-      const rateLimitMiddleware = authorizationRateLimitMiddleware({
-        redis: { ...redisOptions, tlsDisabled: true },
-        keyPrefix: "test-no-identifier",
-        defaultLimiter: {
-          type: "tokenBucket",
-          refillRate: 1,
-          interval: "1m",
-          maxTokens: 1,
-        },
-        pathMatchers: [/^\/api/],
-        // Override supplies a config but no identifier: bucketing stays per-key
-        // (hashed Authorization header), the legacy behavior.
-        limiterConfigOverride: async () => ({
-          config: { type: "tokenBucket", refillRate: 1, interval: "1m", maxTokens: 1 },
-        }),
-      });
+  redisTest("should key per token when no identifier is supplied", async ({ redisOptions }) => {
+    const rateLimitMiddleware = authorizationRateLimitMiddleware({
+      redis: { ...redisOptions, tlsDisabled: true },
+      keyPrefix: "test-no-identifier",
+      defaultLimiter: {
+        type: "tokenBucket",
+        refillRate: 1,
+        interval: "1m",
+        maxTokens: 1,
+      },
+      pathMatchers: [/^\/api/],
+      // Override supplies a config but no identifier: bucketing stays per-key
+      // (hashed Authorization header), the legacy behavior.
+      limiterConfigOverride: async () => ({
+        config: { type: "tokenBucket", refillRate: 1, interval: "1m", maxTokens: 1 },
+      }),
+    });
 
-      app.use(rateLimitMiddleware);
-      app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
+    app.use(rateLimitMiddleware);
+    app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
 
-      const first = await request(app).get("/api/test").set("Authorization", "Bearer token-a");
-      expect(first.status).toBe(200);
+    const first = await request(app).get("/api/test").set("Authorization", "Bearer token-a");
+    expect(first.status).toBe(200);
 
-      // Same token is limited...
-      const firstAgain = await request(app)
-        .get("/api/test")
-        .set("Authorization", "Bearer token-a");
-      expect(firstAgain.status).toBe(429);
+    // Same token is limited...
+    const firstAgain = await request(app).get("/api/test").set("Authorization", "Bearer token-a");
+    expect(firstAgain.status).toBe(429);
 
-      // ...but a different token gets its own bucket.
-      const second = await request(app).get("/api/test").set("Authorization", "Bearer token-b");
-      expect(second.status).toBe(200);
-    }
-  );
+    // ...but a different token gets its own bucket.
+    const second = await request(app).get("/api/test").set("Authorization", "Bearer token-b");
+    expect(second.status).toBe(200);
+  });
 
   describe("Advanced Cases", () => {
     // 1. Test different rate limit configurations