From 7cff78d01c63d8cb5f4c3f35dc33c9cb65746cef Mon Sep 17 00:00:00 2001
From: v1rtl <hi@v1rtl.site>
Date: Wed, 2 Oct 2024 20:34:16 +0300
Subject: [PATCH] add fileSizeLimit

---
 src/index.ts |  9 +++++++--
 test.ts      | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index a6ce6f4..f91471d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -12,7 +12,7 @@ export const hasBody = (method: string) => ['POST', 'PUT', 'PATCH', 'DELETE'].in
 
 const defaultPayloadLimit = 104857600 // 100KB
 
-export type LimitErrorFn = (payloadLimit: number) => Error
+export type LimitErrorFn = (limit: number) => Error
 
 export type ParserOptions = Partial<{
   payloadLimit: number
@@ -91,7 +91,9 @@ const getBoundary = (contentType: string) => {
   return match ? `--${match[1]}` : null
 }
 
-const parseMultipart = (body: string, boundary: string, { fileCountLimit }: MultipartOptions) => {
+const defaultFileSizeLimitErrorFn: LimitErrorFn = (limit) => new Error(`File too large. Limit: ${limit} bytes`)
+
+const parseMultipart = (body: string, boundary: string, { fileCountLimit, fileSizeLimit, fileSizeLimitErrorFn = defaultFileSizeLimitErrorFn }: MultipartOptions) => {
   // Split the body into an array of parts
   const parts = body.split(new RegExp(`${boundary}(--)?`)).filter((part) => !!part && /content-disposition/i.test(part))
   const parsedBody: Record<string, (File | string)[]> = {}
@@ -104,6 +106,8 @@ const parseMultipart = (body: string, boundary: string, { fileCountLimit }: Mult
     const [headers, ...lines] = part.split('\r\n').filter((part) => !!part)
     const data = lines.join('\r\n').trim()
 
+    if (fileSizeLimit && data.length > fileSizeLimit) throw fileSizeLimitErrorFn(fileSizeLimit)
+
     // Extract the name and filename from the headers
     const name = /name="(.+?)"/.exec(headers)![1]
     const filename = /filename="(.+?)"/.exec(headers)
@@ -126,6 +130,7 @@ const parseMultipart = (body: string, boundary: string, { fileCountLimit }: Mult
 type MultipartOptions = Partial<{
   fileCountLimit: number
   fileSizeLimit: number
+  fileSizeLimitErrorFn: LimitErrorFn
 }>
 
 const multipart =
diff --git a/test.ts b/test.ts
index ed540d3..5e9acb0 100644
--- a/test.ts
+++ b/test.ts
@@ -459,3 +459,39 @@ test('should throw multipart if amount of files exceeds limit', async () => {
     method: 'POST',
   }).expect(413, 'Too many files. Limit: 1')
 })
+
+test('should throw multipart if exceeds allowed file size', async () => {
+  const server = createServer(async (req: ReqWithBody, res) => {
+    await multipart({ fileSizeLimit: 10 })(req, res, (err) => {
+      if (err) res.writeHead(413).end(err.message)
+      else res.end(req.body)
+    })
+  })
+
+  const fd = new FormData()
+
+  fd.set('file', new File(['hello world'], 'hello.txt', { type: 'text/plain' }))
+
+  await makeFetch(server)('/', {
+    body: fd,
+    method: 'POST',
+  }).expect(413, 'File too large. Limit: 10 bytes')
+})
+
+test('should throw multipart if exceeds allowed file size with a custom error', async () => {
+  const server = createServer(async (req: ReqWithBody, res) => {
+    await multipart({ fileSizeLimit: 10, fileSizeLimitErrorFn: (limit) => new Error(`File too large. Limit: ${limit / 1024}KB`) })(req, res, (err) => {
+      if (err) res.writeHead(413).end(err.message)
+      else res.end(req.body)
+    })
+  })
+
+  const fd = new FormData()
+
+  fd.set('file', new File(['hello world'], 'hello.txt', { type: 'text/plain' }))
+
+  await makeFetch(server)('/', {
+    body: fd,
+    method: 'POST',
+  }).expect(413, 'File too large. Limit: 0.009765625KB')
+})
\ No newline at end of file