Skip to content

Commit ea974bf

Browse files
YeaseenYeaseen
committed
Close incomplete multipart upload files
1 parent d9cb4c6 commit ea974bf

2 files changed

Lines changed: 115 additions & 4 deletions

File tree

‎python_multipart/multipart.py‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1614,6 +1614,7 @@ def __init__(
16141614
self.boundary = boundary
16151615
self.bytes_received = 0
16161616
self.parser = None
1617+
self._close_files: Callable[[], None] | None = None
16171618

16181619
# Save callbacks.
16191620
self.on_field = on_field
@@ -1647,12 +1648,22 @@ def _on_end() -> None:
16471648

16481649
# Call our callback.
16491650
if on_file:
1650-
on_file(file)
1651+
completed_file = file
1652+
file = None
1653+
on_file(completed_file)
16511654

16521655
# Call the on-end callback.
16531656
if self.on_end is not None:
16541657
self.on_end()
16551658

1659+
def close_current_file() -> None:
1660+
nonlocal file
1661+
if file is not None:
1662+
file.close()
1663+
file = None
1664+
1665+
self._close_files = close_current_file
1666+
16561667
# Instantiate an octet-stream parser
16571668
parser = OctetStreamParser(
16581669
callbacks={"on_start": on_start, "on_data": on_data, "on_end": _on_end},
@@ -1720,6 +1731,7 @@ def _on_end() -> None:
17201731
f_multi: File | Field | None = None
17211732
writer: File | Field | Base64Decoder | QuotedPrintableDecoder | None = None
17221733
is_file = False
1734+
files: list[File] = []
17231735

17241736
def on_part_begin() -> None:
17251737
# Reset headers in case this isn't the first part.
@@ -1739,6 +1751,7 @@ def on_part_end() -> None:
17391751
if is_file:
17401752
if on_file:
17411753
assert isinstance(f_multi, File)
1754+
files.remove(f_multi)
17421755
on_file(f_multi)
17431756
else:
17441757
if on_field:
@@ -1780,6 +1793,7 @@ def on_headers_finished() -> None:
17801793
else:
17811794
f_multi = File(file_name, field_name, config=self.config, content_type=content_type)
17821795
is_file = True
1796+
files.append(f_multi)
17831797

17841798
# Parse the given Content-Transfer-Encoding to determine what
17851799
# we need to do with the incoming data.
@@ -1813,6 +1827,12 @@ def _on_end() -> None:
18131827
if self.on_end is not None:
18141828
self.on_end()
18151829

1830+
def close_files() -> None:
1831+
while files:
1832+
files.pop().close()
1833+
1834+
self._close_files = close_files
1835+
18161836
# Instantiate a multipart parser.
18171837
parser = MultipartParser(
18181838
boundary,
@@ -1859,8 +1879,12 @@ def finalize(self) -> None:
18591879

18601880
def close(self) -> None:
18611881
"""Close the parser."""
1862-
if self.parser is not None and hasattr(self.parser, "close"):
1863-
self.parser.close()
1882+
try:
1883+
if self.parser is not None and hasattr(self.parser, "close"):
1884+
self.parser.close()
1885+
finally:
1886+
if self._close_files is not None:
1887+
self._close_files()
18641888

18651889
def __repr__(self) -> str:
18661890
return f"{self.__class__.__name__}(content_type={self.content_type!r}, parser={self.parser!r})"

‎tests/test_multipart.py‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import unittest
88
from io import BytesIO
99
from typing import TYPE_CHECKING
10-
from unittest.mock import Mock
10+
from unittest.mock import Mock, patch
1111

1212
import pytest
1313
import yaml
@@ -1009,13 +1009,45 @@ def test_upload_delete_tmp_config(self) -> None:
10091009
uploaded_file = self.files[0]
10101010
assert uploaded_file.actual_file_name is not None
10111011
actual_file_name = uploaded_file.actual_file_name.decode(sys.getfilesystemencoding())
1012+
self.f.close()
1013+
1014+
self.assertFalse(uploaded_file.file_object.closed)
10121015
uploaded_file.close()
10131016

10141017
try:
10151018
self.assertTrue(os.path.exists(actual_file_name))
10161019
finally:
10171020
os.unlink(actual_file_name)
10181021

1022+
def test_close_closes_incomplete_upload(self) -> None:
1023+
for transfer_encoding, data in ((b"binary", b"test"), (b"base64", b"dGVzdA=="), (b"quoted-printable", b"test")):
1024+
with (
1025+
self.subTest(transfer_encoding=transfer_encoding),
1026+
tempfile.TemporaryDirectory() as upload_dir,
1027+
patch.object(File, "close", autospec=True, side_effect=File.close) as close,
1028+
):
1029+
self.make(
1030+
"boundary", config={"UPLOAD_DIR": upload_dir, "UPLOAD_DELETE_TMP": False, "MAX_MEMORY_FILE_SIZE": 1}
1031+
)
1032+
body = (
1033+
b"--boundary\r\n"
1034+
b'Content-Disposition: form-data; name="file"; filename="test.txt"\r\n'
1035+
b"Content-Type: text/plain\r\n"
1036+
b"Content-Transfer-Encoding: " + transfer_encoding + b"\r\n\r\n" + data
1037+
)
1038+
1039+
self.f.write(body)
1040+
paths = os.listdir(upload_dir)
1041+
self.assertEqual(len(paths), 1)
1042+
1043+
self.f.close()
1044+
self.f.close()
1045+
1046+
close.assert_called_once()
1047+
self.assertTrue(close.call_args.args[0].file_object.closed)
1048+
self.assertEqual(os.listdir(upload_dir), paths)
1049+
self.assertEqual(self.files, [])
1050+
10191051
@parametrize("param", [t for t in http_tests if t["name"] in single_byte_tests])
10201052
def test_feed_single_bytes(self, param: TestParams) -> None:
10211053
"""
@@ -1294,6 +1326,33 @@ def on_file(f: File) -> None:
12941326
self.assert_file_data(files[0], b"test1234")
12951327
self.assertTrue(on_end.called)
12961328

1329+
f.close()
1330+
self.assertFalse(files[0].file_object.closed)
1331+
files[0].close()
1332+
1333+
def test_close_closes_incomplete_octet_stream_upload(self) -> None:
1334+
with (
1335+
tempfile.TemporaryDirectory() as upload_dir,
1336+
patch.object(File, "close", autospec=True, side_effect=File.close) as close,
1337+
):
1338+
f = FormParser(
1339+
"application/octet-stream",
1340+
None,
1341+
None,
1342+
file_name=b"test.txt",
1343+
config={"UPLOAD_DIR": upload_dir, "UPLOAD_DELETE_TMP": False, "MAX_MEMORY_FILE_SIZE": 1},
1344+
)
1345+
f.write(b"test")
1346+
paths = os.listdir(upload_dir)
1347+
self.assertEqual(len(paths), 1)
1348+
1349+
f.close()
1350+
f.close()
1351+
1352+
close.assert_called_once()
1353+
self.assertTrue(close.call_args.args[0].file_object.closed)
1354+
self.assertEqual(os.listdir(upload_dir), paths)
1355+
12971356
def test_querystring(self) -> None:
12981357
fields: list[Field] = []
12991358

@@ -1390,6 +1449,34 @@ def on_file(f: File) -> None:
13901449
f.finalize()
13911450
self.assert_file_data(files[0], b"Test")
13921451

1452+
def test_bad_content_transfer_encoding_closes_current_file(self) -> None:
1453+
data = (
1454+
b"--boundary\r\n"
1455+
b'Content-Disposition: form-data; name="first"; filename="first.txt"\r\n\r\n'
1456+
b"first\r\n"
1457+
b"--boundary\r\n"
1458+
b'Content-Disposition: form-data; name="second"; filename="second.txt"\r\n'
1459+
b"Content-Transfer-Encoding: badstuff\r\n\r\n"
1460+
)
1461+
files: list[File] = []
1462+
f = FormParser(
1463+
"multipart/form-data", None, files.append, boundary="boundary", config={"UPLOAD_ERROR_ON_BAD_CTE": True}
1464+
)
1465+
1466+
with patch.object(File, "close", autospec=True, side_effect=File.close) as close:
1467+
with self.assertRaises(FormParserError):
1468+
f.write(data)
1469+
1470+
self.assertEqual(len(files), 1)
1471+
f.close()
1472+
1473+
close.assert_called_once()
1474+
self.assertIsNot(close.call_args.args[0], files[0])
1475+
self.assertTrue(close.call_args.args[0].file_object.closed)
1476+
self.assertFalse(files[0].file_object.closed)
1477+
1478+
files[0].close()
1479+
13931480
def test_bad_content_disposition(self) -> None:
13941481
# Field name is required per RFC 7578 §4.2.
13951482
data = (

0 commit comments

Comments
 (0)