UDP stack based on embassy-net

Cargo.toml

@@ -5,7 +5,6 @@ exclude = ["examples/*"]
 
 # For compatibility with ESP IDF
 [patch.crates-io]
-smol = { git = "https://github.com/esp-rs-compat/smol" }
 polling = { git = "https://github.com/esp-rs-compat/polling" }
 socket2 = { git = "https://github.com/esp-rs-compat/socket2" }
 

examples/onoff_light/src/main.rs

@@ -18,7 +18,6 @@
 use core::borrow::Borrow;
 use core::pin::pin;
 
-use embassy_futures::select::select3;
 use log::info;
 use matter::core::{CommissioningData, Matter};
 use matter::data_model::cluster_basic_information::BasicInfoConfig;
@@ -28,13 +27,11 @@ use matter::data_model::objects::*;
 use matter::data_model::root_endpoint;
 use matter::data_model::system_model::descriptor;
 use matter::error::Error;
-use matter::mdns::{DefaultMdns, DefaultMdnsRunner};
+use matter::mdns::MdnsService;
 use matter::persist::FilePsm;
 use matter::secure_channel::spake2p::VerifierData;
 use matter::transport::network::{Ipv4Addr, Ipv6Addr, NetworkStack};
-use matter::transport::runner::{RxBuf, TransportRunner, TxBuf};
+use matter::transport::runner::{AllUdpBuffers, TransportRunner};
-use matter::transport::udp::UdpBuffers;
-use matter::utils::select::EitherUnwrap;
 
 mod dev_att;
 
