Shutdown the WiFi subsystem gracefully dynamically

[renkenono]

Bug description

Discussion initially documented in https://github.com/esp-rs/esp-hal/discussions/4631

To Reproduce

1. See https://github.com/renkenono/esp-hal/blob/fix/wifi-deinit-flow/examples/wifi/embassy_sntp/src/main.rs.
   1. Relying purely on Drop trait, see https://github.com/renkenono/esp-hal/commit/7b6baefc54ebcde29dcb4c8ee9672ee15cd7ee29.
   2. Gracefully shutting down the controller by disconnecting then stopping it, see https://github.com/renkenono/esp-hal/commit/79a032cea40cac0474034795907ea72ab05e7bd7.

Expected behavior

• No crashes.
• WiFi HW is powered off completely.
• WiFi dynamically allocated memory is freed.

Environment

• Target device: [e.g. ESP32-S3] ESP32C6

Chip type:         esp32c6 (revision v0.1)
Crystal frequency: 40 MHz
Flash size:        4MB
Features:          WiFi 6, BT 5
MAC address:       f0:f5:bd:2b:d9:34

Security Information:
=====================
Flags: 0x00000000 (0)
Key Purposes: [0, 0, 0, 0, 0, 0, 12]
Chip ID: 13
API Version: 0
Secure Boot: Disabled
Flash Encryption: Disabled
SPI Boot Crypt Count (SPI_BOOT_CRYPT_CNT): 0x0

• Crate name and version: [e.g. esp-hal 0.20.0] esp-hal v1.0.0

[bugadani]

Thanks for the reproducers. I've tested the Drop one on esp-hal 1.0, and updated it (source below) to current main, both fail the way you've described.

//! Embassy SNTP example
//!
//!
//! Set SSID and PASSWORD env variable before running this example.
//!
//! This gets an ip address via DHCP then performs an SNTP request to update the RTC time with the
//! response. The RTC time is then compared with the received data parsed with jiff.
//! You can change the timezone to your local timezone.

#![no_std]
#![no_main]

use core::net::{IpAddr, SocketAddr};

use embassy_executor::Spawner;
use embassy_net::{
    Runner,
    Stack,
    StackResources,
    dns::DnsQueryType,
    udp::{PacketMetadata, UdpSocket},
};
use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, channel::Channel};
use embassy_time::{Duration, Timer};
use esp_alloc as _;
use esp_backtrace as _;
use esp_hal::{
    clock::CpuClock,
    interrupt::software::SoftwareInterruptControl,
    rng::Rng,
    rtc_cntl::Rtc,
    timer::timg::TimerGroup,
};
use esp_println::println;
use esp_radio::wifi::{
    ModeConfig,
    WifiController,
    WifiDevice,
    WifiEvent,
    WifiStationState,
    scan::ScanConfig,
    sta::StationConfig,
};
use log::{error, info};
use sntpc::{NtpContext, NtpTimestampGenerator, get_time};

esp_bootloader_esp_idf::esp_app_desc!();

// When you are okay with using a nightly compiler it's better to use https://docs.rs/static_cell/2.1.0/static_cell/macro.make_static.html
macro_rules! mk_static {
    ($t:ty,$val:expr) => {{
        static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
        #[deny(unused_attributes)]
        let x = STATIC_CELL.uninit().write(($val));
        x
    }};
}

const SSID: &str = env!("SSID");
const PASSWORD: &str = env!("PASSWORD");
const TIMEZONE: jiff::tz::TimeZone = jiff::tz::get!("UTC");
const NTP_SERVER: &str = "pool.ntp.org";

/// Microseconds in a second
const USEC_IN_SEC: u64 = 1_000_000;

static SNTP_CHANNEL: Channel<CriticalSectionRawMutex, (), 1> = Channel::new();
static SNTP_EVENT_CHANNEL: Channel<CriticalSectionRawMutex, (), 1> = Channel::new();
static CONNECTION_CHANNEL: Channel<CriticalSectionRawMutex, (), 1> = Channel::new();
static CONNECTION_EVENT_CHANNEL: Channel<CriticalSectionRawMutex, (), 1> = Channel::new();

#[derive(Clone, Copy)]
struct Timestamp<'a> {
    rtc: &'a Rtc<'a>,
    current_time_us: u64,
}

impl NtpTimestampGenerator for Timestamp<'_> {
    fn init(&mut self) {
        self.current_time_us = self.rtc.current_time_us();
    }

    fn timestamp_sec(&self) -> u64 {
        self.current_time_us / 1_000_000
    }

    fn timestamp_subsec_micros(&self) -> u32 {
        (self.current_time_us % 1_000_000) as u32
    }
}

