Skip to content

Commit

Permalink
use kv for ratelimiting
Browse files Browse the repository at this point in the history
  • Loading branch information
phughesmcr committed Sep 4, 2024
1 parent 0ab9e62 commit 93f53fa
Showing 1 changed file with 31 additions and 25 deletions.
56 changes: 31 additions & 25 deletions lib/middlewares/ratelimit.ts
Original file line number Diff line number Diff line change
@@ -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) {
Expand All @@ -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;
Expand Down

0 comments on commit 93f53fa

Please sign in to comment.