From 05e76105ad0dd28653701c7900fb70d3928d9b7a Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 5 Nov 2024 17:12:05 +0100
Subject: [PATCH] Improve improv BLE error handling (#129902)

---
 .../components/improv_ble/config_flow.py      | 18 ++++++++++++++----
 tests/components/improv_ble/__init__.py       | 19 +++++++++++++++++++
 .../components/improv_ble/test_config_flow.py | 18 ++++++++++++++++++
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/improv_ble/config_flow.py b/homeassistant/components/improv_ble/config_flow.py
index f38f4830ace..05dd1de449a 100644
--- a/homeassistant/components/improv_ble/config_flow.py
+++ b/homeassistant/components/improv_ble/config_flow.py
@@ -120,12 +120,22 @@ class ImprovBLEConfigFlow(ConfigFlow, domain=DOMAIN):
         assert self._discovery_info is not None
 
         service_data = self._discovery_info.service_data
-        improv_service_data = ImprovServiceData.from_bytes(
-            service_data[SERVICE_DATA_UUID]
-        )
+        try:
+            improv_service_data = ImprovServiceData.from_bytes(
+                service_data[SERVICE_DATA_UUID]
+            )
+        except improv_ble_errors.InvalidCommand as err:
+            _LOGGER.warning(
+                "Aborting improv flow, device %s sent invalid improv data: '%s'",
+                self._discovery_info.address,
+                service_data[SERVICE_DATA_UUID].hex(),
+            )
+            raise AbortFlow("invalid_improv_data") from err
+
         if improv_service_data.state in (State.PROVISIONING, State.PROVISIONED):
             _LOGGER.debug(
-                "Aborting improv flow, device is already provisioned: %s",
+                "Aborting improv flow, device %s is already provisioned: %s",
+                self._discovery_info.address,
                 improv_service_data.state,
             )
             raise AbortFlow("already_provisioned")
diff --git a/tests/components/improv_ble/__init__.py b/tests/components/improv_ble/__init__.py
index 41ea98cda7b..521d0881443 100644
--- a/tests/components/improv_ble/__init__.py
+++ b/tests/components/improv_ble/__init__.py
@@ -25,6 +25,25 @@ IMPROV_BLE_DISCOVERY_INFO = BluetoothServiceInfoBleak(
 )
 
 
+BAD_IMPROV_BLE_DISCOVERY_INFO = BluetoothServiceInfoBleak(
+    name="00123456",
+    address="AA:BB:CC:DD:EE:F0",
+    rssi=-60,
+    manufacturer_data={},
+    service_uuids=[SERVICE_UUID],
+    service_data={SERVICE_DATA_UUID: b"\x00\x00\x00\x00\x00\x00"},
+    source="local",
+    device=generate_ble_device(address="AA:BB:CC:DD:EE:F0", name="00123456"),
+    advertisement=generate_advertisement_data(
+        service_uuids=[SERVICE_UUID],
+        service_data={SERVICE_DATA_UUID: b"\x00\x00\x00\x00\x00\x00"},
+    ),
+    time=0,
+    connectable=True,
+    tx_power=-127,
+)
+
+
 PROVISIONED_IMPROV_BLE_DISCOVERY_INFO = BluetoothServiceInfoBleak(
     name="00123456",
     address="AA:BB:CC:DD:EE:F0",
diff --git a/tests/components/improv_ble/test_config_flow.py b/tests/components/improv_ble/test_config_flow.py
index 640a931bee5..2df4be2ba7d 100644
--- a/tests/components/improv_ble/test_config_flow.py
+++ b/tests/components/improv_ble/test_config_flow.py
@@ -15,6 +15,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult, FlowResultType
 
 from . import (
+    BAD_IMPROV_BLE_DISCOVERY_INFO,
     IMPROV_BLE_DISCOVERY_INFO,
     NOT_IMPROV_BLE_DISCOVERY_INFO,
     PROVISIONED_IMPROV_BLE_DISCOVERY_INFO,
@@ -649,3 +650,20 @@ async def test_provision_retry(hass: HomeAssistant, exc, error) -> None:
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "provision"
     assert result["errors"] == {"base": error}
+
+
+async def test_provision_fails_invalid_data(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test bluetooth flow with error due to invalid data."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+        data=BAD_IMPROV_BLE_DISCOVERY_INFO,
+    )
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "invalid_improv_data"
+    assert (
+        "Aborting improv flow, device AA:BB:CC:DD:EE:F0 sent invalid improv data: '000000000000'"
+        in caplog.text
+    )