diff --git a/python_multipart/multipart.py b/python_multipart/multipart.py index 49cdd8e..0a86607 100644 --- a/python_multipart/multipart.py +++ b/python_multipart/multipart.py @@ -1614,6 +1614,7 @@ def __init__( self.boundary = boundary self.bytes_received = 0 self.parser = None + self._close_files: Callable[[], None] | None = None # Save callbacks. self.on_field = on_field @@ -1647,12 +1648,22 @@ def _on_end() -> None: # Call our callback. if on_file: - on_file(file) + completed_file = file + file = None + on_file(completed_file) # Call the on-end callback. if self.on_end is not None: self.on_end() + def close_current_file() -> None: + nonlocal file + if file is not None: + file.close() + file = None + + self._close_files = close_current_file + # Instantiate an octet-stream parser parser = OctetStreamParser( callbacks={"on_start": on_start, "on_data": on_data, "on_end": _on_end}, @@ -1720,6 +1731,7 @@ def _on_end() -> None: f_multi: File | Field | None = None writer: File | Field | Base64Decoder | QuotedPrintableDecoder | None = None is_file = False + files: list[File] = [] def on_part_begin() -> None: # Reset headers in case this isn't the first part. @@ -1739,6 +1751,7 @@ def on_part_end() -> None: if is_file: if on_file: assert isinstance(f_multi, File) + files.remove(f_multi) on_file(f_multi) else: if on_field: @@ -1780,6 +1793,7 @@ def on_headers_finished() -> None: else: f_multi = File(file_name, field_name, config=self.config, content_type=content_type) is_file = True + files.append(f_multi) # Parse the given Content-Transfer-Encoding to determine what # we need to do with the incoming data. @@ -1813,6 +1827,12 @@ def _on_end() -> None: if self.on_end is not None: self.on_end() + def close_files() -> None: + while files: + files.pop().close() + + self._close_files = close_files + # Instantiate a multipart parser. parser = MultipartParser( boundary, @@ -1859,8 +1879,12 @@ def finalize(self) -> None: def close(self) -> None: """Close the parser.""" - if self.parser is not None and hasattr(self.parser, "close"): - self.parser.close() + try: + if self.parser is not None and hasattr(self.parser, "close"): + self.parser.close() + finally: + if self._close_files is not None: + self._close_files() def __repr__(self) -> str: return f"{self.__class__.__name__}(content_type={self.content_type!r}, parser={self.parser!r})" diff --git a/tests/test_multipart.py b/tests/test_multipart.py index 949d70c..6615c5b 100644 --- a/tests/test_multipart.py +++ b/tests/test_multipart.py @@ -7,7 +7,7 @@ import unittest from io import BytesIO from typing import TYPE_CHECKING -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest import yaml @@ -1009,6 +1009,9 @@ def test_upload_delete_tmp_config(self) -> None: uploaded_file = self.files[0] assert uploaded_file.actual_file_name is not None actual_file_name = uploaded_file.actual_file_name.decode(sys.getfilesystemencoding()) + self.f.close() + + self.assertFalse(uploaded_file.file_object.closed) uploaded_file.close() try: @@ -1016,6 +1019,35 @@ def test_upload_delete_tmp_config(self) -> None: finally: os.unlink(actual_file_name) + def test_close_closes_incomplete_upload(self) -> None: + for transfer_encoding, data in ((b"binary", b"test"), (b"base64", b"dGVzdA=="), (b"quoted-printable", b"test")): + with ( + self.subTest(transfer_encoding=transfer_encoding), + tempfile.TemporaryDirectory() as upload_dir, + patch.object(File, "close", autospec=True, side_effect=File.close) as close, + ): + self.make( + "boundary", config={"UPLOAD_DIR": upload_dir, "UPLOAD_DELETE_TMP": False, "MAX_MEMORY_FILE_SIZE": 1} + ) + body = ( + b"--boundary\r\n" + b'Content-Disposition: form-data; name="file"; filename="test.txt"\r\n' + b"Content-Type: text/plain\r\n" + b"Content-Transfer-Encoding: " + transfer_encoding + b"\r\n\r\n" + data + ) + + self.f.write(body) + paths = os.listdir(upload_dir) + self.assertEqual(len(paths), 1) + + self.f.close() + self.f.close() + + close.assert_called_once() + self.assertTrue(close.call_args.args[0].file_object.closed) + self.assertEqual(os.listdir(upload_dir), paths) + self.assertEqual(self.files, []) + @parametrize("param", [t for t in http_tests if t["name"] in single_byte_tests]) def test_feed_single_bytes(self, param: TestParams) -> None: """ @@ -1294,6 +1326,33 @@ def on_file(f: File) -> None: self.assert_file_data(files[0], b"test1234") self.assertTrue(on_end.called) + f.close() + self.assertFalse(files[0].file_object.closed) + files[0].close() + + def test_close_closes_incomplete_octet_stream_upload(self) -> None: + with ( + tempfile.TemporaryDirectory() as upload_dir, + patch.object(File, "close", autospec=True, side_effect=File.close) as close, + ): + f = FormParser( + "application/octet-stream", + None, + None, + file_name=b"test.txt", + config={"UPLOAD_DIR": upload_dir, "UPLOAD_DELETE_TMP": False, "MAX_MEMORY_FILE_SIZE": 1}, + ) + f.write(b"test") + paths = os.listdir(upload_dir) + self.assertEqual(len(paths), 1) + + f.close() + f.close() + + close.assert_called_once() + self.assertTrue(close.call_args.args[0].file_object.closed) + self.assertEqual(os.listdir(upload_dir), paths) + def test_querystring(self) -> None: fields: list[Field] = [] @@ -1390,6 +1449,34 @@ def on_file(f: File) -> None: f.finalize() self.assert_file_data(files[0], b"Test") + def test_bad_content_transfer_encoding_closes_current_file(self) -> None: + data = ( + b"--boundary\r\n" + b'Content-Disposition: form-data; name="first"; filename="first.txt"\r\n\r\n' + b"first\r\n" + b"--boundary\r\n" + b'Content-Disposition: form-data; name="second"; filename="second.txt"\r\n' + b"Content-Transfer-Encoding: badstuff\r\n\r\n" + ) + files: list[File] = [] + f = FormParser( + "multipart/form-data", None, files.append, boundary="boundary", config={"UPLOAD_ERROR_ON_BAD_CTE": True} + ) + + with patch.object(File, "close", autospec=True, side_effect=File.close) as close: + with self.assertRaises(FormParserError): + f.write(data) + + self.assertEqual(len(files), 1) + f.close() + + close.assert_called_once() + self.assertIsNot(close.call_args.args[0], files[0]) + self.assertTrue(close.call_args.args[0].file_object.closed) + self.assertFalse(files[0].file_object.closed) + + files[0].close() + def test_bad_content_disposition(self) -> None: # Field name is required per RFC 7578 ยง4.2. data = (