Raise GATT errors on read and write etc (#272)

Co-authored-by: J. Nick Koston <nick@koston.org>

aioesphomeapi/api.proto

@@ -1326,3 +1326,31 @@ message BluetoothConnectionsFreeResponse {
   uint32 free = 1;
   uint32 limit = 2;
 }
+
+message BluetoothGATTErrorResponse {
+  option (id) = 82;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_BLUETOOTH_PROXY";
+
+  uint64 address = 1;
+  uint32 handle = 2;
+  int32 error = 3;
+}
+
+message BluetoothGATTWriteResponse {
+  option (id) = 83;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_BLUETOOTH_PROXY";
+
+  uint64 address = 1;
+  uint32 handle = 2;
+}
+
+message BluetoothGATTNotifyResponse {
+  option (id) = 84;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_BLUETOOTH_PROXY";
+
+  uint64 address = 1;
+  uint32 handle = 2;
+}

aioesphomeapi/client.py

@@ -23,16 +23,19 @@ from .api_pb2 import (  # type: ignore
     BluetoothConnectionsFreeResponse,
     BluetoothDeviceConnectionResponse,
     BluetoothDeviceRequest,
+    BluetoothGATTErrorResponse,
     BluetoothGATTGetServicesDoneResponse,
     BluetoothGATTGetServicesRequest,
     BluetoothGATTGetServicesResponse,
     BluetoothGATTNotifyDataResponse,
     BluetoothGATTNotifyRequest,
+    BluetoothGATTNotifyResponse,
     BluetoothGATTReadDescriptorRequest,
     BluetoothGATTReadRequest,
     BluetoothGATTReadResponse,
     BluetoothGATTWriteDescriptorRequest,
     BluetoothGATTWriteRequest,
+    BluetoothGATTWriteResponse,
     BluetoothLEAdvertisementResponse,
     ButtonCommandRequest,
     CameraImageRequest,
@@ -93,7 +96,7 @@ from .api_pb2 import (  # type: ignore
     TextSensorStateResponse,
 )
 from .connection import APIConnection, ConnectionParams
-from .core import APIConnectionError, TimeoutAPIError
+from .core import APIConnectionError, BluetoothGATTAPIError, TimeoutAPIError
 from .host_resolver import ZeroconfInstanceType
 from .model import (
     APIVersion,
@@ -102,6 +105,7 @@ from .model import (
     BluetoothConnectionsFree,
     BluetoothDeviceConnection,
     BluetoothDeviceRequestType,
+    BluetoothGATTError,
     BluetoothGATTRead,
     BluetoothGATTServices,
     BluetoothLEAdvertisement,
@@ -408,6 +412,33 @@ class APIClient:
             SubscribeHomeassistantServicesRequest(), on_msg
         )
 
+    async def _send_bluetooth_message_await_response(
+        self,
+        address: int,
+        handle: int,
+        request: message.Message,
+        response_type: Type[message.Message],
+        timeout: float = 10.0,
+    ) -> message.Message:
+        self._check_authenticated()
+        assert self._connection is not None
+
+        def is_response(msg: message.Message) -> bool:
+            return (
+                isinstance(msg, (BluetoothGATTErrorResponse, response_type))
+                and msg.address == address  # type: ignore[union-attr]
+                and msg.handle == handle  # type: ignore[union-attr]
+            )
+
+        resp = await self._connection.send_message_await_response_complex(
+            request, is_response, is_response, timeout=timeout
+        )
+
+        if isinstance(resp[0], BluetoothGATTErrorResponse):
+            raise BluetoothGATTAPIError(BluetoothGATTError.from_pb(resp[0]))
+
+        return resp[0]
+
     async def subscribe_bluetooth_le_advertisements(
         self, on_bluetooth_le_advertisement: Callable[[BluetoothLEAdvertisement], None]
     ) -> Callable[[], None]:
@@ -509,10 +540,21 @@ class APIClient:
         self._check_authenticated()
 
         def do_append(msg: message.Message) -> bool:
-            return isinstance(msg, BluetoothGATTGetServicesResponse)
+            return (
+                isinstance(
+                    msg, (BluetoothGATTGetServicesResponse, BluetoothGATTErrorResponse)
+                )
+                and msg.address == address
+            )
 
         def do_stop(msg: message.Message) -> bool:
-            return isinstance(msg, BluetoothGATTGetServicesDoneResponse)
+            return (
+                isinstance(
+                    msg,
+                    (BluetoothGATTGetServicesDoneResponse, BluetoothGATTErrorResponse),
+                )
+                and msg.address == address
+            )
 
         assert self._connection is not None
         resp = await self._connection.send_message_await_response_complex(
@@ -523,58 +565,56 @@ class APIClient:
         )
         services = []
         for msg in resp:
+            if isinstance(msg, BluetoothGATTErrorResponse):
+                raise BluetoothGATTAPIError(BluetoothGATTError.from_pb(msg))
+
             services.extend(BluetoothGATTServices.from_pb(msg).services)
+
         return ESPHomeBluetoothGATTServices(address=address, services=services)
 
     async def bluetooth_gatt_read(
         self,
         address: int,
-        characteristic_handle: int,
+        handle: int,
         timeout: float = DEFAULT_BLE_TIMEOUT,
     ) -> bytearray:
-        self._check_authenticated()
-
         req = BluetoothGATTReadRequest()
         req.address = address
-        req.handle = characteristic_handle
+        req.handle = handle
 
-        def is_response(msg: message.Message) -> bool:
-            if isinstance(msg, BluetoothGATTReadResponse):
-                read = BluetoothGATTRead.from_pb(msg)
-
-                if read.address == address and read.handle == characteristic_handle:
-                    return True
-            return False
-
-        assert self._connection is not None
-        resp = await self._connection.send_message_await_response_complex(
-            req, is_response, is_response, timeout=timeout
+        resp = await self._send_bluetooth_message_await_response(
+            address,
+            handle,
+            req,
+            BluetoothGATTReadResponse,
+            timeout=timeout,
         )
 
-        if len(resp) != 1:
-            raise APIConnectionError(f"Expected one result, got {len(resp)}")
-
-        read_response = BluetoothGATTRead.from_pb(resp[0])
+        read_response = BluetoothGATTRead.from_pb(resp)
 
         return bytearray(read_response.data)
 
     async def bluetooth_gatt_write(
         self,
         address: int,
-        characteristic_handle: int,
+        handle: int,
         data: bytes,
         response: bool,
+        timeout: float = 10.0,
     ) -> None:
-        self._check_authenticated()
-
         req = BluetoothGATTWriteRequest()
         req.address = address
-        req.handle = characteristic_handle
+        req.handle = handle
         req.response = response
         req.data = data
 
-        assert self._connection is not None
-        await self._connection.send_message(req)
+        await self._send_bluetooth_message_await_response(
+            address,
+            handle,
+            req,
+            BluetoothGATTWriteResponse,
+            timeout=timeout,
+        )
 
     async def bluetooth_gatt_read_descriptor(
         self,
@@ -582,29 +622,19 @@ class APIClient:
         handle: int,
         timeout: float = DEFAULT_BLE_TIMEOUT,
     ) -> bytearray:
-        self._check_authenticated()
-
         req = BluetoothGATTReadDescriptorRequest()
         req.address = address
         req.handle = handle
 
-        def is_response(msg: message.Message) -> bool:
-            if isinstance(msg, BluetoothGATTReadResponse):
-                read = BluetoothGATTRead.from_pb(msg)
-
-                if read.address == address and read.handle == handle:
-                    return True
-            return False
-
-        assert self._connection is not None
-        resp = await self._connection.send_message_await_response_complex(
-            req, is_response, is_response, timeout=timeout
+        resp = await self._send_bluetooth_message_await_response(
+            address,
+            handle,
+            req,
+            BluetoothGATTReadResponse,
+            timeout=timeout,
         )
 
-        if len(resp) != 1:
-            raise APIConnectionError(f"Expected one result, got {len(resp)}")
-
-        read_response = BluetoothGATTRead.from_pb(resp[0])
+        read_response = BluetoothGATTRead.from_pb(resp)
 
         return bytearray(read_response.data)
 
@@ -613,16 +643,20 @@ class APIClient:
         address: int,
         handle: int,
         data: bytes,
+        timeout: float = 10.0,
     ) -> None:
-        self._check_authenticated()
-
         req = BluetoothGATTWriteDescriptorRequest()
         req.address = address
         req.handle = handle
         req.data = data
 
-        assert self._connection is not None
-        await self._connection.send_message(req)
+        await self._send_bluetooth_message_await_response(
+            address,
+            handle,
+            req,
+            BluetoothGATTWriteResponse,
+            timeout=timeout,
+        )
 
     async def bluetooth_gatt_start_notify(
         self,
@@ -630,7 +664,13 @@ class APIClient:
         handle: int,
         on_bluetooth_gatt_notify: Callable[[int, bytearray], None],
     ) -> Callable[[], Coroutine[Any, Any, None]]:
-        self._check_authenticated()
+
+        await self._send_bluetooth_message_await_response(
+            address,
+            handle,
+            BluetoothGATTNotifyRequest(address=address, handle=handle, enable=True),
+            BluetoothGATTNotifyResponse,
+        )
 
         def on_msg(msg: message.Message) -> None:
             if isinstance(msg, BluetoothGATTNotifyDataResponse):
@@ -639,16 +679,13 @@ class APIClient:
                     on_bluetooth_gatt_notify(handle, bytearray(notify.data))
 
         assert self._connection is not None
-        await self._connection.send_message_callback_response(
-            BluetoothGATTNotifyRequest(address=address, handle=handle, enable=True),
-            on_msg,
-        )
+        remove_callback = self._connection.add_message_callback(on_msg)
 
         async def stop_notify() -> None:
             if self._connection is None:
                 return
 
-            self._connection.remove_message_callback(on_msg)
+            remove_callback()
 
             self._check_authenticated()
 

aioesphomeapi/connection.py

@@ -403,6 +403,17 @@ class APIConnection:
             await self._report_fatal_error(err)
             raise
 
+    def add_message_callback(
+        self, on_message: Callable[[Any], None]
+    ) -> Callable[[], None]:
+        """Add a message callback."""
+        self._message_handlers.append(on_message)
+
+        def unsub() -> None:
+            self._message_handlers.remove(on_message)
+
+        return unsub
+
     def remove_message_callback(self, on_message: Callable[[Any], None]) -> None:
         """Remove a message callback."""
         self._message_handlers.remove(on_message)
@@ -417,10 +428,10 @@ class APIConnection:
     async def send_message_await_response_complex(
         self,
         send_msg: message.Message,
-        do_append: Callable[[Any], bool],
-        do_stop: Callable[[Any], bool],
+        do_append: Callable[[message.Message], bool],
+        do_stop: Callable[[message.Message], bool],
         timeout: float = 10.0,
-    ) -> List[Any]:
+    ) -> List[message.Message]:
         """Send a message to the remote and build up a list response.
 
         :param send_msg: The message (request) to send.

aioesphomeapi/core.py

@@ -1,18 +1,25 @@
+import re
+
+from aioesphomeapi.model import BluetoothGATTError
+
 from .api_pb2 import (  # type: ignore
     BinarySensorStateResponse,
     BluetoothConnectionsFreeResponse,
     BluetoothDeviceConnectionResponse,
     BluetoothDeviceRequest,
+    BluetoothGATTErrorResponse,
     BluetoothGATTGetServicesDoneResponse,
     BluetoothGATTGetServicesRequest,
     BluetoothGATTGetServicesResponse,
     BluetoothGATTNotifyDataResponse,
     BluetoothGATTNotifyRequest,
+    BluetoothGATTNotifyResponse,
     BluetoothGATTReadDescriptorRequest,
     BluetoothGATTReadRequest,
     BluetoothGATTReadResponse,
     BluetoothGATTWriteDescriptorRequest,
     BluetoothGATTWriteRequest,
+    BluetoothGATTWriteResponse,
     BluetoothLEAdvertisementResponse,
     ButtonCommandRequest,
     CameraImageRequest,
@@ -82,6 +89,8 @@ from .api_pb2 import (  # type: ignore
     TextSensorStateResponse,
 )
 
+TWO_CHAR = re.compile(r".{2}")
+
 
 class APIConnectionError(Exception):
     pass
@@ -139,6 +148,19 @@ class ReadFailedAPIError(APIConnectionError):
     pass
 
 
+def to_human_readable_address(address: int) -> str:
+    """Convert a MAC address to a human readable format."""
+    return ":".join(TWO_CHAR.findall(f"{address:012X}"))
+
+
+class BluetoothGATTAPIError(APIConnectionError):
+    def __init__(self, error: BluetoothGATTError) -> None:
+        super().__init__(
+            f"Bluetooth GATT Error address={to_human_readable_address(error.address)} handle={error.handle} error={error.error}"
+        )
+        self.error = error
+
+
 MESSAGE_TYPE_TO_PROTO = {
     1: HelloRequest,
     2: HelloResponse,
@@ -221,4 +243,7 @@ MESSAGE_TYPE_TO_PROTO = {
     79: BluetoothGATTNotifyDataResponse,
     80: SubscribeBluetoothConnectionsFreeRequest,
     81: BluetoothConnectionsFreeResponse,
+    82: BluetoothGATTErrorResponse,
+    83: BluetoothGATTWriteResponse,
+    84: BluetoothGATTNotifyResponse,
 }

aioesphomeapi/model.py

@@ -900,6 +900,13 @@ class BluetoothConnectionsFree(APIModelBase):
     limit: int = 0
 
 
+@dataclass(frozen=True)
+class BluetoothGATTError(APIModelBase):
+    address: int = 0
+    handle: int = 0
+    error: int = 0
+
+
 class BluetoothDeviceRequestType(APIIntEnum):
     CONNECT = 0
     DISCONNECT = 1