From 701f0389718fb698055c1ab2ba5e66306c39ecd0 Mon Sep 17 00:00:00 2001
From: trafficlunar <hello@trafficlunar.net>
Date: Sat, 26 Jul 2025 15:50:20 +0100
Subject: [PATCH] fix: rate limit on images not working properly

---
 src/app/mii/[id]/image/route.ts | 24 +++++++++++-------------
 src/lib/rate-limit.ts           | 14 +++++++++++---
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/app/mii/[id]/image/route.ts b/src/app/mii/[id]/image/route.ts
index e045a1e..c361d82 100644
--- a/src/app/mii/[id]/image/route.ts
+++ b/src/app/mii/[id]/image/route.ts
@@ -1,4 +1,4 @@
-import { NextRequest, NextResponse } from "next/server";
+import { NextRequest } from "next/server";
 
 import fs from "fs/promises";
 import path from "path";
@@ -78,6 +78,8 @@ export async function GET(request: NextRequest, { params }: { params: Promise<{
 		}
 	}
 
+	if (!buffer) return rateLimit.sendResponse({ error: "Image not found" }, 404);
+
 	// Set the file name for the metadata image in the response for SEO
 	if (mii && imageType === "metadata") {
 		const slugify = (str: string) =>
@@ -92,20 +94,16 @@ export async function GET(request: NextRequest, { params }: { params: Promise<{
 
 		const filename = `${name}-mii-${tags}-by-${username}.png`;
 
-		return new NextResponse(buffer, {
-			headers: {
-				"Content-Type": "image/png",
-				"Content-Disposition": `inline; filename="${filename}"`,
-				"Cache-Control": "public, max-age=31536000",
-			},
+		return rateLimit.sendResponse(buffer, 200, {
+			"Content-Type": "image/png",
+			"Content-Disposition": `inline; filename="${filename}"`,
+			"Cache-Control": "public, max-age=31536000",
 		});
 	}
 
-	return new NextResponse(buffer, {
-		headers: {
-			"Content-Type": "image/webp",
-			"X-Robots-Tag": "noindex, noimageindex, nofollow",
-			"Cache-Control": "no-store",
-		},
+	return rateLimit.sendResponse(buffer, 200, {
+		"Content-Type": "image/webp",
+		"X-Robots-Tag": "noindex, noimageindex, nofollow",
+		"Cache-Control": "no-store",
 	});
 }
diff --git a/src/lib/rate-limit.ts b/src/lib/rate-limit.ts
index 47aef3b..c690121 100644
--- a/src/lib/rate-limit.ts
+++ b/src/lib/rate-limit.ts
@@ -64,16 +64,24 @@ export class RateLimit {
 	}
 
 	// Attach rate limit headers to a response
-	sendResponse(message: object, status: number = 200): NextResponse<object> {
-		const response = NextResponse.json(message, { status });
+	sendResponse(body: object | Buffer, status: number = 200, headers?: HeadersInit): NextResponse<object | unknown> {
+		let response: NextResponse;
+
+		if (Buffer.isBuffer(body)) {
+			response = new NextResponse(body, { status, headers });
+		} else {
+			response = NextResponse.json(body, { status, headers });
+		}
+
 		response.headers.set("X-RateLimit-Limit", this.data.limit.toString());
 		response.headers.set("X-RateLimit-Remaining", this.data.remaining.toString());
 		response.headers.set("X-RateLimit-Expires", this.data.expires.toString());
+
 		return response;
 	}
 
 	// Handle both functions above and identifier in one
-	async handle(): Promise<NextResponse<object> | undefined> {
+	async handle(): Promise<NextResponse<object | unknown> | undefined> {
 		const session = await auth();
 		const ip = this.request.headers.get("CF-Connecting-IP") || this.request.headers.get("X-Forwarded-For")?.split(",")[0];
 		const identifier = (session ? session.user.id : ip) ?? "null";