From 4eae8b0997c08c45b35419405274d4a235374043 Mon Sep 17 00:00:00 2001 From: Christoph Li Date: Wed, 9 Sep 2026 15:44:35 -0400 Subject: [PATCH] Close attachment backing on copy/download error paths Attachment.from_file created its SpooledTemporaryFile backing before copying into it, but only restored the source cursor on error and never closed the backing itself, leaking a handle (and an on-disk temp file past the 4 MB rollover) if the copy failed. AttachmentEntry._ensure_attachment had the same issue: the backing buffer created for a download was never closed if the streaming write loop raised mid-transfer. Both now close the backing in an except/raise around the copy, while leaving it open (and handed off to Attachment) on success. Fixes #323 --- src/labapi/entry/attachment.py | 3 ++ src/labapi/entry/entries/attachment.py | 8 +++-- tests/entry/entries/test_attachment.py | 37 ++++++++++++++++++++++ tests/entry/test_attachment.py | 43 ++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 2 deletions(-) diff --git a/src/labapi/entry/attachment.py b/src/labapi/entry/attachment.py index 178f95fe..13df92f3 100644 --- a/src/labapi/entry/attachment.py +++ b/src/labapi/entry/attachment.py @@ -130,6 +130,9 @@ def from_file(file: NamedBinaryIO | PathLike[str] | str) -> Attachment: try: file.seek(0) shutil.copyfileobj(file, backing) + except Exception: + backing.close() + raise finally: file.seek(original_position) backing.seek(0) diff --git a/src/labapi/entry/entries/attachment.py b/src/labapi/entry/entries/attachment.py index def7b80d..7926bb0f 100644 --- a/src/labapi/entry/entries/attachment.py +++ b/src/labapi/entry/entries/attachment.py @@ -117,8 +117,12 @@ def _ensure_attachment(self, use_tempfile: bool) -> None: ) output = _make_backing_io(use_tempfile) - for chunk in attachment_stream: - output.write(chunk) + try: + for chunk in attachment_stream: + output.write(chunk) + except Exception: + output.close() + raise self._filedata = Attachment(output, mime_type, filename, self._data) diff --git a/tests/entry/entries/test_attachment.py b/tests/entry/entries/test_attachment.py index e7cf9cc9..d2347051 100644 --- a/tests/entry/entries/test_attachment.py +++ b/tests/entry/entries/test_attachment.py @@ -376,6 +376,43 @@ def test_attachment_entry_get_attachment_caching(self, client, user: User): assert client.stream_api_get.call_count == 1 assert attachment3.read() == b"Content" + def test_ensure_attachment_closes_backing_on_download_error( + self, client, user: User, monkeypatch + ): + """Test _ensure_attachment closes the backing buffer if the download raises mid-stream.""" + entry = AttachmentEntry("eid_att", "Caption", user) + + created_backings: list[BytesIO] = [] + + def _tracking_make_backing_io(use_tempfile: bool) -> BytesIO: # noqa: ARG001 + io_obj = BytesIO() + created_backings.append(io_obj) + return io_obj + + monkeypatch.setattr( + "labapi.entry.entries.attachment._make_backing_io", + _tracking_make_backing_io, + ) + + def _raising_chunks(): + yield b"partial content" + raise OSError("boom") + + mock_response = Mock() + mock_response.headers = { + "Content-Type": "text/plain", + "Content-Disposition": 'attachment; filename="test.txt"', + } + mock_response.iter_content.return_value = _raising_chunks() + client.stream_api_get = Mock(return_value=StreamingResponse(mock_response)) + + with pytest.raises(OSError, match="boom"): + entry.get_attachment() + + assert entry._filedata is None # pyright: ignore[reportPrivateUsage] + assert len(created_backings) == 1 + assert created_backings[0].closed is True + def test_get_attachment_tempfile_copies_in_chunks(self, client, user: User): """get_attachment(use_tempfile=True) must not read the full payload at once. diff --git a/tests/entry/test_attachment.py b/tests/entry/test_attachment.py index b893ed67..dc1f4741 100644 --- a/tests/entry/test_attachment.py +++ b/tests/entry/test_attachment.py @@ -5,6 +5,7 @@ import tempfile from io import BytesIO from pathlib import Path +from typing import Any import pytest @@ -201,6 +202,48 @@ def test_attachment_getattr_delegation(): assert content == b"Hello" +def test_attachment_from_file_closes_backing_on_copy_error(monkeypatch): + """Test from_file closes the spooled backing buffer if the copy raises.""" + created_backings: list[Any] = [] + real_spooled_temporary_file = tempfile.SpooledTemporaryFile + + class TrackingSpooledFile: + def __init__(self, *args: Any, **kwargs: Any): + self._real = real_spooled_temporary_file(*args, **kwargs) + self.closed = False + created_backings.append(self) + + def close(self) -> None: + self.closed = True + self._real.close() + + def __getattr__(self, name: str) -> Any: + return getattr(self._real, name) + + monkeypatch.setattr( + "labapi.entry.attachment.tempfile.SpooledTemporaryFile", TrackingSpooledFile + ) + + class RaisingFile(BytesIO): + def __init__(self, data: bytes, name: str): + super().__init__(data) + self.name = name + + def read(self, size: int | None = -1, /) -> bytes: # noqa: ARG002 + raise OSError("boom") + + file = RaisingFile(b"Test content", "payload.bin") + file.seek(4) + + with pytest.raises(OSError, match="boom"): + Attachment.from_file(file) + + assert len(created_backings) == 1 + assert created_backings[0].closed is True + # The source cursor is still restored on the error path. + assert file.tell() == 4 + + def test_attachment_seeks_to_beginning(): """Test that Attachment seeks to beginning of seekable backing.""" backing = BytesIO(b"Test")