Fix receive backup streaming (#181012)

Co-authored-by: Erik <erik@montnemery.com>

homeassistant/components/backup/manager.py

@@ -72,8 +72,10 @@ from .store import BackupStore
 from .util import (
     DecryptedBackupStreamer,
     EncryptedBackupStreamer,
+    iter_upload_chunks,
     make_backup_dir,
     read_backup,
+    receive_file,
     validate_password,
     validate_password_stream,
 )
@@ -1004,7 +1006,6 @@ class BackupManager:
         contents: aiohttp.BodyPartReader,
     ) -> str:
         """Receive and store a backup file from upload."""
-        contents.chunk_size = BUF_SIZE
         suggested_filename = contents.filename or "backup.tar"
         safe_filename = PureWindowsPath(suggested_filename).name
         if (
@@ -1022,7 +1023,7 @@ class BackupManager:
         )
         written_backup = await self._reader_writer.async_receive_backup(
             agent_ids=agent_ids,
-            stream=contents,
+            stream=iter_upload_chunks(contents),
             suggested_filename=suggested_filename,
         )
         self.async_on_backup_event(
@@ -1958,6 +1959,47 @@ class CoreBackupReaderWriter(BackupReaderWriter):
             ) from err
         return (tar_file_path, stat_result.st_size)
 
