From 9ee12726d15ccecb269ee99ab90e081c60fa3ee7 Mon Sep 17 00:00:00 2001
From: Aron Carroll <aron@replicate.com>
Date: Fri, 8 Dec 2023 11:19:04 +0000
Subject: [PATCH] Add ratelimiting for api routes

---
 middleware.ts | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)
 create mode 100644 middleware.ts

diff --git a/middleware.ts b/middleware.ts
new file mode 100644
index 0000000..5182850
--- /dev/null
+++ b/middleware.ts
@@ -0,0 +1,35 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { Ratelimit } from '@upstash/ratelimit';
+import { kv } from '@vercel/kv';
+
+const ratelimit = new Ratelimit({
+  redis: kv,
+  // 20 requests from the same IP within a 10 second sliding window
+  limiter: Ratelimit.slidingWindow(20, '10s'),
+  prefix: `v2/zoo/ratelimit/${process.env.VERCEL_ENV ?? 'local'}`,
+});
+
+// Rate limit the /api/predictions/[id] endpoint
+export const config = {
+  matcher: ['/api/predictions/:path+'],
+};
+
+export default async function middleware(request: NextRequest) {
+  if (!process.env.VERCEL_ENV || !process.env.KV_REST_API_URL || !process.env.KV_REST_API_URL) {
+    console.warn('Skipping ratelimiting middleware');
+    return NextResponse.next();
+  }
+
+  const ip = request.ip ?? '127.0.0.1';
+  const { success, limit, remaining, reset } = await ratelimit.limit(ip);
+  const headers = {
+    'X-Ratelimit-Hit': String(!success),
+    'X-Ratelimit-Limit': String(limit),
+    'X-Ratelimit-Remaining': String(remaining),
+    'X-Ratelimit-Reset': String(reset),
+  }
+
+  return success
+    ? NextResponse.next({headers})
+    : NextResponse.json({}, { status: 429, headers });
+}