8 files changed, 368 insertions, 200 deletions

diff --git a/embassy-stm32-wpan/CHANGELOG.md b/embassy-stm32-wpan/CHANGELOG.md
index 7042ad14c..c567fe1de 100644
--- a/embassy-stm32-wpan/CHANGELOG.md
+++ b/embassy-stm32-wpan/CHANGELOG.md
@@ -8,4 +8,5 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 <!-- next-header -->
 ## Unreleased - ReleaseDate
 
+- restructure to allow embassy net driver to work.
 - First release with changelog.
diff --git a/embassy-stm32-wpan/Cargo.toml b/embassy-stm32-wpan/Cargo.toml
index 0802b7328..75d978d1a 100644
--- a/embassy-stm32-wpan/Cargo.toml
+++ b/embassy-stm32-wpan/Cargo.toml
@@ -40,6 +40,7 @@ log = { version = "0.4.17", optional = true }
 cortex-m = "0.7.6"
 heapless = "0.8"
 aligned = "0.4.1"
+critical-section = "1.1"
 
 bit_field = "0.10.2"
 stm32-device-signature = { version = "0.3.3", features = ["stm32wb5x"] }
diff --git a/embassy-stm32-wpan/src/mac/control.rs b/embassy-stm32-wpan/src/mac/control.rs
index e8d2f9f7b..8fb971da3 100644
--- a/embassy-stm32-wpan/src/mac/control.rs
+++ b/embassy-stm32-wpan/src/mac/control.rs
@@ -1,65 +1,184 @@
+use core::cell::RefCell;
 use core::future::Future;
+use core::sync::atomic::{Ordering, compiler_fence};
 use core::task;
 use core::task::Poll;
 
+use embassy_net_driver::LinkState;
+use embassy_sync::blocking_mutex;
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
-use embassy_sync::mutex::MutexGuard;
+use embassy_sync::mutex::Mutex;
 use embassy_sync::signal::Signal;
 use futures_util::FutureExt;
 
-use super::commands::MacCommand;
-use super::event::MacEvent;
-use super::typedefs::MacError;
-use crate::mac::runner::Runner;
+use crate::mac::commands::*;
+use crate::mac::driver::NetworkState;
+use crate::mac::event::MacEvent;
+use crate::mac::runner::ZeroCopyPubSub;
+use crate::mac::typedefs::*;
+use crate::sub::mac::MacTx;
 
 pub struct Control<'a> {