@@ -59,10 +56,11 @@ fn run() -> Result<(), Error> {
     initialize_logger();
 
     info!(
-        "Matter memory: mDNS={}, Matter={}, TransportRunner={}",
+        "Matter memory: mDNS={}, Matter={}, TransportRunner={}, UdpBuffers={}",
-        core::mem::size_of::<DefaultMdns>(),
+        core::mem::size_of::<MdnsService>(),
         core::mem::size_of::<Matter>(),
         core::mem::size_of::<TransportRunner>(),
+        core::mem::size_of::<AllUdpBuffers>(),
     );
 
     let dev_det = BasicInfoConfig {
@@ -83,20 +81,6 @@ fn run() -> Result<(), Error> {
 
     let (ipv4_addr, ipv6_addr, interface) = initialize_network()?;
 
-    let mdns = DefaultMdns::new(
-        0,
-        "matter-demo",
-        ipv4_addr.octets(),
-        Some(ipv6_addr.octets()),
-        interface,
-        &dev_det,
-        matter::MATTER_PORT,
-    );
-
-    let mut mdns_runner = DefaultMdnsRunner::new(&mdns);
-
-    info!("mDNS initialized: {:p}, {:p}", &mdns, &mdns_runner);
-
     let dev_att = dev_att::HardCodedDevAtt::new();
 
     #[cfg(feature = "std")]
@@ -113,6 +97,18 @@ fn run() -> Result<(), Error> {
     #[cfg(not(feature = "std"))]
     let rand = matter::utils::rand::dummy_rand;
 
+    let mdns = MdnsService::new(
+        0,
+        "matter-demo",
+        ipv4_addr.octets(),
+        Some(ipv6_addr.octets()),
+        interface,
+        &dev_det,
+        matter::MATTER_PORT,
+    );
+
+    info!("mDNS initialized: {:p}", &mdns);
+
     let matter = Matter::new(
         // vid/pid should match those in the DAC
         &dev_det,
@@ -125,20 +121,6 @@ fn run() -> Result<(), Error> {
 
     info!("Matter initialized: {:p}", &matter);
 
-    let mut runner = TransportRunner::new(&matter);
-
-    info!("Transport Runner initialized: {:p}", &runner);
-
-    let mut tx_buf = TxBuf::uninit();
-    let mut rx_buf = RxBuf::uninit();
-
-    // NOTE (no_std): If using the `embassy-net` UDP implementation, replace this dummy stack with the `embassy-net` one
-    // When using a custom UDP stack, remove this
-    let stack = NetworkStack::new();
-
-    let mut mdns_udp_buffers = UdpBuffers::new();
-    let mut trans_udp_buffers = UdpBuffers::new();
-
     #[cfg(all(feature = "std", not(target_os = "espidf")))]
     {
         let mut buf = [0; 4096];
@@ -152,62 +134,33 @@ fn run() -> Result<(), Error> {
         }
     }
 
-    let node = Node {
+    let mut runner = TransportRunner::new(&matter);
-        id: 0,
+
-        endpoints: &[
+    info!("Transport runner initialized: {:p}", &runner);
-            root_endpoint::endpoint(0),
-            Endpoint {
-                id: 1,
-                device_type: DEV_TYPE_ON_OFF_LIGHT,
-                clusters: &[descriptor::CLUSTER, cluster_on_off::CLUSTER],
-            },
-        ],
-    };
 
     let handler = HandlerCompat(handler(&matter));
 
-    let matter = &matter;
+    // NOTE (no_std): If using the `embassy-net` UDP implementation, replace this dummy stack with the `embassy-net` one
-    let node = &node;
+    // When using a custom UDP stack, remove this
-    let handler = &handler;
+    let stack = NetworkStack::new();
-    let runner = &mut runner;
-    let tx_buf = &mut tx_buf;
-    let rx_buf = &mut rx_buf;
-    let stack = &stack;
-    let mdns_udp_buffers = &mut mdns_udp_buffers;
-    let trans_udp_buffers = &mut trans_udp_buffers;
 
-    info!(
+    let mut buffers = AllUdpBuffers::new();
-        "About to run wth node {:p}, handler {:p}, transport runner {:p}, mdns_runner {:p}",
-        node, handler, runner, &mdns_runner
-    );
 
-    let mut fut = pin!(async move {
+    let mut fut = pin!(runner.run_udp_all(
-        // NOTE: If using a custom UDP stack, replace `run_udp` with `run`
+        &stack,
-        // and connect the pipes of the `run` method with the custom UDP stack
+        &mdns,
-        let mut transport = pin!(runner.run_udp(
+        &mut buffers,
-            stack,
+        CommissioningData {
-            trans_udp_buffers,
+            // TODO: Hard-coded for now
-            tx_buf,
+            verifier: VerifierData::new_with_pw(123456, *matter.borrow()),
-            rx_buf,
+            discriminator: 250,
-            CommissioningData {
+        },
-                // TODO: Hard-coded for now
+        &handler,
-                verifier: VerifierData::new_with_pw(123456, *matter.borrow()),
+    ));
-                discriminator: 250,
-            },
-            &handler,
-        ));
-
-        // NOTE: If using a custom UDP stack, replace `run_udp` with `run`
-        // and connect the pipes of the `run` method with the custom UDP stack
-        let mut mdns = pin!(mdns_runner.run_udp(stack, mdns_udp_buffers));
-
-        let mut save = pin!(save(matter, &psm));
-        select3(&mut transport, &mut mdns, &mut save).await.unwrap()
-    });
 
     // NOTE: For no_std, replace with your own no_std way of polling the future
     #[cfg(feature = "std")]
-    smol::block_on(&mut fut)?;
+    async_io::block_on(&mut fut)?;
 
     // NOTE (no_std): For no_std, replace with your own more efficient no_std executor,
     // because the executor used below is a simple busy-loop poller

matter/Cargo.toml

@@ -17,7 +17,7 @@ path = "src/lib.rs"
 [features]
 default = ["os", "crypto_rustcrypto"]
 os = ["std", "backtrace", "env_logger", "nix", "critical-section/std", "embassy-sync/std", "embassy-time/std"]
-std = ["alloc", "rand", "qrcode", "async-io", "smol", "esp-idf-sys/std"]
+std = ["alloc", "rand", "qrcode", "async-io", "esp-idf-sys/std"]
 backtrace = []
 alloc = []
 nightly = []
@@ -47,6 +47,8 @@ embassy-time = { version = "0.1.1", features = ["generic-queue-8"] }
 embassy-sync = "0.2"
 critical-section = "1.1.1"
 domain = { version = "0.7.2", default_features = false, features = ["heapless"] }
+
+# embassy-net dependencies
 embassy-net = { version = "0.1", features = ["udp", "igmp", "proto-ipv6", "medium-ethernet", "medium-ip"], optional = true }
 embassy-net-driver = { version = "0.1", optional = true }
 smoltcp = { version = "0.10", default-features = false, optional = true }
@@ -54,7 +56,6 @@ smoltcp = { version = "0.10", default-features = false, optional = true }
 # STD-only dependencies
 rand = { version = "0.8.5", optional = true }
 qrcode = { version = "0.12", default-features = false, optional = true } # Print QR code
-smol = { version = "1.2", optional = true } # =1.2 for compatibility with ESP IDF
 async-io = { version = "=1.12", optional = true } # =1.2 for compatibility with ESP IDF
 
 # crypto

matter/src/mdns.rs

@@ -56,16 +56,18 @@ where
 }
 
 #[cfg(all(feature = "std", target_os = "macos"))]
-pub type DefaultMdns<'a> = astro::Mdns<'a>;
+pub use astro::MdnsRunner;
-
 #[cfg(all(feature = "std", target_os = "macos"))]
-pub type DefaultMdnsRunner<'a> = astro::MdnsRunner<'a>;
+pub use astro::MdnsService;
+#[cfg(all(feature = "std", target_os = "macos"))]
+pub use astro::MdnsUdpBuffers;
 
 #[cfg(not(all(feature = "std", target_os = "macos")))]
-pub type DefaultMdns<'a> = builtin::Mdns<'a>;
+pub use builtin::MdnsRunner;
-
 #[cfg(not(all(feature = "std", target_os = "macos")))]
-pub type DefaultMdnsRunner<'a> = builtin::MdnsRunner<'a>;
+pub use builtin::MdnsService;
+#[cfg(not(all(feature = "std", target_os = "macos")))]
+pub use builtin::MdnsUdpBuffers;
 
 pub struct DummyMdns;
 

matter/src/mdns/astro.rs

@@ -11,13 +11,14 @@ use log::info;
 
 use super::ServiceMode;
 
-pub struct Mdns<'a> {
+pub struct MdnsService<'a> {
     dev_det: &'a BasicInfoConfig<'a>,
     matter_port: u16,
     services: RefCell<HashMap<String, RegisteredDnsService>>,
 }
 
-impl<'a> Mdns<'a> {
+impl<'a> MdnsService<'a> {
+    /// This constructor takes extra parameters for API-compatibility with builtin::MdnsRunner
     pub fn new(
         _id: u16,
         _hostname: &str,
@@ -80,14 +81,37 @@ impl<'a> Mdns<'a> {
     }
 }
 
-pub struct MdnsRunner<'a>(&'a Mdns<'a>);
+/// Only for API-compatibility with builtin::MdnsRunner
+pub struct MdnsUdpBuffers(());
 
+/// Only for API-compatibility with builtin::MdnsRunner
+impl MdnsUdpBuffers {
+    #[inline(always)]
+    pub const fn new() -> Self {
+        Self(())
+    }
+}
+
+impl<'a> super::Mdns for MdnsService<'a> {
+    fn add(&self, service: &str, mode: ServiceMode) -> Result<(), Error> {
+        MdnsService::add(self, service, mode)
+    }
+
+    fn remove(&self, service: &str) -> Result<(), Error> {
+        MdnsService::remove(self, service)
+    }
+}
+
+/// Only for API-compatibility with builtin::MdnsRunner
+pub struct MdnsRunner<'a>(&'a MdnsService<'a>);
+
+/// Only for API-compatibility with builtin::MdnsRunner
 impl<'a> MdnsRunner<'a> {
-    pub const fn new(mdns: &'a Mdns<'a>) -> Self {
+    pub const fn new(mdns: &'a MdnsService<'a>) -> Self {
         Self(mdns)
     }
 
-    pub async fn run_udp(&mut self) -> Result<(), Error> {
+    pub async fn run_udp(&mut self, buffers: &mut MdnsUdpBuffers) -> Result<(), Error> {
         core::future::pending::<Result<(), Error>>().await
     }
 
@@ -95,13 +119,3 @@ impl<'a> MdnsRunner<'a> {
         core::future::pending::<Result<(), Error>>().await
     }
 }
-
-impl<'a> super::Mdns for Mdns<'a> {
-    fn add(&self, service: &str, mode: ServiceMode) -> Result<(), Error> {
-        Mdns::add(self, service, mode)
-    }
-
-    fn remove(&self, service: &str) -> Result<(), Error> {
-        Mdns::remove(self, service)
-    }
-}

matter/src/mdns/builtin.rs

@@ -22,7 +22,7 @@ const IPV6_BROADCAST_ADDR: Ipv6Addr = Ipv6Addr::new(0xff02, 0, 0, 0, 0, 0, 0, 0x
 
 const PORT: u16 = 5353;
 
-pub struct Mdns<'a> {
+pub struct MdnsService<'a> {
     host: Host<'a>,
     #[allow(unused)]
     interface: u32,
@@ -32,7 +32,7 @@ pub struct Mdns<'a> {
     notification: Notification,
 }
 
-impl<'a> Mdns<'a> {
+impl<'a> MdnsService<'a> {
     #[inline(always)]
     pub const fn new(
         id: u16,
@@ -95,10 +95,29 @@ impl<'a> Mdns<'a> {
     }
 }
 
-pub struct MdnsRunner<'a>(&'a Mdns<'a>);
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+pub struct MdnsUdpBuffers {
+    udp: crate::transport::udp::UdpBuffers,
+    tx_buf: core::mem::MaybeUninit<[u8; crate::transport::packet::MAX_TX_BUF_SIZE]>,
+    rx_buf: core::mem::MaybeUninit<[u8; crate::transport::packet::MAX_RX_BUF_SIZE]>,
+}
+
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+impl MdnsUdpBuffers {
+    #[inline(always)]
+    pub const fn new() -> Self {
+        Self {
+            udp: crate::transport::udp::UdpBuffers::new(),
+            tx_buf: core::mem::MaybeUninit::uninit(),
+            rx_buf: core::mem::MaybeUninit::uninit(),
+        }
+    }
+}
+
+pub struct MdnsRunner<'a>(&'a MdnsService<'a>);
 
 impl<'a> MdnsRunner<'a> {
-    pub const fn new(mdns: &'a Mdns<'a>) -> Self {
+    pub const fn new(mdns: &'a MdnsService<'a>) -> Self {
         Self(mdns)
     }
 
@@ -106,31 +125,17 @@ impl<'a> MdnsRunner<'a> {
     pub async fn run_udp<D>(
         &mut self,
         stack: &crate::transport::network::NetworkStack<D>,
-        buffers: &mut crate::transport::udp::UdpBuffers,
+        buffers: &mut MdnsUdpBuffers,
     ) -> Result<(), Error>
     where
         D: crate::transport::network::NetworkStackMulticastDriver
             + crate::transport::network::NetworkStackDriver
             + 'static,
     {
-        let mut tx_buf =
-            core::mem::MaybeUninit::<[u8; crate::transport::packet::MAX_TX_BUF_SIZE]>::uninit();
-        let mut rx_buf =
-            core::mem::MaybeUninit::<[u8; crate::transport::packet::MAX_RX_BUF_SIZE]>::uninit();
-
-        let tx_buf = &mut tx_buf;
-        let rx_buf = &mut rx_buf;
-
-        let tx_pipe = Pipe::new(unsafe { tx_buf.assume_init_mut() });
-        let rx_pipe = Pipe::new(unsafe { rx_buf.assume_init_mut() });
-
-        let tx_pipe = &tx_pipe;
-        let rx_pipe = &rx_pipe;
-
         let mut udp = crate::transport::udp::UdpListener::new(
             stack,
             crate::transport::network::SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), PORT),
-            buffers,
+            &mut buffers.udp,
         )
         .await?;
 
@@ -140,6 +145,11 @@ impl<'a> MdnsRunner<'a> {
             crate::transport::network::Ipv4Addr::from(self.0.host.ip),
         )?;
 
+        let tx_pipe = Pipe::new(unsafe { buffers.tx_buf.assume_init_mut() });
+        let rx_pipe = Pipe::new(unsafe { buffers.rx_buf.assume_init_mut() });
+
+        let tx_pipe = &tx_pipe;
+        let rx_pipe = &rx_pipe;
         let udp = &udp;
 
         let mut tx = pin!(async move {
@@ -295,24 +305,24 @@ impl<'a> MdnsRunner<'a> {
     }
 }
 
-impl<'a> super::Mdns for Mdns<'a> {
+impl<'a> Services for MdnsService<'a> {
-    fn add(&self, service: &str, mode: ServiceMode) -> Result<(), Error> {
-        Mdns::add(self, service, mode)
-    }
-
-    fn remove(&self, service: &str) -> Result<(), Error> {
-        Mdns::remove(self, service)
-    }
-}
-
-impl<'a> Services for Mdns<'a> {
     type Error = crate::error::Error;
 
     fn for_each<F>(&self, callback: F) -> Result<(), Error>
     where
         F: FnMut(&Service) -> Result<(), Error>,
     {
-        Mdns::for_each(self, callback)
+        MdnsService::for_each(self, callback)
+    }
+}
+
+impl<'a> super::Mdns for MdnsService<'a> {
+    fn add(&self, service: &str, mode: ServiceMode) -> Result<(), Error> {
+        MdnsService::add(self, service, mode)
+    }
+
+    fn remove(&self, service: &str) -> Result<(), Error> {
+        MdnsService::remove(self, service)
     }
 }
 

matter/src/transport/network.rs

@@ -60,10 +60,10 @@ impl Debug for Address {
 pub use std_stack::*;
 
 #[cfg(feature = "embassy-net")]
-pub use embassy_stack::*;
+pub use embassy_net_stack::*;
 
-#[cfg(all(feature = "std", not(feature = "embassy-net")))]
+#[cfg(feature = "std")]
-mod std_stack {
+pub mod std_stack {
     pub trait NetworkStackDriver {}
 
     impl NetworkStackDriver for () {}
@@ -82,7 +82,7 @@ mod std_stack {
 }
 
 #[cfg(feature = "embassy-net")]
-mod embassy_stack {
+pub mod embassy_net_stack {
     pub use embassy_net::Stack as NetworkStack;
     pub use embassy_net_driver::Driver as NetworkStackDriver;
     pub use smoltcp::phy::Device as NetworkStackMulticastDriver;

matter/src/transport/runner.rs

@@ -41,8 +41,8 @@ use super::{
     pipe::{Chunk, Pipe},
 };
 
-pub type TxBuf = MaybeUninit<[u8; MAX_TX_BUF_SIZE]>;
+type TxBuf = MaybeUninit<[u8; MAX_TX_BUF_SIZE]>;
-pub type RxBuf = MaybeUninit<[u8; MAX_RX_BUF_SIZE]>;
+type RxBuf = MaybeUninit<[u8; MAX_RX_BUF_SIZE]>;
 type SxBuf = MaybeUninit<[u8; MAX_RX_STATUS_BUF_SIZE]>;
 
 struct PacketPools {
@@ -70,6 +70,42 @@ impl PacketPools {
     }
 }
 
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+pub struct AllUdpBuffers {
+    transport: TransportUdpBuffers,
+    mdns: crate::mdns::MdnsUdpBuffers,
+}
+
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+impl AllUdpBuffers {
+    #[inline(always)]
+    pub const fn new() -> Self {
+        Self {
+            transport: TransportUdpBuffers::new(),
+            mdns: crate::mdns::MdnsUdpBuffers::new(),
+        }
+    }
+}
+
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+pub struct TransportUdpBuffers {
+    udp: crate::transport::udp::UdpBuffers,
+    tx_buf: TxBuf,
+    rx_buf: RxBuf,
+}
+
+#[cfg(any(feature = "std", feature = "embassy-net"))]
+impl TransportUdpBuffers {
+    #[inline(always)]
+    pub const fn new() -> Self {
+        Self {
+            udp: crate::transport::udp::UdpBuffers::new(),
+            tx_buf: core::mem::MaybeUninit::uninit(),
+            rx_buf: core::mem::MaybeUninit::uninit(),
+        }
+    }
+}
+
 /// This struct implements an executor-agnostic option to run the Matter transport stack end-to-end.
 ///
 /// Since it is not possible to use executor tasks spawning in an executor-agnostic way (yet),
@@ -101,13 +137,36 @@ impl<'a> TransportRunner<'a> {
         &self.transport
     }
 
+    #[cfg(any(feature = "std", feature = "embassy-net"))]
+    pub async fn run_udp_all<D, H>(
+        &mut self,
+        stack: &crate::transport::network::NetworkStack<D>,
+        mdns: &crate::mdns::MdnsService<'_>,
+        buffers: &mut AllUdpBuffers,
+        dev_comm: CommissioningData,
+        handler: &H,
+    ) -> Result<(), Error>
+    where
+        D: crate::transport::network::NetworkStackDriver
+            + crate::transport::network::NetworkStackMulticastDriver
+            + 'static,
+        H: DataModelHandler,
+    {
+        let mut mdns_runner = crate::mdns::MdnsRunner::new(mdns);
+
+        let mut mdns = pin!(mdns_runner.run_udp(stack, &mut buffers.mdns));
+        let mut transport = pin!(self.run_udp(stack, &mut buffers.transport, dev_comm, handler));
+
+        embassy_futures::select::select(&mut transport, &mut mdns)
+            .await
+            .unwrap()
+    }
+
     #[cfg(any(feature = "std", feature = "embassy-net"))]
     pub async fn run_udp<D, H>(
         &mut self,
         stack: &crate::transport::network::NetworkStack<D>,
-        buffers: &mut crate::transport::udp::UdpBuffers,
+        buffers: &mut TransportUdpBuffers,
-        tx_buf: &mut TxBuf,
-        rx_buf: &mut RxBuf,
         dev_comm: CommissioningData,
         handler: &H,
     ) -> Result<(), Error>
@@ -123,12 +182,12 @@ impl<'a> TransportRunner<'a> {
                 ),
                 self.transport.matter().port,
             ),
-            buffers,
+            &mut buffers.udp,
         )
         .await?;
 
-        let tx_pipe = Pipe::new(unsafe { tx_buf.assume_init_mut() });
+        let tx_pipe = Pipe::new(unsafe { buffers.tx_buf.assume_init_mut() });
-        let rx_pipe = Pipe::new(unsafe { rx_buf.assume_init_mut() });
+        let rx_pipe = Pipe::new(unsafe { buffers.rx_buf.assume_init_mut() });
 
         let tx_pipe = &tx_pipe;
         let rx_pipe = &rx_pipe;

matter/src/transport/udp.rs

@@ -16,21 +16,25 @@
  */
 
 #[cfg(all(feature = "std", not(feature = "embassy-net")))]
-pub use smol_udp::*;
+pub use async_io::*;
 
 #[cfg(feature = "embassy-net")]
-pub use embassy_udp::*;
+pub use embassy_net::*;
 
-#[cfg(all(feature = "std", not(feature = "embassy-net")))]
+#[cfg(feature = "std")]
-mod smol_udp {
+pub mod async_io {
     use crate::error::*;
-    use log::{debug, info, warn};
-    use smol::net::UdpSocket;
 
-    use crate::transport::network::{
+    use std::net::UdpSocket;
-        Ipv4Addr, Ipv6Addr, NetworkStack, NetworkStackDriver, NetworkStackMulticastDriver,
+
-        SocketAddr,
+    use async_io::Async;
+
+    use log::{debug, info, warn};
+
+    use crate::transport::network::std_stack::{
+        NetworkStack, NetworkStackDriver, NetworkStackMulticastDriver,
     };
+    use crate::transport::network::{Ipv4Addr, Ipv6Addr, SocketAddr};
 
     pub struct UdpBuffers(());
 
@@ -40,7 +44,7 @@ mod smol_udp {
         }
     }
 
-    pub struct UdpListener<'a, D>(UdpSocket, &'a NetworkStack<D>)
+    pub struct UdpListener<'a, D>(Async<UdpSocket>, &'a NetworkStack<D>)
     where
         D: NetworkStackDriver;
 
@@ -53,7 +57,7 @@ mod smol_udp {
             addr: SocketAddr,
             _buffers: &'a mut UdpBuffers,
         ) -> Result<UdpListener<'a, D>, Error> {
-            let listener = UdpListener(UdpSocket::bind((addr.ip(), addr.port())).await?, stack);
+            let listener = UdpListener(Async::<UdpSocket>::bind(addr)?, stack);
 
             info!("Listening on {:?}", addr);
 
@@ -68,7 +72,7 @@ mod smol_udp {
         where
             D: NetworkStackMulticastDriver + 'static,
         {
-            self.0.join_multicast_v6(&multiaddr, interface)?;
+            self.0.get_ref().join_multicast_v6(&multiaddr, interface)?;
 
             info!("Joined IPV6 multicast {}/{}", multiaddr, interface);
 
@@ -84,7 +88,7 @@ mod smol_udp {
             D: NetworkStackMulticastDriver + 'static,
         {
             #[cfg(not(target_os = "espidf"))]
-            self.0.join_multicast_v4(multiaddr, interface)?;
+            self.0.get_ref().join_multicast_v4(&multiaddr, &interface)?;
 
             // join_multicast_v4() is broken for ESP-IDF, most likely due to wrong `ip_mreq` signature in the `libc` crate
             // Note that also most *_multicast_v4 and *_multicast_v6 methods are broken as well in Rust STD for the ESP-IDF
@@ -166,7 +170,7 @@ mod smol_udp {
 }
 
 #[cfg(feature = "embassy-net")]
-mod embassy_udp {
+pub mod embassy_net {
     use core::mem::MaybeUninit;
 
     use embassy_net::udp::{PacketMetadata, UdpSocket};
@@ -177,10 +181,10 @@ mod embassy_udp {
 
     use log::{debug, info, warn};
 
-    use crate::transport::network::{
+    use crate::transport::network::embassy_net_stack::{
-        IpAddr, Ipv4Addr, Ipv6Addr, NetworkStack, NetworkStackDriver, NetworkStackMulticastDriver,
+        NetworkStack, NetworkStackDriver, NetworkStackMulticastDriver,
-        SocketAddr,
     };
+    use crate::transport::network::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
     const RX_BUF_SIZE: usize = 4096;
     const TX_BUF_SIZE: usize = 4096;