Hybrid transport: Switch to event-driven discovery

Author: AlfioEmanueleFresta

libwebauthn/src/transport/ble/btleplug/manager.rs

@@ -2,10 +2,11 @@ use std::collections::HashMap;
 
 use btleplug::api::bleuuid::uuid_from_u16;
 use btleplug::api::{
-    Central as _, Manager as _, Peripheral as _, PeripheralProperties, ScanFilter,
+    Central as _, CentralEvent, Manager as _, Peripheral as _, PeripheralProperties, ScanFilter,
 };
-use btleplug::platform::{Adapter, Manager, Peripheral};
-use tracing::{debug, info, instrument, Level};
+use btleplug::platform::{Adapter, Manager, Peripheral, PeripheralId};
+use futures::{Stream, StreamExt};
+use tracing::{debug, info, instrument, trace, warn, Level};
 use uuid::Uuid;
 
 use super::device::FidoEndpoints;
@@ -50,17 +51,66 @@ impl SupportedRevisions {
     }
 }
 
+async fn on_peripheral_service_data(
+    adapter: &Adapter,
+    id: &PeripheralId,
+    uuids: &[Uuid],
+    service_data: HashMap<Uuid, Vec<u8>>,
+) -> Option<(Peripheral, Vec<u8>)> {
+    for uuid in uuids {
+        if let Some(service_data) = service_data.get(uuid) {
+            trace!(?id, ?service_data, "Found service data");
+            let Ok(peripheral) = adapter.peripheral(id).await else {
+                warn!(?id, "Could not get peripheral");
+                return None;
+            };
+
+            debug!({ ?id, ?service_data }, "Found service data for peripheral");
+            return Some((peripheral, service_data.to_owned()));
+        }
+    }
+
+    trace!(
+        { ?id, ?service_data },
+        "Ignoring periperal as it doesn't have service data for desired UUID"
+    );
+    None
+}
+
 #[instrument(level = Level::DEBUG, skip_all)]
