Replace RateLimiter with RequestCounter

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<string, string>;
 
 interface Env {
-  RATE_LIMITER: DurableObjectNamespace<RateLimiter>;
+  REQUEST_COUNTER: DurableObjectNamespace<RequestCounter>;
 }
 
 
 export const onRequestPost: PagesFunction<Env & StringRecord> = 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<RequestStats> {
+  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<Response> {
   const url = `${env.dockerRunBaseUrl}/run`;
@@ -44,15 +64,18 @@ function run(env: EnvVars, body: ReadableStream): Promise<Response> {
 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 {

glot_cloudflare_rate_limiter/package.json → 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"
   }
-}
+}

glot_cloudflare_request_counter/src/index.ts

@@ -0,0 +1,37 @@
+import { RequestCounter } from './request_counter'
+export { RequestCounter };
+
+
+interface Env {
+	REQUEST_COUNTER: DurableObjectNamespace<RequestCounter>;
+}
+
+export default {
+	async fetch(request, env, ctx): Promise<Response> {
+		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<Env>;
+
+
+
+function errorResponse(status: number, message: string): Response {
+	return new Response(JSON.stringify({ message }), {
+		status,
+		headers: {
+			"Content-Type": "application/json",
+		},
+	});
+}

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<Response> {
+        const stats = await this.increment();
+        return new Response(JSON.stringify(stats), { status: 200 });
+    }
+
+    async increment(): Promise<RequestStats> {
+        const now = Date.now();
+
+        let state = await this.ctx.storage.get<State>("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,
+    };
+}