From 1cb968018641a0203c707cefd730da1272df08d5 Mon Sep 17 00:00:00 2001
From: starlet-dx <15929766099@163.com>
Date: Mon, 14 Aug 2023 17:33:08 +0800
Subject: [PATCH 1/1] limit the maximum number of multipart form parts

Reference:
https://github.com/pallets/werkzeug/commit/517cac5a804e8c4dc4ed038bb20dacd038e7a9f1
---
 src/werkzeug/formparser.py       | 12 +++++++++++-
 src/werkzeug/sansio/multipart.py |  7 +++++++
 src/werkzeug/wrappers/request.py |  8 ++++++++
 tests/test_formparser.py         |  9 +++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/werkzeug/formparser.py b/src/werkzeug/formparser.py
index 6cb758f..92f5b3d 100644
--- a/src/werkzeug/formparser.py
+++ b/src/werkzeug/formparser.py
@@ -181,6 +181,8 @@ class FormDataParser:
     :param cls: an optional dict class to use.  If this is not specified
                        or `None` the default :class:`MultiDict` is used.
     :param silent: If set to False parsing errors will not be caught.
+    :param max_form_parts: The maximum number of parts to be parsed. If this is
+        exceeded, a :exc:`~exceptions.RequestEntityTooLarge` exception is raised.
     """
 
     def __init__(
@@ -192,6 +194,8 @@ class FormDataParser:
         max_content_length: t.Optional[int] = None,
         cls: t.Optional[t.Type[MultiDict]] = None,
         silent: bool = True,
+        *,
+        max_form_parts: t.Optional[int] = None,
     ) -> None:
         if stream_factory is None:
             stream_factory = default_stream_factory
@@ -201,6 +205,7 @@ class FormDataParser:
         self.errors = errors
         self.max_form_memory_size = max_form_memory_size
         self.max_content_length = max_content_length
+        self.max_form_parts = max_form_parts
 
         if cls is None:
             cls = MultiDict
@@ -283,6 +288,7 @@ class FormDataParser:
             self.errors,
             max_form_memory_size=self.max_form_memory_size,
             cls=self.cls,
+            max_form_parts=self.max_form_parts,
         )
         boundary = options.get("boundary", "").encode("ascii")
 
@@ -386,10 +392,12 @@ class MultiPartParser:
         max_form_memory_size: t.Optional[int] = None,
         cls: t.Optional[t.Type[MultiDict]] = None,
         buffer_size: int = 64 * 1024,
+        max_form_parts: t.Optional[int] = None,
     ) -> None:
         self.charset = charset
         self.errors = errors
         self.max_form_memory_size = max_form_memory_size
+        self.max_form_parts = max_form_parts
 
         if stream_factory is None:
             stream_factory = default_stream_factory
@@ -449,7 +457,9 @@ class MultiPartParser:
             [None],
         )
 
-        parser = MultipartDecoder(boundary, self.max_form_memory_size)
+        parser = MultipartDecoder(
+            boundary, self.max_form_memory_size, max_parts=self.max_form_parts
+        )
 
         fields = []
         files = []
diff --git a/src/werkzeug/sansio/multipart.py b/src/werkzeug/sansio/multipart.py
index 2d54422..31a24d0 100644
--- a/src/werkzeug/sansio/multipart.py
+++ b/src/werkzeug/sansio/multipart.py
@@ -83,10 +83,13 @@ class MultipartDecoder:
         self,
         boundary: bytes,
         max_form_memory_size: Optional[int] = None,
+        *,
+        max_parts: Optional[int] = None,
     ) -> None:
         self.buffer = bytearray()
         self.complete = False
         self.max_form_memory_size = max_form_memory_size
+        self.max_parts = max_parts
         self.state = State.PREAMBLE
         self.boundary = boundary
 
@@ -113,6 +116,7 @@ class MultipartDecoder:
             % (LINE_BREAK, re.escape(boundary), LINE_BREAK, LINE_BREAK),
             re.MULTILINE,
         )
+        self._parts_decoded = 0
 
     def last_newline(self) -> int:
         try:
@@ -177,7 +181,10 @@ class MultipartDecoder:
                         name=name,
                     )
                 self.state = State.DATA
+                self._parts_decoded += 1
 
+                if self.max_parts is not None and self._parts_decoded > self.max_parts:
+                    raise RequestEntityTooLarge()
         elif self.state == State.DATA:
             if self.buffer.find(b"--" + self.boundary) == -1:
                 # No complete boundary in the buffer, but there may be
diff --git a/src/werkzeug/wrappers/request.py b/src/werkzeug/wrappers/request.py
index f68dd5a..113cc41 100644
--- a/src/werkzeug/wrappers/request.py
+++ b/src/werkzeug/wrappers/request.py
@@ -81,6 +81,13 @@ class Request(_SansIORequest):
     #: .. versionadded:: 0.5
     max_form_memory_size: t.Optional[int] = None
 
+    #: The maximum number of multipart parts to parse, passed to
+    #: :attr:`form_data_parser_class`. Parsing form data with more than this
+    #: many parts will raise :exc:`~.RequestEntityTooLarge`.
+    #:
+    #: .. versionadded:: 2.2.3
+    max_form_parts = 1000
+
     #: The form data parser that shoud be used.  Can be replaced to customize
     #: the form date parsing.
     form_data_parser_class: t.Type[FormDataParser] = FormDataParser
@@ -265,6 +272,7 @@ class Request(_SansIORequest):
             self.max_form_memory_size,
             self.max_content_length,
             self.parameter_storage_class,
+            max_form_parts=self.max_form_parts,
         )
 
     def _load_form_data(self) -> None:
diff --git a/tests/test_formparser.py b/tests/test_formparser.py
index 18ed1c0..e32657d 100644
--- a/tests/test_formparser.py
+++ b/tests/test_formparser.py
@@ -127,6 +127,15 @@ class TestFormParser:
         req.max_form_memory_size = 400
         assert req.form["foo"] == "Hello World"
 
+        req = Request.from_values(
+            input_stream=io.BytesIO(data),
+            content_length=len(data),
+            content_type="multipart/form-data; boundary=foo",
+            method="POST",
+        )
+        req.max_form_parts = 1
+        pytest.raises(RequestEntityTooLarge, lambda: req.form["foo"])
+
     def test_missing_multipart_boundary(self):
         data = (
             b"--foo\r\nContent-Disposition: form-field; name=foo\r\n\r\n"
-- 
2.30.0