From 04d3cf5ef58c8ac8d28d36ea410fba131f5eff3f Mon Sep 17 00:00:00 2001
From: John Stark <jhnstrk@users.noreply.github.com>
Date: Wed, 11 Dec 2024 16:42:43 +0000
Subject: [PATCH] Handle messages containing only end boundary, fixes #38
 (#142)

---
 python_multipart/multipart.py                 | 24 +++++++++++++++----
 tests/test_data/http/empty_message.http       |  1 +
 tests/test_data/http/empty_message.yaml       |  2 ++
 .../http/empty_message_with_bad_end.http      |  1 +
 .../http/empty_message_with_bad_end.yaml      |  3 +++
 5 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_data/http/empty_message.http
 create mode 100644 tests/test_data/http/empty_message.yaml
 create mode 100644 tests/test_data/http/empty_message_with_bad_end.http
 create mode 100644 tests/test_data/http/empty_message_with_bad_end.yaml

diff --git a/python_multipart/multipart.py b/python_multipart/multipart.py
index a996379..f26a815 100644
--- a/python_multipart/multipart.py
+++ b/python_multipart/multipart.py
@@ -130,7 +130,8 @@ class MultipartState(IntEnum):
     PART_DATA_START = 8
     PART_DATA = 9
     PART_DATA_END = 10
-    END = 11
+    END_BOUNDARY = 11
+    END = 12
 
 
 # Flags for the multipart parser.
@@ -1119,7 +1120,10 @@ def data_callback(name: CallbackName, end_i: int, remaining: bool = False) -> No
                 # Check to ensure that the last 2 characters in our boundary
                 # are CRLF.
                 if index == len(boundary) - 2:
-                    if c != CR:
+                    if c == HYPHEN:
+                        # Potential empty message.
+                        state = MultipartState.END_BOUNDARY
+                    elif c != CR:
                         # Error!
                         msg = "Did not find CR at end of boundary (%d)" % (i,)
                         self.logger.warning(msg)
@@ -1396,6 +1400,18 @@ def data_callback(name: CallbackName, end_i: int, remaining: bool = False) -> No
                     # the start of the boundary itself.
                     i -= 1
 
+            elif state == MultipartState.END_BOUNDARY:
+                if index == len(boundary) - 2 + 1:
+                    if c != HYPHEN:
+                        msg = "Did not find - at end of boundary (%d)" % (i,)
+                        self.logger.warning(msg)
+                        e = MultipartParseError(msg)
+                        e.offset = i
+                        raise e
+                    index += 1
+                    self.callback("end")
+                    state = MultipartState.END
+
             elif state == MultipartState.END:
                 # Don't do anything if chunk ends with CRLF.
                 if c == CR and i + 1 < length and data[i + 1] == LF:
@@ -1707,8 +1723,8 @@ def on_headers_finished() -> None:
 
             def _on_end() -> None:
                 nonlocal writer
-                assert writer is not None
-                writer.finalize()
+                if writer is not None:
+                    writer.finalize()
                 if self.on_end is not None:
                     self.on_end()
 
diff --git a/tests/test_data/http/empty_message.http b/tests/test_data/http/empty_message.http
new file mode 100644
index 0000000..baff7d5
--- /dev/null
+++ b/tests/test_data/http/empty_message.http
@@ -0,0 +1 @@
+----boundary--
diff --git a/tests/test_data/http/empty_message.yaml b/tests/test_data/http/empty_message.yaml
new file mode 100644
index 0000000..ab33940
--- /dev/null
+++ b/tests/test_data/http/empty_message.yaml
@@ -0,0 +1,2 @@
+boundary: --boundary
+expected: []
diff --git a/tests/test_data/http/empty_message_with_bad_end.http b/tests/test_data/http/empty_message_with_bad_end.http
new file mode 100644
index 0000000..a085714
--- /dev/null
+++ b/tests/test_data/http/empty_message_with_bad_end.http
@@ -0,0 +1 @@
+----boundary-X
diff --git a/tests/test_data/http/empty_message_with_bad_end.yaml b/tests/test_data/http/empty_message_with_bad_end.yaml
new file mode 100644
index 0000000..ae920bf
--- /dev/null
+++ b/tests/test_data/http/empty_message_with_bad_end.yaml
@@ -0,0 +1,3 @@
+boundary: --boundary
+expected:
+    error: 13