diff --git a/homeassistant/components/file_upload/__init__.py b/homeassistant/components/file_upload/__init__.py index 9562b77e337a..e884d0362e3c 100644 --- a/homeassistant/components/file_upload/__init__.py +++ b/homeassistant/components/file_upload/__init__.py @@ -118,6 +118,7 @@ async def _receive_file_field( queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = SimpleQueue() def _sync_queue_consumer() -> None: + file_path.parent.mkdir() with file_path.open("wb") as file_handle: while True: if (_chunk_future := queue.get()) is None: @@ -128,6 +129,7 @@ async def _receive_file_field( file_handle.write(_chunk) fut: asyncio.Future[None] | None = None + cancelled: asyncio.CancelledError | None = None try: fut = hass.async_add_executor_job(_sync_queue_consumer) chunks_sent = 0 @@ -143,23 +145,31 @@ async def _receive_file_field( if fut.done(): # The executor job failed break + except asyncio.CancelledError as err: + # Remember a cancellation from the streaming loop so the join below re-raises + # it instead of a later writer error. + cancelled = err + raise finally: # Always terminate the queue consumer, also if the stream raised or the task # was cancelled, otherwise awaiting the consumer future deadlocks. queue.put_nowait(None) if fut is not None: # The executor thread can't be cancelled and is guaranteed to finish once - # it reads the sentinel queued above. Await it even if this task is - # cancelled: awaiting only through a shield keeps the executor future - # itself uncancelled, so the thread is fully done (file written and closed) - # before the caller cleans up. Re-raise any cancellation after. - cancelled: asyncio.CancelledError | None = None + # it reads the sentinel queued above. Wait for it even if this task is + # cancelled: asyncio.wait neither cancels the future nor raises its + # exception, so the thread is fully done (file written and closed) before + # the caller cleans up. The loop re-waits through repeated cancellations. while not fut.done(): try: - await asyncio.shield(fut) + await asyncio.wait({fut}) except asyncio.CancelledError as err: cancelled = err if cancelled is not None: + # A cancellation takes precedence over a writer error; retrieve the + # writer result so its exception isn't flagged as never-retrieved. + if not fut.cancelled(): + fut.exception() raise cancelled fut.result() @@ -216,12 +226,12 @@ class FileUploadView(HomeAssistantView): file_dir = file_upload_data.file_dir(file_id) try: - await hass.async_add_executor_job(file_dir.mkdir) await _receive_file_field(hass, file_field_reader, file_dir / filename) except Exception, asyncio.CancelledError: - # Upload failed: _receive_file_field has joined the writer and closed the - # file, so removing the directory now cannot race the writer. ignore_errors - # covers a failure that happened before the directory was created. + # Upload failed: _receive_file_field has joined the writer, which created + # the directory and closed the file, so removing the directory now cannot + # race the writer. ignore_errors covers a failure that happened before the + # directory was created. await hass.async_add_executor_job( lambda: shutil.rmtree(file_dir, ignore_errors=True) ) diff --git a/tests/components/file_upload/test_init.py b/tests/components/file_upload/test_init.py index 6062a25cb35e..da8a3755ceda 100644 --- a/tests/components/file_upload/test_init.py +++ b/tests/components/file_upload/test_init.py @@ -284,6 +284,10 @@ async def test_upload_cancelled_releases_consumer(hass: HomeAssistant) -> None: with pytest.raises(asyncio.CancelledError): task.result() + # The cancelled upload must not orphan its file directory on disk. + file_upload_data = hass.data[file_upload.DOMAIN] + assert list(file_upload_data.temp_dir.iterdir()) == [] + async def test_receive_file_field_cancelled_while_joining_writer( hass: HomeAssistant, tmp_path: Path @@ -294,7 +298,8 @@ async def test_receive_file_field_cancelled_while_joining_writer( writer (the whole field already streamed, sentinel queued) must not return until the writer thread has finished, so the caller's cleanup cannot race it. """ - file_path = tmp_path / "uploaded.bin" + # Nested under a directory that does not exist yet, so the writer's mkdir runs. + file_path = tmp_path / "upload_dir" / "uploaded.bin" writing_started = asyncio.Event() release_writer = threading.Event() # blocks the writer thread mid-write writes: list[bytes] = [] @@ -347,8 +352,8 @@ async def test_receive_file_field_cancelled_while_joining_writer( task.cancel() for _ in range(10): await asyncio.sleep(0) - # Without the cancellation-safe join the task would finish here (returning - # while the writer thread runs on); the fix keeps it waiting for the writer. + # The task must still be waiting for the writer thread to finish before it + # returns, so the caller's cleanup cannot race the writer. assert not task.done() finally: # Always release the writer so a failed assertion can't leak the blocked @@ -361,3 +366,149 @@ async def test_receive_file_field_cancelled_while_joining_writer( task.result() # The writer finished writing both chunks before the cancellation propagated. assert b"".join(writes) == b"chunk1chunk2" + # The writer created the parent directory as part of the joined job, so cleanup + # cannot race an in-flight mkdir. + assert file_path.parent.is_dir() + + +async def test_receive_file_field_cancel_wins_over_writer_error( + hass: HomeAssistant, tmp_path: Path +) -> None: + """Test a writer error during the join does not mask the cancellation. + + When the task is cancelled while joining the writer and the writer then fails, + the caller must still observe CancelledError, not the writer's exception. + """ + file_path = tmp_path / "upload_dir" / "uploaded.bin" + writing_started = asyncio.Event() + release_writer = threading.Event() # blocks the writer thread mid-write + + class _FailingHandle: + """A file handle whose first write parks the writer, then fails.""" + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> bool: + return False + + def write(self, data: bytes) -> int: + hass.loop.call_soon_threadsafe(writing_started.set) + release_writer.wait() + raise OSError("write failed") + + real_open = Path.open + + def _failing_open(self: Path, *args: object, **kwargs: object) -> object: + if self != file_path: + return real_open(self, *args, **kwargs) + return _FailingHandle() + + chunks = iter([b"chunk1", b"chunk2"]) + + class _Part: + """Fake BodyPartReader yielding two chunks then EOF.""" + + async def read_chunk(self, size: int) -> bytes: + return next(chunks, b"") + + with patch.object(Path, "open", _failing_open): + task = asyncio.create_task( + file_upload._receive_file_field(hass, _Part(), file_path) + ) + try: + # Let the task settle at the join with the writer blocked mid-write, so + # the cancel lands on the join and the writer fails afterwards. + await writing_started.wait() + for _ in range(3): + await asyncio.sleep(0) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + assert not task.done() + finally: + # Always release the writer so a failed assertion can't leak the blocked + # thread and hang teardown. + release_writer.set() + _done, pending = await asyncio.wait({task}, timeout=10) + + assert not pending + # The cancellation wins over the writer's OSError. + assert task.cancelled() + with pytest.raises(asyncio.CancelledError): + task.result() + + +async def test_receive_file_field_cancel_in_stream_wins_over_writer_error( + hass: HomeAssistant, tmp_path: Path +) -> None: + """Test a cancel in the streaming loop is not masked by a writer error. + + When the cancellation lands in read_chunk (not the join) and the writer has + already failed, the caller must still observe CancelledError, not the writer's + exception. + """ + file_path = tmp_path / "upload_dir" / "uploaded.bin" + writer_failed = asyncio.Event() # set once the writer thread has raised + reading_blocked = asyncio.Event() # set when the stream parks on its 2nd read + blocked = asyncio.Event() # never set, so the 2nd read blocks until cancelled + + class _FailingHandle: + """A file handle whose first write fails the writer thread.""" + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> bool: + return False + + def write(self, data: bytes) -> int: + hass.loop.call_soon_threadsafe(writer_failed.set) + raise OSError("write failed") + + real_open = Path.open + + def _failing_open(self: Path, *args: object, **kwargs: object) -> object: + if self != file_path: + return real_open(self, *args, **kwargs) + return _FailingHandle() + + reads = 0 + + class _Part: + """Fake BodyPartReader yielding one chunk, then blocking on the next read.""" + + async def read_chunk(self, size: int) -> bytes: + nonlocal reads + reads += 1 + if reads == 1: + return b"chunk1" + reading_blocked.set() + await blocked.wait() + return b"" + + with patch.object(Path, "open", _failing_open): + task = asyncio.create_task( + file_upload._receive_file_field(hass, _Part(), file_path) + ) + try: + # The stream is parked on its second read and the writer has failed, so + # the cancel lands in the streaming loop with the writer future already + # done with an error. + await reading_blocked.wait() + await writer_failed.wait() + for _ in range(5): + await asyncio.sleep(0) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + finally: + # Always unblock the stream so a failed assertion can't hang teardown. + blocked.set() + _done, pending = await asyncio.wait({task}, timeout=10) + + assert not pending + # The cancellation wins over the writer's OSError. + assert task.cancelled() + with pytest.raises(asyncio.CancelledError): + task.result()