smp: wait for key distribution to be flushed before completing pairing

bumble/device.py

@@ -2097,6 +2097,21 @@ def on_att_mtu_update(self, mtu: int):
     def data_packet_queue(self) -> DataPacketQueue | None:
         return self.device.host.get_data_packet_queue(self.handle)
 
+    async def drain(self) -> None:
+        """Wait until the controller has completed every data packet queued for
+        this connection.
+
+        Returns immediately if this connection has no packet queue, or if
+        nothing has ever been queued for it.
+        """
+        if (data_packet_queue := self.data_packet_queue) is None:
+            return
+        try:
+            await data_packet_queue.drain(self.handle)
+        except ValueError:
+            # Nothing was ever sent on this connection.
+            pass
+
     def cancel_on_disconnection(self, awaitable: Awaitable[_T]) -> Awaitable[_T]:
         """
         Helper method to call `utils.cancel_on_event` for the 'disconnection' event

bumble/smp.py

@@ -1304,6 +1304,13 @@ async def on_pairing(self) -> None:
 
         self.completed = True
 
+        # The keys we have just distributed may still be waiting for the
+        # controller's flow control. Don't report success until they have all
+        # been acknowledged: an application that disconnects as soon as
+        # `pair()` returns would otherwise cut the last key distribution PDU,
+        # and a peer that never receives it discards the incomplete bond.
+        await self.connection.drain()
+
         if self.pairing_result is not None and not self.pairing_result.done():
             self.pairing_result.set_result(None)
 

tests/self_test.py

@@ -361,6 +361,48 @@ def __str__(self):
             await _test_self_smp_with_configs(pairing_config1, pairing_config2)
 
 
+# -----------------------------------------------------------------------------
+@pytest.mark.asyncio
+async def test_self_smp_waits_for_key_distribution():
+    two_devices = await TwoDevices.create_with_connection()
+    for device in two_devices.devices:
+        device.pairing_config_factory = lambda _: PairingConfig(
+            delegate=PairingDelegate(PairingDelegate.IoCapability.NO_OUTPUT_NO_INPUT)
+        )
+
+    connection = two_devices.connections[0]
+    packet_queue = connection.data_packet_queue
+    assert packet_queue is not None
+
+    # Withhold the controller's packet completions, so that everything the
+    # initiator sends stays in flight.
+    withheld: list[tuple[int, int]] = []
+    complete_packets = packet_queue.on_packets_completed
+    packet_queue.on_packets_completed = lambda count, handle: withheld.append(
+        (count, handle)
+    )
+
+    pairing = asyncio.create_task(connection.pair())
+    for _ in range(100):
+        if connection.is_encrypted:
+            break
+        await asyncio.sleep(0.01)
+    await async_barrier()
+
+    # Check that the withholding is in effect, then that pairing has not been
+    # reported as complete while the keys may still be sitting in the queue.
+    assert connection.is_encrypted
+    assert packet_queue.pending > 0
+    assert not pairing.done()
+
+    # Complete the packets: pairing may now finish.
+    packet_queue.on_packets_completed = complete_packets
+    for count, handle in withheld:
+        complete_packets(count, handle)
+    await asyncio.wait_for(pairing, timeout=5.0)
+    assert packet_queue.pending == 0
+
+
 # -----------------------------------------------------------------------------
 @pytest.mark.asyncio
 async def test_self_smp_reject():