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