From ca9419a6fe3259d48fc745d013b23edbfc2f1d5f Mon Sep 17 00:00:00 2001 From: Petter Rasmussen Date: Tue, 23 Jul 2024 21:34:42 +0200 Subject: [PATCH] Replace RateLimiter with RequestCounter --- glot_cloudflare/functions/api/run.ts | 80 ++++++++++--- glot_cloudflare_rate_limiter/src/index.ts | 49 -------- .../src/rate_limiter.ts | 44 ------- .../.editorconfig | 0 .../.gitignore | 0 .../.prettierrc | 0 .../package-lock.json | 0 .../package.json | 4 +- glot_cloudflare_request_counter/src/index.ts | 37 ++++++ .../src/request_counter.ts | 113 ++++++++++++++++++ .../test/index.spec.ts | 0 .../test/tsconfig.json | 0 .../tsconfig.json | 0 .../vitest.config.mts | 0 .../worker-configuration.d.ts | 0 15 files changed, 214 insertions(+), 113 deletions(-) delete mode 100644 glot_cloudflare_rate_limiter/src/index.ts delete mode 100644 glot_cloudflare_rate_limiter/src/rate_limiter.ts rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/.editorconfig (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/.gitignore (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/.prettierrc (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/package-lock.json (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/package.json (93%) create mode 100644 glot_cloudflare_request_counter/src/index.ts create mode 100644 glot_cloudflare_request_counter/src/request_counter.ts rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/test/index.spec.ts (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/test/tsconfig.json (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/tsconfig.json (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/vitest.config.mts (100%) rename {glot_cloudflare_rate_limiter => glot_cloudflare_request_counter}/worker-configuration.d.ts (100%) diff --git a/glot_cloudflare/functions/api/run.ts b/glot_cloudflare/functions/api/run.ts index 6464501..c8bc7c0 100644 --- a/glot_cloudflare/functions/api/run.ts +++ b/glot_cloudflare/functions/api/run.ts @@ -1,32 +1,52 @@ -import { RateLimiter } from "../../../glot_cloudflare_rate_limiter/src/rate_limiter"; +import { RequestCounter, RequestStats } from "../../../glot_cloudflare_request_counter/src/request_counter"; type StringRecord = Record; interface Env { - RATE_LIMITER: DurableObjectNamespace; + REQUEST_COUNTER: DurableObjectNamespace; } export const onRequestPost: PagesFunction = async (context) => { + const envVars = parseEnvVars(context.env); + if (!isAllowed(context.request)) { return errorResponse(403, "Forbidden"); } - const ip = context.request.headers.get("CF-Connecting-IP"); - if (ip === null) { - return errorResponse(400, "Could not determine client IP"); + try { + const ip = getRequestIp(context.request); + const requestStats = await incrementRequestCount(context.env, context.request.clone(), ip); + if (isRatedLimited(envVars, requestStats)) { + return errorResponse(429, "Rate limit exceeded"); + } + + return run(envVars, context.request.body); + } catch (e) { + console.error("Failed to increment request count", e); } +}; - const id = context.env.RATE_LIMITER.idFromName(ip); - const stub = context.env.RATE_LIMITER.get(id); - const response = await stub.fetch(context.request.clone()); - const stats = await response.text(); +function getRequestIp(request: Request): string { + if (request.headers.has("CF-Connecting-IP")) { + return request.headers.get("CF-Connecting-IP") + } else { + return "127.0.0.1" + } +} - console.log(stats) +async function incrementRequestCount(env: Env, request: Request, ip: string): Promise { + const id = env.REQUEST_COUNTER.idFromName(ip); + const stub = env.REQUEST_COUNTER.get(id); + const response = await stub.fetch(request); + return response.json(); +} + +function isRatedLimited(env: EnvVars, stats: RequestStats): boolean { + // TODO + return false +} - const envVars = parseEnvVars(context.env); - return run(envVars, context.request.body); -}; function run(env: EnvVars, body: ReadableStream): Promise { const url = `${env.dockerRunBaseUrl}/run`; @@ -44,15 +64,18 @@ function run(env: EnvVars, body: ReadableStream): Promise { interface EnvVars { dockerRunBaseUrl: string; dockerRunAccessToken: string; + maxRequestsPerMinute: number; + maxRequestsPerHour: number; + maxRequestsPerDay: number; } function parseEnvVars(env: StringRecord): EnvVars { - ensureNotEmpty(env, "DOCKER_RUN_BASE_URL"); - ensureNotEmpty(env, "DOCKER_RUN_ACCESS_TOKEN"); - return { - dockerRunBaseUrl: env.DOCKER_RUN_BASE_URL, - dockerRunAccessToken: env.DOCKER_RUN_ACCESS_TOKEN, + dockerRunBaseUrl: getString(env, "DOCKER_RUN_BASE_URL"), + dockerRunAccessToken: getString(env, "DOCKER_RUN_ACCESS_TOKEN"), + maxRequestsPerMinute: getNumber(env, "MAX_REQUESTS_PER_MINUTE"), + maxRequestsPerHour: getNumber(env, "MAX_REQUESTS_PER_HOUR"), + maxRequestsPerDay: getNumber(env, "MAX_REQUESTS_PER_DAY"), }; } @@ -62,6 +85,27 @@ function ensureNotEmpty(env: StringRecord, field: string) { } } +function ensureInt(env: StringRecord, field: string) { + ensureNotEmpty(env, field); + + const n = parseInt(env[field], 10); + if (isNaN(n)) { + throw new Error(`Invalid number for env var ${field}`); + } +} + +function getString(env: StringRecord, field: string): string { + ensureNotEmpty(env, field); + return env[field]; +} + +function getNumber(env: StringRecord, field: string): number { + ensureNotEmpty(env, field); + ensureInt(env, field); + + return parseInt(env[field], 10); +} + function isAllowed(request: Request): boolean { diff --git a/glot_cloudflare_rate_limiter/src/index.ts b/glot_cloudflare_rate_limiter/src/index.ts deleted file mode 100644 index bb510d9..0000000 --- a/glot_cloudflare_rate_limiter/src/index.ts +++ /dev/null @@ -1,49 +0,0 @@ -import { RateLimiter } from './rate_limiter' -export { RateLimiter }; - -/** - * Welcome to Cloudflare Workers! This is your first worker. - * - * - Run `npm run dev` in your terminal to start a development server - * - Open a browser tab at http://localhost:8787/ to see your worker in action - * - Run `npm run deploy` to publish your worker - * - * Bind resources to your worker in `wrangler.toml`. After adding bindings, a type definition for the - * `Env` object can be regenerated with `npm run cf-typegen`. - * - * Learn more at https://developers.cloudflare.com/workers/ - */ - -interface Env { - RATE_LIMITER: DurableObjectNamespace; -} - -export default { - async fetch(request, env, ctx): Promise { - const ip = request.headers.get("CF-Connecting-IP"); - if (ip === null) { - return errorResponse(400, "Could not determine client IP"); - } - - try { - const id = env.RATE_LIMITER.idFromName(ip); - const stub = env.RATE_LIMITER.get(id); - const stats = await stub.increment({ maxRequests: 10, periodDuration: 60 * 1000 }); - - return new Response(JSON.stringify(stats)); - } catch (e) { - return errorResponse(500, "Internal server error"); - } - }, -} satisfies ExportedHandler; - - - -function errorResponse(status: number, message: string): Response { - return new Response(JSON.stringify({ message }), { - status, - headers: { - "Content-Type": "application/json", - }, - }); -} diff --git a/glot_cloudflare_rate_limiter/src/rate_limiter.ts b/glot_cloudflare_rate_limiter/src/rate_limiter.ts deleted file mode 100644 index e5afa2c..0000000 --- a/glot_cloudflare_rate_limiter/src/rate_limiter.ts +++ /dev/null @@ -1,44 +0,0 @@ -import { DurableObject } from "cloudflare:workers"; - - -export interface RateLimitConfig { - maxRequests: number; - periodDuration: number; -} - -export interface RateLimitStats { - timeUntilReset: number; - requestCount: number; - remainingRequests: number; -} - -export class RateLimiter extends DurableObject { - private startTimestamp: number = 0; - private requestCount: number = 0; - - async fetch(request: Request): Promise { - const stats = await this.increment({ maxRequests: 10, periodDuration: 60 * 1000 }); - return new Response(JSON.stringify(stats), { status: 200 }); - } - - async increment(config: RateLimitConfig): Promise { - const now = Date.now(); - const elapsedTime = now - this.startTimestamp; - const timeUntilReset = config.periodDuration - elapsedTime; - - if (timeUntilReset <= 0) { - this.startTimestamp = now; - this.requestCount = 0; - } - - this.requestCount++; - - const newTimeUntilReset = config.periodDuration - (now - this.startTimestamp); - - return { - timeUntilReset: newTimeUntilReset, - requestCount: this.requestCount, - remainingRequests: Math.max(0, config.maxRequests - this.requestCount), - }; - } -} \ No newline at end of file diff --git a/glot_cloudflare_rate_limiter/.editorconfig b/glot_cloudflare_request_counter/.editorconfig similarity index 100% rename from glot_cloudflare_rate_limiter/.editorconfig rename to glot_cloudflare_request_counter/.editorconfig diff --git a/glot_cloudflare_rate_limiter/.gitignore b/glot_cloudflare_request_counter/.gitignore similarity index 100% rename from glot_cloudflare_rate_limiter/.gitignore rename to glot_cloudflare_request_counter/.gitignore diff --git a/glot_cloudflare_rate_limiter/.prettierrc b/glot_cloudflare_request_counter/.prettierrc similarity index 100% rename from glot_cloudflare_rate_limiter/.prettierrc rename to glot_cloudflare_request_counter/.prettierrc diff --git a/glot_cloudflare_rate_limiter/package-lock.json b/glot_cloudflare_request_counter/package-lock.json similarity index 100% rename from glot_cloudflare_rate_limiter/package-lock.json rename to glot_cloudflare_request_counter/package-lock.json diff --git a/glot_cloudflare_rate_limiter/package.json b/glot_cloudflare_request_counter/package.json similarity index 93% rename from glot_cloudflare_rate_limiter/package.json rename to glot_cloudflare_request_counter/package.json index 6210d36..9a377ac 100644 --- a/glot_cloudflare_rate_limiter/package.json +++ b/glot_cloudflare_request_counter/package.json @@ -1,5 +1,5 @@ { - "name": "rate-limiter", + "name": "request-counter", "version": "0.0.0", "private": true, "scripts": { @@ -16,4 +16,4 @@ "vitest": "1.5.0", "wrangler": "^3.60.3" } -} +} \ No newline at end of file diff --git a/glot_cloudflare_request_counter/src/index.ts b/glot_cloudflare_request_counter/src/index.ts new file mode 100644 index 0000000..fcff037 --- /dev/null +++ b/glot_cloudflare_request_counter/src/index.ts @@ -0,0 +1,37 @@ +import { RequestCounter } from './request_counter' +export { RequestCounter }; + + +interface Env { + REQUEST_COUNTER: DurableObjectNamespace; +} + +export default { + async fetch(request, env, ctx): Promise { + const ip = request.headers.get("CF-Connecting-IP"); + if (ip === null) { + return errorResponse(400, "Could not determine client IP"); + } + + try { + const id = env.REQUEST_COUNTER.idFromName(ip); + const stub = env.REQUEST_COUNTER.get(id); + const stats = await stub.increment(); + + return new Response(JSON.stringify(stats)); + } catch (e) { + return errorResponse(500, "Internal server error"); + } + }, +} satisfies ExportedHandler; + + + +function errorResponse(status: number, message: string): Response { + return new Response(JSON.stringify({ message }), { + status, + headers: { + "Content-Type": "application/json", + }, + }); +} \ No newline at end of file diff --git a/glot_cloudflare_request_counter/src/request_counter.ts b/glot_cloudflare_request_counter/src/request_counter.ts new file mode 100644 index 0000000..eef0b25 --- /dev/null +++ b/glot_cloudflare_request_counter/src/request_counter.ts @@ -0,0 +1,113 @@ +import { DurableObject } from "cloudflare:workers"; + + +interface State { + minutely: Period; + hourly: Period; + daily: Period; + totalCount: number; +} + +function initState(): State { + return { + minutely: minutelyPeriod(), + hourly: hourlyPeriod(), + daily: dailyPeriod(), + totalCount: 0, + }; +} + +export class RequestCounter extends DurableObject { + async fetch(request: Request): Promise { + const stats = await this.increment(); + return new Response(JSON.stringify(stats), { status: 200 }); + } + + async increment(): Promise { + const now = Date.now(); + + let state = await this.ctx.storage.get("state") + if (!state) { + state = initState(); + } + + incrementPeriodRequest(state.minutely, now); + incrementPeriodRequest(state.hourly, now); + incrementPeriodRequest(state.daily, now); + state.totalCount++; + + await this.ctx.storage.put("state", state); + + return { + minutely: getPeriodStats(state.minutely, now), + hourly: getPeriodStats(state.hourly, now), + daily: getPeriodStats(state.daily, now), + totalCount: state.totalCount, + }; + } +} + +interface Period { + startTimestamp: number + count: number; + duration: number; +} + +function minutelyPeriod(): Period { + return { + startTimestamp: 0, + count: 0, + duration: 60 * 1000, + }; +} + +function hourlyPeriod(): Period { + return { + startTimestamp: 0, + count: 0, + duration: 60 * 60 * 1000, + }; +} + +function dailyPeriod(): Period { + return { + startTimestamp: 0, + count: 0, + duration: 24 * 60 * 60 * 1000, + }; +} + +function incrementPeriodRequest(period: Period, now: number): void { + const elapsedTime = now - period.startTimestamp; + const timeUntilReset = period.duration - elapsedTime; + + if (timeUntilReset <= 0) { + period.startTimestamp = now; + period.count = 0; + } + + period.count++; +} + +export interface RequestStats { + minutely: PeriodStats; + hourly: PeriodStats; + daily: PeriodStats; + totalCount: number; +} + +export interface PeriodStats { + count: number; + timeUntilReset: number; +} + + +function getPeriodStats(period: Period, now: number): PeriodStats { + const elapsedTime = now - period.startTimestamp; + const timeUntilReset = period.duration - elapsedTime; + + return { + count: period.count, + timeUntilReset, + }; +} \ No newline at end of file diff --git a/glot_cloudflare_rate_limiter/test/index.spec.ts b/glot_cloudflare_request_counter/test/index.spec.ts similarity index 100% rename from glot_cloudflare_rate_limiter/test/index.spec.ts rename to glot_cloudflare_request_counter/test/index.spec.ts diff --git a/glot_cloudflare_rate_limiter/test/tsconfig.json b/glot_cloudflare_request_counter/test/tsconfig.json similarity index 100% rename from glot_cloudflare_rate_limiter/test/tsconfig.json rename to glot_cloudflare_request_counter/test/tsconfig.json diff --git a/glot_cloudflare_rate_limiter/tsconfig.json b/glot_cloudflare_request_counter/tsconfig.json similarity index 100% rename from glot_cloudflare_rate_limiter/tsconfig.json rename to glot_cloudflare_request_counter/tsconfig.json diff --git a/glot_cloudflare_rate_limiter/vitest.config.mts b/glot_cloudflare_request_counter/vitest.config.mts similarity index 100% rename from glot_cloudflare_rate_limiter/vitest.config.mts rename to glot_cloudflare_request_counter/vitest.config.mts diff --git a/glot_cloudflare_rate_limiter/worker-configuration.d.ts b/glot_cloudflare_request_counter/worker-configuration.d.ts similarity index 100% rename from glot_cloudflare_rate_limiter/worker-configuration.d.ts rename to glot_cloudflare_request_counter/worker-configuration.d.ts