-pub async fn start_discovery(uuids: &[Uuid]) -> Result<(), Error> {
+/// Starts a discovery for devices with specific service data
+pub async fn start_discovery_for_service_data(
+    uuids: Vec<Uuid>,
+) -> Result<impl Stream<Item = (Peripheral, Vec<u8>)>, Error> {
     let adapter = get_adapter().await?;
-    let scan_filter = ScanFilter {
-        services: uuids.to_vec(),
-    };
+    let scan_filter = ScanFilter::default();
+
+    let events = adapter.events().await.or(Err(Error::Unavailable))?;
 
     adapter
         .start_scan(scan_filter)
         .await
-        .or(Err(Error::ConnectionFailed))
+        .or(Err(Error::ConnectionFailed))?;
+
+    let stream = events.filter_map({
+        let adapter = adapter.clone();
+        let uuids = uuids.clone();
+        move |event| {
+            let adapter = adapter.clone();
+            let uuids = uuids.clone();
+            async move {
+                // trace!(?event);
+                match event {
+                    CentralEvent::ServiceDataAdvertisement { id, service_data } => {
+                        on_peripheral_service_data(&adapter, &id, &uuids, service_data).await
+                    }
+                    _ => None,
+                }
+            }
+        }
+    });
+
+    Ok(stream)
 }
 
 /// TODO(#86): Support multiple adapters.
@@ -84,6 +134,7 @@ async fn discover_properties(
             .properties()
             .await
             .or(Err(Error::ConnectionFailed))?;
+        trace!({ ?peripheral, ?properties });
         if let Some(properties) = properties {
             result.push((peripheral, properties));
         }
@@ -117,38 +168,20 @@ pub async fn list_fido_devices() -> Result<Vec<FidoDevice>, Error> {
     Ok(with_properties)
 }
 
-#[instrument(level = Level::DEBUG, skip_all)]
-pub async fn list_devices_with_service_data(
-    service_uuid: Uuid,
-) -> Result<HashMap<FidoDevice, Vec<u8>>, Error> {
-    let adapter = get_adapter().await?;
-    let peripherals = adapter
-        .peripherals()
+pub async fn get_device(peripheral: Peripheral) -> Result<Option<FidoDevice>, Error> {
+    let Some(properties) = peripheral
+        .properties()
         .await
-        .or(Err(Error::ConnectionFailed))?;
-    for peripheral in &peripherals {
-        // TODO: parallelize this
-        peripheral
-            .discover_services()
-            .await
-            .or(Err(Error::ConnectionFailed))?;
-    }
-    let with_properties = discover_properties(peripherals).await?;
-    Ok(with_properties
-        .into_iter()
-        .filter_map(
-            |(peripheral, properties)| match properties.service_data.get(&service_uuid) {
-                Some(service_data) => {
-                    let device = FidoDevice {
-                        peripheral,
-                        properties: properties.to_owned(),
-                    };
-                    Some((device, service_data.to_owned()))
-                }
-                None => None,
-            },
-        )
-        .collect())
+        .or(Err(Error::ConnectionFailed))?
+    else {
+        return Ok(None);
+    };
+
+    let device = FidoDevice {
+        peripheral,
+        properties,
+    };
+    Ok(Some(device))
 }
 
 pub async fn supported_fido_revisions(

libwebauthn/src/transport/ble/btleplug/mod.rs

@@ -8,6 +8,5 @@ pub use connection::Connection;
 pub use device::FidoDevice;
 pub use error::Error;
 pub use manager::{
-    connect, list_devices_with_service_data, list_fido_devices, start_discovery,
-    supported_fido_revisions,
+    connect, list_fido_devices, start_discovery_for_service_data, supported_fido_revisions,
 };

libwebauthn/src/transport/cable/qr_code_device.rs

@@ -1,7 +1,9 @@
 use std::fmt::{Debug, Display};
-use std::time::{Duration, SystemTime};
+use std::pin::pin;
+use std::time::SystemTime;
 
 use async_trait::async_trait;
+use futures::StreamExt;
 use p256::elliptic_curve::sec1::ToEncodedPoint;
 use p256::{NonZeroScalar, SecretKey};
 use rand::rngs::OsRng;
@@ -10,8 +12,7 @@ use serde::Serialize;
 use serde_bytes::ByteArray;
 use serde_indexed::SerializeIndexed;
 use tokio::sync::mpsc;
-use tokio::time::sleep;
-use tracing::{debug, error, trace};
+use tracing::{debug, error, trace, warn};
 use uuid::Uuid;
 
 use super::known_devices::CableKnownDeviceInfoStore;
@@ -28,8 +29,6 @@ use crate::UxUpdate;
 const CABLE_UUID_FIDO: &str = "0000fff9-0000-1000-8000-00805f9b34fb";
 const CABLE_UUID_GOOGLE: &str = "0000fde2-0000-1000-8000-00805f9b34fb";
 
-const ADVERTISEMENT_WAIT_LOOP_MS: u64 = 2000;
-
 #[derive(Debug, Clone, Copy)]
 pub enum QrCodeOperationHint {
     GetAssertionRequest,
@@ -185,50 +184,48 @@ impl CableQrCodeDevice<'_> {
     }
 
     async fn await_advertisement(&self) -> Result<(FidoDevice, DecryptedAdvert), Error> {
-        btleplug::manager::start_discovery(&[
-            Uuid::parse_str(CABLE_UUID_FIDO).unwrap(),
+        let stream = btleplug::manager::start_discovery_for_service_data(vec![
             Uuid::parse_str(CABLE_UUID_GOOGLE).unwrap(),
+            Uuid::parse_str(CABLE_UUID_FIDO).unwrap(),
         ])
         .await
         .or(Err(Error::Transport(TransportError::TransportUnavailable)))?;
 
-        loop {
-            let devices_service_data = btleplug::manager::list_devices_with_service_data(
-                Uuid::parse_str(CABLE_UUID_FIDO).unwrap(),
-            )
-            .await
-            .or(Err(Error::Transport(TransportError::TransportUnavailable)))?;
-            debug!({ ?devices_service_data }, "Found devices with service data");
-
-            let device = devices_service_data
-                .into_iter()
-                .map(|(device, data)| {
-                    let eid_key = derive(&self.qr_code.qr_secret, None, KeyPurpose::EIDKey);
-                    trace!(?device, ?data, ?eid_key);
-                    let decrypted = trial_decrypt_advert(&eid_key, &data);
-                    trace!(?decrypted);
-                    (device, decrypted)
-                })
-                .find(|(_, decrypted)| decrypted.is_some())
-                .map(|(device, decrypted)| {
-                    let decrypted = decrypted.unwrap();
-                    let advert = DecryptedAdvert::from(decrypted.as_slice());
-                    (device, advert)
-                });
-
-            if let Some((device, decrypted)) = device {
-                debug!(
-                    ?device,
-                    ?decrypted,
-                    "Successfully decrypted advertisement from device"
+        let mut stream = pin!(stream);
+        while let Some((peripheral, data)) = stream.as_mut().next().await {
+            debug!({ ?peripheral, ?data }, "Found device with service data");
+
+            let Some(device) = btleplug::manager::get_device(peripheral.clone())
+                .await
+                .or(Err(Error::Transport(TransportError::TransportUnavailable)))?
+            else {
+                warn!(
+                    ?peripheral,
+                    "Unable to fetch peripheral properties, ignoring"
                 );
-
-                return Ok((device, decrypted));
-            }
-
-            debug!("No devices found with matching advertisement, waiting for new advertisement");
-            sleep(Duration::from_millis(ADVERTISEMENT_WAIT_LOOP_MS as u64)).await;
+                continue;
+            };
+
+            let eid_key: Vec<u8> = derive(&self.qr_code.qr_secret, None, KeyPurpose::EIDKey);
+            trace!(?device, ?data, ?eid_key);
+
+            let Some(decrypted) = trial_decrypt_advert(&eid_key, &data) else {
+                warn!(?device, "Trial decrypt failed, ignoring");
+                continue;
+            };
+            trace!(?decrypted);
+
+            let advert = DecryptedAdvert::from(decrypted.as_slice());
+            debug!(
+                ?device,
+                ?decrypted,
+                "Successfully decrypted advertisement from device"
+            );
+
+            return Ok((device, advert));
         }
+
+        unreachable!()
     }
 }