Fix high severity network attack vector in multipart implementation (#57)

Author: abersheeran

baize/asgi/requests.py

@@ -176,6 +176,13 @@ async def json(self) -> Any:
 
         raise UnsupportedMediaType("application/json")
 
+    async def _parse_multipart(self, boundary: bytes, charset: str) -> FormData:
+        return FormData(
+            await parse_multipart(
+                self.stream(), boundary, charset, file_factory=UploadFile
+            )
+        )
+
     @cached_property
     async def form(self) -> FormData:
         """
@@ -193,11 +200,7 @@ async def form(self) -> FormData:
             if "boundary" not in self.content_type.options:
                 raise MalformedMultipart("Missing boundary in header content-type")
             boundary = self.content_type.options["boundary"].encode("latin-1")
-            return FormData(
-                await parse_multipart(
-                    self.stream(), boundary, charset, file_factory=UploadFile
-                )
-            )
+            return await self._parse_multipart(boundary, charset)
         if self.content_type == "application/x-www-form-urlencoded":
             body = (await self.body).decode(
                 encoding=self.content_type.options.get("charset", "latin-1")

baize/exceptions.py

@@ -36,6 +36,17 @@ def abort(
     raise HTTPException(status_code=status_code, headers=headers, content=content)
 
 
+class RequestEntityTooLarge(HTTPException[None]):
+    """
+    413 Request Entity Too Large
+    """
+
+    def __init__(self, retry_after: Optional[int] = None) -> None:
+        super().__init__(
+            413, {"Retry-After": str(retry_after)} if retry_after is not None else None
+        )
+
+
 class UnsupportedMediaType(HTTPException[None]):
     """
     415 Unsupported Media Type

baize/multipart_helper.py

@@ -1,6 +1,7 @@
 from typing import AsyncIterable, Iterable, List, Optional, Tuple, Type, TypeVar, Union
 
 from .datastructures import Headers
+from .exceptions import RequestEntityTooLarge
 from .multipart import (
     Data,
     Epilogue,
@@ -47,6 +48,8 @@ async def parse_async_stream(
     charset: str,
     *,
     file_factory: Type[_AsyncUploadFile],
+    max_form_parts: int = 324,
+    max_form_memory_size: Union[int, None] = None,
 ) -> List[Tuple[str, Union[str, _AsyncUploadFile]]]:
     """
     Parse an asynchronous stream in multipart format
@@ -60,6 +63,8 @@ async def parse_async_stream(
     field_name = ""
     data = bytearray()
     file: Optional[_AsyncUploadFile] = None
+    form_parts_count = 0
+    form_memory_size_count = 0
 
     items: List[Tuple[str, Union[str, _AsyncUploadFile]]] = []
 
@@ -77,6 +82,14 @@ async def parse_async_stream(
             elif isinstance(event, Data):
                 if file is None:
                     data.extend(event.data)
+
+                    # Check if we have exceeded the maximum memory size
+                    form_memory_size_count += len(event.data)
+                    if (
+                        max_form_memory_size is not None
+                        and form_memory_size_count > max_form_memory_size
+                    ):
+                        raise RequestEntityTooLarge()
                 else:
                     await file.awrite(event.data)
 
@@ -88,6 +101,11 @@ async def parse_async_stream(
                         await file.aseek(0)
                         items.append((field_name, file))
                         file = None
+
+                    # Check if we have exceeded the maximum number of form parts
+                    form_parts_count += 1
+                    if form_parts_count > max_form_parts:
+                        raise RequestEntityTooLarge()
     return items
 
 
@@ -97,6 +115,8 @@ def parse_stream(
     charset: str,
     *,
     file_factory: Type[_SyncUploadFile],
+    max_form_parts: int = 324,
+    max_form_memory_size: Union[int, None] = None,
 ) -> List[Tuple[str, Union[str, _SyncUploadFile]]]:
     """
     Parse a synchronous stream in multipart format
@@ -110,6 +130,8 @@ def parse_stream(
     field_name = ""
     data = bytearray()
     file: Optional[_SyncUploadFile] = None
+    form_parts_count = 0
+    form_memory_size_count = 0
 
     items: List[Tuple[str, Union[str, _SyncUploadFile]]] = []
 
@@ -127,6 +149,14 @@ def parse_stream(
             elif isinstance(event, Data):
                 if file is None:
                     data.extend(event.data)
+
+                    # Check if we have exceeded the maximum memory size
+                    form_memory_size_count += len(event.data)
+                    if (
+                        max_form_memory_size is not None
+                        and form_memory_size_count > max_form_memory_size
+                    ):
+                        raise RequestEntityTooLarge()
                 else:
                     file.write(event.data)
 
@@ -138,4 +168,9 @@ def parse_stream(
                         file.seek(0)
                         items.append((field_name, file))
                         file = None
+
+                    # Check if we have exceeded the maximum number of form parts
+                    form_parts_count += 1
+                    if form_parts_count > max_form_parts:
+                        raise RequestEntityTooLarge()
     return items

baize/wsgi/requests.py

@@ -168,6 +168,11 @@ def json(self) -> Any:
 
         raise UnsupportedMediaType("application/json")
 
+    def _parse_multipart(self, boundary: bytes, charset: str) -> FormData:
+        return FormData(
+            parse_multipart(self.stream(), boundary, charset, file_factory=UploadFile)
+        )
+
     @cached_property
     def form(self) -> FormData:
         """
@@ -185,11 +190,7 @@ def form(self) -> FormData:
             if "boundary" not in self.content_type.options:
                 raise MalformedMultipart("Missing boundary in header content-type")
             boundary = self.content_type.options["boundary"].encode("latin-1")
-            return FormData(
-                parse_multipart(
-                    self.stream(), boundary, charset, file_factory=UploadFile
-                )
-            )
+            return self._parse_multipart(boundary, charset)
         if self.content_type == "application/x-www-form-urlencoded":
             body = self.body.decode(
                 encoding=self.content_type.options.get("charset", "latin-1")

tests/test_asgi.py

@@ -32,11 +32,12 @@
     request_response,
     websocket_session,
 )
-from baize.datastructures import UploadFile
+from baize.datastructures import FormData, UploadFile
 from baize.exceptions import (
     HTTPException,
     MalformedJSON,
     MalformedMultipart,
+    RequestEntityTooLarge,
     UnsupportedMediaType,
 )
 from baize.typing import Message, ServerSentEvent
@@ -221,6 +222,44 @@ async def app(scope, receive, send):
             )
 
 
+@pytest.mark.asyncio
+async def test_request_multipart_form_limit():
+    class FRequest(Request):
+        async def _parse_multipart(self, boundary: bytes, charset: str) -> FormData:
+            from baize.multipart_helper import parse_async_stream as parse_multipart
+
+            return FormData(
+                await parse_multipart(
+                    self.stream(),
+                    boundary,
+                    charset,
+                    file_factory=UploadFile,
+                    max_form_parts=2,
+                    max_form_memory_size=1024,
+                )
+            )
+
+    async def app(scope, receive, send):
+        await FRequest(scope, receive).form
+
+    async with httpx.AsyncClient(app=app, base_url="http://testServer/") as client:
+        with pytest.raises(RequestEntityTooLarge):
+            with tempfile.SpooledTemporaryFile(1024) as file:
+                file.write(b"temporary file")
+                file.seek(0, 0)
+                await client.post(
+                    "/", data={"part1": "1", "part2": "2"}, files={"file-key": file}
+                )
+
+        with pytest.raises(RequestEntityTooLarge):
+            with tempfile.SpooledTemporaryFile(1024) as file:
+                file.write(b"temporary file")
+                file.seek(0, 0)
+                await client.post(
+                    "/", data={"abc": "*" * 2048}, files={"file-key": file}
+                )
+
+
 @pytest.mark.asyncio
 async def test_request_body_then_stream():
     async def app(scope, receive, send):

tests/test_wsgi.py

@@ -7,11 +7,12 @@
 import httpx
 import pytest
 
-from baize.datastructures import Address, UploadFile
+from baize.datastructures import Address, FormData, UploadFile
 from baize.exceptions import (
     HTTPException,
     MalformedJSON,
     MalformedMultipart,
+    RequestEntityTooLarge,
     UnsupportedMediaType,
 )
 from baize.typing import ServerSentEvent
@@ -194,6 +195,41 @@ def app(environ, start_response):
             )
 
 
+def test_request_multipart_form_limit():
+    class FRequest(Request):
+        def _parse_multipart(self, boundary: bytes, charset: str) -> FormData:
+            from baize.multipart_helper import parse_stream as parse_multipart
+
+            return FormData(
+                parse_multipart(
+                    self.stream(),
+                    boundary,
+                    charset,
+                    file_factory=UploadFile,
+                    max_form_parts=2,
+                    max_form_memory_size=1024,
+                )
+            )
+
+    def app(environ, start_response):
+        FRequest(environ).form
+
+    with httpx.Client(app=app, base_url="http://testServer/") as client:
+        with pytest.raises(RequestEntityTooLarge):
+            with tempfile.SpooledTemporaryFile(1024) as file:
+                file.write(b"temporary file")
+                file.seek(0, 0)
+                client.post(
+                    "/", data={"part1": "1", "part2": "2"}, files={"file-key": file}
+                )
+
+        with pytest.raises(RequestEntityTooLarge):
+            with tempfile.SpooledTemporaryFile(1024) as file:
+                file.write(b"temporary file")
+                file.seek(0, 0)
+                client.post("/", data={"abc": "*" * 2048}, files={"file-key": file})
+
+
 def test_request_body_then_stream():
     def app(environ, start_response):
         request = Request(environ)