-    runner: &'a Runner<'a>,
+    rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
+    mac_tx: &'a Mutex<CriticalSectionRawMutex, MacTx>,
+    #[allow(unused)]
+    network_state: &'a blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
 }
 
 impl<'a> Control<'a> {
-    pub(crate) fn new(runner: &'a Runner<'a>) -> Self {
-        Self { runner: runner }
+    pub(crate) fn new(
+        rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
+        mac_tx: &'a Mutex<CriticalSectionRawMutex, MacTx>,
+        network_state: &'a blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
+    ) -> Self {
+        Self {
+            rx_event_channel,
+            mac_tx,
+            network_state,
+        }
+    }
+
+    pub async fn init_link(&mut self, short_address: [u8; 2], extended_address: [u8; 8], pan_id: [u8; 2]) {
+        debug!("resetting");
+
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&ResetRequest {
+                set_default_pib: true,
+                ..Default::default()
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        debug!("setting extended address");
+        let extended_address: u64 = u64::from_be_bytes(extended_address);
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&SetRequest {
+                pib_attribute_ptr: &extended_address as *const _ as *const u8,
+                pib_attribute: PibId::ExtendedAddress,
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        debug!("setting short address");
+        let short_address: u16 = u16::from_be_bytes(short_address);
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&SetRequest {
+                pib_attribute_ptr: &short_address as *const _ as *const u8,
+                pib_attribute: PibId::ShortAddress,
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        critical_section::with(|cs| {
+            self.network_state.borrow(cs).borrow_mut().mac_addr = extended_address.to_be_bytes();
+        });
+
+        debug!("setting association permit");
+        let association_permit: bool = true;
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&SetRequest {
+                pib_attribute_ptr: &association_permit as *const _ as *const u8,
+                pib_attribute: PibId::AssociationPermit,
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        debug!("setting TX power");
+        let transmit_power: i8 = 2;
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&SetRequest {
+                pib_attribute_ptr: &transmit_power as *const _ as *const u8,
+                pib_attribute: PibId::TransmitPower,
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        debug!("starting FFD device");
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&StartRequest {
+                pan_id: PanId(pan_id),
+                channel_number: MacChannel::Channel16,
+                beacon_order: 0x0F,
+                superframe_order: 0x0F,
+                pan_coordinator: true,
+                battery_life_extension: false,
+                ..Default::default()
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        debug!("setting RX on when idle");
+        let rx_on_while_idle: bool = true;
+        debug!(
+            "{:#x}",
+            self.send_command_and_get_response(&SetRequest {
+                pib_attribute_ptr: &rx_on_while_idle as *const _ as *const u8,
+                pib_attribute: PibId::RxOnWhenIdle,
+            })
+            .await
+            .unwrap()
+            .await
+        );
+
+        critical_section::with(|cs| {
+            let mut network_state = self.network_state.borrow(cs).borrow_mut();
+
+            network_state.link_state = LinkState::Up;
+            network_state.link_waker.wake();
+        });
     }
 
     pub async fn send_command<T>(&self, cmd: &T) -> Result<(), MacError>
     where
         T: MacCommand,
     {
-        let _wm = self.runner.write_mutex.lock().await;
-
-        self.runner.mac_subsystem.send_command(cmd).await
+        self.mac_tx.lock().await.send_command(cmd).await
     }
 
     pub async fn send_command_and_get_response<T>(&self, cmd: &T) -> Result<EventToken<'a>, MacError>
     where
         T: MacCommand,
     {
-        let rm = self.runner.read_mutex.lock().await;
-        let _wm = self.runner.write_mutex.lock().await;
-        let token = EventToken::new(self.runner, rm);
+        let token = EventToken::new(self.rx_event_channel);
+
+        compiler_fence(Ordering::Release);
 
-        self.runner.mac_subsystem.send_command(cmd).await?;
+        self.mac_tx.lock().await.send_command(cmd).await?;
 
         Ok(token)
     }
 }
 
 pub struct EventToken<'a> {
-    runner: &'a Runner<'a>,
-    _mutex_guard: MutexGuard<'a, CriticalSectionRawMutex, ()>,
+    rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
 }
 
 impl<'a> EventToken<'a> {
-    pub(crate) fn new(runner: &'a Runner<'a>, mutex_guard: MutexGuard<'a, CriticalSectionRawMutex, ()>) -> Self {
+    pub(crate) fn new(rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>) -> Self {
         // Enable event receiving
-        runner.rx_event_channel.lock(|s| {
+        rx_event_channel.lock(|s| {
             *s.borrow_mut() = Some(Signal::new());
         });
 
-        Self {
-            runner: runner,
-            _mutex_guard: mutex_guard,
-        }
+        Self { rx_event_channel }
     }
 }
 
@@ -67,7 +186,7 @@ impl<'a> Future for EventToken<'a> {
     type Output = MacEvent<'a>;
 
     fn poll(self: core::pin::Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
-        self.get_mut().runner.rx_event_channel.lock(|s| {
+        self.rx_event_channel.lock(|s| {
             let signal = s.borrow_mut();
             let signal = match &*signal {
                 Some(s) => s,
@@ -88,7 +207,7 @@ impl<'a> Drop for EventToken<'a> {
     fn drop(&mut self) {
         // Disable event receiving
         // This will also drop the contained event, if it exists, and will free up receiving the next event
-        self.runner.rx_event_channel.lock(|s| {
+        self.rx_event_channel.lock(|s| {
             *s.borrow_mut() = None;
         });
     }
diff --git a/embassy-stm32-wpan/src/mac/driver.rs b/embassy-stm32-wpan/src/mac/driver.rs
index 480ac3790..819299b48 100644
--- a/embassy-stm32-wpan/src/mac/driver.rs
+++ b/embassy-stm32-wpan/src/mac/driver.rs
@@ -1,22 +1,97 @@
 #![deny(unused_must_use)]
 
+use core::cell::RefCell;
 use core::task::Context;
 
 use embassy_net_driver::{Capabilities, HardwareAddress, LinkState};
+use embassy_sync::blocking_mutex;
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 use embassy_sync::channel::Channel;
+use embassy_sync::mutex::Mutex;
+use embassy_sync::waitqueue::AtomicWaker;
 
-use crate::mac::MTU;
 use crate::mac::event::MacEvent;
-use crate::mac::runner::Runner;
+use crate::mac::runner::{BUF_SIZE, ZeroCopyPubSub};
+use crate::mac::{Control, MTU, Runner};
+use crate::sub::mac::{Mac, MacRx, MacTx};
+
+pub struct NetworkState {
+    pub mac_addr: [u8; 8],
+    pub link_state: LinkState,
+    pub link_waker: AtomicWaker,
+}
+
+impl NetworkState {
+    pub const fn new() -> Self {
+        Self {
+            mac_addr: [0u8; 8],
+            link_state: LinkState::Down,
+            link_waker: AtomicWaker::new(),
+        }
+    }
+}
+
+pub struct DriverState<'d> {
+    pub mac_tx: Mutex<CriticalSectionRawMutex, MacTx>,
+    pub mac_rx: MacRx,
+    pub rx_event_channel: ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'d>>,
+    pub rx_data_channel: Channel<CriticalSectionRawMutex, MacEvent<'d>, 1>,
+    pub tx_data_channel: Channel<CriticalSectionRawMutex, (&'d mut [u8; MTU], usize), BUF_SIZE>,
+    pub tx_buf_channel: Channel<CriticalSectionRawMutex, &'d mut [u8; MTU], BUF_SIZE>,
+    pub tx_buf_queue: [[u8; MTU]; BUF_SIZE],
+    pub network_state: blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
+}
+
+impl<'d> DriverState<'d> {
+    pub const fn new(mac: Mac) -> Self {
+        let (mac_rx, mac_tx) = mac.split();
+        let mac_tx = Mutex::new(mac_tx);
+
+        Self {
+            mac_tx,
+            mac_rx,
+            rx_event_channel: ZeroCopyPubSub::new(RefCell::new(None)),
+            rx_data_channel: Channel::new(),
+            tx_data_channel: Channel::new(),
+            tx_buf_channel: Channel::new(),
+            tx_buf_queue: [[0u8; MTU]; BUF_SIZE],
+            network_state: blocking_mutex::Mutex::new(RefCell::new(NetworkState::new())),
+        }
+    }
+}
 
 pub struct Driver<'d> {
-    runner: &'d Runner<'d>,
+    tx_data_channel: &'d Channel<CriticalSectionRawMutex, (&'d mut [u8; MTU], usize), BUF_SIZE>,
+    tx_buf_channel: &'d Channel<CriticalSectionRawMutex, &'d mut [u8; MTU], BUF_SIZE>,
+    rx_data_channel: &'d Channel<CriticalSectionRawMutex, MacEvent<'d>, 1>,
+    network_state: &'d blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
 }
 
 impl<'d> Driver<'d> {
-    pub(crate) fn new(runner: &'d Runner<'d>) -> Self {
-        Self { runner: runner }
+    pub fn new(driver_state: &'d mut DriverState<'d>) -> (Self, Runner<'d>, Control<'d>) {
+        (
+            Self {
+                tx_data_channel: &driver_state.tx_data_channel,
+                tx_buf_channel: &driver_state.tx_buf_channel,
+                rx_data_channel: &driver_state.rx_data_channel,
+                network_state: &driver_state.network_state,
+            },
+            Runner::new(
+                &driver_state.rx_event_channel,
+                &driver_state.rx_data_channel,
+                &mut driver_state.mac_rx,
+                &driver_state.tx_data_channel,
+                &driver_state.tx_buf_channel,
+                &driver_state.mac_tx,
+                &mut driver_state.tx_buf_queue,
+                &driver_state.network_state,
+            ),
+            Control::new(
+                &driver_state.rx_event_channel,
+                &driver_state.mac_tx,
+                &driver_state.network_state,
+            ),
+        )
     }
 }
 
@@ -33,16 +108,16 @@ impl<'d> embassy_net_driver::Driver for Driver<'d> {
         Self: 'a;
 
     fn receive(&mut self, cx: &mut Context) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
-        if self.runner.rx_channel.poll_ready_to_receive(cx).is_ready()
-            && self.runner.tx_buf_channel.poll_ready_to_receive(cx).is_ready()
+        if self.rx_data_channel.poll_ready_to_receive(cx).is_ready()
+            && self.tx_buf_channel.poll_ready_to_receive(cx).is_ready()
         {
             Some((
                 RxToken {
-                    rx: &self.runner.rx_channel,
+                    rx: self.rx_data_channel,
                 },
                 TxToken {
-                    tx: &self.runner.tx_channel,
-                    tx_buf: &self.runner.tx_buf_channel,
+                    tx: self.tx_data_channel,
+                    tx_buf: self.tx_buf_channel,
                 },
             ))
         } else {
@@ -51,10 +126,10 @@ impl<'d> embassy_net_driver::Driver for Driver<'d> {
     }
 
     fn transmit(&mut self, cx: &mut Context) -> Option<Self::TxToken<'_>> {
-        if self.runner.tx_buf_channel.poll_ready_to_receive(cx).is_ready() {
+        if self.tx_buf_channel.poll_ready_to_receive(cx).is_ready() {
             Some(TxToken {
-                tx: &self.runner.tx_channel,
-                tx_buf: &self.runner.tx_buf_channel,
+                tx: self.tx_data_channel,
+                tx_buf: self.tx_buf_channel,
             })
         } else {
             None
@@ -68,13 +143,20 @@ impl<'d> embassy_net_driver::Driver for Driver<'d> {
         caps
     }
 
-    fn link_state(&mut self, _cx: &mut Context) -> LinkState {
-        LinkState::Down
+    fn link_state(&mut self, cx: &mut Context) -> LinkState {
+        critical_section::with(|cs| {
+            let network_state = self.network_state.borrow(cs).borrow_mut();
+
+            // Unconditionally register the waker to avoid a race
+            network_state.link_waker.register(cx.waker());
+            network_state.link_state
+        })
     }
 
     fn hardware_address(&self) -> HardwareAddress {
-        // self.mac_addr
-        HardwareAddress::Ieee802154([0; 8])
+        HardwareAddress::Ieee802154(critical_section::with(|cs| {
+            self.network_state.borrow(cs).borrow().mac_addr
+        }))
     }
 }
 
@@ -99,8 +181,8 @@ impl<'d> embassy_net_driver::RxToken for RxToken<'d> {
 }
 
 pub struct TxToken<'d> {
-    tx: &'d Channel<CriticalSectionRawMutex, (&'d mut [u8; MTU], usize), 5>,
-    tx_buf: &'d Channel<CriticalSectionRawMutex, &'d mut [u8; MTU], 5>,
+    tx: &'d Channel<CriticalSectionRawMutex, (&'d mut [u8; MTU], usize), BUF_SIZE>,
+    tx_buf: &'d Channel<CriticalSectionRawMutex, &'d mut [u8; MTU], BUF_SIZE>,
 }
 
 impl<'d> embassy_net_driver::TxToken for TxToken<'d> {
diff --git a/embassy-stm32-wpan/src/mac/mod.rs b/embassy-stm32-wpan/src/mac/mod.rs
index c847a5cca..ac50a6b29 100644
--- a/embassy-stm32-wpan/src/mac/mod.rs
+++ b/embassy-stm32-wpan/src/mac/mod.rs
@@ -11,11 +11,7 @@ pub mod runner;
 pub mod typedefs;
 
 pub use crate::mac::control::Control;
-use crate::mac::driver::Driver;
+pub use crate::mac::driver::{Driver, DriverState};
 pub use crate::mac::runner::Runner;
 
 const MTU: usize = 127;
-
-pub async fn new<'a>(runner: &'a Runner<'a>) -> (Control<'a>, Driver<'a>) {
-    (Control::new(runner), Driver::new(runner))
-}
diff --git a/embassy-stm32-wpan/src/mac/runner.rs b/embassy-stm32-wpan/src/mac/runner.rs
index 2409f994d..26fdf23e0 100644
--- a/embassy-stm32-wpan/src/mac/runner.rs
+++ b/embassy-stm32-wpan/src/mac/runner.rs
@@ -8,52 +8,65 @@ use embassy_sync::mutex::Mutex;
 use embassy_sync::signal::Signal;
 
 use crate::mac::MTU;
-use crate::mac::commands::DataRequest;
+use crate::mac::commands::*;
+use crate::mac::driver::NetworkState;
 use crate::mac::event::MacEvent;
-use crate::mac::typedefs::{AddressMode, MacAddress, PanId, SecurityLevel};
-use crate::sub::mac::Mac;
+use crate::mac::typedefs::*;
+use crate::sub::mac::{MacRx, MacTx};
 
-type ZeroCopyPubSub<M, T> = blocking_mutex::Mutex<M, RefCell<Option<Signal<NoopRawMutex, T>>>>;
+pub type ZeroCopyPubSub<M, T> = blocking_mutex::Mutex<M, RefCell<Option<Signal<NoopRawMutex, T>>>>;
+
+pub const BUF_SIZE: usize = 3;
 
 pub struct Runner<'a> {
-    pub(crate) mac_subsystem: Mac,
     // rx event backpressure is already provided through the MacEvent drop mechanism
     // therefore, we don't need to worry about overwriting events
-    pub(crate) rx_event_channel: ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
-    pub(crate) read_mutex: Mutex<CriticalSectionRawMutex, ()>,
-    pub(crate) write_mutex: Mutex<CriticalSectionRawMutex, ()>,
-    pub(crate) rx_channel: Channel<CriticalSectionRawMutex, MacEvent<'a>, 1>,
-    pub(crate) tx_channel: Channel<CriticalSectionRawMutex, (&'a mut [u8; MTU], usize), 5>,
-    pub(crate) tx_buf_channel: Channel<CriticalSectionRawMutex, &'a mut [u8; MTU], 5>,
+    rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
+    rx_data_channel: &'a Channel<CriticalSectionRawMutex, MacEvent<'a>, 1>,
+    mac_rx: &'a mut MacRx,
+
+    tx_data_channel: &'a Channel<CriticalSectionRawMutex, (&'a mut [u8; MTU], usize), BUF_SIZE>,
+    tx_buf_channel: &'a Channel<CriticalSectionRawMutex, &'a mut [u8; MTU], BUF_SIZE>,
+    mac_tx: &'a Mutex<CriticalSectionRawMutex, MacTx>,
+
+    #[allow(unused)]
+    network_state: &'a blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
 }
 
 impl<'a> Runner<'a> {
-    pub fn new(mac: Mac, tx_buf_queue: [&'a mut [u8; MTU]; 5]) -> Self {
-        let this = Self {
-            mac_subsystem: mac,
-            rx_event_channel: blocking_mutex::Mutex::new(RefCell::new(None)),
-            read_mutex: Mutex::new(()),
-            write_mutex: Mutex::new(()),
-            rx_channel: Channel::new(),
-            tx_channel: Channel::new(),
-            tx_buf_channel: Channel::new(),
-        };
-
+    pub fn new(
+        rx_event_channel: &'a ZeroCopyPubSub<CriticalSectionRawMutex, MacEvent<'a>>,
+        rx_data_channel: &'a Channel<CriticalSectionRawMutex, MacEvent<'a>, 1>,
+        mac_rx: &'a mut MacRx,
+        tx_data_channel: &'a Channel<CriticalSectionRawMutex, (&'a mut [u8; MTU], usize), BUF_SIZE>,
+        tx_buf_channel: &'a Channel<CriticalSectionRawMutex, &'a mut [u8; MTU], BUF_SIZE>,
+        mac_tx: &'a Mutex<CriticalSectionRawMutex, MacTx>,
+        tx_buf_queue: &'a mut [[u8; MTU]; BUF_SIZE],
+        network_state: &'a blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<NetworkState>>,
+    ) -> Self {
         for buf in tx_buf_queue {
-            this.tx_buf_channel.try_send(buf).unwrap();
+            tx_buf_channel.try_send(buf).unwrap();
         }
 
-        this
+        Self {
+            rx_event_channel,
+            rx_data_channel,
+            mac_rx,
+            tx_data_channel,
+            tx_buf_channel,
+            mac_tx,
+            network_state,
+        }
     }
 
     pub async fn run(&'a self) -> ! {
         join::join(
             async {
                 loop {
-                    if let Ok(mac_event) = self.mac_subsystem.read().await {
+                    if let Ok(mac_event) = self.mac_rx.read().await {
                         match mac_event {
                             MacEvent::McpsDataInd(_) => {
-                                self.rx_channel.send(mac_event).await;
+                                self.rx_data_channel.send(mac_event).await;
                             }
                             _ => {
                                 self.rx_event_channel.lock(|s| {
@@ -73,11 +86,13 @@ impl<'a> Runner<'a> {
                 let mut msdu_handle = 0x02;
 
                 loop {
-                    let (buf, len) = self.tx_channel.receive().await;
-                    let _wm = self.write_mutex.lock().await;
+                    let (buf, len) = self.tx_data_channel.receive().await;
+                    let mac_tx = self.mac_tx.lock().await;
+
+                    // TODO: skip this if the link state is down
 
                     // The mutex should be dropped on the next loop iteration
-                    self.mac_subsystem
+                    mac_tx
                         .send_command(
                             DataRequest {
                                 src_addr_mode: AddressMode::Short,
diff --git a/embassy-stm32-wpan/src/sub/mac.rs b/embassy-stm32-wpan/src/sub/mac.rs
index baf4da979..93cafbc72 100644
--- a/embassy-stm32-wpan/src/sub/mac.rs
+++ b/embassy-stm32-wpan/src/sub/mac.rs
@@ -28,32 +28,39 @@ impl Mac {
         Self { _private: () }
     }
 
-    /// `HW_IPCC_MAC_802_15_4_EvtNot`
-    ///
-    /// This function will stall if the previous `EvtBox` has not been dropped
-    pub async fn tl_read(&self) -> EvtBox<Self> {
-        // Wait for the last event box to be dropped
-        poll_fn(|cx| {
-            MAC_WAKER.register(cx.waker());
-            if MAC_EVT_OUT.load(Ordering::SeqCst) {
-                Poll::Pending
-            } else {
-                Poll::Ready(())
-            }
-        })
-        .await;
+    pub const fn split(self) -> (MacRx, MacTx) {
+        (MacRx { _private: () }, MacTx { _private: () })
+    }
 
-        // Return a new event box
-        Ipcc::receive(channels::cpu2::IPCC_MAC_802_15_4_NOTIFICATION_ACK_CHANNEL, || unsafe {
-            // The closure is not async, therefore the closure must execute to completion (cannot be dropped)
-            // Therefore, the event box is guaranteed to be cleaned up if it's not leaked
-            MAC_EVT_OUT.store(true, Ordering::SeqCst);
+    pub async fn tl_write_and_get_response(&self, opcode: u16, payload: &[u8]) -> u8 {
+        MacTx { _private: () }.tl_write_and_get_response(opcode, payload).await
+    }
 
-            Some(EvtBox::new(MAC_802_15_4_NOTIF_RSP_EVT_BUFFER.as_mut_ptr() as *mut _))
-        })
-        .await
+    pub async fn tl_write(&self, opcode: u16, payload: &[u8]) {
+        MacTx { _private: () }.tl_write(opcode, payload).await
     }
 
+    pub async fn send_command<T>(&self, cmd: &T) -> Result<(), MacError>
+    where
+        T: MacCommand,
+    {
+        MacTx { _private: () }.send_command(cmd).await
+    }
+
+    pub async fn tl_read(&self) -> EvtBox<Mac> {
+        MacRx { _private: () }.tl_read().await
+    }
+
+    pub async fn read(&self) -> Result<MacEvent<'_>, ()> {
+        MacRx { _private: () }.read().await
+    }
+}
+
+pub struct MacTx {
+    _private: (),
+}
+
+impl MacTx {
     /// `HW_IPCC_MAC_802_15_4_CmdEvtNot`
     pub async fn tl_write_and_get_response(&self, opcode: u16, payload: &[u8]) -> u8 {
         self.tl_write(opcode, payload).await;
@@ -92,6 +99,38 @@ impl Mac {
             Err(MacError::from(response))
         }
     }
+}
+
+pub struct MacRx {
+    _private: (),
+}
+
+impl MacRx {
+    /// `HW_IPCC_MAC_802_15_4_EvtNot`
+    ///
+    /// This function will stall if the previous `EvtBox` has not been dropped
+    pub async fn tl_read(&self) -> EvtBox<Mac> {
+        // Wait for the last event box to be dropped
+        poll_fn(|cx| {
+            MAC_WAKER.register(cx.waker());
+            if MAC_EVT_OUT.load(Ordering::SeqCst) {
+                Poll::Pending
+            } else {
+                Poll::Ready(())
+            }
+        })
+        .await;
+
+        // Return a new event box
+        Ipcc::receive(channels::cpu2::IPCC_MAC_802_15_4_NOTIFICATION_ACK_CHANNEL, || unsafe {
+            // The closure is not async, therefore the closure must execute to completion (cannot be dropped)
+            // Therefore, the event box is guaranteed to be cleaned up if it's not leaked
+            MAC_EVT_OUT.store(true, Ordering::SeqCst);
+
+            Some(EvtBox::new(MAC_802_15_4_NOTIF_RSP_EVT_BUFFER.as_mut_ptr() as *mut _))
+        })
+        .await
+    }
 
     pub async fn read(&self) -> Result<MacEvent<'_>, ()> {
         MacEvent::new(self.tl_read().await)
diff --git a/examples/stm32wb/src/bin/mac_ffd_net.rs b/examples/stm32wb/src/bin/mac_ffd_net.rs
index 5296943a1..9b705dda9 100644
--- a/examples/stm32wb/src/bin/mac_ffd_net.rs
+++ b/examples/stm32wb/src/bin/mac_ffd_net.rs
@@ -7,9 +7,7 @@ use embassy_stm32::bind_interrupts;
 use embassy_stm32::ipcc::{Config, ReceiveInterruptHandler, TransmitInterruptHandler};
 use embassy_stm32::rcc::WPAN_DEFAULT;
 use embassy_stm32_wpan::TlMbox;
-use embassy_stm32_wpan::mac::commands::{ResetRequest, SetRequest, StartRequest};
-use embassy_stm32_wpan::mac::typedefs::{MacChannel, PanId, PibId};
-use embassy_stm32_wpan::mac::{self, Runner};
+use embassy_stm32_wpan::mac::{Driver, DriverState, Runner};
 use embassy_stm32_wpan::sub::mm;
 use static_cell::StaticCell;
 use {defmt_rtt as _, panic_probe as _};
@@ -72,106 +70,23 @@ async fn main(spawner: Spawner) {
     let result = mbox.sys_subsystem.shci_c2_mac_802_15_4_init().await;
     info!("initialized mac: {}", result);
 
-    info!("resetting");
-    mbox.mac_subsystem
-        .send_command(&ResetRequest {
-            set_default_pib: true,
-            ..Default::default()
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("setting extended address");
-    let extended_address: u64 = 0xACDE480000000001;
-    mbox.mac_subsystem
-        .send_command(&SetRequest {
-            pib_attribute_ptr: &extended_address as *const _ as *const u8,
-            pib_attribute: PibId::ExtendedAddress,
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("setting short address");
-    let short_address: u16 = 0x1122;
-    mbox.mac_subsystem
-        .send_command(&SetRequest {
-            pib_attribute_ptr: &short_address as *const _ as *const u8,
-            pib_attribute: PibId::ShortAddress,
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("setting association permit");
-    let association_permit: bool = true;
-    mbox.mac_subsystem
-        .send_command(&SetRequest {
-            pib_attribute_ptr: &association_permit as *const _ as *const u8,
-            pib_attribute: PibId::AssociationPermit,
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("setting TX power");
-    let transmit_power: i8 = 2;
-    mbox.mac_subsystem
-        .send_command(&SetRequest {
-            pib_attribute_ptr: &transmit_power as *const _ as *const u8,
-            pib_attribute: PibId::TransmitPower,
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("starting FFD device");
-    mbox.mac_subsystem
-        .send_command(&StartRequest {
-            pan_id: PanId([0x1A, 0xAA]),
-            channel_number: MacChannel::Channel16,
-            beacon_order: 0x0F,
-            superframe_order: 0x0F,
-            pan_coordinator: true,
-            battery_life_extension: false,
-            ..Default::default()
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    info!("setting RX on when idle");
-    let rx_on_while_idle: bool = true;
-    mbox.mac_subsystem
-        .send_command(&SetRequest {
-            pib_attribute_ptr: &rx_on_while_idle as *const _ as *const u8,
-            pib_attribute: PibId::RxOnWhenIdle,
-        })
-        .await
-        .unwrap();
-    defmt::info!("{:#x}", mbox.mac_subsystem.read().await.unwrap());
-
-    static TX1: StaticCell<[u8; 127]> = StaticCell::new();
-    static TX2: StaticCell<[u8; 127]> = StaticCell::new();
-    static TX3: StaticCell<[u8; 127]> = StaticCell::new();
-    static TX4: StaticCell<[u8; 127]> = StaticCell::new();
-    static TX5: StaticCell<[u8; 127]> = StaticCell::new();
-    let tx_queue = [
-        TX1.init([0u8; 127]),
-        TX2.init([0u8; 127]),
-        TX3.init([0u8; 127]),
-        TX4.init([0u8; 127]),
-        TX5.init([0u8; 127]),
-    ];
-
+    static DRIVER_STATE: StaticCell<DriverState> = StaticCell::new();
     static RUNNER: StaticCell<Runner> = StaticCell::new();
-    let runner = RUNNER.init(Runner::new(mbox.mac_subsystem, tx_queue));
 
-    spawner.spawn(run_mac(runner).unwrap());
+    let driver_state = DRIVER_STATE.init(DriverState::new(mbox.mac_subsystem));
+    let (driver, runner, mut control) = Driver::new(driver_state);
+
+    spawner.spawn(run_mac(RUNNER.init(runner)).unwrap());
+
+    control
+        .init_link(
+            0x1122u16.to_be_bytes().try_into().unwrap(),
+            0xACDE480000000001u64.to_be_bytes().try_into().unwrap(),
+            [0x1A, 0xAA],
+        )
+        .await;
 
-    let (driver, control) = mac::new(runner).await;
+    cortex_m::asm::bkpt();
 
     let _ = driver;
-    let _ = control;
 }