#[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);
    let rtc = Rtc::new(peripherals.LPWR);

    esp_alloc::heap_allocator!(size: 72 * 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 (controller, interfaces) =
        esp_radio::wifi::new(peripherals.WIFI, Default::default()).unwrap();

    let wifi_interface = interfaces.station;

    let config = embassy_net::Config::dhcpv4(Default::default());

    let rng = Rng::new();
    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    // Init network stack
    let (stack, runner) = embassy_net::new(
        wifi_interface,
        config,
        mk_static!(StackResources<3>, StackResources::<3>::new()),
        seed,
    );

    spawner.must_spawn(sntp(stack, rtc));
    spawner.must_spawn(connection(controller));
    spawner.must_spawn(net_task(runner));

    let _ = CONNECTION_EVENT_CHANNEL.receive().await;
    println!("connection is up");
    Timer::after_secs(5).await;

    SNTP_CHANNEL.send(()).await;
    let _ = SNTP_EVENT_CHANNEL.receive().await;
    println!("sntp task done");

    CONNECTION_CHANNEL.send(()).await;
    println!("shutting down wifi subsystem");

    // Busy working...
    Timer::after_secs(600).await;

    loop {}
}

#[embassy_executor::task]
async fn sntp(stack: Stack<'static>, rtc: Rtc<'static>) {
    let mut rx_meta = [PacketMetadata::EMPTY; 16];
    let mut rx_buffer = [0; 4096];
    let mut tx_meta = [PacketMetadata::EMPTY; 16];
    let mut tx_buffer = [0; 4096];

    stack.wait_link_up().await;
    println!("Waiting to get IP address...");

    stack.wait_config_up().await;
    if let Some(config) = stack.config_v4() {
        println!("Got IP: {}", config.address);
    }

    let ntp_addrs = stack.dns_query(NTP_SERVER, DnsQueryType::A).await.unwrap();

    if ntp_addrs.is_empty() {
        panic!("Failed to resolve DNS. Empty result");
    }

    let mut socket = UdpSocket::new(
        stack,
        &mut rx_meta,
        &mut rx_buffer,
        &mut tx_meta,
        &mut tx_buffer,
    );

    socket.bind(123).unwrap();

    // Display initial Rtc time before synchronization
    let now = jiff::Timestamp::from_microsecond(rtc.current_time_us() as i64).unwrap();
    info!("Rtc: {now}");

    let addr: IpAddr = ntp_addrs[0].into();
    let result = get_time(
        SocketAddr::from((addr, 123)),
        &socket,
        NtpContext::new(Timestamp {
            rtc: &rtc,
            current_time_us: 0,
        }),
    )
    .await;

    match result {
        Ok(time) => {
            // Set time immediately after receiving to reduce time offset.
            rtc.set_current_time_us(
                (time.sec() as u64 * USEC_IN_SEC)
                    + ((time.sec_fraction() as u64 * USEC_IN_SEC) >> 32),
            );

            // Compare RTC to parsed time
            info!(
                "Response: {:?}\nTime: {}\nRtc : {}",
                time,
                // Create a Jiff Timestamp from seconds and nanoseconds
                jiff::Timestamp::from_second(time.sec() as i64)
                    .unwrap()
                    .checked_add(
                        jiff::Span::new()
                            .nanoseconds((time.seconds_fraction as i64 * 1_000_000_000) >> 32),
                    )
                    .unwrap()
                    .to_zoned(TIMEZONE),
                jiff::Timestamp::from_microsecond(rtc.current_time_us() as i64)
                    .unwrap()
                    .to_zoned(TIMEZONE)
            );
        }
        Err(e) => {
            error!("Error getting time: {e:?}");
        }
    }

    Timer::after(Duration::from_secs(1)).await;
    SNTP_EVENT_CHANNEL.send(()).await;
}

#[embassy_executor::task]
async fn connection(mut controller: WifiController<'static>) {
    println!("start connection task: {} - {}", SSID, PASSWORD);
    println!("Device capabilities: {:?}", controller.capabilities());
    let client_config = ModeConfig::Station(
        StationConfig::default()
            .with_ssid(SSID.into())
            .with_password(PASSWORD.into()),
    );
    controller.set_config(&client_config).unwrap();
    println!("Starting wifi");
    controller.start_async().await.unwrap();
    println!("Wifi started!");

    println!("Scan");
    let scan_config = ScanConfig::default().with_max(10).with_ssid(SSID);
    let result = controller
        .scan_with_config_async(scan_config)
        .await
        .unwrap();
    for ap in result {
        println!("{:?}", ap);
    }

    println!("About to connect...");

    match controller.connect_async().await {
        Ok(_) => println!("Wifi connected!"),
        Err(e) => {
            println!("Failed to connect to wifi: {e:?}");
            Timer::after(Duration::from_millis(5000)).await
        }
    }

    CONNECTION_EVENT_CHANNEL.send(()).await;

    let _ = CONNECTION_CHANNEL.receive().await;

    println!("connection task end");
}

#[embassy_executor::task]
async fn net_task(mut runner: Runner<'static, WifiDevice<'static>>) {
    runner.run().await
}

The crash happens because we don't check the state of the wifi stack in esp_wifi_send_data, but unconditionally call esp_wifi_internal_tx. We shouldn't provide a TxToken if the stack is not active.