fix(upload): derive R2 key extension from content type, not filename (#77)

* fix(upload): derive R2 key extension from content type, not filename

getPresignedUploadUrl trusted the extension off the user-supplied filename even though contentType is already validated against ALLOWED_TYPES in the presign route. A hostile filename could put an arbitrary extension on a public object key. Map the validated content type to an extension instead, and keep the original filename only as sanitized object metadata.

* style: prettier

---------

Co-authored-by: royalpinto007 <royalpinto007@gmail.com>

Author: Renji

lib/r2/upload.test.ts

@@ -0,0 +1,71 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const { putObjectCommand, getSignedUrl } = vi.hoisted(() => ({
+  putObjectCommand: vi.fn((input: unknown) => ({ input })),
+  getSignedUrl: vi.fn(async () => "https://upload.example/signed"),
+}));
+
+vi.mock("@aws-sdk/client-s3", () => ({
+  S3Client: vi.fn(() => ({})),
+  PutObjectCommand: putObjectCommand,
+  GetObjectCommand: vi.fn((input: unknown) => ({ input })),
+}));
+
+vi.mock("@aws-sdk/s3-request-presigner", () => ({
+  getSignedUrl,
+}));
+
+import { getPresignedUploadUrl } from "./upload";
+
+describe("getPresignedUploadUrl", () => {
+  const originalEnv = { ...process.env };
+
+  beforeEach(() => {
+    process.env.R2_ACCOUNT_ID = "test-account";
+    process.env.R2_ACCESS_KEY_ID = "test-key";
+    process.env.R2_SECRET_ACCESS_KEY = "test-secret";
+    process.env.R2_BUCKET_NAME = "test-bucket";
+    process.env.R2_PUBLIC_URL = "https://cdn.example.com";
+    putObjectCommand.mockClear();
+  });
+
+  afterEach(() => {
+    process.env = { ...originalEnv };
+  });
+
+  it.each([
+    ["image/jpeg", "jpg"],
+    ["image/png", "png"],
+    ["image/webp", "webp"],
+    ["image/gif", "gif"],
+  ])("maps %s to a .%s key", async (contentType, ext) => {
+    const { key } = await getPresignedUploadUrl("whatever.html", contentType);
+
+    expect(key).toMatch(new RegExp(`^screenshots/[^/]+\\.${ext}$`));
+  });
+
+  it("ignores a hostile filename with a disallowed extension", async () => {
+    const { key } = await getPresignedUploadUrl("payload.html", "image/png");
+
+    expect(key.endsWith(".html")).toBe(false);
+    expect(key.endsWith(".png")).toBe(true);
+  });
+
+  it("falls back to .bin for an unrecognized content type", async () => {
+    const { key } = await getPresignedUploadUrl(
+      "file",
+      "application/octet-stream",
+    );
+
+    expect(key).toMatch(/^screenshots\/[^/]+\.bin$/);
+  });
+
+  it("passes the sanitized original filename as object metadata, not the key", async () => {
+    await getPresignedUploadUrl("my photo #1!.png", "image/png");
+
+    const command = putObjectCommand.mock.calls[0][0] as {
+      Metadata?: Record<string, string>;
+    };
+    expect(command.Metadata?.["original-filename"]).toBe("my_photo__1_.png");
+  });
+});

lib/r2/upload.ts

@@ -9,6 +9,27 @@ import { getR2PublicBaseUrl } from "@/lib/utils/urls";
 
 const R2_ENDPOINT = `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com`;
 
+const EXT_BY_CONTENT_TYPE: Record<string, string> = {
+  "image/jpeg": "jpg",
+  "image/png": "png",
+  "image/webp": "webp",
+  "image/gif": "gif",
+};
+
+/**
+ * Maps a validated content type to a storage extension. Never derive the
+ * extension from a user-supplied filename — it isn't checked against the
+ * allowlist the route validates `contentType` with.
+ */
+function extensionForContentType(contentType: string): string {
+  return EXT_BY_CONTENT_TYPE[contentType] ?? "bin";
+}
+
+/** Strips the filename down to characters safe for an S3 metadata header. */
+function sanitizeFilenameForMetadata(filename: string): string {
+  return filename.replace(/[^\w.-]/g, "_").slice(0, 255);
+}
+
 function getR2Client() {
   return new S3Client({
     region: "auto",
@@ -38,14 +59,15 @@ export async function getPresignedUploadUrl(
   contentType: string,
   folder = "screenshots",
 ): Promise<PresignUploadResult> {
-  const ext = filename.split(".").pop() ?? "bin";
+  const ext = extensionForContentType(contentType);
   const key = `${folder}/${randomUUID()}.${ext}`;
 
   const client = getR2Client();
   const command = new PutObjectCommand({
     Bucket: process.env.R2_BUCKET_NAME!,
     Key: key,
     ContentType: contentType,
+    Metadata: { "original-filename": sanitizeFilenameForMetadata(filename) },
   });
 
   const uploadUrl = await getSignedUrl(client, command, { expiresIn: 300 });