From d27446b021e9663e58e080156a9f638453e43a79 Mon Sep 17 00:00:00 2001 From: Kavish Shah Date: Wed, 7 May 2025 23:44:26 -0700 Subject: [PATCH] feat: add per-route rate-limit middleware --- middleware.ts | 127 ++++++++++++++++++++++++++++++++-------------- package-lock.json | 35 +++++++++++++ package.json | 2 + 3 files changed, 126 insertions(+), 38 deletions(-) diff --git a/middleware.ts b/middleware.ts index 56fccad..587ed70 100644 --- a/middleware.ts +++ b/middleware.ts @@ -1,59 +1,110 @@ -import { NextResponse } from "next/server"; +// middleware.ts +import { NextResponse, type NextRequest } from "next/server"; +import { Ratelimit } from "@upstash/ratelimit"; +import { Redis } from "@upstash/redis"; +import { getToken } from "next-auth/jwt"; import { AUTH_ROUTES, DEFAULT_LOGIN_REDIRECT, PROTECTED_BASE_ROUTES, PROTECTED_ROUTES, } from "./routes"; -import authConfig from "./auth.config"; -import NextAuth from "next-auth"; -// 2. Wrapped middleware option -const { auth } = NextAuth(authConfig); +// ─── Upstash Redis REST client (Edge) ─────────────────────────────── +const redis = new Redis({ + url: process.env.UPSTASH_REDIS_REST_URL!, + token: process.env.UPSTASH_REDIS_REST_TOKEN!, +}); -export default auth((req) => { - // Your custom middleware logic goes here - const currentPathname = req.nextUrl.pathname; - // !! converts the value into its boolean equivalent - const isLoggedIn = !!req.auth; +// ─── Define one Ratelimit instance per “task” ──────────────────────── +const pageLimiter = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(100, "60 s"), // 100 page‐views/min +}); - const isProtectedBaseRoute = PROTECTED_BASE_ROUTES.some((el) => - currentPathname.startsWith(el), - ); +const apiLimiter = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(60, "60 s"), // 60 general API calls/min +}); - if (isProtectedBaseRoute && !isLoggedIn) { - console.log( - "Access denied for not logged-in users trying to access a protected base route", - ); - return NextResponse.redirect(new URL("/auth/login", req.url)); +const chatLimiter = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(20, "60 s"), // 20 chat messages/min +}); + +const researchLimiter = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(10, "60 s"), // 10 research calls/min +}); + +const uploadLimiter = new Ratelimit({ + redis, + limiter: Ratelimit.slidingWindow(5, "60 s"), // 5 uploads/min +}); + +export async function middleware(req: NextRequest) { + // — Identify the client (by IP) — + const forwarded = req.headers.get("x-forwarded-for") ?? ""; + const ip = forwarded.split(",")[0] || "unknown"; + const path = req.nextUrl.pathname; + + // — Pick the right limiter based on path — + let limiter = pageLimiter; + if (path.startsWith("/api/chat")) { + limiter = chatLimiter; + } else if (path.startsWith("/api/research")) { + limiter = researchLimiter; + } else if ( + path.startsWith("/api/upload") || + path.startsWith("/api/bulk-upload") + ) { + limiter = uploadLimiter; + } else if (path.startsWith("/api/")) { + limiter = apiLimiter; } - if (AUTH_ROUTES.includes(currentPathname)) { - // we are accessing an auth route - if (isLoggedIn) { - console.log( - "access denied for accessing auth routes for logged in users", - ); - // we are already logged in so we cant access the auth routes anymore - return NextResponse.redirect(new URL(DEFAULT_LOGIN_REDIRECT, req.url)); - } + // — Apply rate limit — + const { success, limit, remaining } = await limiter.limit(ip); + if (!success) { + return new NextResponse("Too Many Requests", { + status: 429, + headers: { + "Retry-After": process.env.RATE_LIMIT_WINDOW || "60", + "X-RateLimit-Limit": String(limit), + "X-RateLimit-Remaining": String(remaining), + }, + }); } - if (PROTECTED_ROUTES.includes(currentPathname)) { - // we are accessing a protected routes - // check for valid sessions - // redirect unauthorized users + // — Push headers on every successful pass — + const response = NextResponse.next(); + response.headers.set("X-RateLimit-Limit", String(limit)); + response.headers.set("X-RateLimit-Remaining", String(remaining)); + + // — Authentication / Protected‐route logic (unchanged) — + const token = await getToken({ req, secret: process.env.NEXTAUTH_SECRET }); + const isLoggedIn = !!token; - if (!isLoggedIn) { - console.log("access denied for not logged in users"); - return NextResponse.redirect(new URL("/auth/login", req.url)); - } + if ( + PROTECTED_BASE_ROUTES.some((r) => path.startsWith(r)) && + !isLoggedIn + ) { + return NextResponse.redirect(new URL("/auth/login", req.url)); } - return NextResponse.next(); -}); + if (AUTH_ROUTES.includes(path) && isLoggedIn) { + return NextResponse.redirect( + new URL(DEFAULT_LOGIN_REDIRECT, req.url) + ); + } + + if (PROTECTED_ROUTES.includes(path) && !isLoggedIn) { + return NextResponse.redirect(new URL("/auth/login", req.url)); + } + + return response; +} -// run for all routes export const config = { matcher: ["/((?!.*\\..*|_next).*)", "/", "/(api|trpc)(.*)"], }; diff --git a/package-lock.json b/package-lock.json index 1cba5d7..5c6931b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -45,6 +45,8 @@ "@tanstack/react-table": "^8.20.6", "@types/axios": "^0.9.36", "@types/pdf-parse": "^1.1.4", + "@upstash/ratelimit": "^2.0.5", + "@upstash/redis": "^1.34.9", "@vercel/blob": "^0.26.0", "ai": "^4.3.10", "axios": "^1.8.4", @@ -3845,6 +3847,39 @@ "resolved": "https://registry.npmjs.org/@ungap/structured-clone/-/structured-clone-1.2.0.tgz", "integrity": "sha512-zuVdFrMJiuCDQUMCzQaD6KL28MjnqqN8XnAqiEq9PNm/hCPTSGfrXCOfwj1ow4LFb/tNymJPwsNbVePc1xFqrQ==" }, + "node_modules/@upstash/core-analytics": { + "version": "0.0.10", + "resolved": "https://registry.npmjs.org/@upstash/core-analytics/-/core-analytics-0.0.10.tgz", + "integrity": "sha512-7qJHGxpQgQr9/vmeS1PktEwvNAF7TI4iJDi8Pu2CFZ9YUGHZH4fOP5TfYlZ4aVxfopnELiE4BS4FBjyK7V1/xQ==", + "license": "MIT", + "dependencies": { + "@upstash/redis": "^1.28.3" + }, + "engines": { + "node": ">=16.0.0" + } + }, + "node_modules/@upstash/ratelimit": { + "version": "2.0.5", + "resolved": "https://registry.npmjs.org/@upstash/ratelimit/-/ratelimit-2.0.5.tgz", + "integrity": "sha512-1FRv0cs3ZlBjCNOCpCmKYmt9BYGIJf0J0R3pucOPE88R21rL7jNjXG+I+rN/BVOvYJhI9niRAS/JaSNjiSICxA==", + "license": "MIT", + "dependencies": { + "@upstash/core-analytics": "^0.0.10" + }, + "peerDependencies": { + "@upstash/redis": "^1.34.3" + } + }, + "node_modules/@upstash/redis": { + "version": "1.34.9", + "resolved": "https://registry.npmjs.org/@upstash/redis/-/redis-1.34.9.tgz", + "integrity": "sha512-7qzzF2FQP5VxR2YUNjemWs+hl/8VzJJ6fOkT7O7kt9Ct8olEVzb1g6/ik6B8Pb8W7ZmYv81SdlVV9F6O8bh/gw==", + "license": "MIT", + "dependencies": { + "crypto-js": "^4.2.0" + } + }, "node_modules/@vercel/blob": { "version": "0.26.0", "resolved": "https://registry.npmjs.org/@vercel/blob/-/blob-0.26.0.tgz", diff --git a/package.json b/package.json index 0ace3ef..3aef2bc 100644 --- a/package.json +++ b/package.json @@ -53,6 +53,8 @@ "@tanstack/react-table": "^8.20.6", "@types/axios": "^0.9.36", "@types/pdf-parse": "^1.1.4", + "@upstash/ratelimit": "^2.0.5", + "@upstash/redis": "^1.34.9", "@vercel/blob": "^0.26.0", "ai": "^4.3.10", "axios": "^1.8.4",