feat(webapp): share rate limit bucket across additional API keys per environment (#4508)
## What
Rate-limit the API by **environment** rather than per API key.
Previously the limiter keyed its bucket on the hash of the full
`Authorization` header — one bucket per key. With additional environment
API keys (`tr_*_sk_*`), an environment can mint many keys and each got
its own full bucket, so more keys = higher effective rate limit. This
collapses all of an environment's keys onto a single shared
per-environment bucket, so the ceiling is exactly the configured limit
regardless of key mix.
## How
- `authorizationRateLimitMiddleware` now lets the override return `{
config?, identifier? }`. `identifier`, when present, is the rate limit
bucket key; otherwise it falls back to the hashed `Authorization` header
(unchanged legacy behavior, still used by `engineRateLimiter` and any
unauthenticated fallthrough).
- `apiRateLimiter`'s override resolves the environment id and uses it as
the identifier:
- **Additional keys** (`isAdditionalApiKey`) resolve via a new
`resolveAdditionalApiKeyRateLimitScope()` — a **scope-agnostic** keyHash
→ (environmentId, org limiter config) lookup. It is deliberately
permissive (restricted keys resolve too) because it's used **only for
bucketing, never as an auth decision** — request auth still goes through
the RBAC bearer controller, which enforces scopes. Revoked/expired keys
are excluded so they can't hold a bucket warm.
- **Root/legacy keys** reuse the environment already resolved by
`authenticateAuthorizationHeader` and key on `environment.id` too.
- The identifier is always the stable environment id, never the secret
key (which can rotate and would split the bucket).
- The whole override result is cached per key by the existing SWR cache,
so **no extra per-request lookup and no separate Redis mapping** is
added.
## Behavior notes
- Root + additional keys of the same environment now share one bucket
(ceiling = configured limit, not a multiple of it). Restricted
additional keys are included — they were the biggest gap, since they
authenticate via the RBAC controller and previously fell back to per-key
buckets.
- **Public JWTs** keep their existing fixed-window, per-token bucketing.
- One-time bucket reset on deploy (bucket keys change); harmless.
## Tests
- New: two tokens resolving to the same identifier share one bucket.
- New: with no identifier, bucketing stays per-key (legacy behavior
preserved).
- Updated existing override tests to the new `{ config }` return shape.
Base: `feat/multi-keys-surface`. Closes TRI-12888.
This commit is contained in:
@@ -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.
|
||||
@@ -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"]) {
|
||||
|
||||
@@ -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<PrivateApiKeyRateLimitScope | null> {
|
||||
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`).
|
||||
|
||||
@@ -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<LimitsResult> {
|
||||
// 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<number | null> {
|
||||
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<number>
|
||||
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", {
|
||||
|
||||
-1
@@ -78,7 +78,6 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
|
||||
projectId: project.id,
|
||||
environmentId: environment.id,
|
||||
environmentType: environment.type,
|
||||
environmentApiKey: environment.apiKey,
|
||||
})
|
||||
);
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -52,7 +52,14 @@ export const RateLimiterConfig = z.discriminatedUnion("type", [
|
||||
|
||||
export type RateLimiterConfig = z.infer<typeof RateLimiterConfig>;
|
||||
|
||||
type LimitConfigOverrideFunction = (authorizationValue: string) => Promise<unknown>;
|
||||
type RateLimitOverride = {
|
||||
config?: unknown;
|
||||
identifier?: string;
|
||||
};
|
||||
|
||||
type LimitConfigOverrideFunction = (
|
||||
authorizationValue: string
|
||||
) => Promise<RateLimitOverride | undefined>;
|
||||
|
||||
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<RateLimiterConfig> {
|
||||
): Promise<ResolvedRateLimit> {
|
||||
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<RateLimiterConfig>(ctx, {
|
||||
limiter: new Namespace<ResolvedRateLimit>(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
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user