diff --git a/src/middleware.ts b/src/middleware.ts new file mode 100644 index 0000000..615e699 --- /dev/null +++ b/src/middleware.ts @@ -0,0 +1,199 @@ +import { getToken } from "next-auth/jwt"; +import { NextRequest, NextResponse } from "next/server"; + +const WINDOW_SECONDS = 60; +const AUTHENTICATED_LIMIT = 60; +const ANONYMOUS_LIMIT = 10; +const memoryBuckets = new Map(); + +type RateLimitResult = { + allowed: boolean; + limit: number; + remaining: number; + reset: number; +}; + +function getIp(req: NextRequest) { + return ( + req.ip ?? + req.headers.get("x-forwarded-for")?.split(",")[0]?.trim() ?? + req.headers.get("x-real-ip") ?? + "unknown" + ); +} + +function buildHeaders(result: RateLimitResult) { + const headers = new Headers(); + headers.set("X-RateLimit-Limit", String(result.limit)); + headers.set("X-RateLimit-Remaining", String(result.remaining)); + headers.set("X-RateLimit-Reset", String(result.reset)); + + if (!result.allowed) { + headers.set( + "Retry-After", + String(Math.max(result.reset - Math.floor(Date.now() / 1000), 1)) + ); + } + + return headers; +} + +function pruneMemoryBuckets(now: number) { + if (memoryBuckets.size < 500) { + return; + } + + const cutoff = now - WINDOW_SECONDS * 1000; + for (const [key, values] of Array.from(memoryBuckets.entries())) { + const active = values.filter((timestamp: number) => timestamp > cutoff); + if (active.length === 0) { + memoryBuckets.delete(key); + } else { + memoryBuckets.set(key, active); + } + } +} + +function checkMemoryLimit( + key: string, + limit: number, + now: number +): RateLimitResult { + pruneMemoryBuckets(now); + + const cutoff = now - WINDOW_SECONDS * 1000; + const active = (memoryBuckets.get(key) ?? []).filter( + (timestamp) => timestamp > cutoff + ); + const reset = Math.ceil( + ((active[0] ?? now) + WINDOW_SECONDS * 1000) / 1000 + ); + + if (active.length >= limit) { + memoryBuckets.set(key, active); + return { + allowed: false, + limit, + remaining: 0, + reset, + }; + } + + active.push(now); + memoryBuckets.set(key, active); + + return { + allowed: true, + limit, + remaining: Math.max(limit - active.length, 0), + reset, + }; +} + +async function checkUpstashLimit( + key: string, + limit: number, + now: number +): Promise { + const url = process.env.UPSTASH_REDIS_REST_URL; + const token = process.env.UPSTASH_REDIS_REST_TOKEN; + + if (!url || !token) { + return null; + } + + const cutoff = now - WINDOW_SECONDS * 1000; + const reset = Math.ceil((now + WINDOW_SECONDS * 1000) / 1000); + + try { + const response = await fetch(`${url}/pipeline`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify([ + ["ZREMRANGEBYSCORE", key, 0, cutoff], + ["ZCARD", key], + ]), + cache: "no-store", + }); + + if (!response.ok) { + return null; + } + + const pipeline = (await response.json()) as Array<{ result?: number }>; + const previousCount = Number(pipeline[1]?.result ?? 0); + + if (previousCount >= limit) { + return { + allowed: false, + limit, + remaining: 0, + reset, + }; + } + + await fetch(`${url}/pipeline`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify([ + ["ZADD", key, now, `${now}:${Math.random().toString(36).slice(2)}`], + ["EXPIRE", key, WINDOW_SECONDS], + ]), + cache: "no-store", + }); + + return { + allowed: true, + limit, + remaining: Math.max(limit - previousCount - 1, 0), + reset, + }; + } catch { + return null; + } +} + +async function checkRateLimit(identifier: string, limit: number) { + const now = Date.now(); + const key = `metrics-rate-limit:${identifier}`; + return ( + (await checkUpstashLimit(key, limit, now)) ?? + checkMemoryLimit(key, limit, now) + ); +} + +export async function middleware(req: NextRequest) { + const token = await getToken({ req, secret: process.env.NEXTAUTH_SECRET }); + const githubId = typeof token?.githubId === "string" ? token.githubId : null; + const identifier = githubId ? `user:${githubId}` : `ip:${getIp(req)}`; + const limit = githubId ? AUTHENTICATED_LIMIT : ANONYMOUS_LIMIT; + const result = await checkRateLimit(identifier, limit); + const headers = buildHeaders(result); + + if (!result.allowed) { + console.warn("metrics_rate_limit_hit", { + identifier, + path: req.nextUrl.pathname, + limit, + }); + + return NextResponse.json( + { error: "Too many metrics requests. Please retry shortly." }, + { status: 429, headers } + ); + } + + const response = NextResponse.next(); + headers.forEach((value, key) => response.headers.set(key, value)); + return response; +} + +export const config = { + matcher: "/api/metrics/:path*", +};