From 7278eafc0c772489ebb97bb369e12b794132fd1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Thu, 30 Oct 2025 10:48:29 +0800 Subject: [PATCH] Clean up processing scan results --- esp-radio/src/wifi/mod.rs | 94 ++++++-------------------------- esp-radio/src/wifi/scan.rs | 107 +++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 78 deletions(-) create mode 100644 esp-radio/src/wifi/scan.rs diff --git a/esp-radio/src/wifi/mod.rs b/esp-radio/src/wifi/mod.rs index ae20d17a29..b2fbc8fdb3 100644 --- a/esp-radio/src/wifi/mod.rs +++ b/esp-radio/src/wifi/mod.rs @@ -5,16 +5,10 @@ pub mod event; mod internal; pub(crate) mod os_adapter; +mod scan; pub(crate) mod state; -use alloc::{collections::vec_deque::VecDeque, string::String}; -use core::{ - fmt::Debug, - marker::PhantomData, - mem::MaybeUninit, - ptr::addr_of, - task::Poll, - time::Duration, -}; +use alloc::{collections::vec_deque::VecDeque, string::String, vec::Vec}; +use core::{fmt::Debug, marker::PhantomData, ptr::addr_of, task::Poll, time::Duration}; use enumset::{EnumSet, EnumSetType}; use esp_config::esp_config_int; @@ -84,7 +78,10 @@ use crate::{ wifi_scan_channel_bitmap_t, }, }, - wifi::private::PacketBuffer, + wifi::{ + private::PacketBuffer, + scan::{FreeApListOnDrop, ScanResults}, + }, }; const MTU: usize = esp_config_int!(usize, "ESP_RADIO_CONFIG_WIFI_MTU"); @@ -1779,7 +1776,7 @@ pub(crate) fn wifi_start_scan( }; let mut ssid_buf = ssid.map(|m| { - let mut buf = alloc::vec::Vec::from_iter(m.bytes()); + let mut buf = Vec::from_iter(m.bytes()); buf.push(b'\0'); buf }); @@ -1978,33 +1975,6 @@ impl WifiDevice<'_> { } } -fn convert_ap_info(record: &include::wifi_ap_record_t) -> AccessPointInfo { - let str_len = record - .ssid - .iter() - .position(|&c| c == 0) - .unwrap_or(record.ssid.len()); - let ssid_ref = unsafe { core::str::from_utf8_unchecked(&record.ssid[..str_len]) }; - - let mut ssid = String::new(); - ssid.push_str(ssid_ref); - - AccessPointInfo { - ssid, - bssid: record.bssid, - channel: record.primary, - secondary_channel: match record.second { - include::wifi_second_chan_t_WIFI_SECOND_CHAN_NONE => SecondaryChannel::None, - include::wifi_second_chan_t_WIFI_SECOND_CHAN_ABOVE => SecondaryChannel::Above, - include::wifi_second_chan_t_WIFI_SECOND_CHAN_BELOW => SecondaryChannel::Below, - _ => panic!(), - }, - signal_strength: record.rssi, - auth_method: Some(AuthMethod::from_raw(record.authmode)), - country: Country::try_from_c(&record.country), - } -} - /// The radio metadata header of the received packet, which is the common header /// at the beginning of all RX callback buffers in promiscuous mode. #[cfg(not(esp32c6))] @@ -2564,21 +2534,6 @@ pub(crate) fn apply_power_saving(ps: PowerSaveMode) -> Result<(), WifiError> { Ok(()) } -struct FreeApListOnDrop; -impl FreeApListOnDrop { - pub fn defuse(self) { - core::mem::forget(self); - } -} - -impl Drop for FreeApListOnDrop { - fn drop(&mut self) { - unsafe { - include::esp_wifi_clear_ap_list(); - } - } -} - /// Represents the Wi-Fi controller and its associated interfaces. #[non_exhaustive] pub struct Interfaces<'d> { @@ -3049,33 +3004,17 @@ impl WifiController<'_> { pub fn scan_with_config( &mut self, config: ScanConfig<'_>, - ) -> Result, WifiError> { + ) -> Result, WifiError> { esp_wifi_result!(crate::wifi::wifi_start_scan(true, config))?; self.scan_results(config.max.unwrap_or(usize::MAX)) } - fn scan_results(&mut self, max: usize) -> Result, WifiError> { - let mut scanned = alloc::vec::Vec::::new(); - let mut bss_total: u16 = max as u16; - - // Prevents memory leak on error - let guard = FreeApListOnDrop; - - unsafe { esp_wifi_result!(include::esp_wifi_scan_get_ap_num(&mut bss_total))? }; - - guard.defuse(); - - let mut record: MaybeUninit = MaybeUninit::uninit(); - for _ in 0..usize::min(bss_total as usize, max) { - unsafe { esp_wifi_result!(include::esp_wifi_scan_get_ap_record(record.as_mut_ptr()))? }; - let record = unsafe { MaybeUninit::assume_init_ref(&record) }; - let ap_info = convert_ap_info(record); - scanned.push(ap_info); - } - - unsafe { esp_wifi_result!(include::esp_wifi_clear_ap_list())? }; + fn scan_results_iter(&mut self) -> Result, WifiError> { + ScanResults::new(self) + } - Ok(scanned) + fn scan_results(&mut self, max: usize) -> Result, WifiError> { + Ok(self.scan_results_iter()?.take(max).collect::>()) } /// Starts the Wi-Fi controller. @@ -3282,14 +3221,13 @@ impl WifiController<'_> { pub async fn scan_with_config_async( &mut self, config: ScanConfig<'_>, - ) -> Result, WifiError> { + ) -> Result, WifiError> { Self::clear_events(WifiEvent::ScanDone); esp_wifi_result!(wifi_start_scan(false, config))?; - // Prevents memory leak if `scan_n`'s future is dropped. + // Prevents memory leak if `scan_with_config_async`'s future is dropped. let guard = FreeApListOnDrop; WifiEventFuture::new(WifiEvent::ScanDone).await; - guard.defuse(); let result = self.scan_results(config.max.unwrap_or(usize::MAX))?; diff --git a/esp-radio/src/wifi/scan.rs b/esp-radio/src/wifi/scan.rs new file mode 100644 index 0000000000..43b657e190 --- /dev/null +++ b/esp-radio/src/wifi/scan.rs @@ -0,0 +1,107 @@ +use core::{marker::PhantomData, mem::MaybeUninit}; + +use crate::{ + esp_wifi_result, + sys::include, + wifi::{ + AccessPointInfo, + AuthMethod, + AuthMethodExt as _, + Country, + SecondaryChannel, + WifiController, + WifiError, + }, +}; + +pub struct ScanResults<'d> { + /// Number of APs to return + remaining: usize, + + /// Ensures the result list is free'd when this struct is dropped. + _drop_guard: FreeApListOnDrop, + + /// Hold a lifetime to ensure the scan list is freed before a new scan is started. + _marker: PhantomData<&'d mut ()>, +} + +impl<'d> ScanResults<'d> { + pub fn new(_controller: &'d mut WifiController<'_>) -> Result { + // Construct Self first. This ensures we'll free the result list even if `get_ap_num` + // returns an error. + let mut this = Self { + remaining: 0, + _drop_guard: FreeApListOnDrop, + _marker: PhantomData, + }; + + let mut bss_total = 0; + unsafe { esp_wifi_result!(include::esp_wifi_scan_get_ap_num(&mut bss_total))? }; + + this.remaining = bss_total as usize; + + Ok(this) + } +} + +impl Iterator for ScanResults<'_> { + type Item = AccessPointInfo; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + + self.remaining -= 1; + + let mut record: MaybeUninit = MaybeUninit::uninit(); + + // We could detect ESP_FAIL to see if we've exhausted the list, but we know the number of + // results. Reading the number of results also ensures we're in the correct state, so + // unwrapping here should never fail. + unwrap!(unsafe { + esp_wifi_result!(include::esp_wifi_scan_get_ap_record(record.as_mut_ptr())) + }); + + Some(convert_ap_info(unsafe { record.assume_init_ref() })) + } +} + +fn convert_ap_info(record: &include::wifi_ap_record_t) -> AccessPointInfo { + let str_len = record + .ssid + .iter() + .position(|&c| c == 0) + .unwrap_or(record.ssid.len()); + let ssid = alloc::string::String::from_utf8_lossy(&record.ssid[..str_len]).into_owned(); + + AccessPointInfo { + ssid, + bssid: record.bssid, + channel: record.primary, + secondary_channel: match record.second { + include::wifi_second_chan_t_WIFI_SECOND_CHAN_NONE => SecondaryChannel::None, + include::wifi_second_chan_t_WIFI_SECOND_CHAN_ABOVE => SecondaryChannel::Above, + include::wifi_second_chan_t_WIFI_SECOND_CHAN_BELOW => SecondaryChannel::Below, + _ => panic!(), + }, + signal_strength: record.rssi, + auth_method: Some(AuthMethod::from_raw(record.authmode)), + country: Country::try_from_c(&record.country), + } +} + +pub struct FreeApListOnDrop; +impl FreeApListOnDrop { + pub fn defuse(self) { + core::mem::forget(self); + } +} + +impl Drop for FreeApListOnDrop { + fn drop(&mut self) { + unsafe { + include::esp_wifi_clear_ap_list(); + } + } +}