From 936ccc7d8c88ef2d3b94c2e6179aaeab9f8a10e5 Mon Sep 17 00:00:00 2001
From: RinYato <chearithorn@gmail.com>
Date: Sun, 31 Mar 2024 09:25:39 +0700
Subject: [PATCH] feat: better file type detection (#63)

* chore: update dep

* feat: better file type detection
---
 package-lock.json           |  6 ++++++
 package.json                |  1 +
 src/app/api/upload/route.ts | 15 ++++++++++-----
 src/lib/r2.ts               |  2 +-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 503ebd12..dcb65b63 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -55,6 +55,7 @@
         "eslint-plugin-jest": "^27.6.3",
         "lucia": "^3.1.1",
         "lucide-react": "^0.309.0",
+        "magic-bytes.js": "^1.10.0",
         "next": "14.0.4",
         "oslo": "^1.1.3",
         "react": "^18",
@@ -12232,6 +12233,11 @@
         "lz-string": "bin/bin.js"
       }
     },
+    "node_modules/magic-bytes.js": {
+      "version": "1.10.0",
+      "resolved": "https://registry.npmjs.org/magic-bytes.js/-/magic-bytes.js-1.10.0.tgz",
+      "integrity": "sha512-/k20Lg2q8LE5xiaaSkMXk4sfvI+9EGEykFS4b0CHHGWqDYU0bGUFSwchNOMA56D7TCs9GwVTkqe9als1/ns8UQ=="
+    },
     "node_modules/make-dir": {
       "version": "4.0.0",
       "resolved": "https://registry.npmjs.org/make-dir/-/make-dir-4.0.0.tgz",
diff --git a/package.json b/package.json
index 773c968a..f379f0ce 100644
--- a/package.json
+++ b/package.json
@@ -67,6 +67,7 @@
     "eslint-plugin-jest": "^27.6.3",
     "lucia": "^3.1.1",
     "lucide-react": "^0.309.0",
+    "magic-bytes.js": "^1.10.0",
     "next": "14.0.4",
     "oslo": "^1.1.3",
     "react": "^18",
diff --git a/src/app/api/upload/route.ts b/src/app/api/upload/route.ts
index b07ee88a..44e23d52 100644
--- a/src/app/api/upload/route.ts
+++ b/src/app/api/upload/route.ts
@@ -11,6 +11,7 @@ import { generateId } from "lucia";
 import { ApiError } from "@/lib/api-error";
 import { HttpStatus } from "@/constants/http-status";
 import { env } from "@/env";
+import { filetypeinfo } from "magic-bytes.js";
 
 export const runtime = "edge";
 
@@ -61,20 +62,24 @@ export const POST = withUser(async ({ req, user }) => {
     });
   }
 
+  // we can do compression or resizing here and return a new buffer
+  const buffer = new Uint8Array(await file.arrayBuffer());
+
   // e.g. image/png -> png
-  const fileExtenstion = file.type.split("/")[1] || "";
+  const fallbackFileType = file.type.split("/")[1] || "";
+  const fallbackExtension = file.name.split(".").pop() || fallbackFileType;
 
-  // we can do compression or resizing here and return a new buffer
-  const buffer = Buffer.from(await file.arrayBuffer());
+  const { extension = fallbackExtension, mime = file.type } =
+    filetypeinfo(buffer)[0];
 
   // hash the content and append the file extension
-  const hashedFilename = concat(await hash(buffer), ".", fileExtenstion);
+  const hashedFilename = concat(await hash(buffer), ".", extension);
 
   const uploadToR2 = await R2.upload({
     buffer,
     userId: user.id,
     filename: hashedFilename,
-    contentType: file.type,
+    contentType: mime,
   })
     .then(ok)
     .catch(err);
diff --git a/src/lib/r2.ts b/src/lib/r2.ts
index 04723760..365bf1c2 100644
--- a/src/lib/r2.ts
+++ b/src/lib/r2.ts
@@ -16,7 +16,7 @@ async function upload({
   filename,
   contentType,
 }: {
-  buffer: Buffer;
+  buffer: Buffer | Uint8Array;
   userId: string;
   filename: string;
   contentType: string;