fix: rate limit on images not working properly

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, {
+		return rateLimit.sendResponse(buffer, 200, {
-			headers: {
 			"Content-Type": "image/png",
 			"Content-Disposition": `inline; filename="${filename}"`,
 			"Cache-Control": "public, max-age=31536000",
-			},
 		});
 	}
 
-	return new NextResponse(buffer, {
+	return rateLimit.sendResponse(buffer, 200, {
-		headers: {
 		"Content-Type": "image/webp",
 		"X-Robots-Tag": "noindex, noimageindex, nofollow",
 		"Cache-Control": "no-store",
-		},
 	});
 }

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> {
+	sendResponse(body: object | Buffer, status: number = 200, headers?: HeadersInit): NextResponse<object | unknown> {
-		const response = NextResponse.json(message, { status });
+		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";