088f68b373
## 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.
138 lines
5.1 KiB
TypeScript
138 lines
5.1 KiB
TypeScript
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";
|
|
import type { Duration } from "./rateLimiter.server";
|
|
|
|
const BATCH_STREAM_ITEMS_PATH = /^\/api\/v3\/batches\/([^/]+)\/items$/;
|
|
|
|
export const apiRateLimiter = authorizationRateLimitMiddleware({
|
|
redis: {
|
|
port: env.RATE_LIMIT_REDIS_PORT,
|
|
host: env.RATE_LIMIT_REDIS_HOST,
|
|
username: env.RATE_LIMIT_REDIS_USERNAME,
|
|
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
|
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
|
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
|
},
|
|
keyPrefix: "api",
|
|
defaultLimiter: {
|
|
type: "tokenBucket",
|
|
refillRate: env.API_RATE_LIMIT_REFILL_RATE,
|
|
interval: env.API_RATE_LIMIT_REFILL_INTERVAL as Duration,
|
|
maxTokens: env.API_RATE_LIMIT_MAX,
|
|
},
|
|
limiterCache: {
|
|
fresh: 60_000 * 10, // Data is fresh for 10 minutes
|
|
stale: 60_000 * 20, // Date is stale after 20 minutes
|
|
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,
|
|
});
|
|
|
|
if (!authenticatedEnv || !authenticatedEnv.ok) {
|
|
return;
|
|
}
|
|
|
|
if (authenticatedEnv.type === "PUBLIC_JWT") {
|
|
return {
|
|
config: {
|
|
type: "fixedWindow",
|
|
window: env.API_RATE_LIMIT_JWT_WINDOW,
|
|
tokens: env.API_RATE_LIMIT_JWT_TOKENS,
|
|
},
|
|
};
|
|
}
|
|
|
|
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
|
|
pathWhiteList: [
|
|
"/api/internal/stripe_webhooks",
|
|
// Keep allowlisted: these CLI endpoints are intentionally unauthenticated,
|
|
// so this Authorization-header-keyed limiter would 401 them. They are
|
|
// throttled separately by authCodeRateLimiter.server.ts.
|
|
"/api/v1/authorization-code",
|
|
"/api/v1/token",
|
|
"/api/v1/usage/ingest",
|
|
"/api/v1/plain/customer-cards",
|
|
/^\/api\/v1\/tasks\/[^/]+\/callback\/[^/]+$/, // /api/v1/tasks/$id/callback/$secret
|
|
/^\/api\/v1\/runs\/[^/]+\/tasks\/[^/]+\/callback\/[^/]+$/, // /api/v1/runs/$runId/tasks/$id/callback/$secret
|
|
/^\/api\/v1\/http-endpoints\/[^/]+\/env\/[^/]+\/[^/]+$/, // /api/v1/http-endpoints/$httpEndpointId/env/$envType/$shortcode
|
|
/^\/api\/v1\/sources\/http\/[^/]+$/, // /api/v1/sources/http/$id
|
|
/^\/api\/v1\/endpoints\/[^/]+\/[^/]+\/index\/[^/]+$/, // /api/v1/endpoints/$environmentId/$endpointSlug/index/$indexHookIdentifier
|
|
"/api/v1/timezones",
|
|
"/api/v1/usage/ingest",
|
|
"/api/v1/auth/jwt/claims",
|
|
/^\/api\/v1\/runs\/[^/]+\/attempts$/, // /api/v1/runs/$runFriendlyId/attempts
|
|
/^\/api\/v1\/waitpoints\/tokens\/[^/]+\/callback\/[^/]+$/, // /api/v1/waitpoints/tokens/$waitpointFriendlyId/callback/$hash
|
|
/^\/api\/v\d+\/deployments/, // /api/v{1,2,3,n}/deployments/*
|
|
// Internal SDK plumbing — packets are presigned-URL handshakes for
|
|
// payload uploads (v2 PUT) and downloads (v1 GET), authenticated via
|
|
// run-scoped JWT, called once per task/turn boundary by the runtime.
|
|
// Same shape as `/api/v1/runs/$runFriendlyId/attempts` above; not a
|
|
// customer-facing surface so customer rate limits shouldn't apply.
|
|
/^\/api\/v1\/packets\//,
|
|
/^\/api\/v2\/packets\//,
|
|
/^\/api\/v1\/sessions\/[^/]+\/snapshot-url$/,
|
|
],
|
|
bypass: async (req) => {
|
|
const match = BATCH_STREAM_ITEMS_PATH.exec(req.path);
|
|
|
|
if (!match) {
|
|
return false;
|
|
}
|
|
|
|
const batchFriendlyId = match[1];
|
|
const authorizationValue = req.headers.authorization;
|
|
|
|
if (!batchFriendlyId || !authorizationValue) {
|
|
return false;
|
|
}
|
|
|
|
const [authError, authenticated] = await tryCatch(
|
|
authenticateAuthorizationHeader(authorizationValue, {
|
|
allowPublicKey: true,
|
|
})
|
|
);
|
|
|
|
if (authError || !authenticated || !authenticated.ok) {
|
|
return false;
|
|
}
|
|
|
|
return batchStreamGrants.spend(authenticated.environment.id, batchFriendlyId);
|
|
},
|
|
log: {
|
|
rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
|
|
requests: env.API_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1",
|
|
limiter: env.API_RATE_LIMIT_LIMITER_LOGS_ENABLED === "1",
|
|
},
|
|
});
|
|
|
|
export type RateLimitMiddleware = ReturnType<typeof authorizationRateLimitMiddleware>;
|