From 67186ef25ab866ff065bdfcb7d32c79930807bbc Mon Sep 17 00:00:00 2001 From: vimtor Date: Thu, 21 May 2026 19:36:53 +0200 Subject: [PATCH] add redis-backed ip rate limiting with database fallback --- bun.lock | 3 + infra/console.ts | 2 + infra/secret.ts | 2 + packages/console/app/package.json | 1 + .../app/src/routes/zen/util/ipRateLimiter.ts | 80 ++++++++++++++++--- .../app/src/routes/zen/util/rateLimit.ts | 13 +++ .../console/app/src/routes/zen/util/redis.ts | 12 +++ packages/console/app/test/rateLimiter.test.ts | 2 +- sst-env.d.ts | 8 ++ 9 files changed, 113 insertions(+), 10 deletions(-) create mode 100644 packages/console/app/src/routes/zen/util/rateLimit.ts create mode 100644 packages/console/app/src/routes/zen/util/redis.ts diff --git a/bun.lock b/bun.lock index 0d7612922a..95c5b9fdc7 100644 --- a/bun.lock +++ b/bun.lock @@ -101,6 +101,7 @@ "@solidjs/router": "catalog:", "@solidjs/start": "catalog:", "@stripe/stripe-js": "8.6.1", + "@upstash/redis": "1.38.0", "chart.js": "4.5.1", "nitro": "3.0.1-alpha.1", "solid-js": "catalog:", @@ -2467,6 +2468,8 @@ "@ungap/structured-clone": ["@ungap/structured-clone@1.3.0", "", {}, "sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g=="], + "@upstash/redis": ["@upstash/redis@1.38.0", "", { "dependencies": { "uncrypto": "^0.1.3" } }, "sha512-wu+dZBptlLy0+MCUEoHmzrY/TnmgDey3+c7EbIGwrLqAvkP8yi5MWZHYGIFtAygmL4Bkz2TdFu+eU0vFPncIcg=="], + "@valibot/to-json-schema": ["@valibot/to-json-schema@1.6.0", "", { "peerDependencies": { "valibot": "^1.3.0" } }, "sha512-d6rYyK5KVa2XdqamWgZ4/Nr+cXhxjy7lmpe6Iajw15J/jmU+gyxl2IEd1Otg1d7Rl3gOQL5reulnSypzBtYy1A=="], "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="], diff --git a/infra/console.ts b/infra/console.ts index 29e473de37..6f5fd0dd85 100644 --- a/infra/console.ts +++ b/infra/console.ts @@ -250,6 +250,8 @@ new sst.cloudflare.x.SolidStart("Console", { bucket, bucketNew, database, + SECRET.UpstashRedisRestUrl, + SECRET.UpstashRedisRestToken, AUTH_API_URL, STRIPE_WEBHOOK_SECRET, DISCORD_INCIDENT_WEBHOOK_URL, diff --git a/infra/secret.ts b/infra/secret.ts index d4e8b148fc..eafbd91ed2 100644 --- a/infra/secret.ts +++ b/infra/secret.ts @@ -8,4 +8,6 @@ export const SECRET = { R2AccessKey: new sst.Secret("R2AccessKey", "unknown"), R2SecretKey: new sst.Secret("R2SecretKey", "unknown"), HoneycombWebhookSecret: new random.RandomPassword("HoneycombWebhookSecret", { length: 24 }), + UpstashRedisRestUrl: new sst.Secret("UpstashRedisRestUrl"), + UpstashRedisRestToken: new sst.Secret("UpstashRedisRestToken"), } diff --git a/packages/console/app/package.json b/packages/console/app/package.json index bab3e8a491..ca8e589170 100644 --- a/packages/console/app/package.json +++ b/packages/console/app/package.json @@ -26,6 +26,7 @@ "@solidjs/router": "catalog:", "@solidjs/start": "catalog:", "@stripe/stripe-js": "8.6.1", + "@upstash/redis": "1.38.0", "chart.js": "4.5.1", "nitro": "3.0.1-alpha.1", "solid-js": "catalog:", diff --git a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts index 03eca06c47..be63d6c0fe 100644 --- a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts +++ b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts @@ -2,7 +2,9 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" import { FreeUsageLimitError } from "./error" import { logger } from "./logger" -import { i18n } from "~/i18n" +import { getRetryAfterDay, type RateLimiter, type RateLimiterState } from "./rateLimit" +import { buildRateLimitKey, redis } from "./redis" +import { i18n, type Dict } from "~/i18n" import { localeFromRequest } from "~/lib/language" import { Subscription } from "@opencode-ai/console-core/subscription.js" @@ -18,10 +20,23 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined const ip = !rawIp.length ? "unknown" : rawIp const now = Date.now() - const lifetimeInterval = "" const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) + const retryAfter = getRetryAfterDay(now) + const state = { isNew: false, fallbackDatabase: false } + const databaseLimiter = createDatabaseRateLimiter(ip, dailyInterval, dailyLimit, isDefaultModel, dict, retryAfter, state) + return createUpstashRateLimiter(ip, dailyInterval, dailyLimit, isDefaultModel, dict, retryAfter, state, databaseLimiter) +} - let _isNew: boolean +function createDatabaseRateLimiter( + ip: string, + dailyInterval: string, + dailyLimit: number, + isDefaultModel: boolean, + dict: Dict, + retryAfter: number, + state: RateLimiterState, +): RateLimiter { + const lifetimeInterval = "" return { check: async () => { @@ -42,10 +57,10 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) - _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + state.isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 - if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit)) - throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now)) + if ((state.isNew && dailyCount >= dailyLimit * 2) || (!state.isNew && dailyCount >= dailyLimit)) + throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter) }, track: async () => { await Database.use((tx) => @@ -53,7 +68,7 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined .insert(IpRateLimitTable) .values([ { ip, interval: dailyInterval, count: 1 }, - ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), + ...(state.isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), ]) .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), ) @@ -61,8 +76,55 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined } } -export function getRetryAfterDay(now: number) { - return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000) +function createUpstashRateLimiter( + ip: string, + dailyInterval: string, + dailyLimit: number, + isDefaultModel: boolean, + dict: Dict, + retryAfter: number, + state: RateLimiterState, + databaseLimiter: RateLimiter, +): RateLimiter { + const lifetimeInterval = "" + const lifetimeKey = buildRateLimitKey("ip", ip, lifetimeInterval) + const dailyKey = buildRateLimitKey("ip", ip, dailyInterval) + + return { + check: async () => { + try { + const keys = isDefaultModel + ? [lifetimeKey, dailyKey] + : [dailyKey] + const counts = await redis.mget<(string | number | null)[]>(keys) + const lifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0 + const dailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0) + logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) + + state.isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + if ((state.isNew && dailyCount >= dailyLimit * 2) || (!state.isNew && dailyCount >= dailyLimit)) + throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter) + } catch (error) { + if (error instanceof FreeUsageLimitError) throw error + + state.fallbackDatabase = true + await databaseLimiter.check() + } + }, + track: async () => { + if (state.fallbackDatabase) return databaseLimiter.track() + + try { + const pipeline = redis.pipeline() + pipeline.incr(dailyKey) + pipeline.expire(dailyKey, retryAfter) + if (state.isNew) pipeline.incr(lifetimeKey) + await pipeline.exec() + } catch { + await databaseLimiter.track() + } + }, + } } function buildYYYYMMDD(timestamp: number) { diff --git a/packages/console/app/src/routes/zen/util/rateLimit.ts b/packages/console/app/src/routes/zen/util/rateLimit.ts new file mode 100644 index 0000000000..27eef8aa92 --- /dev/null +++ b/packages/console/app/src/routes/zen/util/rateLimit.ts @@ -0,0 +1,13 @@ +export type RateLimiter = { + check: () => Promise + track: () => Promise +} + +export type RateLimiterState = { + isNew: boolean + fallbackDatabase: boolean +} + +export function getRetryAfterDay(now: number) { + return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000) +} diff --git a/packages/console/app/src/routes/zen/util/redis.ts b/packages/console/app/src/routes/zen/util/redis.ts new file mode 100644 index 0000000000..5fb77816d9 --- /dev/null +++ b/packages/console/app/src/routes/zen/util/redis.ts @@ -0,0 +1,12 @@ +import { Resource } from "@opencode-ai/console-resource" +import { Redis } from "@upstash/redis/cloudflare" + +export const redis = new Redis({ + url: Resource.UpstashRedisRestUrl.value, + token: Resource.UpstashRedisRestToken.value, + enableTelemetry: false, +}) + +export function buildRateLimitKey(kind: string, identifier: string, interval: string) { + return `zen:${Resource.App.stage}:ratelimit:${kind}:${identifier}:${interval || "lifetime"}` +} diff --git a/packages/console/app/test/rateLimiter.test.ts b/packages/console/app/test/rateLimiter.test.ts index 6c96226278..d89e9cb53e 100644 --- a/packages/console/app/test/rateLimiter.test.ts +++ b/packages/console/app/test/rateLimiter.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test" -import { getRetryAfterDay } from "../src/routes/zen/util/ipRateLimiter" +import { getRetryAfterDay } from "../src/routes/zen/util/rateLimit" describe("getRetryAfterDay", () => { test("returns full day at midnight UTC", () => { diff --git a/sst-env.d.ts b/sst-env.d.ts index 1eaebd1e59..aa79ec87d3 100644 --- a/sst-env.d.ts +++ b/sst-env.d.ts @@ -137,6 +137,14 @@ declare module "sst" { "type": "sst.cloudflare.SolidStart" "url": string } + "UpstashRedisRestToken": { + "type": "sst.sst.Secret" + "value": string + } + "UpstashRedisRestUrl": { + "type": "sst.sst.Secret" + "value": string + } "Web": { "type": "sst.cloudflare.Astro" "url": string