From 391b8c7f41b8bf1933a545a2f434e71e3ce3cf06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1niel=20Buga?= Date: Mon, 9 Oct 2023 10:30:29 +0200 Subject: [PATCH] Move socket collection into the ap/sta implementations --- src/board/wifi/ap.rs | 18 +++----- src/board/wifi/mod.rs | 96 ++++++++++++++++++++++++++----------------- src/board/wifi/sta.rs | 20 ++++----- 3 files changed, 72 insertions(+), 62 deletions(-) diff --git a/src/board/wifi/ap.rs b/src/board/wifi/ap.rs index 7719cde7..f8b21908 100644 --- a/src/board/wifi/ap.rs +++ b/src/board/wifi/ap.rs @@ -5,13 +5,13 @@ use gui::widgets::wifi::WifiState; use crate::{ board::{ hal::{radio::Wifi, Rng}, - wifi::{net_task, STACK_SOCKET_COUNT}, + wifi::{net_task, StackWrapper}, }, task_control::{TaskControlToken, TaskController}, }; use embassy_executor::Spawner; use embassy_futures::join::join; -use embassy_net::{Config, Stack, StackResources}; +use embassy_net::{Config, Stack}; use embedded_svc::wifi::{AccessPointConfiguration, Configuration, Wifi as _}; use esp_wifi::{ wifi::{WifiController, WifiDevice, WifiEvent, WifiMode, WifiState as WifiStackState}, @@ -36,7 +36,7 @@ impl From for WifiState { #[derive(Clone)] pub struct Ap { - stack: Rc>>, + stack: Rc, client_count: Rc, } @@ -65,7 +65,7 @@ impl Ap { pub(super) struct ApState { init: EspWifiInitialization, controller: Option>>, - stack: Rc>>, + stack: Rc, connection_task_control: Option>, net_task_control: TaskController, client_count: Rc, @@ -76,23 +76,17 @@ impl ApState { init: EspWifiInitialization, config: Config, wifi: &'static mut Wifi, - resources: &'static mut StackResources, - mut rng: Rng, + rng: Rng, ) -> Self { info!("Configuring AP"); let (wifi_interface, controller) = unwrap!(esp_wifi::wifi::new_with_mode(&init, wifi, WifiMode::Ap)); - let lower = rng.random() as u64; - let upper = rng.random() as u64; - - let random_seed = upper << 32 | lower; - Self { init, controller: Some(Box::new(controller)), - stack: Rc::new(Stack::new(wifi_interface, config, resources, random_seed)), + stack: Rc::new(StackWrapper::new(wifi_interface, config, rng)), connection_task_control: None, net_task_control: TaskController::new(), client_count: Rc::new(AtomicU32::new(0)), diff --git a/src/board/wifi/mod.rs b/src/board/wifi/mod.rs index 508058c0..9ff9bce3 100644 --- a/src/board/wifi/mod.rs +++ b/src/board/wifi/mod.rs @@ -1,6 +1,8 @@ use core::{ hint::unreachable_unchecked, mem::{self, MaybeUninit}, + ops::Deref, + ptr::NonNull, }; use crate::{ @@ -29,11 +31,50 @@ pub unsafe fn as_static_mut(what: &mut T) -> &'static mut T { mem::transmute(what) } +const STACK_SOCKET_COUNT: usize = 3; + +struct StackWrapper( + NonNull>, + Stack>, +); + +impl StackWrapper { + fn new(wifi_interface: WifiDevice<'static>, config: Config, mut rng: Rng) -> StackWrapper { + let lower = rng.random() as u64; + let upper = rng.random() as u64; + + let random_seed = upper << 32 | lower; + + let resources = Box::new(StackResources::new()); + let resources_ref = Box::leak(resources); + + Self( + NonNull::from(&mut *resources_ref), + Stack::new(wifi_interface, config, resources_ref, random_seed), + ) + } +} + +impl Deref for StackWrapper { + type Target = Stack>; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + &self.1 + } +} + +impl Drop for StackWrapper { + fn drop(&mut self) { + unsafe { + _ = Box::from_raw(self.0.as_mut()); + } + } +} + pub mod ap; pub mod sta; -const STACK_SOCKET_COUNT: usize = 3; - #[derive(Default, PartialEq)] pub enum GenericConnectionState { Sta(sta::ConnectionState), @@ -72,14 +113,8 @@ impl WifiHandle { enum WifiDriverState { Uninitialized(WifiInitResources), Initialized(EspWifiInitialization), - Ap( - MaybeUninit, - Box>, - ), - Sta( - MaybeUninit, - Box>, - ), + Ap(MaybeUninit), + Sta(MaybeUninit), } impl WifiDriverState { @@ -105,7 +140,7 @@ impl WifiDriverState { async fn uninit_mode(&mut self) { match self { - WifiDriverState::Sta(sta, _) => { + WifiDriverState::Sta(sta) => { { let sta = unsafe { // Preinit is only called immediately before initialization, which means we @@ -123,7 +158,7 @@ impl WifiDriverState { }); } - WifiDriverState::Ap(ap, _) => { + WifiDriverState::Ap(ap) => { { let ap = unsafe { // Preinit is only called immediately before initialization, which means we @@ -147,25 +182,17 @@ impl WifiDriverState { fn replace_with(&mut self, mode: WifiMode) -> Self { match mode { - WifiMode::Ap => mem::replace( - self, - Self::Ap(MaybeUninit::uninit(), Box::new(StackResources::new())), - ), - WifiMode::Sta => mem::replace( - self, - Self::Sta(MaybeUninit::uninit(), Box::new(StackResources::new())), - ), + WifiMode::Ap => mem::replace(self, Self::Ap(MaybeUninit::uninit())), + WifiMode::Sta => mem::replace(self, Self::Sta(MaybeUninit::uninit())), } } fn handle(&self) -> Option { match self { - WifiDriverState::Sta(sta, _) => unsafe { + WifiDriverState::Sta(sta) => unsafe { sta.assume_init_ref().handle().map(WifiHandle::Sta) }, - WifiDriverState::Ap(ap, _) => unsafe { - ap.assume_init_ref().handle().map(WifiHandle::Ap) - }, + WifiDriverState::Ap(ap) => unsafe { ap.assume_init_ref().handle().map(WifiHandle::Ap) }, _ => None, } } @@ -203,8 +230,8 @@ impl WifiDriver { pub fn wifi_mode(&self) -> Option { match self.state { - WifiDriverState::Ap(_, _) => Some(WifiMode::Ap), - WifiDriverState::Sta(_, _) => Some(WifiMode::Sta), + WifiDriverState::Ap(_) => Some(WifiMode::Ap), + WifiDriverState::Sta(_) => Some(WifiMode::Sta), _ => None, } } @@ -229,14 +256,13 @@ impl WifiDriver { // Init AP mode match &mut self.state { - WifiDriverState::Ap(ap, resources) => { + WifiDriverState::Ap(ap) => { // Initialize the memory if we need to if let Some(init) = init { ap.write(ApState::init( init, config, unsafe { as_static_mut(&mut self.wifi) }, - unsafe { as_static_mut(resources) }, self.rng, )); } @@ -247,7 +273,7 @@ impl WifiDriver { WifiDriverState::Uninitialized { .. } | WifiDriverState::Initialized { .. } - | WifiDriverState::Sta(_, _) => { + | WifiDriverState::Sta(_) => { unreachable!() } } @@ -259,14 +285,13 @@ impl WifiDriver { // Init STA mode match &mut self.state { - WifiDriverState::Sta(sta, resources) => { + WifiDriverState::Sta(sta) => { // Initialize the memory if we need to if let Some(init) = init { sta.write(StaState::init( init, config, unsafe { as_static_mut(&mut self.wifi) }, - unsafe { as_static_mut(resources) }, self.rng, )); } @@ -277,7 +302,7 @@ impl WifiDriver { WifiDriverState::Uninitialized { .. } | WifiDriverState::Initialized { .. } - | WifiDriverState::Ap(_, _) => { + | WifiDriverState::Ap(_) => { unreachable!() } } @@ -288,7 +313,7 @@ impl WifiDriver { } pub fn ap_running(&self) -> bool { - if let WifiDriverState::Ap(ap, _) = &self.state { + if let WifiDriverState::Ap(ap) = &self.state { let ap = unsafe { ap.assume_init_ref() }; ap.is_running() } else { @@ -308,9 +333,6 @@ impl WifiDriver { } #[cardio::task] -pub async fn net_task( - stack: Rc>>, - mut task_control: TaskControlToken, -) { +async fn net_task(stack: Rc, mut task_control: TaskControlToken) { task_control.run_cancellable(|_| stack.run()).await; } diff --git a/src/board/wifi/sta.rs b/src/board/wifi/sta.rs index e6fe223b..7c9c1d9a 100644 --- a/src/board/wifi/sta.rs +++ b/src/board/wifi/sta.rs @@ -4,7 +4,7 @@ use crate::{ board::{ hal::{radio::Wifi, Rng}, initialized::Board, - wifi::{net_task, STACK_SOCKET_COUNT}, + wifi::{net_task, StackWrapper}, }, states::display_message, task_control::{TaskControlToken, TaskController}, @@ -18,7 +18,7 @@ use embassy_futures::{ join::join, select::{select, Either}, }; -use embassy_net::{dns::DnsSocket, Config, Stack, StackResources}; +use embassy_net::{dns::DnsSocket, Config}; use embassy_sync::{ blocking_mutex::raw::NoopRawMutex, mutex::{Mutex, MutexGuard}, @@ -124,7 +124,7 @@ impl From for ConnectionState { #[derive(Clone)] pub struct Sta { - stack: Rc>>, + stack: Rc, networks: Shared>, known_networks: Shared>, state: Rc, @@ -267,7 +267,7 @@ impl<'a> HttpsClientResources<'a> { pub(super) struct StaState { init: EspWifiInitialization, controller: Option>>, - stack: Rc>>, + stack: Rc, networks: Shared>, known_networks: Shared>, state: Rc, @@ -281,23 +281,17 @@ impl StaState { init: EspWifiInitialization, config: Config, wifi: &'static mut Wifi, - resources: &'static mut StackResources, - mut rng: Rng, + rng: Rng, ) -> Self { info!("Configuring STA"); let (wifi_interface, controller) = unwrap!(esp_wifi::wifi::new_with_mode(&init, wifi, WifiMode::Sta)); - let lower = rng.random() as u64; - let upper = rng.random() as u64; - - let random_seed = upper << 32 | lower; - Self { init, controller: Some(Box::new(controller)), - stack: Rc::new(Stack::new(wifi_interface, config, resources, random_seed)), + stack: Rc::new(StackWrapper::new(wifi_interface, config, rng)), networks: Rc::new(Mutex::new(heapless::Vec::new())), known_networks: Rc::new(Mutex::new(Vec::new())), state: Rc::new(State::new(InternalConnectionState::NotConnected)), @@ -380,7 +374,7 @@ async fn sta_task( networks: Shared>, known_networks: Shared>, state: Rc, - stack: Rc>>, + stack: Rc, mut task_control: TaskControlToken<(), StaTaskResources>, ) { const SCAN_PERIOD: Duration = Duration::from_secs(5);