Update cloud backup agent to use calculate_b64md5 from lib (#138391)

* Update cloud backup agent to use calculate_b64md5 from lib

* Catch error, add test

* Address review comments

* Update tests/components/cloud/test_backup.py

Co-authored-by: Abílio Costa <abmantis@users.noreply.github.com>

---------

Co-authored-by: Abílio Costa <abmantis@users.noreply.github.com>

homeassistant/components/cloud/backup.py

@@ -3,9 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-import base64
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
-import hashlib
 import logging
 import random
 from typing import Any
@@ -14,7 +12,7 @@ from aiohttp import ClientError
 from hass_nabucasa import Cloud, CloudError
 from hass_nabucasa.api import CloudApiNonRetryableError
 from hass_nabucasa.cloud_api import async_files_delete_file, async_files_list
-from hass_nabucasa.files import StorageType
+from hass_nabucasa.files import FilesError, StorageType, calculate_b64md5
 
 from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError
 from homeassistant.core import HomeAssistant, callback
@@ -30,14 +28,6 @@ _RETRY_SECONDS_MIN = 60
 _RETRY_SECONDS_MAX = 600
 
 
-async def _b64md5(stream: AsyncIterator[bytes]) -> str:
-    """Calculate the MD5 hash of a file."""
-    file_hash = hashlib.md5()
-    async for chunk in stream:
-        file_hash.update(chunk)
-    return base64.b64encode(file_hash.digest()).decode()
-
-
 async def async_get_backup_agents(
     hass: HomeAssistant,
     **kwargs: Any,
@@ -129,10 +119,13 @@ class CloudBackupAgent(BackupAgent):
         if not backup.protected:
             raise BackupAgentError("Cloud backups must be protected")
 
-        base64md5hash = await _b64md5(await open_stream())
+        size = backup.size
+        try:
+            base64md5hash = await calculate_b64md5(open_stream, size)
+        except FilesError as err:
+            raise BackupAgentError(err) from err
         filename = self._get_backup_filename()
         metadata = backup.as_dict()
-        size = backup.size
 
         tries = 1
         while tries <= _RETRY_LIMIT:

tests/components/cloud/test_backup.py

@@ -285,6 +285,7 @@ async def test_agents_upload(
 ) -> None:
     """Test agent upload backup."""
     client = await hass_client()
+    backup_data = "test"
     backup_id = "test-backup"
     test_backup = AgentBackup(
         addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
@@ -297,7 +298,7 @@ async def test_agents_upload(
         homeassistant_version="2024.12.0",
         name="Test",
         protected=True,
-        size=0,
+        size=len(backup_data),
     )
     with (
         patch(
@@ -309,11 +310,11 @@ async def test_agents_upload(
         ),
         patch("pathlib.Path.open") as mocked_open,
     ):
-        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        mocked_open.return_value.read = Mock(side_effect=[backup_data.encode(), b""])
         fetch_backup.return_value = test_backup
         resp = await client.post(
             "/api/backup/upload?agent_id=cloud.cloud",
-            data={"file": StringIO("test")},
+            data={"file": StringIO(backup_data)},
         )
 
     assert len(cloud.files.upload.mock_calls) == 1
@@ -336,6 +337,7 @@ async def test_agents_upload_fail(
 ) -> None:
     """Test agent upload backup fails."""
     client = await hass_client()
+    backup_data = "test"
     backup_id = "test-backup"
     test_backup = AgentBackup(
         addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
@@ -348,7 +350,7 @@ async def test_agents_upload_fail(
         homeassistant_version="2024.12.0",
         name="Test",
         protected=True,
-        size=0,
+        size=len(backup_data),
     )
 
     cloud.files.upload.side_effect = side_effect
@@ -366,11 +368,11 @@ async def test_agents_upload_fail(
         patch("homeassistant.components.cloud.backup.random.randint", return_value=60),
         patch("homeassistant.components.cloud.backup._RETRY_LIMIT", 2),
     ):
-        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        mocked_open.return_value.read = Mock(side_effect=[backup_data.encode(), b""])
         fetch_backup.return_value = test_backup
         resp = await client.post(
             "/api/backup/upload?agent_id=cloud.cloud",
-            data={"file": StringIO("test")},
+            data={"file": StringIO(backup_data)},
         )
         await hass.async_block_till_done()
 
@@ -409,6 +411,7 @@ async def test_agents_upload_fail_non_retryable(
 ) -> None:
     """Test agent upload backup fails with non-retryable error."""
     client = await hass_client()
+    backup_data = "test"
     backup_id = "test-backup"
     test_backup = AgentBackup(
         addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
@@ -435,12 +438,13 @@ async def test_agents_upload_fail_non_retryable(
             return_value=test_backup,
         ),
         patch("pathlib.Path.open") as mocked_open,
+        patch("homeassistant.components.cloud.backup.calculate_b64md5"),
     ):
-        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        mocked_open.return_value.read = Mock(side_effect=[backup_data.encode(), b""])
         fetch_backup.return_value = test_backup
         resp = await client.post(
             "/api/backup/upload?agent_id=cloud.cloud",
-            data={"file": StringIO("test")},
+            data={"file": StringIO(backup_data)},
         )
         await hass.async_block_till_done()
 
@@ -461,6 +465,7 @@ async def test_agents_upload_not_protected(
 ) -> None:
     """Test agent upload backup, when cloud user is logged in."""
     client = await hass_client()
+    backup_data = "test"
     backup_id = "test-backup"
     test_backup = AgentBackup(
         addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
@@ -473,7 +478,7 @@ async def test_agents_upload_not_protected(
         homeassistant_version="2024.12.0",
         name="Test",
         protected=False,
-        size=0,
+        size=len(backup_data),
     )
     with (
         patch("pathlib.Path.open"),
@@ -484,7 +489,7 @@ async def test_agents_upload_not_protected(
     ):
         resp = await client.post(
             "/api/backup/upload?agent_id=cloud.cloud",
-            data={"file": StringIO("test")},
+            data={"file": StringIO(backup_data)},
         )
         await hass.async_block_till_done()
 
@@ -496,6 +501,53 @@ async def test_agents_upload_not_protected(
     assert stored_backup["failed_agent_ids"] == ["cloud.cloud"]
 
 
+@pytest.mark.usefixtures("cloud_logged_in", "mock_list_files")
+async def test_agents_upload_wrong_size(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    caplog: pytest.LogCaptureFixture,
+    cloud: Mock,
+) -> None:
+    """Test agent upload backup with the wrong size."""
+    client = await hass_client()
+    backup_data = "test"
+    backup_id = "test-backup"
+    test_backup = AgentBackup(
+        addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
+        backup_id=backup_id,
+        database_included=True,
+        date="1970-01-01T00:00:00.000Z",
+        extra_metadata={},
+        folders=[Folder.MEDIA, Folder.SHARE],
+        homeassistant_included=True,
+        homeassistant_version="2024.12.0",
+        name="Test",
+        protected=True,
+        size=len(backup_data) - 1,
+    )
+    with (
+        patch(
+            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
+        ) as fetch_backup,
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            return_value=test_backup,
+        ),
+        patch("pathlib.Path.open") as mocked_open,
+    ):
+        mocked_open.return_value.read = Mock(side_effect=[backup_data.encode(), b""])
+        fetch_backup.return_value = test_backup
+        resp = await client.post(
+            "/api/backup/upload?agent_id=cloud.cloud",
+            data={"file": StringIO(backup_data)},
+        )
+
+    assert len(cloud.files.upload.mock_calls) == 0
+
+    assert resp.status == 201
+    assert "Upload failed for cloud.cloud" in caplog.text
+
+
 @pytest.mark.usefixtures("cloud_logged_in", "mock_list_files")
 async def test_agents_delete(
     hass: HomeAssistant,