diff --git a/homeassistant/components/backup/manager.py b/homeassistant/components/backup/manager.py index bff31d7b50be..f345178dce9b 100644 --- a/homeassistant/components/backup/manager.py +++ b/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") diff --git a/homeassistant/components/backup/util.py b/homeassistant/components/backup/util.py index 953b4ab4c102..f5a26da4e45d 100644 --- a/homeassistant/components/backup/util.py +++ b/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: diff --git a/tests/components/backup/test_manager.py b/tests/components/backup/test_manager.py index 62df097b75f7..3b63d8b25738 100644 --- a/tests/components/backup/test_manager.py +++ b/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,