diff --git a/.server-changes/additional-api-key-rate-limit-bucket.md b/.server-changes/additional-api-key-rate-limit-bucket.md new file mode 100644 index 0000000000..d67a3b8e2a --- /dev/null +++ b/.server-changes/additional-api-key-rate-limit-bucket.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +API rate limits now apply per environment, so creating extra API keys no longer increases how many requests an environment can make. diff --git a/apps/webapp/app/models/api-key.server.ts b/apps/webapp/app/models/api-key.server.ts index 7bd283344b..64564a5894 100644 --- a/apps/webapp/app/models/api-key.server.ts +++ b/apps/webapp/app/models/api-key.server.ts @@ -1,6 +1,5 @@ import type { PrismaClient, RuntimeEnvironment } from "@trigger.dev/database"; import type { HostRbacController } from "@trigger.dev/rbac"; -import { trail } from "agentcrumbs"; // @crumbs import { customAlphabet } from "nanoid"; import { MAX_API_KEY_TASK_IDENTIFIERS } from "~/consts"; import { prisma } from "~/db.server"; @@ -11,8 +10,6 @@ import { rbac } from "~/services/rbac.server"; import { generateAdditionalApiKey, generateRootApiKey } from "~/utils/apiKeys"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; -const crumb = trail("webapp"); // @crumbs - const apiKeyId = customAlphabet( "1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", 12 @@ -220,12 +217,6 @@ export async function createEnvironmentApiKey( })(); telemetryRecorder.recordOperation("create", "success"); - crumb("environment API key created", { - apiKeyId: apiKey.id, - environmentId, - presetId: apiKey.presetId, - }); // @crumbs - return { apiKey, plaintext: generated.apiKey }; } @@ -267,7 +258,6 @@ export async function revokeEnvironmentApiKey( } telemetryRecorder.recordOperation("revoke", "success"); - crumb("environment API key revoked", { apiKeyId, environmentId }); // @crumbs } export function createApiKeyForEnv(envType: RuntimeEnvironment["type"]) { diff --git a/apps/webapp/app/models/runtimeEnvironment.server.ts b/apps/webapp/app/models/runtimeEnvironment.server.ts index 84d2979bb1..8dc5a68b63 100644 --- a/apps/webapp/app/models/runtimeEnvironment.server.ts +++ b/apps/webapp/app/models/runtimeEnvironment.server.ts @@ -301,6 +301,89 @@ export async function findEnvironmentByApiKeyWithResolution( return resolveEnvironmentByApiKey(apiKey, branchName, tx, additionalApiKeyLookupEnabled); } +export type PrivateApiKeyRateLimitScope = { + environmentId: string; + apiRateLimiterConfig: unknown; +}; + +export async function resolvePrivateApiKeyRateLimitScope( + apiKey: string, + tx: PrismaClientOrTransaction = $replica +): Promise { + const now = new Date(); + + if (isAdditionalApiKey(apiKey)) { + const match = await tx.apiKey.findFirst({ + where: { + keyHash: hashApiKey(apiKey), + revokedAt: null, + OR: [{ expiresAt: null }, { expiresAt: { gt: now } }], + }, + select: { + runtimeEnvironment: { + select: { + id: true, + project: { select: { deletedAt: true } }, + organization: { select: { apiRateLimiterConfig: true } }, + }, + }, + }, + }); + + if (!match?.runtimeEnvironment || match.runtimeEnvironment.project.deletedAt) { + return null; + } + + return { + environmentId: match.runtimeEnvironment.id, + apiRateLimiterConfig: match.runtimeEnvironment.organization.apiRateLimiterConfig, + }; + } + + const environment = await tx.runtimeEnvironment.findFirst({ + where: { apiKey }, + select: { + id: true, + project: { select: { deletedAt: true } }, + organization: { select: { apiRateLimiterConfig: true } }, + }, + }); + + if (environment) { + if (environment.project.deletedAt) { + return null; + } + + return { + environmentId: environment.id, + apiRateLimiterConfig: environment.organization.apiRateLimiterConfig, + }; + } + + const revokedApiKey = await tx.revokedApiKey.findFirst({ + where: { apiKey, expiresAt: { gt: now } }, + select: { + runtimeEnvironment: { + select: { + id: true, + project: { select: { deletedAt: true } }, + organization: { select: { apiRateLimiterConfig: true } }, + }, + }, + }, + }); + + const revokedEnvironment = revokedApiKey?.runtimeEnvironment; + if (!revokedEnvironment || revokedEnvironment.project.deletedAt) { + return null; + } + + return { + environmentId: revokedEnvironment.id, + apiRateLimiterConfig: revokedEnvironment.organization.apiRateLimiterConfig, + }; +} + /** * @deprecated We don't use public API keys (`pk_*` tokens) anymore — public * access goes through public JWTs (see `isPublicJWT` / `validatePublicJwtKey`). diff --git a/apps/webapp/app/presenters/v3/LimitsPresenter.server.ts b/apps/webapp/app/presenters/v3/LimitsPresenter.server.ts index e468efd921..099a2e3553 100644 --- a/apps/webapp/app/presenters/v3/LimitsPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/LimitsPresenter.server.ts @@ -1,6 +1,5 @@ import { Ratelimit } from "@upstash/ratelimit"; import type { RuntimeEnvironmentType } from "@trigger.dev/database"; -import { createHash } from "node:crypto"; import { env } from "~/env.server"; import { getCurrentPlan } from "~/services/platform.v3.server"; import { @@ -90,13 +89,11 @@ export class LimitsPresenter extends BasePresenter { projectId, environmentId, environmentType, - environmentApiKey, }: { organizationId: string; projectId: string; environmentId: string; environmentType: RuntimeEnvironmentType; - environmentApiKey: string; }): Promise { // Get organization with all limit-related fields const organization = await this._replica.organization.findFirstOrThrow({ @@ -168,10 +165,21 @@ export class LimitsPresenter extends BasePresenter { where: { organizationId }, }); - // Get current rate limit tokens for this environment's API key + const runtimeEnv = await this._replica.runtimeEnvironment.findFirst({ + where: { id: environmentId }, + select: { + id: true, + parentEnvironmentId: true, + maximumConcurrencyLimit: true, + concurrencyLimitBurstFactor: true, + }, + }); + const apiRateLimitEnvironmentId = runtimeEnv?.parentEnvironmentId ?? environmentId; + + // Get current rate limit tokens for this environment's API bucket const apiRateLimitTokens = await getRateLimitRemainingTokens( "api", - environmentApiKey, + apiRateLimitEnvironmentId, apiRateLimitConfig ); // Batch rate limiter uses environment ID directly (not hashed) with a different key prefix @@ -181,15 +189,6 @@ export class LimitsPresenter extends BasePresenter { ); // Get current queue size for this environment - // We need the runtime environment fields for the engine query - const runtimeEnv = await this._replica.runtimeEnvironment.findFirst({ - where: { id: environmentId }, - select: { - id: true, - maximumConcurrencyLimit: true, - concurrencyLimitBurstFactor: true, - }, - }); let currentQueueSize = 0; if (runtimeEnv) { @@ -454,20 +453,14 @@ function resolveBatchConcurrencyConfig(batchConcurrencyConfig?: unknown): { /** * Query the current remaining tokens for a rate limiter using the Upstash getRemaining method. - * This uses the same configuration and hashing logic as the rate limit middleware. + * The API limiter uses the environment ID as the bucket identifier for private API keys. */ async function getRateLimitRemainingTokens( keyPrefix: string, - apiKey: string, + identifier: string, config: RateLimiterConfig ): Promise { try { - // Hash the authorization header the same way the rate limiter does - const authorizationValue = `Bearer ${apiKey}`; - const hash = createHash("sha256"); - hash.update(authorizationValue); - const hashedKey = hash.digest("hex"); - // Create a Ratelimit instance with the same configuration const limiter = createLimiterFromConfig(config); const ratelimit = new Ratelimit({ @@ -478,9 +471,9 @@ async function getRateLimitRemainingTokens( prefix: `ratelimit:${keyPrefix}`, }); - // Use the getRemaining method to get the current remaining tokens + // Use the same identifier as the API rate-limit middleware. // getRemaining returns a Promise - const remaining = await ratelimit.getRemaining(hashedKey); + const remaining = await ratelimit.getRemaining(identifier); return remaining; } catch (error) { logger.warn("Failed to get rate limit remaining tokens", { diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx index 833dffdec6..a9f8100a6e 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx @@ -78,7 +78,6 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { projectId: project.id, environmentId: environment.id, environmentType: environment.type, - environmentApiKey: environment.apiKey, }) ); diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts index 1b8d8a3ed1..435c82f606 100644 --- a/apps/webapp/app/services/apiRateLimit.server.ts +++ b/apps/webapp/app/services/apiRateLimit.server.ts @@ -1,5 +1,6 @@ import { tryCatch } from "@trigger.dev/core/v3"; import { env } from "~/env.server"; +import { resolvePrivateApiKeyRateLimitScope } from "~/models/runtimeEnvironment.server"; import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server"; import { authenticateAuthorizationHeader } from "./apiAuth.server"; import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server"; @@ -29,6 +30,21 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({ maxItems: 1000, }, limiterConfigOverride: async (authorizationValue) => { + const rawApiKey = authorizationValue.replace(/^Bearer /, ""); + + if (rawApiKey.startsWith("tr_")) { + const scope = await resolvePrivateApiKeyRateLimitScope(rawApiKey); + + if (!scope) { + return; + } + + return { + config: scope.apiRateLimiterConfig, + identifier: scope.environmentId, + }; + } + const authenticatedEnv = await authenticateAuthorizationHeader(authorizationValue, { allowPublicKey: true, allowJWT: true, @@ -40,13 +56,19 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({ if (authenticatedEnv.type === "PUBLIC_JWT") { return { - type: "fixedWindow", - window: env.API_RATE_LIMIT_JWT_WINDOW, - tokens: env.API_RATE_LIMIT_JWT_TOKENS, + config: { + type: "fixedWindow", + window: env.API_RATE_LIMIT_JWT_WINDOW, + tokens: env.API_RATE_LIMIT_JWT_TOKENS, + }, }; - } else { - return authenticatedEnv.environment.organization.apiRateLimiterConfig; } + + return { + config: authenticatedEnv.environment.organization.apiRateLimiterConfig, + // Public keys are browser-distributed, so keep them on per-key buckets. + identifier: authenticatedEnv.type === "PRIVATE" ? authenticatedEnv.environment.id : undefined, + }; }, pathMatchers: [/^\/api/], // Allow /api/v1/tasks/:id/callback/:secret diff --git a/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts b/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts index 5fcd6eb545..ce0b8b50d2 100644 --- a/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts +++ b/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts @@ -52,7 +52,14 @@ export const RateLimiterConfig = z.discriminatedUnion("type", [ export type RateLimiterConfig = z.infer; -type LimitConfigOverrideFunction = (authorizationValue: string) => Promise; +type RateLimitOverride = { + config?: unknown; + identifier?: string; +}; + +type LimitConfigOverrideFunction = ( + authorizationValue: string +) => Promise; type Options = { redis: RedisWithClusterOptions; @@ -80,16 +87,22 @@ type Options = { }; }; -async function resolveLimitConfig( +type ResolvedRateLimit = { + config: RateLimiterConfig; + // Bucket key to use, or undefined to fall back to the hashed Authorization header. + identifier?: string; +}; + +async function resolveRateLimit( authorizationValue: string, hashedAuthorizationValue: string, defaultLimiter: RateLimiterConfig, - cache: UnkeyCache<{ limiter: RateLimiterConfig }>, + cache: UnkeyCache<{ limiter: ResolvedRateLimit }>, logsEnabled: boolean, limiterConfigOverride?: LimitConfigOverrideFunction -): Promise { +): Promise { if (!limiterConfigOverride) { - return defaultLimiter; + return { config: defaultLimiter }; } if (logsEnabled) { @@ -110,10 +123,16 @@ async function resolveLimitConfig( }); } - return defaultLimiter; + return { config: defaultLimiter } satisfies ResolvedRateLimit; } - const parsedOverride = RateLimiterConfig.safeParse(override); + const identifier = override.identifier; + + if (!override.config) { + return { config: defaultLimiter, identifier } satisfies ResolvedRateLimit; + } + + const parsedOverride = RateLimiterConfig.safeParse(override.config); if (!parsedOverride.success) { logger.error("Error parsing rate limiter override", { @@ -121,7 +140,7 @@ async function resolveLimitConfig( errors: parsedOverride.error.errors, }); - return defaultLimiter; + return { config: defaultLimiter, identifier } satisfies ResolvedRateLimit; } if (logsEnabled && parsedOverride.data) { @@ -132,10 +151,22 @@ async function resolveLimitConfig( }); } - return parsedOverride.data; + return { config: parsedOverride.data, identifier } satisfies ResolvedRateLimit; }); - return cacheResult.val ?? defaultLimiter; + // Defensive read: the cache is keyed on a shared Redis namespace, so during a + // deploy an entry could have been written by a server running a different + // code version (a different stored shape). Re-validate here so a stale/foreign + // entry can never reach createLimiterFromConfig with an undefined config and + // throw. The cache key is also versioned (see RedisCacheStore keyPrefix), so + // this is belt-and-suspenders. + const cached = cacheResult.val; + const parsedConfig = RateLimiterConfig.safeParse(cached?.config); + + return { + config: parsedConfig.success ? parsedConfig.data : defaultLimiter, + identifier: typeof cached?.identifier === "string" ? cached.identifier : undefined, + }; } /** @@ -169,14 +200,17 @@ export function authorizationRateLimitMiddleware({ const memory = createLRUMemoryStore(limiterCache?.maxItems ?? 1000); const redisCacheStore = new RedisCacheStore({ connection: { - keyPrefix: `cache:${keyPrefix}:rate-limit-cache:`, + // Versioned namespace: the cached value shape is part of this key. Bump + // the version whenever ResolvedRateLimit changes so a rolling deploy never + // reads entries written in a previous shape (and vice versa). + keyPrefix: `cache:${keyPrefix}:rate-limit-cache:v2:`, ...redis, }, }); // This cache holds the rate limit configuration for each org, so we don't have to fetch it every request const cache = createCache({ - limiter: new Namespace(ctx, { + limiter: new Namespace(ctx, { stores: [memory, redisCacheStore], fresh: limiterCache?.fresh ?? 30_000, stale: limiterCache?.stale ?? 60_000, @@ -269,7 +303,7 @@ export function authorizationRateLimitMiddleware({ hash.update(authorizationValue); const hashedAuthorizationValue = hash.digest("hex"); - const limiterConfig = await resolveLimitConfig( + const { config: limiterConfig, identifier } = await resolveRateLimit( authorizationValue, hashedAuthorizationValue, defaultLimiter, @@ -278,6 +312,8 @@ export function authorizationRateLimitMiddleware({ limiterConfigOverride ); + const rateLimitIdentifier = identifier ?? hashedAuthorizationValue; + const limiter = createLimiterFromConfig(limiterConfig); const rateLimiter = new RateLimiter({ @@ -288,7 +324,7 @@ export function authorizationRateLimitMiddleware({ logFailure: log.rejections, }); - const { success, limit, reset, remaining } = await rateLimiter.limit(hashedAuthorizationValue); + const { success, limit, reset, remaining } = await rateLimiter.limit(rateLimitIdentifier); const $remaining = Math.max(0, remaining); // remaining can be negative if the user has exceeded the limit, so clamp it to 0 diff --git a/apps/webapp/test/authorizationRateLimitMiddleware.test.ts b/apps/webapp/test/authorizationRateLimitMiddleware.test.ts index b6076cef0d..29318c319e 100644 --- a/apps/webapp/test/authorizationRateLimitMiddleware.test.ts +++ b/apps/webapp/test/authorizationRateLimitMiddleware.test.ts @@ -150,10 +150,12 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", limiterConfigOverride: async (authorizationValue) => { if (authorizationValue === "Bearer premium-token") { return { - type: "tokenBucket", - refillRate: 10, - interval: "1m", - maxTokens: 100, + config: { + type: "tokenBucket", + refillRate: 10, + interval: "1m", + maxTokens: 100, + }, }; } return undefined; @@ -184,6 +186,75 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", } ); + redisTest( + "should share a bucket across tokens that resolve to the same identifier", + async ({ redisOptions }) => { + const rateLimitMiddleware = authorizationRateLimitMiddleware({ + redis: { ...redisOptions, tlsDisabled: true }, + keyPrefix: "test-identifier", + defaultLimiter: { + type: "tokenBucket", + refillRate: 1, + interval: "1m", + maxTokens: 1, + }, + pathMatchers: [/^\/api/], + // Both tokens map to the same environment identifier, so they should + // consume from a single shared bucket rather than one bucket each. + limiterConfigOverride: async () => ({ identifier: "env_shared" }), + }); + + app.use(rateLimitMiddleware); + app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" })); + + // First token uses the single token in the shared bucket. + const first = await request(app) + .get("/api/test") + .set("Authorization", "Bearer tr_prod_sk_aaaaaaaaaaaaaaaaaaaaaaaa"); + expect(first.status).toBe(200); + + // A different token that resolves to the same identifier is limited, + // because the bucket is shared rather than per-key. + const second = await request(app) + .get("/api/test") + .set("Authorization", "Bearer tr_prod_sk_bbbbbbbbbbbbbbbbbbbbbbbb"); + expect(second.status).toBe(429); + } + ); + + 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" })); + + 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); + + // ...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 redisTest("should enforce fixed window rate limiting", async ({ redisOptions }) => { @@ -375,10 +446,12 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", configOverrideCalls++; if (authorizationValue === "Bearer premium-token") { return { - type: "tokenBucket", - refillRate: 10, - interval: "1m", - maxTokens: 100, + config: { + type: "tokenBucket", + refillRate: 10, + interval: "1m", + maxTokens: 100, + }, }; } return undefined; diff --git a/apps/webapp/test/findEnvironmentByApiKey.test.ts b/apps/webapp/test/findEnvironmentByApiKey.test.ts index c301a02b20..10425ecbfe 100644 --- a/apps/webapp/test/findEnvironmentByApiKey.test.ts +++ b/apps/webapp/test/findEnvironmentByApiKey.test.ts @@ -1,7 +1,10 @@ import { postgresTest } from "@internal/testcontainers"; import { type PrismaClient } from "@trigger.dev/database"; import { describe, expect, it, vi } from "vitest"; -import { findEnvironmentByApiKey } from "~/models/runtimeEnvironment.server"; +import { + findEnvironmentByApiKey, + resolvePrivateApiKeyRateLimitScope, +} from "~/models/runtimeEnvironment.server"; import { generateAdditionalApiKey, hashApiKey } from "~/utils/apiKeys"; import { createTestOrgProjectWithMember, uniqueId } from "./fixtures/environmentVariablesFixtures"; @@ -143,6 +146,36 @@ describe("findEnvironmentByApiKey — PREVIEW (regression guard)", () => { expect(resolved?.apiKey).toBe(previewParent.apiKey); } ); + + postgresTest( + "rate limit scope resolves root and additional keys to the preview parent", + async ({ prisma }) => { + const { organization, project, user } = await createTestOrgProjectWithMember(prisma); + const previewParent = await createEnv(prisma, project.id, organization.id, { + type: "PREVIEW", + isBranchableEnvironment: true, + }); + const additional = generateAdditionalApiKey("PREVIEW").apiKey; + + await prisma.apiKey.create({ + data: { + name: "Preview integration", + keyHash: hashApiKey(additional), + lastFour: additional.slice(-4), + runtimeEnvironmentId: previewParent.id, + createdByUserId: user.id, + presetId: null, + scopes: ["admin"], + }, + }); + + const rootScope = await resolvePrivateApiKeyRateLimitScope(previewParent.apiKey, prisma); + const additionalScope = await resolvePrivateApiKeyRateLimitScope(additional, prisma); + + expect(rootScope?.environmentId).toBe(previewParent.id); + expect(additionalScope?.environmentId).toBe(previewParent.id); + } + ); }); describe("findEnvironmentByApiKey — non-branchable", () => { @@ -366,4 +399,30 @@ describe("findEnvironmentByApiKey — additional and disabled keys", () => { ).resolves.toMatchObject({ id: environment.id }); } ); + + postgresTest("does not resolve additional keys for deleted projects", async ({ prisma }) => { + const { organization, project, user } = await createTestOrgProjectWithMember(prisma); + const environment = await createEnv(prisma, project.id, organization.id, { + type: "PRODUCTION", + }); + const additional = generateAdditionalApiKey("PRODUCTION").apiKey; + + await prisma.apiKey.create({ + data: { + name: "Deleted project key", + keyHash: hashApiKey(additional), + lastFour: additional.slice(-4), + runtimeEnvironmentId: environment.id, + createdByUserId: user.id, + presetId: null, + scopes: ["admin"], + }, + }); + await prisma.project.update({ + where: { id: project.id }, + data: { deletedAt: new Date() }, + }); + + await expect(resolvePrivateApiKeyRateLimitScope(additional, prisma)).resolves.toBeNull(); + }); });