+    async def _receive_and_move_backup(
+        self,
+        *,
+        agent_ids: list[str],
+        stream: AsyncIterator[bytes],
+        temp_file: Path,
+    ) -> tuple[AgentBackup, Path]:
+        """Receive the upload into temp_file, validate it, and move it into place.
+
+        Remove temp_file on any failure, including cancellation from a client
+        disconnect, so a partial or unparsable upload does not orphan a
+        potentially large temp file.
+        """
+        async_add_executor_job = self._hass.async_add_executor_job
+        try:
+            await receive_file(self._hass, stream, temp_file)
+            try:
+                backup = await async_add_executor_job(read_backup, temp_file)
+            except (
+                OSError,
+                tarfile.TarError,
+                json.JSONDecodeError,
+                KeyError,
+                InvalidBackupFilename,
+            ) as err:
+                LOGGER.warning("Unable to parse backup %s: %s", temp_file, err)
+                raise
+
+            manager = self._hass.data[DATA_MANAGER]
+            if self._local_agent_id in agent_ids:
+                local_agent = manager.local_backup_agents[self._local_agent_id]
+                tar_file_path = local_agent.get_new_backup_path(backup)
+                await async_add_executor_job(make_backup_dir, tar_file_path.parent)
+                await async_add_executor_job(shutil.move, temp_file, tar_file_path)
+            else:
+                tar_file_path = temp_file
+        except Exception, asyncio.CancelledError:
+            await async_add_executor_job(temp_file.unlink, True)
+            raise
+        return backup, tar_file_path
+
     @override
     async def async_receive_backup(
         self,
@@ -1971,33 +2013,9 @@ class CoreBackupReaderWriter(BackupReaderWriter):
 
         async_add_executor_job = self._hass.async_add_executor_job
         await async_add_executor_job(make_backup_dir, self.temp_backup_dir)
-        f = await async_add_executor_job(temp_file.open, "wb")
-        try:
-            async for chunk in stream:
-                await async_add_executor_job(f.write, chunk)
-        finally:
-            await async_add_executor_job(f.close)
-
-        try:
-            backup = await async_add_executor_job(read_backup, temp_file)
-        except (
-            OSError,
-            tarfile.TarError,
-            json.JSONDecodeError,
-            KeyError,
-            InvalidBackupFilename,
-        ) as err:
-            LOGGER.warning("Unable to parse backup %s: %s", temp_file, err)
-            raise
-
-        manager = self._hass.data[DATA_MANAGER]
-        if self._local_agent_id in agent_ids:
-            local_agent = manager.local_backup_agents[self._local_agent_id]
-            tar_file_path = local_agent.get_new_backup_path(backup)
-            await async_add_executor_job(make_backup_dir, tar_file_path.parent)
-            await async_add_executor_job(shutil.move, temp_file, tar_file_path)
-        else:
-            tar_file_path = temp_file
+        backup, tar_file_path = await self._receive_and_move_backup(
+            agent_ids=agent_ids, stream=stream, temp_file=temp_file
+        )
 
         async def send_backup() -> AsyncIterator[bytes]:
             f = await async_add_executor_job(tar_file_path.open, "rb")

homeassistant/components/backup/util.py

@@ -507,6 +507,16 @@ class EncryptedBackupStreamer(_CipherBackupStreamer):
         return replace(self._backup, protected=True, size=self.size())
 
 
+async def iter_upload_chunks(contents: aiohttp.BodyPartReader) -> AsyncIterator[bytes]:
+    """Yield chunks of an uploaded file.
+
+    Iterating a BodyPartReader reads the whole part into memory and enforces the
+    request's client_max_size limit; reading it in chunks does neither.
+    """
+    while chunk := await contents.read_chunk(BUF_SIZE):
+        yield chunk
+
+
 async def receive_file(
     hass: HomeAssistant, contents: aiohttp.BodyPartReader, path: Path
 ) -> None:

tests/components/backup/test_manager.py

@@ -4,7 +4,7 @@ import asyncio
 from collections.abc import Callable, Generator
 from dataclasses import replace
 from datetime import timedelta
-from io import StringIO
+from io import BytesIO, StringIO
 import json
 from pathlib import Path
 import re
@@ -34,7 +34,7 @@ from homeassistant.components.backup import (
     LocalBackupAgent,
 )
 from homeassistant.components.backup.agent import BackupAgentError
-from homeassistant.components.backup.const import DATA_MANAGER
+from homeassistant.components.backup.const import BUF_SIZE, DATA_MANAGER
 from homeassistant.components.backup.manager import (
     AddonErrorData,
     AddonInfo,
@@ -50,6 +50,7 @@ from homeassistant.components.backup.manager import (
     UploadBackupEvent,
     WrittenBackup,
 )
+from homeassistant.components.http.server import MAX_CLIENT_SIZE
 from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STARTED
 from homeassistant.core import CoreState, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -2017,6 +2018,141 @@ async def test_receive_backup(
     assert unlink_mock.call_count == temp_file_unlink_call_count
 
 
+@pytest.mark.parametrize(
+    ("upload_size", "min_chunk_count"),
+    [
+        pytest.param(1024, 1, id="small"),
+        pytest.param(MAX_CLIENT_SIZE + 1024, 2, id="above_max_client_size"),
+    ],
+)
+async def test_receive_large_backup(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    upload_size: int,
+    min_chunk_count: int,
+) -> None:
+    """Test receiving a backup larger than the max request body size."""
+    await setup_backup_integration(hass)
+    # Make sure we wait for Platform.EVENT and Platform.SENSOR to be fully processed,
+    # to avoid interference with the Path.open patching below which is used to verify
+    # that the file is written to the expected location.
+    await hass.async_block_till_done(True)
+    client = await hass_client()
+    open_mock = mock_open()
+
+    with (
+        patch("pathlib.Path.open", open_mock),
+        patch("homeassistant.components.backup.manager.make_backup_dir"),
+        patch("shutil.move"),
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            return_value=TEST_BACKUP_ABC123,
+        ),
+    ):
+        data = FormData(quote_fields=False)
+        data.add_field(
+            "file",
+            BytesIO(b"\0" * upload_size),
+            filename="backup.tar",
+            content_type="application/octet-stream",
+        )
+        resp = await client.post("/api/backup/upload?agent_id=backup.local", data=data)
+        await hass.async_block_till_done()
+
+    assert resp.status == 201
+    assert await resp.json() == {"backup_id": TEST_BACKUP_ABC123.backup_id}
+    written_chunks = [
+        call.args[0] for call in open_mock.return_value.write.call_args_list
+    ]
+    assert sum(len(chunk) for chunk in written_chunks) == upload_size
+    # The file must be written in bounded chunks, not buffered into memory whole
+    assert len(written_chunks) >= min_chunk_count
+    assert max(len(chunk) for chunk in written_chunks) <= BUF_SIZE
+
+
+class _DisconnectingBodyPartReader:
+    """Minimal BodyPartReader whose read_chunk raises after the first chunk.
+
+    Models a client disconnect mid-upload: aiohttp sets a ConnectionResetError on
+    the request payload and cancels the handler when the connection is lost.
+    """
+
+    filename = "backup.tar"
+
+    def __init__(self, chunk: bytes, error: Exception) -> None:
+        """Initialize the reader."""
+        self._chunk = chunk
+        self._error = error
+        self._sent = False
+
+    async def read_chunk(self, size: int) -> bytes:
+        """Return one chunk, then raise on the next read."""
+        if self._sent:
+            raise self._error
+        self._sent = True
+        return self._chunk
+
+
+async def test_receive_backup_stream_error_resets_state(hass: HomeAssistant) -> None:
+    """Test the backup manager returns to IDLE when the upload stream fails.
+
+    A client disconnect (or other stream error) mid-upload must not wedge the
+    manager in a busy state; this drives the manager with a stream that raises
+    rather than a real socket disconnect, which the test client can't produce.
+    """
+    await setup_backup_integration(hass)
+    await hass.async_block_till_done(True)
+    manager = hass.data[DATA_MANAGER]
+    contents = _DisconnectingBodyPartReader(
+        b"\0" * 1024, ConnectionResetError("Connection lost")
+    )
+
+    with (
+        patch("pathlib.Path.open", mock_open()),
+        patch("homeassistant.components.backup.manager.make_backup_dir"),
+        patch("pathlib.Path.unlink") as unlink_mock,
+    ):
+        # Bound the await so a reintroduced deadlock fails fast instead of hanging.
+        async with asyncio.timeout(10):
+            with pytest.raises(ConnectionResetError):
+                await manager.async_receive_backup(
+                    agent_ids=["backup.local"], contents=contents
+                )
+
+    assert manager.state is BackupManagerState.IDLE
+    # The partially written temp file is removed on the failed upload.
+    assert unlink_mock.call_count == 1
+
+
+async def test_receive_backup_unparsable_file_removed(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+) -> None:
+    """Test the temp file is removed when the uploaded backup can't be parsed."""
+    await setup_backup_integration(hass)
+    await hass.async_block_till_done(True)
+    client = await hass_client()
+
+    with (
+        patch("pathlib.Path.open", mock_open(read_data=b"test")),
+        patch("homeassistant.components.backup.manager.make_backup_dir"),
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            side_effect=OSError("Boom"),
+        ),
+        patch("pathlib.Path.unlink") as unlink_mock,
+    ):
+        resp = await client.post(
+            "/api/backup/upload?agent_id=backup.local",
+            data={"file": StringIO("test")},
+        )
+        await hass.async_block_till_done()
+
+    assert resp.status == 500
+    # The unparsable temp file is removed exactly once.
+    assert unlink_mock.call_count == 1
+
+
 async def test_receive_backup_valid_filename(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,