smp: save pairing keys to keystore before draining connection

bumble/smp.py

@@ -1304,16 +1304,6 @@ 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)
-
         # Use the peer address from the pairing protocol or the connection
         if self.peer_bd_addr is not None:
             peer_address = self.peer_bd_addr
@@ -1366,6 +1356,16 @@ async def on_pairing(self) -> None:
             )
         await self.manager.on_pairing(self, peer_address, keys)
 
+        # 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)
+
     def on_pairing_failure(self, reason: ErrorCode) -> None:
         logger.warning('pairing failure (%s)', reason.name)
 

tests/self_test.py

@@ -389,10 +389,18 @@ async def test_self_smp_waits_for_key_distribution():
     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.
+    # reported as complete while the keys may still be sitting in the queue,
+    # even though the keys have already been saved to the keystore.
     assert connection.is_encrypted
     assert packet_queue.pending > 0
     assert not pairing.done()
+    assert two_devices.devices[0].keystore is not None
+    assert (
+        await two_devices.devices[0].keystore.get(
+            str(two_devices.devices[1].public_address)
+        )
+        is not None
+    )
 
     # Complete the packets: pairing may now finish.
     packet_queue.on_packets_completed = complete_packets