Use multicast to wait for new scan results

Author: Sh3Rm4n

examples/nl80211_trigger_scan.rs

@@ -2,10 +2,21 @@
 
 use std::env::args;
 
-use anyhow::{bail, Context, Error};
-use futures::stream::TryStreamExt;
+use anyhow::{anyhow, bail, Context, Result};
+use futures::{stream::TryStreamExt, StreamExt};
+use netlink_packet_core::{NetlinkMessage, NetlinkPayload, NLM_F_REQUEST};
+use netlink_packet_generic::{
+    ctrl::{
+        nlas::{GenlCtrlAttrs, McastGrpAttrs},
+        GenlCtrl, GenlCtrlCmd,
+    },
+    GenlMessage,
+};
+use netlink_packet_utils::ParseableParametrized;
+use netlink_sys::AsyncSocket;
+use wl_nl80211::{Nl80211Attr, Nl80211Command, Nl80211Message};
 
-fn main() -> Result<(), Error> {
+fn main() -> Result<()> {
     let argv: Vec<_> = args().collect();
 
     if argv.len() < 2 {
@@ -21,13 +32,20 @@ fn main() -> Result<(), Error> {
         .enable_time()
         .build()
         .unwrap();
-    rt.block_on(dump_scan(index));
+    rt.block_on(dump_scan(index))?;
 
     Ok(())
 }
 
-async fn dump_scan(if_index: u32) {
-    let (connection, handle, _) = wl_nl80211::new_connection().unwrap();
+async fn dump_scan(if_index: u32) -> Result<()> {
+    let (mut connection, handle, mut messages) = wl_nl80211::new_connection()?;
+
+    // Attach the connection socket to the multicast scan group to find out,
+    // when the scan is finished.
+    let socket = connection.socket_mut().socket_mut();
+    socket.bind_auto()?;
+    socket.add_membership(get_scan_multicast_id().await?)?;
+
     tokio::spawn(connection);
 
     let attrs = wl_nl80211::Nl80211Scan::new(if_index)
@@ -41,7 +59,23 @@ async fn dump_scan(if_index: u32) {
     while let Some(msg) = scan_handle.try_next().await.unwrap() {
         msgs.push(msg);
     }
-    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+
+    while let Some((message, _)) = messages.next().await {
+        match message.payload {
+            NetlinkPayload::InnerMessage(msg) => {
+                let msg = Nl80211Message::parse_with_param(
+                    msg.payload.as_slice(),
+                    msg.header,
+                )?;
+                if msg.cmd == Nl80211Command::NewScanResults
+                    && msg.attributes.contains(&Nl80211Attr::IfIndex(if_index))
+                {
+                    break;
+                }
+            }
+            _ => continue,
+        }
+    }
 
     let mut dump = handle.scan().dump(if_index).execute().await;
     let mut msgs = Vec::new();
@@ -52,4 +86,65 @@ async fn dump_scan(if_index: u32) {
     for msg in msgs {
         println!("{msg:?}");
     }
+
+    Ok(())
+}
+
+async fn get_scan_multicast_id() -> Result<u32> {
+    let (conn, mut handle, _) = wl_nl80211::new_connection()?;
+    tokio::spawn(conn);
+
+    let mut nl_msg =
+        NetlinkMessage::from(GenlMessage::from_payload(GenlCtrl {
+            cmd: GenlCtrlCmd::GetFamily,
+            nlas: vec![GenlCtrlAttrs::FamilyName("nl80211".to_owned())],
+        }));
+
+    // To get the mcast groups for the nl80211 family, we must also set the
+    // message type id
+    nl_msg.header.message_type =
+        handle.handle.resolve_family_id::<Nl80211Message>().await?;
+    // This is a request, but not a dump. Which means, the family name has to be
+    // specified, to obtain it's information.
+    nl_msg.header.flags = NLM_F_REQUEST;
+
+    let responses = handle.handle.request(nl_msg).await?;
+    let nl80211_family: Vec<Vec<GenlCtrlAttrs>> = responses
+        .try_filter_map(|msg| async move {
+            match msg.payload {
+                NetlinkPayload::InnerMessage(genlmsg)
+                    if genlmsg.payload.cmd == GenlCtrlCmd::NewFamily
+                        && genlmsg.payload.nlas.contains(
+                            &GenlCtrlAttrs::FamilyName("nl80211".to_owned()),
+                        ) =>
+                {
+                    Ok(Some(genlmsg.payload.nlas.clone()))
+                }
+                _ => Ok(None),
+            }
+        })
+        .try_collect()
+        .await?;
+
+    // Now get the mcid for "nl80211" "scan" group
+    let scan_multicast_id = nl80211_family
+        .first()
+        .ok_or_else(|| anyhow!("Missing \"nl80211\" family"))?
+        .iter()
+        .find_map(|attr| match attr {
+            GenlCtrlAttrs::McastGroups(mcast_groups) => Some(mcast_groups),
+            _ => None,
+        })
+        .ok_or_else(|| anyhow!("Missing McastGroup attribute"))?
+        .iter()
+        .find(|grp| grp.contains(&McastGrpAttrs::Name("scan".to_owned())))
+        .ok_or_else(|| anyhow!("Missing scan group"))?
+        .iter()
+        .find_map(|grp_attr| match grp_attr {
+            McastGrpAttrs::Id(id) => Some(*id),
+            _ => None,
+        })
+        .ok_or_else(|| anyhow!("No multicast id defined for scan group"))?;
+
+    Ok(scan_multicast_id)
 }