Avoid crashing after shutting down WiFi

esp-radio-rtos-driver/CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Renamed `Scheduler` to `SchedulerImplementation` and `scheduler_impl!` to `register_scheduler_implementation!` (#4559)
 - `current_task`, `task_create` and `schedule_task_deletion` functions now work with `ThreadPtr` instead of `*mut c_void` (#4559)
 - `schedule_task_deletion` now takes an `Option` (#4559)
+- `SemaphoreHandle` is now `Send` (#4761)
 
 ### Fixed
 

esp-radio-rtos-driver/src/semaphore.rs

@@ -446,6 +446,8 @@ impl Drop for SemaphoreHandle {
     }
 }
 
+unsafe impl Send for SemaphoreHandle {}
+
 #[cfg(feature = "ipc-implementations")]
 mod implementation {
     use alloc::boxed::Box;

esp-radio/CHANGELOG.md

@@ -40,6 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed a linker error (about missing symbols) when the `wifi` feature is selected but the code doesn't use it (#4513)
 - `Controller::stop_async()` now returns `WifiError::NotStarted` when the `Controller` has not been started (#4504)
 - ESP32-C2: Disable BLE controller before deinitializing the stack (#4606)
+- Fix a crash after shutting down WiFi (#4761)
 
 ### Removed
 

esp-radio/src/wifi/mod.rs

@@ -1025,6 +1025,20 @@ impl WifiDeviceMode {
         }
 
         if self.can_send() {
+            // even checking for !Uninitialized would be enough to not crash
+            match self {
+                WifiDeviceMode::Station => {
+                    if !matches!(station_state(), WifiStationState::Connected) {
+                        return None;
+                    }
+                }
+                WifiDeviceMode::AccessPoint => {
+                    if !matches!(access_point_state(), WifiAccessPointState::Started) {
+                        return None;
+                    }
+                }
+            }
+
             Some(WifiTxToken { mode: *self })
         } else {
             None
@@ -1436,19 +1450,32 @@ impl TxToken for WifiTxToken {
 // though in reality `esp_wifi_internal_tx` copies the buffer into its own
 // memory and does not modify
 pub(crate) fn esp_wifi_send_data(interface: wifi_interface_t, data: &mut [u8]) {
-    trace!("sending... {} bytes", data.len());
-    dump_packet_info(data);
+    // `esp_wifi_internal_tx` will crash if wifi is uninitialized or de-inited
+
+    state::locked(|| {
+        // even checking for !Uninitialized would be enough to not crash
+        if (interface == wifi_interface_t_WIFI_IF_STA
+            && !matches!(station_state(), WifiStationState::Connected))
+            || (interface == wifi_interface_t_WIFI_IF_AP
+                && !matches!(access_point_state(), WifiAccessPointState::Started))
+        {
+            return;
+        }
 
-    let len = data.len() as u16;
-    let ptr = data.as_mut_ptr().cast();
+        trace!("sending... {} bytes", data.len());
+        dump_packet_info(data);
 
-    let res = unsafe { esp_wifi_result!(esp_wifi_internal_tx(interface, ptr, len)) };
+        let len = data.len() as u16;
+        let ptr = data.as_mut_ptr().cast();
 
-    if res.is_err() {
-        decrement_inflight_counter();
-    } else {
-        trace!("esp_wifi_internal_tx ok");
-    }
+        let res = unsafe { esp_wifi_result!(esp_wifi_internal_tx(interface, ptr, len)) };
+
+        if res.is_err() {
+            decrement_inflight_counter();
+        } else {
+            trace!("esp_wifi_internal_tx ok");
+        }
+    })
 }
 
 fn dump_packet_info(_buffer: &mut [u8]) {
@@ -1960,16 +1987,18 @@ pub struct WifiController<'d> {
 
 impl Drop for WifiController<'_> {
     fn drop(&mut self) {
-        if let Err(e) = crate::wifi::wifi_deinit() {
-            warn!("Failed to cleanly deinit wifi: {:?}", e);
-        }
+        state::locked(|| {
+            if let Err(e) = crate::wifi::wifi_deinit() {
+                warn!("Failed to cleanly deinit wifi: {:?}", e);
+            }
 
-        set_access_point_state(WifiAccessPointState::Uninitialized);
-        set_station_state(WifiStationState::Uninitialized);
+            set_access_point_state(WifiAccessPointState::Uninitialized);
+            set_station_state(WifiStationState::Uninitialized);
 
-        esp_hal::rng::TrngSource::decrease_entropy_source_counter(unsafe {
-            esp_hal::Internal::conjure()
-        });
+            esp_hal::rng::TrngSource::decrease_entropy_source_counter(unsafe {
+                esp_hal::Internal::conjure()
+            });
+        })
     }
 }
 

esp-radio/src/wifi/state.rs

@@ -1,56 +1,53 @@
 use core::sync::atomic::Ordering;
 
-use private::{AtomicWifiAccessPointState, AtomicWifiStationState};
-pub(crate) use private::{WifiAccessPointState, WifiStationState};
+use esp_radio_rtos_driver::semaphore::{SemaphoreHandle, SemaphoreKind};
+use esp_sync::NonReentrantMutex;
+use portable_atomic_enum::atomic_enum;
 
 use super::WifiEvent;
 
-mod private {
-    use portable_atomic_enum::atomic_enum;
-
-    /// Wi-Fi interface for station state.
-    #[atomic_enum]
-    #[derive(PartialEq, Debug, Clone, Copy, Hash)]
-    #[cfg_attr(feature = "defmt", derive(defmt::Format))]
-    #[non_exhaustive]
-    pub enum WifiStationState {
-        /// Start initiated.
-        Starting,
-        /// Station started.
-        Started,
-        /// Connect initiated.
-        Connecting,
-        /// Station connected.
-        Connected,
-        /// Disconnect initiated.
-        Disconnecting,
-        /// Station disconnected.
-        Disconnected,
-        /// Stop initiated.
-        Stopping,
-        /// Station stopped
-        Stopped,
-        /// Uninitialized state.
-        Uninitialized,
-    }
+/// Wi-Fi interface for station state.
+#[atomic_enum]
+#[derive(PartialEq, Debug, Clone, Copy, Hash)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+#[non_exhaustive]
+pub(crate) enum WifiStationState {
+    /// Start initiated.
+    Starting,
+    /// Station started.
+    Started,
+    /// Connect initiated.
+    Connecting,
+    /// Station connected.
+    Connected,
+    /// Disconnect initiated.
+    Disconnecting,
+    /// Station disconnected.
+    Disconnected,
+    /// Stop initiated.
+    Stopping,
+    /// Station stopped
+    Stopped,
+    /// Uninitialized state.
+    Uninitialized,
+}
 
-    /// Wi-Fi interface for access point state.
-    #[atomic_enum]
-    #[derive(PartialEq, Debug, Clone, Copy, Hash)]
-    #[cfg_attr(feature = "defmt", derive(defmt::Format))]
-    #[non_exhaustive]
-    pub enum WifiAccessPointState {
-        /// Start initiated.
-        Starting,
-        /// Access point started.
-        Started,
-        /// Stop initiated.
-        Stopping,
-        /// Access point stopped.
-        Stopped,
-        /// Uninitialized state.
-        Uninitialized,
-    }
+/// Wi-Fi interface for access point state.
+#[atomic_enum]
+#[derive(PartialEq, Debug, Clone, Copy, Hash)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+#[non_exhaustive]
+pub(crate) enum WifiAccessPointState {
+    /// Start initiated.
+    Starting,
+    /// Access point started.
+    Started,
+    /// Stop initiated.
+    Stopping,
+    /// Access point stopped.
+    Stopped,
+    /// Uninitialized state.
+    Uninitialized,
 }
 
 impl From<WifiEvent> for WifiStationState {
@@ -75,9 +72,9 @@ impl From<WifiEvent> for WifiAccessPointState {
     }
 }
 
-pub(crate) static STATION_STATE: AtomicWifiStationState =
+static STATION_STATE: AtomicWifiStationState =
     AtomicWifiStationState::new(WifiStationState::Uninitialized);
-pub(crate) static ACCESS_POINT_STATE: AtomicWifiAccessPointState =
+static ACCESS_POINT_STATE: AtomicWifiAccessPointState =
     AtomicWifiAccessPointState::new(WifiAccessPointState::Uninitialized);
 
 /// Get the current state of the access point.
@@ -118,3 +115,21 @@ pub(crate) fn set_access_point_state(state: WifiAccessPointState) {
 pub(crate) fn set_station_state(state: WifiStationState) {
     STATION_STATE.store(state, Ordering::Relaxed)
 }
+
+pub(crate) fn locked<R>(f: impl FnOnce() -> R) -> R {
+    static LOCK: NonReentrantMutex<Option<SemaphoreHandle>> = NonReentrantMutex::new(None);
+
+    LOCK.with(|sem| {
+        sem.get_or_insert_with(|| SemaphoreHandle::new(SemaphoreKind::Mutex))
+            .take(None)
+    });
+
+    let res: R = f();
+
+    LOCK.with(|sem| {
+        sem.get_or_insert_with(|| SemaphoreHandle::new(SemaphoreKind::Mutex))
+            .give()
+    });
+
+    res
+}

qa-test/src/bin/embassy_wifi_drop.rs

@@ -0,0 +1,121 @@
+//! Check that trying to send after drop doesn't crash
+
+//% FEATURES: esp-radio esp-radio/wifi esp-radio/unstable esp-hal/unstable
+//% CHIPS: esp32 esp32c2 esp32c3 esp32c6 esp32s2 esp32s3
+
+#![no_std]
+#![no_main]
+
+use embassy_executor::Spawner;
+use esp_alloc as _;
+use esp_backtrace as _;
+use esp_hal::{
+    clock::CpuClock,
+    interrupt::software::SoftwareInterruptControl,
+    ram,
+    timer::timg::TimerGroup,
+};
+use esp_println::println;
+use esp_radio::wifi::{ModeConfig, sta::StationConfig};
+use smoltcp::phy::Device;
+
+esp_bootloader_esp_idf::esp_app_desc!();
+
+const SSID: &str = env!("SSID");
+const PASSWORD: &str = env!("PASSWORD");
+
+#[esp_rtos::main]
+async fn main(_spawner: Spawner) {
+    esp_println::logger::init_logger_from_env();
+    let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
+    let peripherals = esp_hal::init(config);
+
+    esp_alloc::heap_allocator!(size: 32 * 1024);
+    // add some more RAM
+    esp_alloc::heap_allocator!(#[ram(reclaimed)] size: 64 * 1024);
+
+    let timg0 = TimerGroup::new(peripherals.TIMG0);
+    let sw_int = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
+    esp_rtos::start(timg0.timer0, sw_int.software_interrupt0);
+
+    let mut wifi = peripherals.WIFI;
+    {
+        let (mut controller, interfaces) =
+            esp_radio::wifi::new(wifi.reborrow(), Default::default()).unwrap();
+
+        let mut wifi_interface = interfaces.station;
+
+        let token = wifi_interface.transmit(smoltcp::time::Instant::from_millis(
+            esp_hal::time::Instant::now()
+                .duration_since_epoch()
+                .as_millis() as i64,
+        ));
+        assert!(matches!(token, None));
+
+        let station_config = ModeConfig::Station(
+            StationConfig::default()
+                .with_ssid(SSID.into())
+                .with_password(PASSWORD.into()),
+        );
+        controller.set_config(&station_config).unwrap();
+        controller.start_async().await.unwrap();
+        controller.connect_async().await.unwrap();
+
+        let token = wifi_interface.transmit(smoltcp::time::Instant::from_millis(
+            esp_hal::time::Instant::now()
+                .duration_since_epoch()
+                .as_millis() as i64,
+        ));
+
+        println!("got token {}", token.is_some());
+
+        core::mem::drop(controller);
+
+        if let Some(token) = token {
+            token.consume_token(10, |tx| tx.fill(0));
+            println!("survived consumer_token");
+        } else {
+            println!("no token");
+        }
+
+        let token = wifi_interface.transmit(smoltcp::time::Instant::from_millis(
+            esp_hal::time::Instant::now()
+                .duration_since_epoch()
+                .as_millis() as i64,
+        ));
+        assert!(matches!(token, None));
+    }
+
+    {
+        let (mut controller, interfaces) =
+            esp_radio::wifi::new(wifi.reborrow(), Default::default()).unwrap();
+
+        let station_config = ModeConfig::Station(
+            StationConfig::default()
+                .with_ssid(SSID.into())
+                .with_password(PASSWORD.into()),
+        );
+        controller.set_config(&station_config).unwrap();
+        controller.start_async().await.unwrap();
+        controller.connect_async().await.unwrap();
+
+        let mut wifi_interface = interfaces.station;
+
+        let token = wifi_interface.transmit(smoltcp::time::Instant::from_millis(
+            esp_hal::time::Instant::now()
+                .duration_since_epoch()
+                .as_millis() as i64,
+        ));
+
+        println!("got token {}", token.is_some());
+
+        if let Some(token) = token {
+            token.consume_token(10, |tx| tx.fill(0));
+            println!("tx done");
+        } else {
+            println!("no token");
+        }
+    }
+
+    println!("Done");
+}