implement wifi event handling with data (#2453)

* implement event handling for apsta{,dis}connect and probe

add to wifi_access_point example

* hide internal `Option` to simplify api

* move handler call logic to event.rs

* update_handler_leak

* - Add comments

- match simpler api from `std::panic::update_hook`

- do not assume size_of in prelude

- make all events handleable

- box static instead of leak

* update changelog

* elide lifetime on default handler

* newtypes for all event types.

* add doc to newtypes

* fix previous doc example

* - `get_handler` -> `handler`
- pass critical section to event handlers
- comment on perf of Box<ZST>
- don't pass `prev` in `update_handler`, instead call previous handler first unconditionally

* pass cs to dispatch_event_handler

* don't print "unhandled event" for handled events.

esp-wifi/CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `serde` support through the `serde` feature (#2346)
 - Added `PowerSaveMode` and `set_power_saving` methods on `EspNowManager` & `WifiController` (#2446)
 - Added CSI support (#2422)
+- Enable setting event handlers for wifi events (#2453)
 
 ### Changed
 

esp-wifi/src/wifi/event.rs

@@ -0,0 +1,348 @@
+use alloc::boxed::Box;
+use core::cell::RefCell;
+
+use critical_section::Mutex;
+
+use super::WifiEvent;
+
+pub(crate) mod sealed {
+    use super::*;
+
+    pub trait Event {
+        /// Get the static reference to the handler for this event.
+        fn handler() -> &'static Mutex<RefCell<Option<Box<Handler<Self>>>>>;
+        /// # Safety
+        /// `ptr` must be a valid for casting to this event's inner event data.
+        unsafe fn from_raw_event_data(ptr: *mut crate::binary::c_types::c_void) -> Self;
+    }
+}
+/// The type of handlers of events.
+pub type Handler<T> = dyn FnMut(critical_section::CriticalSection<'_>, &T) + Sync + Send;
+
+fn default_handler<Event: 'static>() -> Box<Handler<Event>> {
+    fn drop_ref<T>(_: critical_section::CriticalSection<'_>, _: &T) {}
+    // perf: `drop_ref` is a ZST [function item](https://doc.rust-lang.org/reference/types/function-item.html)
+    // so this doesn't actually allocate.
+    Box::new(drop_ref)
+}
+
+/// Extension trait for setting handlers for an event.
+///
+/// Register a new event handler like:
+/// ```
+/// # use esp_wifi::wifi::event::{self, *};
+/// # fn new_handler(_: critical_section::CriticalSection<'_>, _: &ApStaconnected) {}
+/// event::ApStaconnected::update_handler(|_cs, event| {
+///     new_handler(event);
+/// })
+/// ```
+// Implemented like this instead of free functions because the syntax would be
+// ```
+// event::update_handler::<event::ApStaconnected, _>(...)
+// ```
+pub trait EventExt: sealed::Event + Sized + 'static {
+    /// Get the handler for this event, replacing it with the default handler.
+    fn take_handler() -> Box<Handler<Self>> {
+        critical_section::with(|cs| {
+            Self::handler()
+                .borrow_ref_mut(cs)
+                .take()
+                .unwrap_or_else(default_handler::<Self>)
+        })
+    }
+    /// Set the handler for this event, returning the old handler.
+    fn replace_handler<
+        F: FnMut(critical_section::CriticalSection<'_>, &Self) + Sync + Send + 'static,
+    >(
+        f: F,
+    ) -> Box<Handler<Self>> {
+        critical_section::with(|cs| {
+            Self::handler()
+                .borrow_ref_mut(cs)
+                .replace(Box::new(f))
+                .unwrap_or_else(default_handler::<Self>)
+        })
+    }
+    /// Atomic combination of [`take_handler`] and [`replace_handler`]. Use this
+    /// to add a new handler which runs after the previously registered
+    /// handlers.
+    fn update_handler<
+        F: FnMut(critical_section::CriticalSection<'_>, &Self) + Sync + Send + 'static,
+    >(
+        mut f: F,
+    ) {
+        critical_section::with(move |cs| {
+            let mut handler: Box<Handler<Self>> = Self::handler()
+                .borrow_ref_mut(cs)
+                .take()
+                .unwrap_or_else(default_handler::<Self>);
+            Self::handler().borrow_ref_mut(cs).replace(Box::new(
+                move |cs: critical_section::CriticalSection<'_>, event| {
+                    handler(cs, event);
+                    f(cs, event)
+                },
+            ));
+        });
+    }
+}
+impl<T: sealed::Event + 'static> EventExt for T {}
+
+macro_rules! impl_wifi_event {
+    // no data
+    ($newtype:ident) => {
+        /// See [`WifiEvent`].
+        #[derive(Copy, Clone)]
+        pub struct $newtype;
+        impl sealed::Event for $newtype {
+            unsafe fn from_raw_event_data(_: *mut crate::binary::c_types::c_void) -> Self {
+                Self
+            }
+            fn handler() -> &'static Mutex<RefCell<Option<Box<Handler<Self>>>>> {
+                static HANDLE: Mutex<RefCell<Option<Box<Handler<$newtype>>>>> =
+                    Mutex::new(RefCell::new(None));
+                &HANDLE
+            }
+        }
+    };
+    // data
+    ($newtype:ident, $data:ident) => {
+        pub use esp_wifi_sys::include::$data;
+        /// See [`WifiEvent`].
+        #[derive(Copy, Clone)]
+        pub struct $newtype(pub $data);
+        impl sealed::Event for $newtype {
+            unsafe fn from_raw_event_data(ptr: *mut crate::binary::c_types::c_void) -> Self {
+                Self(unsafe { *ptr.cast() })
+            }
+            fn handler() -> &'static Mutex<RefCell<Option<Box<Handler<Self>>>>> {
+                static HANDLE: Mutex<RefCell<Option<Box<Handler<$newtype>>>>> =
+                    Mutex::new(RefCell::new(None));
+                &HANDLE
+            }
+        }
+    };
+}
+
+impl_wifi_event!(WifiReady);
+impl_wifi_event!(ScanDone);
+impl_wifi_event!(StaStart);
+impl_wifi_event!(StaStop);
+impl_wifi_event!(StaConnected, wifi_event_sta_connected_t);
+impl_wifi_event!(StaDisconnected, wifi_event_sta_disconnected_t);
+impl_wifi_event!(StaAuthmodeChange, wifi_event_sta_authmode_change_t);
+impl_wifi_event!(StaWpsErSuccess, wifi_event_sta_wps_er_success_t);
+impl_wifi_event!(StaWpsErFailed);
+impl_wifi_event!(StaWpsErTimeout);
+impl_wifi_event!(StaWpsErPin, wifi_event_sta_wps_er_pin_t);
+impl_wifi_event!(StaWpsErPbcOverlap);
+impl_wifi_event!(ApStart);
+impl_wifi_event!(ApStop);
+impl_wifi_event!(ApStaconnected, wifi_event_ap_staconnected_t);
+impl_wifi_event!(ApStadisconnected, wifi_event_ap_stadisconnected_t);
+impl_wifi_event!(ApProbereqrecved, wifi_event_ap_probe_req_rx_t);
+impl_wifi_event!(FtmReport, wifi_event_ftm_report_t);
+impl_wifi_event!(StaBssRssiLow, wifi_event_bss_rssi_low_t);
+impl_wifi_event!(ActionTxStatus, wifi_event_action_tx_status_t);
+impl_wifi_event!(RocDone, wifi_event_roc_done_t);
+impl_wifi_event!(StaBeaconTimeout);
+impl_wifi_event!(ConnectionlessModuleWakeIntervalStart);
+impl_wifi_event!(ApWpsRgSuccess, wifi_event_ap_wps_rg_success_t);
+impl_wifi_event!(ApWpsRgFailed, wifi_event_ap_wps_rg_fail_reason_t);
+impl_wifi_event!(ApWpsRgTimeout);
+impl_wifi_event!(ApWpsRgPin, wifi_event_ap_wps_rg_pin_t);
+impl_wifi_event!(ApWpsRgPbcOverlap);
+impl_wifi_event!(ItwtSetup);
+impl_wifi_event!(ItwtTeardown);
+impl_wifi_event!(ItwtProbe);
+impl_wifi_event!(ItwtSuspend);
+impl_wifi_event!(TwtWakeup);
+impl_wifi_event!(BtwtSetup);
+impl_wifi_event!(BtwtTeardown);
+impl_wifi_event!(NanStarted);
+impl_wifi_event!(NanStopped);
+impl_wifi_event!(NanSvcMatch, wifi_event_nan_svc_match_t);
+impl_wifi_event!(NanReplied, wifi_event_nan_replied_t);
+impl_wifi_event!(NanReceive, wifi_event_nan_receive_t);
+impl_wifi_event!(NdpIndication, wifi_event_ndp_indication_t);
+impl_wifi_event!(NdpConfirm, wifi_event_ndp_confirm_t);
+impl_wifi_event!(NdpTerminated, wifi_event_ndp_terminated_t);
+impl_wifi_event!(HomeChannelChange, wifi_event_home_channel_change_t);
+impl_wifi_event!(StaNeighborRep, wifi_event_neighbor_report_t);
+
+/// Handle the given event using the registered event handlers.
+pub fn handle<Event: EventExt>(
+    cs: critical_section::CriticalSection<'_>,
+    event_data: &Event,
+) -> bool {
+    if let Some(handler) = &mut *Event::handler().borrow_ref_mut(cs) {
+        handler(cs, event_data);
+        true
+    } else {
+        false
+    }
+}
+
+/// Handle an event given the raw pointers.
+/// # Safety
+/// The pointer should be valid to cast to `Event`'s inner type (if it has one)
+pub(crate) unsafe fn handle_raw<Event: EventExt>(
+    cs: critical_section::CriticalSection<'_>,
+    event_data: *mut crate::binary::c_types::c_void,
+    event_data_size: usize,
+) -> bool {
+    debug_assert_eq!(
+        event_data_size,
+        core::mem::size_of::<Event>(),
+        "wrong size event data"
+    );
+    handle::<Event>(cs, unsafe { &Event::from_raw_event_data(event_data) })
+}
+
+/// Handle event regardless of its type.
+/// # Safety
+/// Arguments should be self-consistent.
+#[rustfmt::skip]
+pub(crate) unsafe fn dispatch_event_handler(
+    cs: critical_section::CriticalSection<'_>,
+    event: WifiEvent,
+    event_data: *mut crate::binary::c_types::c_void,
+    event_data_size: usize,
+) -> bool {
+    match event {
+        WifiEvent::WifiReady => {
+            handle_raw::<WifiReady>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ScanDone => {
+            handle_raw::<ScanDone>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaStart => {
+            handle_raw::<StaStart>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaStop => {
+            handle_raw::<StaStop>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaConnected => {
+            handle_raw::<StaConnected>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaDisconnected => {
+            handle_raw::<StaDisconnected>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaAuthmodeChange => {
+            handle_raw::<StaAuthmodeChange>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaWpsErSuccess => {
+            handle_raw::<StaWpsErSuccess>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaWpsErFailed => {
+            handle_raw::<StaWpsErFailed>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaWpsErTimeout => {
+            handle_raw::<StaWpsErTimeout>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaWpsErPin => {
+            handle_raw::<StaWpsErPin>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaWpsErPbcOverlap => {
+            handle_raw::<StaWpsErPbcOverlap>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApStart => {
+            handle_raw::<ApStart>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApStop => {
+            handle_raw::<ApStop>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApStaconnected => {
+            handle_raw::<ApStaconnected>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApStadisconnected => {
+            handle_raw::<ApStadisconnected>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApProbereqrecved => {
+            handle_raw::<ApProbereqrecved>(cs, event_data, event_data_size)
+        }
+        WifiEvent::FtmReport => {
+            handle_raw::<FtmReport>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaBssRssiLow => {
+            handle_raw::<StaBssRssiLow>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ActionTxStatus => {
+            handle_raw::<ActionTxStatus>(cs, event_data, event_data_size)
+        }
+        WifiEvent::RocDone => {
+            handle_raw::<RocDone>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaBeaconTimeout => {
+            handle_raw::<StaBeaconTimeout>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ConnectionlessModuleWakeIntervalStart => {
+            handle_raw::<ConnectionlessModuleWakeIntervalStart>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApWpsRgSuccess => {
+            handle_raw::<ApWpsRgSuccess>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApWpsRgFailed => {
+            handle_raw::<ApWpsRgFailed>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApWpsRgTimeout => {
+            handle_raw::<ApWpsRgTimeout>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApWpsRgPin => {
+            handle_raw::<ApWpsRgPin>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ApWpsRgPbcOverlap => {
+            handle_raw::<ApWpsRgPbcOverlap>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ItwtSetup => {
+            handle_raw::<ItwtSetup>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ItwtTeardown => {
+            handle_raw::<ItwtTeardown>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ItwtProbe => {
+            handle_raw::<ItwtProbe>(cs, event_data, event_data_size)
+        }
+        WifiEvent::ItwtSuspend => {
+            handle_raw::<ItwtSuspend>(cs, event_data, event_data_size)
+        }
+        WifiEvent::TwtWakeup => {
+            handle_raw::<TwtWakeup>(cs, event_data, event_data_size)
+        }
+        WifiEvent::BtwtSetup => {
+            handle_raw::<BtwtSetup>(cs, event_data, event_data_size)
+        }
+        WifiEvent::BtwtTeardown => {
+            handle_raw::<BtwtTeardown>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NanStarted => {
+            handle_raw::<NanStarted>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NanStopped => {
+            handle_raw::<NanStopped>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NanSvcMatch => {
+            handle_raw::<NanSvcMatch>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NanReplied => {
+            handle_raw::<NanReplied>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NanReceive => {
+            handle_raw::<NanReceive>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NdpIndication => {
+            handle_raw::<NdpIndication>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NdpConfirm => {
+            handle_raw::<NdpConfirm>(cs, event_data, event_data_size)
+        }
+        WifiEvent::NdpTerminated => {
+            handle_raw::<NdpTerminated>(cs, event_data, event_data_size)
+        }
+        WifiEvent::HomeChannelChange => {
+            handle_raw::<HomeChannelChange>(cs, event_data, event_data_size)
+        }
+        WifiEvent::StaNeighborRep => {
+            handle_raw::<StaNeighborRep>(cs, event_data, event_data_size)
+        }
+    }
+}

esp-wifi/src/wifi/mod.rs

@@ -1,5 +1,6 @@
 //! WiFi
 
+pub mod event;
 pub(crate) mod os_adapter;
 pub(crate) mod state;
 

esp-wifi/src/wifi/os_adapter.rs

@@ -866,9 +866,14 @@ pub unsafe extern "C" fn event_post(
 
     let event = unwrap!(WifiEvent::from_i32(event_id));
     trace!("EVENT: {:?}", event);
-    critical_section::with(|cs| WIFI_EVENTS.borrow_ref_mut(cs).insert(event));
 
-    super::state::update_state(event);
+    let mut handled = false;
+    critical_section::with(|cs| {
+        WIFI_EVENTS.borrow_ref_mut(cs).insert(event);
+        handled = super::event::dispatch_event_handler(cs, event, event_data, event_data_size);
+    });
+
+    super::state::update_state(event, handled);
 
     event.waker().wake();
 

esp-wifi/src/wifi/state.rs

@@ -47,7 +47,7 @@ pub fn get_sta_state() -> WifiState {
     STA_STATE.load(Ordering::Relaxed)
 }
 
-pub(crate) fn update_state(event: WifiEvent) {
+pub(crate) fn update_state(event: WifiEvent, handled: bool) {
     match event {
         WifiEvent::StaConnected
         | WifiEvent::StaDisconnected
@@ -58,7 +58,11 @@ pub(crate) fn update_state(event: WifiEvent) {
             AP_STATE.store(WifiState::from(event), Ordering::Relaxed)
         }
 
-        other => debug!("Unhandled event: {:?}", other),
+        other => {
+            if !handled {
+                debug!("Unhandled event: {:?}", other)
+            }
+        }
     }
 }
 

examples/src/bin/wifi_access_point.rs

@@ -28,6 +28,7 @@ use esp_println::{print, println};
 use esp_wifi::{
     init,
     wifi::{
+        event::{self, EventExt},
         utils::create_network_interface,
         AccessPointConfiguration,
         Configuration,
@@ -49,6 +50,24 @@ fn main() -> ! {
 
     let timg0 = TimerGroup::new(peripherals.TIMG0);
 
+    // Set event handlers for wifi before init to avoid missing any.
+    let mut connections = 0u32;
+    _ = event::ApStart::replace_handler(|_, _| esp_println::println!("ap start event"));
+    event::ApStaconnected::update_handler(move |_, event| {
+        connections += 1;
+        esp_println::println!("connected {}, mac: {:?}", connections, event.0.mac);
+    });
+    event::ApStaconnected::update_handler(|_, event| {
+        esp_println::println!("connected aid: {}", event.0.aid);
+    });
+    event::ApStadisconnected::update_handler(|_, event| {
+        esp_println::println!(
+            "disconnected mac: {:?}, reason: {:?}",
+            event.0.mac,
+            event.0.reason
+        );
+    });
+
     let mut rng = Rng::new(peripherals.RNG);
 
     let init = init(timg0.timer0, rng.clone(), peripherals.RADIO_CLK).unwrap();