Skip to content

Commit

Permalink
Move socket collection into the ap/sta implementations
Browse files Browse the repository at this point in the history
  • Loading branch information
@bugadani
bugadani committed Oct 9, 2023
1 parent f75d5f3 commit 391b8c7
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 62 deletions.
18 changes: 6 additions & 12 deletions src/board/wifi/ap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -36,7 +36,7 @@ impl From<ApConnectionState> for WifiState {

#[derive(Clone)]
pub struct Ap {
stack: Rc<Stack<WifiDevice<'static>>>,
stack: Rc<StackWrapper>,
client_count: Rc<AtomicU32>,
}

Expand Down Expand Up @@ -65,7 +65,7 @@ impl Ap {
pub(super) struct ApState {
init: EspWifiInitialization,
controller: Option<Box<WifiController<'static>>>,
stack: Rc<Stack<WifiDevice<'static>>>,
stack: Rc<StackWrapper>,
connection_task_control: Option<TaskController<(), ApTaskResources>>,
net_task_control: TaskController<!>,
client_count: Rc<AtomicU32>,
Expand All @@ -76,23 +76,17 @@ impl ApState {
init: EspWifiInitialization,
config: Config,
wifi: &'static mut Wifi,
resources: &'static mut StackResources<STACK_SOCKET_COUNT>,
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)),
Expand Down
96 changes: 59 additions & 37 deletions src/board/wifi/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use core::{
hint::unreachable_unchecked,
mem::{self, MaybeUninit},
ops::Deref,
ptr::NonNull,
};

use crate::{
Expand Down Expand Up @@ -29,11 +31,50 @@ pub unsafe fn as_static_mut<T>(what: &mut T) -> &'static mut T {
mem::transmute(what)
}

const STACK_SOCKET_COUNT: usize = 3;

struct StackWrapper(
NonNull<StackResources<STACK_SOCKET_COUNT>>,
Stack<WifiDevice<'static>>,
);

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<WifiDevice<'static>>;

#[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),
Expand Down Expand Up @@ -72,14 +113,8 @@ impl WifiHandle {
enum WifiDriverState {
Uninitialized(WifiInitResources),
Initialized(EspWifiInitialization),
Ap(
MaybeUninit<ApState>,
Box<StackResources<STACK_SOCKET_COUNT>>,
),
Sta(
MaybeUninit<StaState>,
Box<StackResources<STACK_SOCKET_COUNT>>,
),
Ap(MaybeUninit<ApState>),
Sta(MaybeUninit<StaState>),
}

impl WifiDriverState {
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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<WifiHandle> {
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,
}
}
Expand Down Expand Up @@ -203,8 +230,8 @@ impl WifiDriver {

pub fn wifi_mode(&self) -> Option<WifiMode> {
match self.state {
WifiDriverState::Ap(_, _) => Some(WifiMode::Ap),
WifiDriverState::Sta(_, _) => Some(WifiMode::Sta),
WifiDriverState::Ap(_) => Some(WifiMode::Ap),
WifiDriverState::Sta(_) => Some(WifiMode::Sta),
_ => None,
}
}
Expand All @@ -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,
));
}
Expand All @@ -247,7 +273,7 @@ impl WifiDriver {

WifiDriverState::Uninitialized { .. }
| WifiDriverState::Initialized { .. }
| WifiDriverState::Sta(_, _) => {
| WifiDriverState::Sta(_) => {
unreachable!()
}
}
Expand All @@ -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,
));
}
Expand All @@ -277,7 +302,7 @@ impl WifiDriver {

WifiDriverState::Uninitialized { .. }
| WifiDriverState::Initialized { .. }
| WifiDriverState::Ap(_, _) => {
| WifiDriverState::Ap(_) => {
unreachable!()
}
}
Expand All @@ -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 {
Expand All @@ -308,9 +333,6 @@ impl WifiDriver {
}

#[cardio::task]
pub async fn net_task(
stack: Rc<Stack<WifiDevice<'static>>>,
mut task_control: TaskControlToken<!>,
) {
async fn net_task(stack: Rc<StackWrapper>, mut task_control: TaskControlToken<!>) {
task_control.run_cancellable(|_| stack.run()).await;
}
20 changes: 7 additions & 13 deletions src/board/wifi/sta.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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},
Expand Down Expand Up @@ -124,7 +124,7 @@ impl From<InternalConnectionState> for ConnectionState {

#[derive(Clone)]
pub struct Sta {
stack: Rc<Stack<WifiDevice<'static>>>,
stack: Rc<StackWrapper>,
networks: Shared<heapless::Vec<AccessPointInfo, SCAN_RESULTS>>,
known_networks: Shared<Vec<KnownNetwork>>,
state: Rc<State>,
Expand Down Expand Up @@ -267,7 +267,7 @@ impl<'a> HttpsClientResources<'a> {
pub(super) struct StaState {
init: EspWifiInitialization,
controller: Option<Box<WifiController<'static>>>,
stack: Rc<Stack<WifiDevice<'static>>>,
stack: Rc<StackWrapper>,
networks: Shared<heapless::Vec<AccessPointInfo, SCAN_RESULTS>>,
known_networks: Shared<Vec<KnownNetwork>>,
state: Rc<State>,
Expand All @@ -281,23 +281,17 @@ impl StaState {
init: EspWifiInitialization,
config: Config,
wifi: &'static mut Wifi,
resources: &'static mut StackResources<STACK_SOCKET_COUNT>,
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)),
Expand Down Expand Up @@ -380,7 +374,7 @@ async fn sta_task(
networks: Shared<heapless::Vec<AccessPointInfo, SCAN_RESULTS>>,
known_networks: Shared<Vec<KnownNetwork>>,
state: Rc<State>,
stack: Rc<Stack<WifiDevice<'static>>>,
stack: Rc<StackWrapper>,
mut task_control: TaskControlToken<(), StaTaskResources>,
) {
const SCAN_PERIOD: Duration = Duration::from_secs(5);
Expand Down

0 comments on commit 391b8c7

Please sign in to comment.