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