From 3a869a4b3780cf43a064b28f61169e6e95e7a258 Mon Sep 17 00:00:00 2001 From: AlCalzone Date: Sat, 31 Aug 2024 20:09:50 +0200 Subject: [PATCH] feat: use native rate-limiter (#139) --- package.json | 1 - src/durable_objects/RateLimiter.ts | 99 ------------------------------ src/durable_objects/itty.d.ts | 9 --- src/lib/apiKeys.ts | 4 +- src/maintenance/upload.ts | 2 +- src/routes/admin.ts | 32 ---------- src/routes/api.ts | 47 +++++++------- src/worker.ts | 5 +- wrangler.toml | 35 +++++++++-- yarn.lock | 11 ---- 10 files changed, 62 insertions(+), 183 deletions(-) delete mode 100644 src/durable_objects/RateLimiter.ts delete mode 100644 src/durable_objects/itty.d.ts diff --git a/package.json b/package.json index 0a2fbc6..5c6b597 100644 --- a/package.json +++ b/package.json @@ -16,7 +16,6 @@ "upload": "yarn tsx src/maintenance/upload.ts" }, "dependencies": { - "itty-durable": "^1.2.0", "itty-router": "^2.6.1", "itty-router-extras": "^0.4.2", "json-logic-js": "^2.0.2", diff --git a/src/durable_objects/RateLimiter.ts b/src/durable_objects/RateLimiter.ts deleted file mode 100644 index 8986e2a..0000000 --- a/src/durable_objects/RateLimiter.ts +++ /dev/null @@ -1,99 +0,0 @@ -import { createDurable } from "itty-durable"; - -interface RateLimit { - resetDate: number; - remaining: number; -} - -const PERSIST_INTERVAL_MS = 20 * 1000; - -export class RateLimiter extends createDurable() { - private resetDate: number = 0; - private remaining: number = 0; - - private lastPersist: number = Date.now(); - private persistTimeout: number | undefined; - - // Avoid persisting the values related to throttling the persist calls - public getPersistable(): any { - const { persistTimeout, lastPersist, ...persistable } = this; - return persistable; - } - - // private getStorage(): DurableObjectStorage { - // return (this.state as any).storage; - // } - - /** - * Checks if the current request is allowed. If not, returns the rate limit information. - */ - public async request( - maxPerHour: number - ): Promise { - const now = Date.now(); - - // // Destroy the durable object if it hasn't been used for 2 hours - // await this.getStorage().deleteAlarm(); - // await this.getStorage().setAlarm(now + 2 * 60 * 60 * 1000); - - if (this.resetDate < now) { - // Not initialized yet or the rate limit has expired - this.resetDate = now + 3600000; - this.remaining = maxPerHour; - } - - let limitExceeded = false; - if (this.remaining > 0) { - this.remaining = Math.max(0, this.remaining - 1); - } else { - limitExceeded = true; - } - - await this.throttlePersist(); - - return { - resetDate: this.resetDate, - remaining: this.remaining, - limitExceeded, - }; - } - - private async throttlePersist(): Promise { - // Persisting the state on every request is expensive. Instead we persist regularly after changes. - // This has the risk of data loss if the worker crashes, or the DO gets evicted early, but that's acceptable for this use case. - - // (Re-)schedule a persist for later, so we automatically persist if nothing happens for a while - if (this.persistTimeout) clearTimeout(this.persistTimeout); - // eslint-disable-next-line @typescript-eslint/ban-ts-comment - // @ts-ignore - For some reason, the Node.JS globals end up in this file - this.persistTimeout = setTimeout(() => { - void this.doPersist(); - }, PERSIST_INTERVAL_MS); - - // Also make sure to persist busy objects at least every PERSIST_INTERVAL_MS - if (Date.now() - this.lastPersist > PERSIST_INTERVAL_MS) { - // We haven't persisted in a while, so persist now - await this.doPersist(); - } - } - - private async doPersist(): Promise { - this.lastPersist = Date.now(); - await this.persist(); - } - - /** Refreshes the rate limiter and sets its remaining requests to the given value */ - public async setTo(maxPerHour: number): Promise { - this.resetDate = Date.now() + 3600000; - this.remaining = maxPerHour; - await this.doPersist(); - } - - // public async alarm(): Promise { - // await this.destroy(); - // } -} - -export type RateLimiterProps = { - RateLimiter: IttyDurableObjectNamespace; -}; diff --git a/src/durable_objects/itty.d.ts b/src/durable_objects/itty.d.ts deleted file mode 100644 index b7a7e15..0000000 --- a/src/durable_objects/itty.d.ts +++ /dev/null @@ -1,9 +0,0 @@ -type PromisifyPublicFunctions = { - [K in keyof T]: T[K] extends (...args: any[]) => any - ? (...args: Parameters) => Promise>> - : never; -}; - -interface IttyDurableObjectNamespace { - get(id: string | DurableObjectId): PromisifyPublicFunctions; -} diff --git a/src/lib/apiKeys.ts b/src/lib/apiKeys.ts index eb1e47e..7e3ea64 100644 --- a/src/lib/apiKeys.ts +++ b/src/lib/apiKeys.ts @@ -5,8 +5,10 @@ const AUTH_TAG_LEN = 8; export interface APIKey { id: number; - /** The maximum number of requests per hour */ + /** @deprecated The maximum number of requests per hour */ rateLimit: number; + /** Which ratelimiting bucket to use */ + bucket?: string; } // Encoded API key format: diff --git a/src/maintenance/upload.ts b/src/maintenance/upload.ts index 98e3070..888b50c 100644 --- a/src/maintenance/upload.ts +++ b/src/maintenance/upload.ts @@ -72,7 +72,7 @@ void (async () => { })), }; console.log( - `Uplaoding files ${cursor + 1}...${ + `Uploading files ${cursor + 1}...${ cursor + currentBatch.length } of ${files.length}...` ); diff --git a/src/routes/admin.ts b/src/routes/admin.ts index 2a880aa..9914d30 100644 --- a/src/routes/admin.ts +++ b/src/routes/admin.ts @@ -1,4 +1,3 @@ -import { withDurables } from "itty-durable"; import { json, missing, @@ -6,7 +5,6 @@ import { withParams, type ThrowableRouter, } from "itty-router-extras"; -import type { RateLimiterProps } from "../durable_objects/RateLimiter"; import { encryptAPIKey } from "../lib/apiKeys"; import { createCachedR2FS, @@ -161,34 +159,4 @@ export default function register(router: ThrowableRouter): void { return text(ret || ""); } ); - - router.post( - "/admin/resetRateLimit/:id/:limit", - withParams, - withDurables({ parse: true }), - async ( - req: RequestWithProps< - [{ params: { id: string; limit: string } }, RateLimiterProps] - >, - env: CloudflareEnvironment - ) => { - const id = parseInt(req.params.id); - const limit = parseInt(req.params.limit); - if ( - Number.isNaN(id) || - id < 1 || - Number.isNaN(limit) || - limit < 1 - ) { - console.error("Usage: /admin/resetRateLimit/:id/:limit"); - return clientError("Invalid id or limit"); - } - - const objId = env.RateLimiter.idFromName(req.params.id); - const RateLimiter = req.RateLimiter.get(objId); - await RateLimiter.setTo(limit); - - return json({ ok: true }); - } - ); } diff --git a/src/routes/api.ts b/src/routes/api.ts index ad84f23..fe8b269 100644 --- a/src/routes/api.ts +++ b/src/routes/api.ts @@ -1,4 +1,4 @@ -import { withDurables } from "itty-durable"; +import type { RateLimit } from "@cloudflare/workers-types/experimental"; import { json, type ThrowableRouter } from "itty-router-extras"; import { compare } from "semver"; import { @@ -8,7 +8,6 @@ import { APIv3_RequestSchema, APIv3_Response, } from "../apiDefinitions"; -import type { RateLimiterProps } from "../durable_objects/RateLimiter"; import { withCache } from "../lib/cache"; import { lookupConfig } from "../lib/config"; import type { UpgradeInfo } from "../lib/configSchema"; @@ -132,35 +131,41 @@ async function handleUpdateRequest( ); } +function getBucket(rateLimit: number): string { + if (rateLimit === 9) return "TEST"; + if (rateLimit <= 100) return "FREE"; + if (rateLimit <= 1000) return "1k"; + if (rateLimit <= 10000) return "10k"; + if (rateLimit <= 100000) return "100k"; + return "FREE"; +} + export default function register(router: ThrowableRouter): void { // Check API keys and apply rate limiter - router.all("/api/*", withAPIKey, withDurables({ parse: true }), (async ( - req: RequestWithProps<[APIKeyProps, RateLimiterProps]>, + router.all("/api/*", withAPIKey, (async ( + req: RequestWithProps<[APIKeyProps]>, env: CloudflareEnvironment ) => { - const objId = env.RateLimiter.idFromName( - (req.apiKey?.id ?? 0).toString() - ); - const RateLimiter = req.RateLimiter.get(objId); + const bucket = + req.apiKey?.bucket ?? + (req.apiKey?.rateLimit && getBucket(req.apiKey.rateLimit)) ?? + "FREE"; + const rateLimiter = ((env as any)[`RL_${bucket}`] ?? + env.RL_FREE) as RateLimit; + const apiKeyId = req.apiKey?.id ?? 0; - const maxPerHour = req.apiKey?.rateLimit ?? 10000; - const result = await RateLimiter.request(maxPerHour); + const { success } = await rateLimiter.limit({ + key: apiKeyId.toString(), + }); - env.responseHeaders = { - ...env.responseHeaders, - "x-ratelimit-limit": maxPerHour.toString(), - "x-ratelimit-remaining": result.remaining.toString(), - "x-ratelimit-reset": Math.ceil(result.resetDate / 1000).toString(), - }; - - if (result.limitExceeded) { + if (!success) { // Rate limit exceeded + return new Response(undefined, { status: 429, headers: { - "retry-after": Math.ceil( - (result.resetDate - Date.now()) / 1000 - ).toString(), + // Right now we have no way to know when the rate limit will reset + "retry-after": "60", }, }); } diff --git a/src/worker.ts b/src/worker.ts index f80ef2d..5b68f1f 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -7,17 +7,14 @@ export interface CloudflareEnvironment { CONFIG_FILES: R2Bucket; - RateLimiter: DurableObjectNamespace; - API_KEYS: KVNamespace; + RL_FREE: RateLimit; responseHeaders: Record; } const router = build(); -export { RateLimiter as RateLimiterDurableObject } from "./durable_objects/RateLimiter"; - export default { async fetch( request: Request, diff --git a/wrangler.toml b/wrangler.toml index 55edfda..69e7117 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -11,11 +11,38 @@ kv_namespaces = [ tag = "v1" # Should be unique for each entry new_classes = ["RateLimiterDurableObject"] -[durable_objects] -bindings = [ - {name = "RateLimiter", class_name = "RateLimiterDurableObject"}, -] +[[migrations]] +tag = "v2" +deleted_classes = ["RateLimiterDurableObject"] [[r2_buckets]] binding = "CONFIG_FILES" bucket_name = "zwave-js-firmware-updates--config-files" + +[[unsafe.bindings]] +type = "ratelimit" +# Free ratelimit bucket - 120 requests per hour +name = "RL_FREE" +namespace_id = "1001" +simple = { limit = 2, period = 60 } # requests per seconds + +[[unsafe.bindings]] +type = "ratelimit" +# Ratelimit bucket: 1200 requests per hour +name = "RL_1k" +namespace_id = "1002" +simple = { limit = 20, period = 60 } # requests per seconds + +[[unsafe.bindings]] +type = "ratelimit" +# Ratelimit bucket: 12000 requests per hour +name = "RL_10k" +namespace_id = "1003" +simple = { limit = 200, period = 60 } # requests per seconds + +[[unsafe.bindings]] +type = "ratelimit" +# Ratelimit bucket: 120000 requests per hour +name = "RL_100k" +namespace_id = "1004" +simple = { limit = 2000, period = 60 } # requests per seconds diff --git a/yarn.lock b/yarn.lock index 385dbdf..5763ad9 100644 --- a/yarn.lock +++ b/yarn.lock @@ -990,7 +990,6 @@ __metadata: eslint: ^8.23.1 eslint-config-prettier: ^8.5.0 eslint-plugin-prettier: ^4.2.1 - itty-durable: ^1.2.0 itty-router: ^2.6.1 itty-router-extras: ^0.4.2 json-logic-js: ^2.0.2 @@ -3154,16 +3153,6 @@ __metadata: languageName: node linkType: hard -"itty-durable@npm:^1.2.0": - version: 1.2.0 - resolution: "itty-durable@npm:1.2.0" - dependencies: - itty-router: ^2.6.1 - itty-router-extras: ^0.4.2 - checksum: 11e945a19b56c5e5ab896a7eaf042f4d2f8e760af528cac7d21468aac7e3c47a604365518720ffcda719dc2016b6babb5864f4ed21124594d892b11a7cab1e64 - languageName: node - linkType: hard - "itty-router-extras@npm:^0.4.2": version: 0.4.2 resolution: "itty-router-extras@npm:0.4.2"