Fix potential deadlock in receive_file backup util (#181070)

homeassistant/components/backup/util.py

@@ -12,7 +12,6 @@ import tarfile
 import threading
 from typing import IO, Any, cast
 
-import aiohttp
 from securetar import (
     InvalidPasswordError,
     SecureTarArchive,
@@ -518,7 +517,7 @@ async def iter_upload_chunks(contents: aiohttp.BodyPartReader) -> AsyncIterator[
 
 
 async def receive_file(
-    hass: HomeAssistant, contents: aiohttp.BodyPartReader, path: Path
+    hass: HomeAssistant, stream: AsyncIterator[bytes], path: Path
 ) -> None:
     """Receive a file from a stream and write it to a file."""
     queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = SimpleQueue()
@@ -536,10 +535,10 @@ async def receive_file(
     fut: asyncio.Future[None] | None = None
     try:
         fut = hass.async_add_executor_job(_sync_queue_consumer)
-        megabytes_sending = 0
-        while chunk := await contents.read_chunk(BUF_SIZE):
-            megabytes_sending += 1
-            if megabytes_sending % 5 != 0:
+        chunks_sent = 0
+        async for chunk in stream:
+            chunks_sent += 1
+            if chunks_sent % 5 != 0:
                 queue.put_nowait((chunk, None))
                 continue
 
@@ -552,8 +551,9 @@ async def receive_file(
             if fut.done():
                 # The executor job failed
                 break
-
-        queue.put_nowait(None)  # terminate queue consumer
     finally:
+        # Always terminate the queue consumer, also if the stream raised or the
+        # task was cancelled.
+        queue.put_nowait(None)
         if fut is not None:
             await fut

tests/components/backup/test_util.py

@@ -19,6 +19,7 @@ from homeassistant.components.backup.util import (
     DecryptedBackupStreamer,
     EncryptedBackupStreamer,
     read_backup,
+    receive_file,
     suggested_filename,
     validate_password,
 )
@@ -784,3 +785,87 @@ def test_suggested_filename(name: str, resulting_filename: str) -> None:
         size=1234,
     )
     assert suggested_filename(backup) == resulting_filename
+
+
+# Bound receive_file awaits so a reintroduced deadlock fails fast instead of
+# hanging the test run.
+_RECEIVE_FILE_TIMEOUT = 10
+
+
+async def _stream_chunks(
+    chunks: list[bytes], error: Exception | None = None
+) -> AsyncIterator[bytes]:
+    """Yield chunks, then optionally raise to simulate a broken upload stream."""
+    for chunk in chunks:
+        yield chunk
+    if error is not None:
+        raise error
+
+
+@pytest.mark.parametrize(
+    "chunks",
+    [
+        pytest.param([], id="empty"),
+        pytest.param([b"single"], id="single_chunk"),
+        pytest.param([b"chunk1", b"chunk2", b"chunk3"], id="multi_chunk"),
+        # >5 chunks crosses the backpressure checkpoint (every 5th chunk).
+        pytest.param([f"chunk{i}".encode() for i in range(12)], id="many_chunks"),
+    ],
+)
+async def test_receive_file(
+    hass: HomeAssistant, tmp_path: Path, chunks: list[bytes]
+) -> None:
+    """Test receiving a stream and writing it to a file."""
+    path = tmp_path / "received.bin"
+    async with asyncio.timeout(_RECEIVE_FILE_TIMEOUT):
+        await receive_file(hass, _stream_chunks(chunks), path)
+    assert path.read_bytes() == b"".join(chunks)
+
+
+async def test_receive_file_writer_error(hass: HomeAssistant, tmp_path: Path) -> None:
+    """Test an OSError from the file writer propagates without deadlocking."""
+    # The parent directory does not exist, so opening the file for writing fails.
+    path = tmp_path / "missing" / "received.bin"
+    async with asyncio.timeout(_RECEIVE_FILE_TIMEOUT):
+        with pytest.raises(FileNotFoundError):
+            await receive_file(hass, _stream_chunks([b"data"] * 10), path)
+
+
+async def test_receive_file_stream_error(hass: HomeAssistant, tmp_path: Path) -> None:
+    """Test a stream error propagates and the consumer still terminates."""
+    path = tmp_path / "received.bin"
+    stream = _stream_chunks(
+        [b"chunk1", b"chunk2"], ConnectionResetError("Connection lost")
+    )
+    async with asyncio.timeout(_RECEIVE_FILE_TIMEOUT):
+        with pytest.raises(ConnectionResetError):
+            await receive_file(hass, stream, path)
+    # The consumer drained the queued chunks and closed the file before the error
+    # surfaced, proving it terminated rather than deadlocked.
+    assert path.read_bytes() == b"chunk1chunk2"
+
+
+async def test_receive_file_cancelled(hass: HomeAssistant, tmp_path: Path) -> None:
+    """Test cancelling mid-transfer propagates CancelledError without deadlocking."""
+    path = tmp_path / "received.bin"
+    first_chunk_sent = asyncio.Event()
+    blocked = asyncio.Event()  # never set, so the stream blocks until cancelled
+
+    async def _blocking_stream() -> AsyncIterator[bytes]:
+        yield b"chunk1"
+        first_chunk_sent.set()
+        await blocked.wait()
+
+    task = asyncio.create_task(receive_file(hass, _blocking_stream(), path))
+    await first_chunk_sent.wait()
+    task.cancel()
+
+    # asyncio.wait does not cancel on timeout, so a deadlocked task stays pending
+    # and the assertion fails fast instead of the run hanging.
+    _done, pending = await asyncio.wait({task}, timeout=_RECEIVE_FILE_TIMEOUT)
+    assert not pending
+    with pytest.raises(asyncio.CancelledError):
+        task.result()
+    # The first chunk was flushed and the file closed before cancellation
+    # completed, proving the consumer terminated rather than deadlocked.
+    assert path.read_bytes() == b"chunk1"