Skip to main content

embedded_wg_display/wifi/
mod.rs

1//! WiFi station mode driver.
2//!
3//! Call [`Wifi::start_station`] once from `main()` to configure the radio,
4//! spawn the connection and network tasks, then await
5//! [`Wifi::wait_for_connection`] before starting any network-dependent tasks.
6use crate::util::globals;
7use alloc::string::ToString;
8use core::net::Ipv4Addr;
9use core::str::FromStr;
10use defmt::{info, warn};
11use embassy_executor::Spawner;
12use embassy_net::{Ipv4Cidr, Runner, Stack, StackResources, StaticConfigV4};
13use embassy_time::{Duration, Timer};
14use esp_alloc as _;
15use esp_hal::rng::Rng;
16use esp_radio::wifi::{
17    AccessPointConfig, AuthMethod, ClientConfig, ModeConfig, WifiController, WifiDevice, WifiEvent,
18    WifiStaState,
19};
20
21const AP_GATEWAY_IP: &str = "192.168.2.1";
22const AP_SSID: &str = "WG Display AP";
23const AP_PASSWORD: &str = "wgdisplay123";
24
25#[allow(
26    clippy::large_stack_frames,
27    reason = "wifi module is allowd to have large stack frames"
28)]
29// 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
30macro_rules! mk_static {
31    ($t:ty,$val:expr) => {{
32        static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
33        #[deny(unused_attributes)]
34        let x = STATIC_CELL.uninit().write($val);
35        x
36    }};
37}
38
39/// Struct to hold relevenat states for other modules (http server/client)
40pub struct Wifi {
41    stack: Stack<'static>,
42    tls_seed: u64,
43}
44
45impl Wifi {
46    /// Initialises the WiFi radio.
47    ///
48    /// With `use_ap = false`, the device will attempt to connect to the specified WiFi network as a station.
49    /// If it fails to connect after a number of attempts, it will switch to AP mode and reboot to allow the user to connect and configure WiFi credentials via UI.
50    ///
51    /// Call [`wait_for_connection`](Self::wait_for_connection) to wait until connection is fully established is in station mode.
52    pub fn start_station(
53        wifi_peripheral: esp_hal::peripherals::WIFI<'static>,
54        spawner: &Spawner,
55        ssid: alloc::string::String,
56        password: alloc::string::String,
57        use_ap: bool,
58    ) -> Self {
59        // init radio wifi
60        let radio_init = &*mk_static!(
61            esp_radio::Controller<'static>,
62            esp_radio::init().expect("Failed to initialize Wi-Fi/BLE controller")
63        );
64        let (mut controller, interfaces) =
65            esp_radio::wifi::new(radio_init, wifi_peripheral, Default::default()).unwrap();
66
67        let rng = Rng::new();
68        let net_seed = rng.random() as u64 | ((rng.random() as u64) << 32);
69        let tls_seed = rng.random() as u64 | ((rng.random() as u64) << 32);
70
71        let (wifi_interface, mode_config, net_config, log_msg) = if use_ap {
72            // Access Point mode
73            (
74                interfaces.ap,
75                ModeConfig::AccessPoint(
76                    AccessPointConfig::default()
77                        .with_ssid(AP_SSID.to_string())
78                        .with_auth_method(AuthMethod::Wpa2Personal)
79                        .with_password(AP_PASSWORD.to_string()),
80                ),
81                embassy_net::Config::ipv4_static(StaticConfigV4 {
82                    address: Ipv4Cidr::new(Ipv4Addr::new(192, 168, 2, 1), 24),
83                    gateway: Some(Ipv4Addr::new(192, 168, 2, 1)),
84                    dns_servers: Default::default(),
85                }),
86                "WiFi configured in AP mode",
87            )
88        } else {
89            // Station mode
90            (
91                interfaces.sta,
92                ModeConfig::Client(
93                    ClientConfig::default()
94                        .with_ssid(ssid.clone())
95                        .with_password(password),
96                ),
97                embassy_net::Config::dhcpv4(Default::default()),
98                "WiFi configured in station mode",
99            )
100        };
101
102        // Init network stack
103        let (stack, runner) = embassy_net::new(
104            wifi_interface,
105            net_config,
106            mk_static!(StackResources<8>, StackResources::<8>::new()),
107            net_seed,
108        );
109
110        // Configure WiFi with selected mode
111        controller.set_config(&mode_config).unwrap();
112        info!("{}", log_msg);
113
114        // Spawn wifi connection tasks
115        if use_ap {
116            spawner.spawn(connection_ap(controller)).ok();
117            spawner.spawn(run_dhcp(stack, AP_GATEWAY_IP)).ok();
118        } else {
119            spawner.spawn(connection(controller)).ok();
120        }
121        spawner.spawn(net_task(runner)).ok();
122
123        Self { stack, tls_seed }
124    }
125
126    /// Returns the embassy-net stack. Pass to [`EspHttpClient::new`](crate::http_client::EspHttpClient::new).
127    pub fn stack(&self) -> Stack<'static> {
128        self.stack
129    }
130
131    /// Returns the TLS seed. Pass to [`EspHttpClient::new`](crate::http_client::EspHttpClient::new).
132    pub fn tls_seed(&self) -> u64 {
133        self.tls_seed
134    }
135
136    /// Waits asynchronously until the WiFi link is up.
137    pub async fn wait_for_connection(&self) -> Ipv4Cidr {
138        info!("Waiting for link to be up");
139        loop {
140            if self.stack.is_link_up() {
141                break;
142            }
143            Timer::after(Duration::from_millis(500)).await;
144        }
145
146        info!("Waiting to get IP address...");
147        loop {
148            if let Some(config) = self.stack.config_v4() {
149                info!("Got IP: {}", config.address);
150                return config.address;
151            }
152            Timer::after(Duration::from_millis(500)).await;
153        }
154    }
155}
156
157// Task to handle Wifi connection and reconnection
158#[embassy_executor::task]
159async fn connection(mut controller: WifiController<'static>) {
160    info!("start connection task");
161
162    loop {
163        if esp_radio::wifi::sta_state() == WifiStaState::Connected {
164            // wait until we're no longer connected
165            controller.wait_for_event(WifiEvent::StaDisconnected).await;
166            Timer::after(Duration::from_millis(5000)).await
167        }
168        if !matches!(controller.is_started(), Ok(true)) {
169            info!("Starting wifi");
170            controller.start_async().await.unwrap();
171            info!("Wifi started!");
172        }
173        info!("About to connect...");
174
175        globals::console_println("Connecting to WiFi").await;
176
177        match controller.connect_async().await {
178            Ok(_) => {
179                info!("Wifi connected!");
180            }
181            Err(e) => {
182                warn!("Failed to connect to wifi: {:?}", e);
183
184                Timer::after(Duration::from_millis(5000)).await
185            }
186        }
187    }
188}
189
190#[embassy_executor::task]
191async fn run_dhcp(stack: Stack<'static>, gw_ip_addr: &'static str) {
192    use core::net::{Ipv4Addr, SocketAddrV4};
193
194    use edge_dhcp::{
195        io::{self, DEFAULT_SERVER_PORT},
196        server::{Server, ServerOptions},
197    };
198    use edge_nal::UdpBind;
199    use edge_nal_embassy::{Udp, UdpBuffers};
200
201    let ip = Ipv4Addr::from_str(gw_ip_addr).expect("dhcp task failed to parse gw ip");
202
203    let mut buf = [0u8; 1500];
204
205    let mut gw_buf = [Ipv4Addr::UNSPECIFIED];
206
207    let buffers = UdpBuffers::<3, 1024, 1024, 10>::new();
208    let unbound_socket = Udp::new(stack, &buffers);
209    let mut bound_socket = unbound_socket
210        .bind(core::net::SocketAddr::V4(SocketAddrV4::new(
211            Ipv4Addr::UNSPECIFIED,
212            DEFAULT_SERVER_PORT,
213        )))
214        .await
215        .unwrap();
216
217    loop {
218        _ = io::server::run(
219            &mut Server::<_, 64>::new_with_et(ip),
220            &ServerOptions::new(ip, Some(&mut gw_buf)),
221            &mut bound_socket,
222            &mut buf,
223        )
224        .await
225        .inspect_err(|_| warn!("DHCP server error"));
226        Timer::after(Duration::from_millis(500)).await;
227    }
228}
229
230#[embassy_executor::task]
231async fn connection_ap(mut controller: WifiController<'static>) {
232    info!("start AP connection task");
233
234    if !matches!(controller.is_started(), Ok(true)) {
235        info!("Starting WiFi AP");
236        controller.start_async().await.unwrap();
237        info!("WiFi AP started");
238    }
239
240    loop {
241        let events = controller
242            .wait_for_events(
243                WifiEvent::ApStaConnected | WifiEvent::ApStaDisconnected,
244                true,
245            )
246            .await;
247
248        if events.contains(WifiEvent::ApStaConnected) {
249            info!("Station connected to AP");
250        }
251
252        if events.contains(WifiEvent::ApStaDisconnected) {
253            info!("Station disconnected from AP");
254        }
255
256        Timer::after(Duration::from_millis(200)).await;
257    }
258}
259
260#[embassy_executor::task]
261async fn net_task(mut runner: Runner<'static, WifiDevice<'static>>) {
262    runner.run().await
263}