From 93f53fac3a8c3b73bbe25304dea9764badef7785 Mon Sep 17 00:00:00 2001
From: Peter Hughes <github@phugh.es>
Date: Wed, 4 Sep 2024 20:34:54 +0100
Subject: [PATCH] use kv for ratelimiting

---
 lib/middlewares/ratelimit.ts | 56 ++++++++++++++++++++----------------
 1 file changed, 31 insertions(+), 25 deletions(-)

diff --git a/lib/middlewares/ratelimit.ts b/lib/middlewares/ratelimit.ts
index 4b8eeef..52bd134 100644
--- a/lib/middlewares/ratelimit.ts
+++ b/lib/middlewares/ratelimit.ts
@@ -1,48 +1,53 @@
 import type { FreshContext } from "$fresh/server.ts";
+import { kv } from "lib/kv.ts";
 
 const WINDOW_SIZE = 60 * 1000; // 1 minute in milliseconds
 const MAX_REQUESTS = 100; // Maximum requests per minute
 
-interface RateLimitEntry {
-  timestamps: number[];
-  lastCleanup: number;
+function getClientIP(req: Request, addr: Deno.NetAddr): string {
+  const forwardedFor = req.headers.get("x-forwarded-for");
+  if (forwardedFor) {
+    // Take the first IP in the list
+    return forwardedFor.split(',')[0].trim();
+  }
+  // Fall back to the direct connection IP
+  return addr.hostname;
 }
 
-interface RateLimitStore {
-  [ip: string]: RateLimitEntry;
+async function getRateLimitEntry(ip: string): Promise<{ timestamps: number[], lastCleanup: number }> {
+  const entry = await kv.get<{ timestamps: number[], lastCleanup: number }>(['ratelimit', ip]);
+  return entry.value || { timestamps: [], lastCleanup: Date.now() };
 }
-const store: RateLimitStore = {};
 
-function cleanupStoreEntry(ip: string, now: number): void {
-  if (!store[ip]) return;
+async function setRateLimitEntry(ip: string, entry: { timestamps: number[], lastCleanup: number }): Promise<void> {
+  await kv.set(['ratelimit', ip], entry);
+}
 
-  store[ip].timestamps = store[ip].timestamps.filter((timestamp) => now - timestamp < WINDOW_SIZE);
-  store[ip].lastCleanup = now;
+async function cleanupStoreEntry(ip: string, now: number): Promise<void> {
+  const entry = await getRateLimitEntry(ip);
+  entry.timestamps = entry.timestamps.filter((timestamp) => now - timestamp < WINDOW_SIZE);
+  entry.lastCleanup = now;
 
-  if (store[ip].timestamps.length === 0) {
-    delete store[ip];
+  if (entry.timestamps.length === 0) {
+    await kv.delete(['ratelimit', ip]);
+  } else {
+    await setRateLimitEntry(ip, entry);
   }
 }
 
 export default async function handler(req: Request, ctx: FreshContext) {
-  const ip = req.headers.get("x-forwarded-for") || "unknown";
+  const ip = getClientIP(req, ctx.remoteAddr);
   const now = Date.now();
 
-  if (!store[ip]) {
-    store[ip] = { timestamps: [], lastCleanup: now };
-  }
+  let entry = await getRateLimitEntry(ip);
 
   // Perform cleanup if it's been more than WINDOW_SIZE since last cleanup
-  if (now - store[ip].lastCleanup >= WINDOW_SIZE) {
-    cleanupStoreEntry(ip, now);
-  }
-
-  // Ensure the entry exists after cleanup
-  if (!store[ip]) {
-    store[ip] = { timestamps: [], lastCleanup: now };
+  if (now - entry.lastCleanup >= WINDOW_SIZE) {
+    await cleanupStoreEntry(ip, now);
+    entry = await getRateLimitEntry(ip);
   }
 
-  const remaining = MAX_REQUESTS - store[ip].timestamps.length;
+  const remaining = MAX_REQUESTS - entry.timestamps.length;
   const isRateLimited = remaining <= 0;
 
   if (isRateLimited) {
@@ -61,7 +66,8 @@ export default async function handler(req: Request, ctx: FreshContext) {
     );
   }
 
-  store[ip].timestamps.push(now);
+  entry.timestamps.push(now);
+  await setRateLimitEntry(ip, entry);
 
   const resp = await ctx.next();
   const headers = resp.headers;