From 4fc13585cc7aaceed2fc9493d31a0668dd52c013 Mon Sep 17 00:00:00 2001
From: Josh Wu <joshwu@google.com>
Date: Fri, 17 Mar 2023 22:58:03 +0800
Subject: [PATCH] Handle BR/EDR connection roles

---
 bumble/device.py | 10 +++++++++-
 bumble/host.py   | 29 +++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/bumble/device.py b/bumble/device.py
index 512bb1d..eec3685 100644
--- a/bumble/device.py
+++ b/bumble/device.py
@@ -615,7 +615,9 @@ class Connection(CompositeEventEmitter):
         assert self.transport == BT_BR_EDR_TRANSPORT
         self.handle = handle
         self.peer_resolvable_address = peer_resolvable_address
-        self.role = role
+        # Quirk: role might be known before complete
+        if self.role is None:
+            self.role = role
         self.parameters = parameters
 
     @property
@@ -2905,6 +2907,12 @@ class Device(CompositeEventEmitter):
         )
         connection.emit('connection_data_length_change')
 
+    # [Classic only]
+    @host_event_handler
+    @with_connection_from_address
+    def on_role_change(self, connection, new_role):
+        connection.role = new_role
+
     @with_connection_from_handle
     def on_pairing_start(self, connection):
         connection.emit('pairing_start')
diff --git a/bumble/host.py b/bumble/host.py
index 87ec610..96391a1 100644
--- a/bumble/host.py
+++ b/bumble/host.py
@@ -24,7 +24,10 @@ from bumble.colors import color
 from bumble.l2cap import L2CAP_PDU
 from bumble.snoop import Snooper
 
+from typing import Optional
+
 from .hci import (
+    Address,
     HCI_ACL_DATA_PACKET,
     HCI_COMMAND_COMPLETE_EVENT,
     HCI_COMMAND_PACKET,
@@ -142,6 +145,24 @@ class Host(AbortableEventEmitter):
         if controller_sink:
             self.set_packet_sink(controller_sink)
 
+    def find_connection_by_bd_addr(
+        self,
+        bd_addr: Address,
+        transport: Optional[int] = None,
+        check_address_type: bool = False,
+    ) -> Optional[Connection]:
+        for connection in self.connections.values():
+            if connection.peer_address.to_bytes() == bd_addr.to_bytes():
+                if (
+                    check_address_type
+                    and connection.peer_address.address_type != bd_addr.address_type
+                ):
+                    continue
+                if transport is None or connection.transport == transport:
+                    return connection
+
+        return None
+
     async def flush(self) -> None:
         # Make sure no command is pending
         await self.command_semaphore.acquire()
@@ -582,7 +603,7 @@ class Host(AbortableEventEmitter):
                 BT_BR_EDR_TRANSPORT,
                 event.bd_addr,
                 None,
-                BT_CENTRAL_ROLE,
+                role,
                 None,
             )
         else:
@@ -719,7 +740,11 @@ class Host(AbortableEventEmitter):
                 f'role change for {event.bd_addr}: '
                 f'{HCI_Constant.role_name(event.new_role)}'
             )
-            # TODO: lookup the connection and update the role
+            if connection := self.find_connection_by_bd_addr(
+                event.bd_addr, BT_BR_EDR_TRANSPORT
+            ):
+                connection.role = event.new_role
+             self.emit('role_change', event.bd_addr, event.new_role)
         else:
             logger.debug(
                 f'role change for {event.bd_addr} failed: '