Merge remote-tracking branch 'origin/main' into nrf-pdm

8 files changed, 1269 insertions, 749 deletions

diff --git a/embassy-net/src/device.rs b/embassy-net/src/device.rs
index c183bd58a..4513c86d3 100644
--- a/embassy-net/src/device.rs
+++ b/embassy-net/src/device.rs
@@ -1,129 +1,106 @@
-use core::task::Waker;
-
-use smoltcp::phy::{Device as SmolDevice, DeviceCapabilities};
-use smoltcp::time::Instant as SmolInstant;
-
-use crate::packet_pool::PacketBoxExt;
-use crate::{Packet, PacketBox, PacketBuf};
-
-#[derive(PartialEq, Eq, Clone, Copy)]
-pub enum LinkState {
-    Down,
-    Up,
-}
-
-// 'static required due to the "fake GAT" in smoltcp::phy::Device.
-// https://github.com/smoltcp-rs/smoltcp/pull/572
-pub trait Device {
-    fn is_transmit_ready(&mut self) -> bool;
-    fn transmit(&mut self, pkt: PacketBuf);
-    fn receive(&mut self) -> Option<PacketBuf>;
-
-    fn register_waker(&mut self, waker: &Waker);
-    fn capabilities(&self) -> DeviceCapabilities;
-    fn link_state(&mut self) -> LinkState;
-    fn ethernet_address(&self) -> [u8; 6];
+use core::task::Context;
+
+use embassy_net_driver::{Capabilities, Checksum, Driver, Medium, RxToken, TxToken};
+use smoltcp::phy;
+use smoltcp::time::Instant;
+
+pub(crate) struct DriverAdapter<'d, 'c, T>
+where
+    T: Driver,
+{
+    // must be Some when actually using this to rx/tx
+    pub cx: Option<&'d mut Context<'c>>,
+    pub inner: &'d mut T,
 }
 
-impl<T: ?Sized + Device> Device for &'static mut T {
-    fn is_transmit_ready(&mut self) -> bool {
-        T::is_transmit_ready(self)
-    }
-    fn transmit(&mut self, pkt: PacketBuf) {
-        T::transmit(self, pkt)
-    }
-    fn receive(&mut self) -> Option<PacketBuf> {
-        T::receive(self)
+impl<'d, 'c, T> phy::Device for DriverAdapter<'d, 'c, T>
+where
+    T: Driver,
+{
+    type RxToken<'a> = RxTokenAdapter<T::RxToken<'a>> where Self: 'a;
+    type TxToken<'a> = TxTokenAdapter<T::TxToken<'a>> where Self: 'a;
+
+    fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
+        self.inner
+            .receive(self.cx.as_deref_mut().unwrap())
+            .map(|(rx, tx)| (RxTokenAdapter(rx), TxTokenAdapter(tx)))
     }
-    fn register_waker(&mut self, waker: &Waker) {
-        T::register_waker(self, waker)
-    }
-    fn capabilities(&self) -> DeviceCapabilities {
-        T::capabilities(self)
-    }
-    fn link_state(&mut self) -> LinkState {
-        T::link_state(self)
-    }
-    fn ethernet_address(&self) -> [u8; 6] {
-        T::ethernet_address(self)
-    }
-}
-
-pub struct DeviceAdapter<D: Device> {
-    pub device: D,
-    caps: DeviceCapabilities,
-}
 
-impl<D: Device> DeviceAdapter<D> {
-    pub(crate) fn new(device: D) -> Self {
-        Self {
-            caps: device.capabilities(),
-            device,
-        }
+    /// Construct a transmit token.
+    fn transmit(&mut self, _timestamp: Instant) -> Option<Self::TxToken<'_>> {
+        self.inner.transmit(self.cx.as_deref_mut().unwrap()).map(TxTokenAdapter)
     }
-}
 
-impl<'a, D: Device + 'static> SmolDevice<'a> for DeviceAdapter<D> {
-    type RxToken = RxToken;
-    type TxToken = TxToken<'a, D>;
-
-    fn receive(&'a mut self) -> Option<(Self::RxToken, Self::TxToken)> {
-        let tx_pkt = PacketBox::new(Packet::new())?;
-        let rx_pkt = self.device.receive()?;
-        let rx_token = RxToken { pkt: rx_pkt };
-        let tx_token = TxToken {
-            device: &mut self.device,
-            pkt: tx_pkt,
+    /// Get a description of device capabilities.
+    fn capabilities(&self) -> phy::DeviceCapabilities {
+        fn convert(c: Checksum) -> phy::Checksum {
+            match c {
+                Checksum::Both => phy::Checksum::Both,
+                Checksum::Tx => phy::Checksum::Tx,
+                Checksum::Rx => phy::Checksum::Rx,
+                Checksum::None => phy::Checksum::None,
+            }
+        }
+        let caps: Capabilities = self.inner.capabilities();
+        let mut smolcaps = phy::DeviceCapabilities::default();
+
+        smolcaps.max_transmission_unit = caps.max_transmission_unit;
+        smolcaps.max_burst_size = caps.max_burst_size;
+        smolcaps.medium = match caps.medium {
+            #[cfg(feature = "medium-ethernet")]
+            Medium::Ethernet => phy::Medium::Ethernet,
+            #[cfg(feature = "medium-ip")]
+            Medium::Ip => phy::Medium::Ip,
+            #[allow(unreachable_patterns)]
+            _ => panic!(
+                "Unsupported medium {:?}. Make sure to enable it in embassy-net's Cargo features.",
+                caps.medium
+            ),
         };
-
-        Some((rx_token, tx_token))
-    }
-
-    /// Construct a transmit token.
-    fn transmit(&'a mut self) -> Option<Self::TxToken> {
-        if !self.device.is_transmit_ready() {
-            return None;
+        smolcaps.checksum.ipv4 = convert(caps.checksum.ipv4);
+        smolcaps.checksum.tcp = convert(caps.checksum.tcp);
+        smolcaps.checksum.udp = convert(caps.checksum.udp);
+        #[cfg(feature = "proto-ipv4")]
+        {
+            smolcaps.checksum.icmpv4 = convert(caps.checksum.icmpv4);
+        }
+        #[cfg(feature = "proto-ipv6")]
+        {
+            smolcaps.checksum.icmpv6 = convert(caps.checksum.icmpv6);
         }
 
-        let tx_pkt = PacketBox::new(Packet::new())?;
-        Some(TxToken {
-            device: &mut self.device,
-            pkt: tx_pkt,
-        })
-    }
-
-    /// Get a description of device capabilities.
-    fn capabilities(&self) -> DeviceCapabilities {
-        self.caps.clone()
+        smolcaps
     }
 }
 
-pub struct RxToken {
-    pkt: PacketBuf,
-}
+pub(crate) struct RxTokenAdapter<T>(T)
+where
+    T: RxToken;
 
-impl smoltcp::phy::RxToken for RxToken {
-    fn consume<R, F>(mut self, _timestamp: SmolInstant, f: F) -> smoltcp::Result<R>
+impl<T> phy::RxToken for RxTokenAdapter<T>
+where
+    T: RxToken,
+{
+    fn consume<R, F>(self, f: F) -> R
     where
-        F: FnOnce(&mut [u8]) -> smoltcp::Result<R>,
+        F: FnOnce(&mut [u8]) -> R,
     {
-        f(&mut self.pkt)
+        self.0.consume(|buf| f(buf))
     }
 }
 
-pub struct TxToken<'a, D: Device> {
-    device: &'a mut D,
-    pkt: PacketBox,
-}
+pub(crate) struct TxTokenAdapter<T>(T)
+where
+    T: TxToken;
 
-impl<'a, D: Device> smoltcp::phy::TxToken for TxToken<'a, D> {
-    fn consume<R, F>(self, _timestamp: SmolInstant, len: usize, f: F) -> smoltcp::Result<R>
+impl<T> phy::TxToken for TxTokenAdapter<T>
+where
+    T: TxToken,
+{
+    fn consume<R, F>(self, len: usize, f: F) -> R
     where
-        F: FnOnce(&mut [u8]) -> smoltcp::Result<R>,
+        F: FnOnce(&mut [u8]) -> R,
     {
-        let mut buf = self.pkt.slice(0..len);
-        let r = f(&mut buf)?;
-        self.device.transmit(buf);
-        Ok(r)
+        self.0.consume(len, |buf| f(buf))
     }
 }
diff --git a/embassy-net/src/dns.rs b/embassy-net/src/dns.rs
new file mode 100644
index 000000000..94f75f108
--- /dev/null
+++ b/embassy-net/src/dns.rs
@@ -0,0 +1,107 @@
+//! DNS client compatible with the `embedded-nal-async` traits.
+//!
+//! This exists only for compatibility with crates that use `embedded-nal-async`.
+//! Prefer using [`Stack::dns_query`](crate::Stack::dns_query) directly if you're
+//! not using `embedded-nal-async`.
+
+use heapless::Vec;
+pub use smoltcp::socket::dns::{DnsQuery, Socket};
+pub(crate) use smoltcp::socket::dns::{GetQueryResultError, StartQueryError};
+pub use smoltcp::wire::{DnsQueryType, IpAddress};
+
+use crate::{Driver, Stack};
+
+/// Errors returned by DnsSocket.
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+pub enum Error {
+    /// Invalid name
+    InvalidName,
+    /// Name too long
+    NameTooLong,
+    /// Name lookup failed
+    Failed,
+}
+
+impl From<GetQueryResultError> for Error {
+    fn from(_: GetQueryResultError) -> Self {
+        Self::Failed
+    }
+}
+
+impl From<StartQueryError> for Error {
+    fn from(e: StartQueryError) -> Self {
+        match e {
+            StartQueryError::NoFreeSlot => Self::Failed,
+            StartQueryError::InvalidName => Self::InvalidName,
+            StartQueryError::NameTooLong => Self::NameTooLong,
+        }
+    }
+}
+
+/// DNS client compatible with the `embedded-nal-async` traits.
+///
+/// This exists only for compatibility with crates that use `embedded-nal-async`.
+/// Prefer using [`Stack::dns_query`](crate::Stack::dns_query) directly if you're
+/// not using `embedded-nal-async`.
+pub struct DnsSocket<'a, D>
+where
+    D: Driver + 'static,
+{
+    stack: &'a Stack<D>,
+}
+
+impl<'a, D> DnsSocket<'a, D>
+where
+    D: Driver + 'static,
+{
+    /// Create a new DNS socket using the provided stack.
+    ///
+    /// NOTE: If using DHCP, make sure it has reconfigured the stack to ensure the DNS servers are updated.
+    pub fn new(stack: &'a Stack<D>) -> Self {
+        Self { stack }
+    }
+
+    /// Make a query for a given name and return the corresponding IP addresses.
+    pub async fn query(&self, name: &str, qtype: DnsQueryType) -> Result<Vec<IpAddress, 1>, Error> {
+        self.stack.dns_query(name, qtype).await
+    }
+}
+
+#[cfg(all(feature = "unstable-traits", feature = "nightly"))]
+impl<'a, D> embedded_nal_async::Dns for DnsSocket<'a, D>
+where
+    D: Driver + 'static,
+{
+    type Error = Error;
+
+    async fn get_host_by_name(
+        &self,
+        host: &str,
+        addr_type: embedded_nal_async::AddrType,
+    ) -> Result<embedded_nal_async::IpAddr, Self::Error> {
+        use embedded_nal_async::{AddrType, IpAddr};
+        let qtype = match addr_type {
+            AddrType::IPv6 => DnsQueryType::Aaaa,
+            _ => DnsQueryType::A,
+        };
+        let addrs = self.query(host, qtype).await?;
+        if let Some(first) = addrs.get(0) {
+            Ok(match first {
+                #[cfg(feature = "proto-ipv4")]
+                IpAddress::Ipv4(addr) => IpAddr::V4(addr.0.into()),
+                #[cfg(feature = "proto-ipv6")]
+                IpAddress::Ipv6(addr) => IpAddr::V6(addr.0.into()),
+            })
+        } else {
+            Err(Error::Failed)
+        }
+    }
+
+    async fn get_host_by_address(
+        &self,
+        _addr: embedded_nal_async::IpAddr,
+    ) -> Result<heapless::String<256>, Self::Error> {
+        todo!()
+    }
+}
diff --git a/embassy-net/src/lib.rs b/embassy-net/src/lib.rs
index 83d364715..0d0a986f6 100644
--- a/embassy-net/src/lib.rs
+++ b/embassy-net/src/lib.rs
@@ -1,31 +1,726 @@
 #![cfg_attr(not(feature = "std"), no_std)]
-#![allow(clippy::new_without_default)]
-#![feature(generic_associated_types, type_alias_impl_trait)]
+#![cfg_attr(feature = "nightly", feature(async_fn_in_trait, impl_trait_projections))]
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
 
 // This mod MUST go first, so that the others see its macros.
 pub(crate) mod fmt;
 
 mod device;
-mod packet_pool;
-mod stack;
-
-pub use device::{Device, LinkState};
-pub use packet_pool::{Packet, PacketBox, PacketBoxExt, PacketBuf, MTU};
-pub use stack::{Config, ConfigStrategy, Stack, StackResources};
-
+#[cfg(feature = "dns")]
+pub mod dns;
 #[cfg(feature = "tcp")]
 pub mod tcp;
-
+mod time;
 #[cfg(feature = "udp")]
 pub mod udp;
 
-// smoltcp reexports
-pub use smoltcp::phy::{DeviceCapabilities, Medium};
-pub use smoltcp::time::{Duration as SmolDuration, Instant as SmolInstant};
+use core::cell::RefCell;
+use core::future::{poll_fn, Future};
+use core::task::{Context, Poll};
+
+pub use embassy_net_driver as driver;
+use embassy_net_driver::{Driver, LinkState, Medium};
+use embassy_sync::waitqueue::WakerRegistration;
+use embassy_time::{Instant, Timer};
+use futures::pin_mut;
+use heapless::Vec;
+#[cfg(feature = "igmp")]
+pub use smoltcp::iface::MulticastError;
+use smoltcp::iface::{Interface, SocketHandle, SocketSet, SocketStorage};
+#[cfg(feature = "dhcpv4")]
+use smoltcp::socket::dhcpv4::{self, RetryConfig};
+#[cfg(feature = "udp")]
+pub use smoltcp::wire::IpListenEndpoint;
 #[cfg(feature = "medium-ethernet")]
 pub use smoltcp::wire::{EthernetAddress, HardwareAddress};
-pub use smoltcp::wire::{IpAddress, IpCidr, Ipv4Address, Ipv4Cidr};
+pub use smoltcp::wire::{IpAddress, IpCidr, IpEndpoint};
+#[cfg(feature = "proto-ipv4")]
+pub use smoltcp::wire::{Ipv4Address, Ipv4Cidr};
 #[cfg(feature = "proto-ipv6")]
 pub use smoltcp::wire::{Ipv6Address, Ipv6Cidr};
-#[cfg(feature = "udp")]
-pub use smoltcp::{socket::udp::PacketMetadata, wire::IpListenEndpoint};
+
+use crate::device::DriverAdapter;
+use crate::time::{instant_from_smoltcp, instant_to_smoltcp};
+
+const LOCAL_PORT_MIN: u16 = 1025;
+const LOCAL_PORT_MAX: u16 = 65535;
+#[cfg(feature = "dns")]
+const MAX_QUERIES: usize = 4;
+
+/// Memory resources needed for a network stack.
+pub struct StackResources<const SOCK: usize> {
+    sockets: [SocketStorage<'static>; SOCK],
+    #[cfg(feature = "dns")]
+    queries: [Option<dns::DnsQuery>; MAX_QUERIES],
+}
+
+impl<const SOCK: usize> StackResources<SOCK> {
+    /// Create a new set of stack resources.
+    pub const fn new() -> Self {
+        #[cfg(feature = "dns")]
+        const INIT: Option<dns::DnsQuery> = None;
+        Self {
+            sockets: [SocketStorage::EMPTY; SOCK],
+            #[cfg(feature = "dns")]
+            queries: [INIT; MAX_QUERIES],
+        }
+    }
+}
+
+/// Static IP address configuration.
+#[cfg(feature = "proto-ipv4")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StaticConfigV4 {
+    /// IP address and subnet mask.
+    pub address: Ipv4Cidr,
+    /// Default gateway.
+    pub gateway: Option<Ipv4Address>,
+    /// DNS servers.
+    pub dns_servers: Vec<Ipv4Address, 3>,
+}
+
+/// Static IPv6 address configuration
+#[cfg(feature = "proto-ipv6")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StaticConfigV6 {
+    /// IP address and subnet mask.
+    pub address: Ipv6Cidr,
+    /// Default gateway.
+    pub gateway: Option<Ipv6Address>,
+    /// DNS servers.
+    pub dns_servers: Vec<Ipv6Address, 3>,
+}
+
+/// DHCP configuration.
+#[cfg(feature = "dhcpv4")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DhcpConfig {
+    /// Maximum lease duration.
+    ///
+    /// If not set, the lease duration specified by the server will be used.
+    /// If set, the lease duration will be capped at this value.
+    pub max_lease_duration: Option<embassy_time::Duration>,
+    /// Retry configuration.
+    pub retry_config: RetryConfig,
+    /// Ignore NAKs from DHCP servers.
+    ///
+    /// This is not compliant with the DHCP RFCs, since theoretically we must stop using the assigned IP when receiving a NAK. This can increase reliability on broken networks with buggy routers or rogue DHCP servers, however.
+    pub ignore_naks: bool,
+    /// Server port. This is almost always 67. Do not change unless you know what you're doing.
+    pub server_port: u16,
+    /// Client port. This is almost always 68. Do not change unless you know what you're doing.
+    pub client_port: u16,
+}
+
+#[cfg(feature = "dhcpv4")]
+impl Default for DhcpConfig {
+    fn default() -> Self {
+        Self {
+            max_lease_duration: Default::default(),
+            retry_config: Default::default(),
+            ignore_naks: Default::default(),
+            server_port: smoltcp::wire::DHCP_SERVER_PORT,
+            client_port: smoltcp::wire::DHCP_CLIENT_PORT,
+        }
+    }
+}
+
+/// Network stack configuration.
+pub struct Config {
+    /// IPv4 configuration
+    #[cfg(feature = "proto-ipv4")]
+    pub ipv4: ConfigV4,
+    /// IPv6 configuration
+    #[cfg(feature = "proto-ipv6")]
+    pub ipv6: ConfigV6,
+}
+
+impl Config {
+    /// IPv4 configuration with static addressing.
+    #[cfg(feature = "proto-ipv4")]
+    pub fn ipv4_static(config: StaticConfigV4) -> Self {
+        Self {
+            ipv4: ConfigV4::Static(config),
+            #[cfg(feature = "proto-ipv6")]
+            ipv6: ConfigV6::None,
+        }
+    }
+
+    /// IPv6 configuration with static addressing.
+    #[cfg(feature = "proto-ipv6")]
+    pub fn ipv6_static(config: StaticConfigV6) -> Self {
+        Self {
+            #[cfg(feature = "proto-ipv4")]
+            ipv4: ConfigV4::None,
+            ipv6: ConfigV6::Static(config),
+        }
+    }
+
+    /// IPv6 configuration with dynamic addressing.
+    ///
+    /// # Example
+    /// ```rust
+    /// let _cfg = Config::dhcpv4(Default::default());
+    /// ```
+    #[cfg(feature = "dhcpv4")]
+    pub fn dhcpv4(config: DhcpConfig) -> Self {
+        Self {
+            ipv4: ConfigV4::Dhcp(config),
+            #[cfg(feature = "proto-ipv6")]
+            ipv6: ConfigV6::None,
+        }
+    }
+}
+
+/// Network stack IPv4 configuration.
+#[cfg(feature = "proto-ipv4")]
+pub enum ConfigV4 {
+    /// Use a static IPv4 address configuration.
+    Static(StaticConfigV4),
+    /// Use DHCP to obtain an IP address configuration.
+    #[cfg(feature = "dhcpv4")]
+    Dhcp(DhcpConfig),
+    /// Do not configure IPv6.
+    None,
+}
+
+/// Network stack IPv6 configuration.
+#[cfg(feature = "proto-ipv6")]
+pub enum ConfigV6 {
+    /// Use a static IPv6 address configuration.
+    Static(StaticConfigV6),
+    /// Do not configure IPv6.
+    None,
+}
+
+/// A network stack.
+///
+/// This is the main entry point for the network stack.
+pub struct Stack<D: Driver> {
+    pub(crate) socket: RefCell<SocketStack>,
+    inner: RefCell<Inner<D>>,
+}
+
+struct Inner<D: Driver> {
+    device: D,
+    link_up: bool,
+    #[cfg(feature = "proto-ipv4")]
+    static_v4: Option<StaticConfigV4>,
+    #[cfg(feature = "proto-ipv6")]
+    static_v6: Option<StaticConfigV6>,
+    #[cfg(feature = "dhcpv4")]
+    dhcp_socket: Option<SocketHandle>,
+    #[cfg(feature = "dns")]
+    dns_socket: SocketHandle,
+    #[cfg(feature = "dns")]
+    dns_waker: WakerRegistration,
+}
+
+pub(crate) struct SocketStack {
+    pub(crate) sockets: SocketSet<'static>,
+    pub(crate) iface: Interface,
+    pub(crate) waker: WakerRegistration,
+    next_local_port: u16,
+}
+
+impl<D: Driver + 'static> Stack<D> {
+    /// Create a new network stack.
+    pub fn new<const SOCK: usize>(
+        mut device: D,
+        config: Config,
+        resources: &'static mut StackResources<SOCK>,
+        random_seed: u64,
+    ) -> Self {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = device.capabilities().medium;
+
+        let hardware_addr = match medium {
+            #[cfg(feature = "medium-ethernet")]
+            Medium::Ethernet => HardwareAddress::Ethernet(EthernetAddress(device.ethernet_address())),
+            #[cfg(feature = "medium-ip")]
+            Medium::Ip => HardwareAddress::Ip,
+            #[allow(unreachable_patterns)]
+            _ => panic!(
+                "Unsupported medium {:?}. Make sure to enable it in embassy-net's Cargo features.",
+                medium
+            ),
+        };
+        let mut iface_cfg = smoltcp::iface::Config::new(hardware_addr);
+        iface_cfg.random_seed = random_seed;
+
+        let iface = Interface::new(
+            iface_cfg,
+            &mut DriverAdapter {
+                inner: &mut device,
+                cx: None,
+            },
+            instant_to_smoltcp(Instant::now()),
+        );
+
+        let sockets = SocketSet::new(&mut resources.sockets[..]);
+
+        let next_local_port = (random_seed % (LOCAL_PORT_MAX - LOCAL_PORT_MIN) as u64) as u16 + LOCAL_PORT_MIN;
+
+        let mut socket = SocketStack {
+            sockets,
+            iface,
+            waker: WakerRegistration::new(),
+            next_local_port,
+        };
+
+        let mut inner = Inner {
+            device,
+            link_up: false,
+            #[cfg(feature = "proto-ipv4")]
+            static_v4: None,
+            #[cfg(feature = "proto-ipv6")]
+            static_v6: None,
+            #[cfg(feature = "dhcpv4")]
+            dhcp_socket: None,
+            #[cfg(feature = "dns")]
+            dns_socket: socket.sockets.add(dns::Socket::new(
+                &[],
+                managed::ManagedSlice::Borrowed(&mut resources.queries),
+            )),
+            #[cfg(feature = "dns")]
+            dns_waker: WakerRegistration::new(),
+        };
+
+        #[cfg(feature = "proto-ipv4")]
+        match config.ipv4 {
+            ConfigV4::Static(config) => {
+                inner.apply_config_v4(&mut socket, config);
+            }
+            #[cfg(feature = "dhcpv4")]
+            ConfigV4::Dhcp(config) => {
+                let mut dhcp_socket = smoltcp::socket::dhcpv4::Socket::new();
+                inner.apply_dhcp_config(&mut dhcp_socket, config);
+                let handle = socket.sockets.add(dhcp_socket);
+                inner.dhcp_socket = Some(handle);
+            }
+            ConfigV4::None => {}
+        }
+        #[cfg(feature = "proto-ipv6")]
+        match config.ipv6 {
+            ConfigV6::Static(config) => {
+                inner.apply_config_v6(&mut socket, config);
+            }
+            ConfigV6::None => {}
+        }
+
+        Self {
+            socket: RefCell::new(socket),
+            inner: RefCell::new(inner),
+        }
+    }
+
+    fn with<R>(&self, f: impl FnOnce(&SocketStack, &Inner<D>) -> R) -> R {
+        f(&*self.socket.borrow(), &*self.inner.borrow())
+    }
+
+    fn with_mut<R>(&self, f: impl FnOnce(&mut SocketStack, &mut Inner<D>) -> R) -> R {
+        f(&mut *self.socket.borrow_mut(), &mut *self.inner.borrow_mut())
+    }
+
+    /// Get the MAC address of the network interface.
+    pub fn ethernet_address(&self) -> [u8; 6] {
+        self.with(|_s, i| i.device.ethernet_address())
+    }
+
+    /// Get whether the link is up.
+    pub fn is_link_up(&self) -> bool {
+        self.with(|_s, i| i.link_up)
+    }
+
+    /// Get whether the network stack has a valid IP configuration.
+    /// This is true if the network stack has a static IP configuration or if DHCP has completed
+    pub fn is_config_up(&self) -> bool {
+        let v4_up;
+        let v6_up;
+
+        #[cfg(feature = "proto-ipv4")]
+        {
+            v4_up = self.config_v4().is_some();
+        }
+        #[cfg(not(feature = "proto-ipv4"))]
+        {
+            v4_up = false;
+        }
+
+        #[cfg(feature = "proto-ipv6")]
+        {
+            v6_up = self.config_v6().is_some();
+        }
+        #[cfg(not(feature = "proto-ipv6"))]
+        {
+            v6_up = false;
+        }
+
+        v4_up || v6_up
+    }
+
+    /// Get the current IPv4 configuration.
+    #[cfg(feature = "proto-ipv4")]
+    pub fn config_v4(&self) -> Option<StaticConfigV4> {
+        self.with(|_s, i| i.static_v4.clone())
+    }
+
+    /// Get the current IPv6 configuration.
+    #[cfg(feature = "proto-ipv6")]
+    pub fn config_v6(&self) -> Option<StaticConfigV6> {
+        self.with(|_s, i| i.static_v6.clone())
+    }
+
+    /// Run the network stack.
+    ///
+    /// You must call this in a background task, to process network events.
+    pub async fn run(&self) -> ! {
+        poll_fn(|cx| {
+            self.with_mut(|s, i| i.poll(cx, s));
+            Poll::<()>::Pending
+        })
+        .await;
+        unreachable!()
+    }
+
+    /// Make a query for a given name and return the corresponding IP addresses.
+    #[cfg(feature = "dns")]
+    pub async fn dns_query(&self, name: &str, qtype: dns::DnsQueryType) -> Result<Vec<IpAddress, 1>, dns::Error> {
+        // For A and AAAA queries we try detect whether `name` is just an IP address
+        match qtype {
+            #[cfg(feature = "proto-ipv4")]
+            dns::DnsQueryType::A => {
+                if let Ok(ip) = name.parse().map(IpAddress::Ipv4) {
+                    return Ok([ip].into_iter().collect());
+                }
+            }
+            #[cfg(feature = "proto-ipv6")]
+            dns::DnsQueryType::Aaaa => {
+                if let Ok(ip) = name.parse().map(IpAddress::Ipv6) {
+                    return Ok([ip].into_iter().collect());
+                }
+            }
+            _ => {}
+        }
+
+        let query = poll_fn(|cx| {
+            self.with_mut(|s, i| {
+                let socket = s.sockets.get_mut::<dns::Socket>(i.dns_socket);
+                match socket.start_query(s.iface.context(), name, qtype) {
+                    Ok(handle) => Poll::Ready(Ok(handle)),
+                    Err(dns::StartQueryError::NoFreeSlot) => {
+                        i.dns_waker.register(cx.waker());
+                        Poll::Pending
+                    }
+                    Err(e) => Poll::Ready(Err(e)),
+                }
+            })
+        })
+        .await?;
+
+        #[must_use = "to delay the drop handler invocation to the end of the scope"]
+        struct OnDrop<F: FnOnce()> {
+            f: core::mem::MaybeUninit<F>,
+        }
+
+        impl<F: FnOnce()> OnDrop<F> {
+            fn new(f: F) -> Self {
+                Self {
+                    f: core::mem::MaybeUninit::new(f),
+                }
+            }
+
+            fn defuse(self) {
+                core::mem::forget(self)
+            }
+        }
+
+        impl<F: FnOnce()> Drop for OnDrop<F> {
+            fn drop(&mut self) {
+                unsafe { self.f.as_ptr().read()() }
+            }
+        }
+
+        let drop = OnDrop::new(|| {
+            self.with_mut(|s, i| {
+                let socket = s.sockets.get_mut::<dns::Socket>(i.dns_socket);
+                socket.cancel_query(query);
+                s.waker.wake();
+                i.dns_waker.wake();
+            })
+        });
+
+        let res = poll_fn(|cx| {
+            self.with_mut(|s, i| {
+                let socket = s.sockets.get_mut::<dns::Socket>(i.dns_socket);
+                match socket.get_query_result(query) {
+                    Ok(addrs) => {
+                        i.dns_waker.wake();
+                        Poll::Ready(Ok(addrs))
+                    }
+                    Err(dns::GetQueryResultError::Pending) => {
+                        socket.register_query_waker(query, cx.waker());
+                        Poll::Pending
+                    }
+                    Err(e) => {
+                        i.dns_waker.wake();
+                        Poll::Ready(Err(e.into()))
+                    }
+                }
+            })
+        })
+        .await;
+
+        drop.defuse();
+
+        res
+    }
+}
+
+#[cfg(feature = "igmp")]
+impl<D: Driver + smoltcp::phy::Device + 'static> Stack<D> {
+    /// Join a multicast group.
+    pub fn join_multicast_group<T>(&self, addr: T) -> Result<bool, MulticastError>
+    where
+        T: Into<IpAddress>,
+    {
+        let addr = addr.into();
+
+        self.with_mut(|s, i| {
+            s.iface
+                .join_multicast_group(&mut i.device, addr, instant_to_smoltcp(Instant::now()))
+        })
+    }
+
+    /// Leave a multicast group.
+    pub fn leave_multicast_group<T>(&self, addr: T) -> Result<bool, MulticastError>
+    where
+        T: Into<IpAddress>,
+    {
+        let addr = addr.into();
+
+        self.with_mut(|s, i| {
+            s.iface
+                .leave_multicast_group(&mut i.device, addr, instant_to_smoltcp(Instant::now()))
+        })
+    }
+
+    /// Get whether the network stack has joined the given multicast group.
+    pub fn has_multicast_group<T: Into<IpAddress>>(&self, addr: T) -> bool {
+        self.socket.borrow().iface.has_multicast_group(addr)
+    }
+}
+
+impl SocketStack {
+    #[allow(clippy::absurd_extreme_comparisons, dead_code)]
+    pub fn get_local_port(&mut self) -> u16 {
+        let res = self.next_local_port;
+        self.next_local_port = if res >= LOCAL_PORT_MAX { LOCAL_PORT_MIN } else { res + 1 };
+        res
+    }
+}
+
+impl<D: Driver + 'static> Inner<D> {
+    #[cfg(feature = "proto-ipv4")]
+    fn apply_config_v4(&mut self, s: &mut SocketStack, config: StaticConfigV4) {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = self.device.capabilities().medium;
+
+        debug!("Acquired IP configuration:");
+
+        debug!("   IP address:      {}", config.address);
+        s.iface.update_ip_addrs(|addrs| {
+            if addrs.is_empty() {
+                addrs.push(IpCidr::Ipv4(config.address)).unwrap();
+            } else {
+                addrs[0] = IpCidr::Ipv4(config.address);
+            }
+        });
+
+        #[cfg(feature = "medium-ethernet")]
+        if medium == Medium::Ethernet {
+            if let Some(gateway) = config.gateway {
+                debug!("   Default gateway: {}", gateway);
+                s.iface.routes_mut().add_default_ipv4_route(gateway).unwrap();
+            } else {
+                debug!("   Default gateway: None");
+                s.iface.routes_mut().remove_default_ipv4_route();
+            }
+        }
+        for (i, s) in config.dns_servers.iter().enumerate() {
+            debug!("   DNS server {}:    {}", i, s);
+        }
+
+        self.static_v4 = Some(config);
+
+        #[cfg(feature = "dns")]
+        {
+            self.update_dns_servers(s)
+        }
+    }
+
+    /// Replaces the current IPv6 static configuration with a newly supplied config.
+    #[cfg(feature = "proto-ipv6")]
+    fn apply_config_v6(&mut self, s: &mut SocketStack, config: StaticConfigV6) {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = self.device.capabilities().medium;
+
+        debug!("Acquired IPv6 configuration:");
+
+        debug!("   IP address:      {}", config.address);
+        s.iface.update_ip_addrs(|addrs| {
+            if addrs.is_empty() {
+                addrs.push(IpCidr::Ipv6(config.address)).unwrap();
+            } else {
+                addrs[0] = IpCidr::Ipv6(config.address);
+            }
+        });
+
+        #[cfg(feature = "medium-ethernet")]
+        if Medium::Ethernet == medium {
+            if let Some(gateway) = config.gateway {
+                debug!("   Default gateway: {}", gateway);
+                s.iface.routes_mut().add_default_ipv6_route(gateway).unwrap();
+            } else {
+                debug!("   Default gateway: None");
+                s.iface.routes_mut().remove_default_ipv6_route();
+            }
+        }
+        for (i, s) in config.dns_servers.iter().enumerate() {
+            debug!("   DNS server {}:    {}", i, s);
+        }
+
+        self.static_v6 = Some(config);
+
+        #[cfg(feature = "dns")]
+        {
+            self.update_dns_servers(s)
+        }
+    }
+
+    #[cfg(feature = "dns")]
+    fn update_dns_servers(&mut self, s: &mut SocketStack) {
+        let socket = s.sockets.get_mut::<smoltcp::socket::dns::Socket>(self.dns_socket);
+
+        let servers_v4;
+        #[cfg(feature = "proto-ipv4")]
+        {
+            servers_v4 = self
+                .static_v4
+                .iter()
+                .flat_map(|cfg| cfg.dns_servers.iter().map(|c| IpAddress::Ipv4(*c)));
+        };
+        #[cfg(not(feature = "proto-ipv4"))]
+        {
+            servers_v4 = core::iter::empty();
+        }
+
+        let servers_v6;
+        #[cfg(feature = "proto-ipv6")]
+        {
+            servers_v6 = self
+                .static_v6
+                .iter()
+                .flat_map(|cfg| cfg.dns_servers.iter().map(|c| IpAddress::Ipv6(*c)));
+        }
+        #[cfg(not(feature = "proto-ipv6"))]
+        {
+            servers_v6 = core::iter::empty();
+        }
+
+        // Prefer the v6 DNS servers over the v4 servers
+        let servers: Vec<IpAddress, 6> = servers_v6.chain(servers_v4).collect();
+        socket.update_servers(&servers[..]);
+    }
+
+    #[cfg(feature = "dhcpv4")]
+    fn apply_dhcp_config(&self, socket: &mut smoltcp::socket::dhcpv4::Socket, config: DhcpConfig) {
+        socket.set_ignore_naks(config.ignore_naks);
+        socket.set_max_lease_duration(config.max_lease_duration.map(crate::time::duration_to_smoltcp));
+        socket.set_ports(config.server_port, config.client_port);
+        socket.set_retry_config(config.retry_config);
+    }
+
+    #[allow(unused)] // used only with dhcp
+    fn unapply_config(&mut self, s: &mut SocketStack) {
+        #[cfg(feature = "medium-ethernet")]
+        let medium = self.device.capabilities().medium;
+
+        debug!("Lost IP configuration");
+        s.iface.update_ip_addrs(|ip_addrs| ip_addrs.clear());
+        #[cfg(feature = "medium-ethernet")]
+        if medium == Medium::Ethernet {
+            #[cfg(feature = "proto-ipv4")]
+            {
+                s.iface.routes_mut().remove_default_ipv4_route();
+            }
+        }
+        #[cfg(feature = "proto-ipv4")]
+        {
+            self.static_v4 = None
+        }
+    }
+
+    fn poll(&mut self, cx: &mut Context<'_>, s: &mut SocketStack) {
+        s.waker.register(cx.waker());
+
+        #[cfg(feature = "medium-ethernet")]
+        if self.device.capabilities().medium == Medium::Ethernet {
+            s.iface.set_hardware_addr(HardwareAddress::Ethernet(EthernetAddress(
+                self.device.ethernet_address(),
+            )));
+        }
+
+        let timestamp = instant_to_smoltcp(Instant::now());
+        let mut smoldev = DriverAdapter {
+            cx: Some(cx),
+            inner: &mut self.device,
+        };
+        s.iface.poll(timestamp, &mut smoldev, &mut s.sockets);
+
+        // Update link up
+        let old_link_up = self.link_up;
+        self.link_up = self.device.link_state(cx) == LinkState::Up;
+
+        // Print when changed
+        if old_link_up != self.link_up {
+            info!("link_up = {:?}", self.link_up);
+        }
+
+        #[cfg(feature = "dhcpv4")]
+        if let Some(dhcp_handle) = self.dhcp_socket {
+            let socket = s.sockets.get_mut::<dhcpv4::Socket>(dhcp_handle);
+
+            if self.link_up {
+                match socket.poll() {
+                    None => {}
+                    Some(dhcpv4::Event::Deconfigured) => self.unapply_config(s),
+                    Some(dhcpv4::Event::Configured(config)) => {
+                        let config = StaticConfigV4 {
+                            address: config.address,
+                            gateway: config.router,
+                            dns_servers: config.dns_servers,
+                        };
+                        self.apply_config_v4(s, config)
+                    }
+                }
+            } else if old_link_up {
+                socket.reset();
+                self.unapply_config(s);
+            }
+        }
+        //if old_link_up || self.link_up {
+        //    self.poll_configurator(timestamp)
+        //}
+        //
+
+        if let Some(poll_at) = s.iface.poll_at(timestamp, &mut s.sockets) {
+            let t = Timer::at(instant_from_smoltcp(poll_at));
+            pin_mut!(t);
+            if t.poll(cx).is_ready() {
+                cx.waker().wake_by_ref();
+            }
+        }
+    }
+}
diff --git a/embassy-net/src/packet_pool.rs b/embassy-net/src/packet_pool.rs
deleted file mode 100644
index cb8a1316c..000000000
--- a/embassy-net/src/packet_pool.rs
+++ /dev/null
@@ -1,107 +0,0 @@
-use core::ops::{Deref, DerefMut, Range};
-
-use as_slice::{AsMutSlice, AsSlice};
-use atomic_pool::{pool, Box};
-
-pub const MTU: usize = 1516;
-
-#[cfg(feature = "pool-4")]
-pub const PACKET_POOL_SIZE: usize = 4;
-
-#[cfg(feature = "pool-8")]
-pub const PACKET_POOL_SIZE: usize = 8;
-
-#[cfg(feature = "pool-16")]
-pub const PACKET_POOL_SIZE: usize = 16;
-
-#[cfg(feature = "pool-32")]
-pub const PACKET_POOL_SIZE: usize = 32;
-
-#[cfg(feature = "pool-64")]
-pub const PACKET_POOL_SIZE: usize = 64;
-
-#[cfg(feature = "pool-128")]
-pub const PACKET_POOL_SIZE: usize = 128;
-
-pool!(pub PacketPool: [Packet; PACKET_POOL_SIZE]);
-pub type PacketBox = Box<PacketPool>;
-
-#[repr(align(4))]
-pub struct Packet(pub [u8; MTU]);
-
-impl Packet {
-    pub const fn new() -> Self {
-        Self([0; MTU])
-    }
-}
-
-pub trait PacketBoxExt {
-    fn slice(self, range: Range<usize>) -> PacketBuf;
-}
-
-impl PacketBoxExt for PacketBox {
-    fn slice(self, range: Range<usize>) -> PacketBuf {
-        PacketBuf { packet: self, range }
-    }
-}
-
-impl AsSlice for Packet {
-    type Element = u8;
-
-    fn as_slice(&self) -> &[Self::Element] {
-        &self.deref()[..]
-    }
-}
-
-impl AsMutSlice for Packet {
-    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
-        &mut self.deref_mut()[..]
-    }
-}
-
-impl Deref for Packet {
-    type Target = [u8; MTU];
-
-    fn deref(&self) -> &[u8; MTU] {
-        &self.0
-    }
-}
-
-impl DerefMut for Packet {
-    fn deref_mut(&mut self) -> &mut [u8; MTU] {
-        &mut self.0
-    }
-}
-
-pub struct PacketBuf {
-    packet: PacketBox,
-    range: Range<usize>,
-}
-
-impl AsSlice for PacketBuf {
-    type Element = u8;
-
-    fn as_slice(&self) -> &[Self::Element] {
-        &self.packet[self.range.clone()]
-    }
-}
-
-impl AsMutSlice for PacketBuf {
-    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
-        &mut self.packet[self.range.clone()]
-    }
-}
-
-impl Deref for PacketBuf {
-    type Target = [u8];
-
-    fn deref(&self) -> &[u8] {
-        &self.packet[self.range.clone()]
-    }
-}
-
-impl DerefMut for PacketBuf {
-    fn deref_mut(&mut self) -> &mut [u8] {
-        &mut self.packet[self.range.clone()]
-    }
-}
diff --git a/embassy-net/src/stack.rs b/embassy-net/src/stack.rs
deleted file mode 100644
index 8d2dd4bca..000000000
--- a/embassy-net/src/stack.rs
+++ /dev/null
@@ -1,316 +0,0 @@
-use core::cell::UnsafeCell;
-use core::future::Future;
-use core::task::{Context, Poll};
-
-use embassy_sync::waitqueue::WakerRegistration;
-use embassy_time::{Instant, Timer};
-use futures::future::poll_fn;
-use futures::pin_mut;
-use heapless::Vec;
-#[cfg(feature = "dhcpv4")]
-use smoltcp::iface::SocketHandle;
-use smoltcp::iface::{Interface, InterfaceBuilder, SocketSet, SocketStorage};
-#[cfg(feature = "medium-ethernet")]
-use smoltcp::iface::{Neighbor, NeighborCache, Route, Routes};
-#[cfg(feature = "medium-ethernet")]
-use smoltcp::phy::{Device as _, Medium};
-#[cfg(feature = "dhcpv4")]
-use smoltcp::socket::dhcpv4;
-use smoltcp::time::Instant as SmolInstant;
-#[cfg(feature = "medium-ethernet")]
-use smoltcp::wire::{EthernetAddress, HardwareAddress, IpAddress};
-use smoltcp::wire::{IpCidr, Ipv4Address, Ipv4Cidr};
-
-use crate::device::{Device, DeviceAdapter, LinkState};
-
-const LOCAL_PORT_MIN: u16 = 1025;
-const LOCAL_PORT_MAX: u16 = 65535;
-
-pub struct StackResources<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> {
-    addresses: [IpCidr; ADDR],
-    sockets: [SocketStorage<'static>; SOCK],
-
-    #[cfg(feature = "medium-ethernet")]
-    routes: [Option<(IpCidr, Route)>; 1],
-    #[cfg(feature = "medium-ethernet")]
-    neighbor_cache: [Option<(IpAddress, Neighbor)>; NEIGHBOR],
-}
-
-impl<const ADDR: usize, const SOCK: usize, const NEIGHBOR: usize> StackResources<ADDR, SOCK, NEIGHBOR> {
-    pub fn new() -> Self {
-        Self {
-            addresses: [IpCidr::new(Ipv4Address::UNSPECIFIED.into(), 32); ADDR],
-            sockets: [SocketStorage::EMPTY; SOCK],
-            #[cfg(feature = "medium-ethernet")]
-            routes: [None; 1],
-            #[cfg(feature = "medium-ethernet")]
-            neighbor_cache: [None; NEIGHBOR],
-        }
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct Config {
-    pub address: Ipv4Cidr,
-    pub gateway: Option<Ipv4Address>,
-    pub dns_servers: Vec<Ipv4Address, 3>,
-}
-
-pub enum ConfigStrategy {
-    Static(Config),
-    #[cfg(feature = "dhcpv4")]
-    Dhcp,
-}
-
-pub struct Stack<D: Device> {
-    pub(crate) socket: UnsafeCell<SocketStack>,
-    inner: UnsafeCell<Inner<D>>,
-}
-
-struct Inner<D: Device> {
-    device: DeviceAdapter<D>,
-    link_up: bool,
-    config: Option<Config>,
-    #[cfg(feature = "dhcpv4")]
-    dhcp_socket: Option<SocketHandle>,
-}
-
-pub(crate) struct SocketStack {
-    pub(crate) sockets: SocketSet<'static>,
-    pub(crate) iface: Interface<'static>,
-    pub(crate) waker: WakerRegistration,
-    next_local_port: u16,
-}
-
-unsafe impl<D: Device> Send for Stack<D> {}
-
-impl<D: Device + 'static> Stack<D> {
-    pub fn new<const ADDR: usize, const SOCK: usize, const NEIGH: usize>(
-        device: D,
-        config: ConfigStrategy,
-        resources: &'static mut StackResources<ADDR, SOCK, NEIGH>,
-        random_seed: u64,
-    ) -> Self {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = device.capabilities().medium;
-
-        #[cfg(feature = "medium-ethernet")]
-        let ethernet_addr = if medium == Medium::Ethernet {
-            device.ethernet_address()
-        } else {
-            [0, 0, 0, 0, 0, 0]
-        };
-
-        let mut device = DeviceAdapter::new(device);
-
-        let mut b = InterfaceBuilder::new();
-        b = b.ip_addrs(&mut resources.addresses[..]);
-        b = b.random_seed(random_seed);
-
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            b = b.hardware_addr(HardwareAddress::Ethernet(EthernetAddress(ethernet_addr)));
-            b = b.neighbor_cache(NeighborCache::new(&mut resources.neighbor_cache[..]));
-            b = b.routes(Routes::new(&mut resources.routes[..]));
-        }
-
-        let iface = b.finalize(&mut device);
-
-        let sockets = SocketSet::new(&mut resources.sockets[..]);
-
-        let next_local_port = (random_seed % (LOCAL_PORT_MAX - LOCAL_PORT_MIN) as u64) as u16 + LOCAL_PORT_MIN;
-
-        let mut inner = Inner {
-            device,
-            link_up: false,
-            config: None,
-            #[cfg(feature = "dhcpv4")]
-            dhcp_socket: None,
-        };
-        let mut socket = SocketStack {
-            sockets,
-            iface,
-            waker: WakerRegistration::new(),
-            next_local_port,
-        };
-
-        match config {
-            ConfigStrategy::Static(config) => inner.apply_config(&mut socket, config),
-            #[cfg(feature = "dhcpv4")]
-            ConfigStrategy::Dhcp => {
-                let handle = socket.sockets.add(smoltcp::socket::dhcpv4::Socket::new());
-                inner.dhcp_socket = Some(handle);
-            }
-        }
-
-        Self {
-            socket: UnsafeCell::new(socket),
-            inner: UnsafeCell::new(inner),
-        }
-    }
-
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with<R>(&self, f: impl FnOnce(&SocketStack, &Inner<D>) -> R) -> R {
-        f(&*self.socket.get(), &*self.inner.get())
-    }
-
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with_mut<R>(&self, f: impl FnOnce(&mut SocketStack, &mut Inner<D>) -> R) -> R {
-        f(&mut *self.socket.get(), &mut *self.inner.get())
-    }
-
-    pub fn ethernet_address(&self) -> [u8; 6] {
-        unsafe { self.with(|_s, i| i.device.device.ethernet_address()) }
-    }
-
-    pub fn is_link_up(&self) -> bool {
-        unsafe { self.with(|_s, i| i.link_up) }
-    }
-
-    pub fn is_config_up(&self) -> bool {
-        unsafe { self.with(|_s, i| i.config.is_some()) }
-    }
-
-    pub fn config(&self) -> Option<Config> {
-        unsafe { self.with(|_s, i| i.config.clone()) }
-    }
-
-    pub async fn run(&self) -> ! {
-        poll_fn(|cx| {
-            unsafe { self.with_mut(|s, i| i.poll(cx, s)) }
-            Poll::<()>::Pending
-        })
-        .await;
-        unreachable!()
-    }
-}
-
-impl SocketStack {
-    #[allow(clippy::absurd_extreme_comparisons)]
-    pub fn get_local_port(&mut self) -> u16 {
-        let res = self.next_local_port;
-        self.next_local_port = if res >= LOCAL_PORT_MAX { LOCAL_PORT_MIN } else { res + 1 };
-        res
-    }
-}
-
-impl<D: Device + 'static> Inner<D> {
-    fn apply_config(&mut self, s: &mut SocketStack, config: Config) {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = self.device.capabilities().medium;
-
-        debug!("Acquired IP configuration:");
-
-        debug!("   IP address:      {}", config.address);
-        self.set_ipv4_addr(s, config.address);
-
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            if let Some(gateway) = config.gateway {
-                debug!("   Default gateway: {}", gateway);
-                s.iface.routes_mut().add_default_ipv4_route(gateway).unwrap();
-            } else {
-                debug!("   Default gateway: None");
-                s.iface.routes_mut().remove_default_ipv4_route();
-            }
-        }
-        for (i, s) in config.dns_servers.iter().enumerate() {
-            debug!("   DNS server {}:    {}", i, s);
-        }
-
-        self.config = Some(config)
-    }
-
-    #[allow(unused)] // used only with dhcp
-    fn unapply_config(&mut self, s: &mut SocketStack) {
-        #[cfg(feature = "medium-ethernet")]
-        let medium = self.device.capabilities().medium;
-
-        debug!("Lost IP configuration");
-        self.set_ipv4_addr(s, Ipv4Cidr::new(Ipv4Address::UNSPECIFIED, 0));
-        #[cfg(feature = "medium-ethernet")]
-        if medium == Medium::Ethernet {
-            s.iface.routes_mut().remove_default_ipv4_route();
-        }
-        self.config = None
-    }
-
-    fn set_ipv4_addr(&mut self, s: &mut SocketStack, cidr: Ipv4Cidr) {
-        s.iface.update_ip_addrs(|addrs| {
-            let dest = addrs.iter_mut().next().unwrap();
-            *dest = IpCidr::Ipv4(cidr);
-        });
-    }
-
-    fn poll(&mut self, cx: &mut Context<'_>, s: &mut SocketStack) {
-        self.device.device.register_waker(cx.waker());
-        s.waker.register(cx.waker());
-
-        let timestamp = instant_to_smoltcp(Instant::now());
-        if s.iface.poll(timestamp, &mut self.device, &mut s.sockets).is_err() {
-            // If poll() returns error, it may not be done yet, so poll again later.
-            cx.waker().wake_by_ref();
-            return;
-        }
-
-        // Update link up
-        let old_link_up = self.link_up;
-        self.link_up = self.device.device.link_state() == LinkState::Up;
-
-        // Print when changed
-        if old_link_up != self.link_up {
-            info!("link_up = {:?}", self.link_up);
-        }
-
-        #[cfg(feature = "dhcpv4")]
-        if let Some(dhcp_handle) = self.dhcp_socket {
-            let socket = s.sockets.get_mut::<dhcpv4::Socket>(dhcp_handle);
-
-            if self.link_up {
-                match socket.poll() {
-                    None => {}
-                    Some(dhcpv4::Event::Deconfigured) => self.unapply_config(s),
-                    Some(dhcpv4::Event::Configured(config)) => {
-                        let mut dns_servers = Vec::new();
-                        for s in &config.dns_servers {
-                            if let Some(addr) = s {
-                                dns_servers.push(addr.clone()).unwrap();
-                            }
-                        }
-
-                        self.apply_config(
-                            s,
-                            Config {
-                                address: config.address,
-                                gateway: config.router,
-                                dns_servers,
-                            },
-                        )
-                    }
-                }
-            } else if old_link_up {
-                socket.reset();
-                self.unapply_config(s);
-            }
-        }
-        //if old_link_up || self.link_up {
-        //    self.poll_configurator(timestamp)
-        //}
-
-        if let Some(poll_at) = s.iface.poll_at(timestamp, &mut s.sockets) {
-            let t = Timer::at(instant_from_smoltcp(poll_at));
-            pin_mut!(t);
-            if t.poll(cx).is_ready() {
-                cx.waker().wake_by_ref();
-            }
-        }
-    }
-}
-
-fn instant_to_smoltcp(instant: Instant) -> SmolInstant {
-    SmolInstant::from_millis(instant.as_millis() as i64)
-}
-
-fn instant_from_smoltcp(instant: SmolInstant) -> Instant {
-    Instant::from_millis(instant.total_millis() as u64)
-}
diff --git a/embassy-net/src/tcp.rs b/embassy-net/src/tcp.rs
index 910772c7d..367675b13 100644
--- a/embassy-net/src/tcp.rs
+++ b/embassy-net/src/tcp.rs
@@ -1,24 +1,39 @@
-use core::cell::UnsafeCell;
-use core::future::Future;
+//! TCP sockets.
+//!
+//! # Listening
+//!
+//! `embassy-net` does not have a `TcpListener`. Instead, individual `TcpSocket`s can be put into
+//! listening mode by calling [`TcpSocket::accept`].
+//!
+//! Incoming connections when no socket is listening are rejected. To accept many incoming
+//! connections, create many sockets and put them all into listening mode.
+
+use core::cell::RefCell;
+use core::future::poll_fn;
 use core::mem;
 use core::task::Poll;
 
-use futures::future::poll_fn;
+use embassy_net_driver::Driver;
+use embassy_time::Duration;
 use smoltcp::iface::{Interface, SocketHandle};
 use smoltcp::socket::tcp;
-use smoltcp::time::Duration;
+pub use smoltcp::socket::tcp::State;
 use smoltcp::wire::{IpEndpoint, IpListenEndpoint};
 
-use super::stack::Stack;
-use crate::stack::SocketStack;
-use crate::Device;
+use crate::time::duration_to_smoltcp;
+use crate::{SocketStack, Stack};
 
+/// Error returned by TcpSocket read/write functions.
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum Error {
+    /// The connection was reset.
+    ///
+    /// This can happen on receiving a RST packet, or on timeout.
     ConnectionReset,
 }
 
+/// Error returned by [`TcpSocket::connect`].
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum ConnectError {
@@ -32,6 +47,7 @@ pub enum ConnectError {
     NoRoute,
 }
 
+/// Error returned by [`TcpSocket::accept`].
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum AcceptError {
@@ -43,22 +59,53 @@ pub enum AcceptError {
     ConnectionReset,
 }
 
+/// A TCP socket.
 pub struct TcpSocket<'a> {
     io: TcpIo<'a>,
 }
 
+/// The reader half of a TCP socket.
 pub struct TcpReader<'a> {
     io: TcpIo<'a>,
 }
 
+/// The writer half of a TCP socket.
 pub struct TcpWriter<'a> {
     io: TcpIo<'a>,
 }
 
+impl<'a> TcpReader<'a> {
+    /// Read data from the socket.
+    ///
+    /// Returns how many bytes were read, or an error. If no data is available, it waits
+    /// until there is at least one byte available.
+    pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
+        self.io.read(buf).await
+    }
+}
+
+impl<'a> TcpWriter<'a> {
+    /// Write data to the socket.
+    ///
+    /// Returns how many bytes were written, or an error. If the socket is not ready to
+    /// accept data, it waits until it is.
+    pub async fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
+        self.io.write(buf).await
+    }
+
+    /// Flushes the written data to the socket.
+    ///
+    /// This waits until all data has been sent, and ACKed by the remote host. For a connection
+    /// closed with [`abort()`](TcpSocket::abort) it will wait for the TCP RST packet to be sent.
+    pub async fn flush(&mut self) -> Result<(), Error> {
+        self.io.flush().await
+    }
+}
+
 impl<'a> TcpSocket<'a> {
-    pub fn new<D: Device>(stack: &'a Stack<D>, rx_buffer: &'a mut [u8], tx_buffer: &'a mut [u8]) -> Self {
-        // safety: not accessed reentrantly.
-        let s = unsafe { &mut *stack.socket.get() };
+    /// Create a new TCP socket on the given stack, with the given buffers.
+    pub fn new<D: Driver>(stack: &'a Stack<D>, rx_buffer: &'a mut [u8], tx_buffer: &'a mut [u8]) -> Self {
+        let s = &mut *stack.socket.borrow_mut();
         let rx_buffer: &'static mut [u8] = unsafe { mem::transmute(rx_buffer) };
         let tx_buffer: &'static mut [u8] = unsafe { mem::transmute(tx_buffer) };
         let handle = s.sockets.add(tcp::Socket::new(
@@ -74,25 +121,28 @@ impl<'a> TcpSocket<'a> {
         }
     }
 
+    /// Split the socket into reader and a writer halves.
     pub fn split(&mut self) -> (TcpReader<'_>, TcpWriter<'_>) {
         (TcpReader { io: self.io }, TcpWriter { io: self.io })
     }
 
+    /// Connect to a remote host.
     pub async fn connect<T>(&mut self, remote_endpoint: T) -> Result<(), ConnectError>
     where
         T: Into<IpEndpoint>,
     {
-        // safety: not accessed reentrantly.
-        let local_port = unsafe { &mut *self.io.stack.get() }.get_local_port();
+        let local_port = self.io.stack.borrow_mut().get_local_port();
 
-        // safety: not accessed reentrantly.
-        match unsafe { self.io.with_mut(|s, i| s.connect(i, remote_endpoint, local_port)) } {
+        match {
+            self.io
+                .with_mut(|s, i| s.connect(i.context(), remote_endpoint, local_port))
+        } {
             Ok(()) => {}
             Err(tcp::ConnectError::InvalidState) => return Err(ConnectError::InvalidState),
             Err(tcp::ConnectError::Unaddressable) => return Err(ConnectError::NoRoute),
         }
 
-        futures::future::poll_fn(|cx| unsafe {
+        poll_fn(|cx| {
             self.io.with_mut(|s, _| match s.state() {
                 tcp::State::Closed | tcp::State::TimeWait => Poll::Ready(Err(ConnectError::ConnectionReset)),
                 tcp::State::Listen => unreachable!(),
@@ -106,18 +156,20 @@ impl<'a> TcpSocket<'a> {
         .await
     }
 
+    /// Accept a connection from a remote host.
+    ///
+    /// This function puts the socket in listening mode, and waits until a connection is received.
     pub async fn accept<T>(&mut self, local_endpoint: T) -> Result<(), AcceptError>
     where
         T: Into<IpListenEndpoint>,
     {
-        // safety: not accessed reentrantly.
-        match unsafe { self.io.with_mut(|s, _| s.listen(local_endpoint)) } {
+        match self.io.with_mut(|s, _| s.listen(local_endpoint)) {
             Ok(()) => {}
             Err(tcp::ListenError::InvalidState) => return Err(AcceptError::InvalidState),
             Err(tcp::ListenError::Unaddressable) => return Err(AcceptError::InvalidPort),
         }
 
-        futures::future::poll_fn(|cx| unsafe {
+        poll_fn(|cx| {
             self.io.with_mut(|s, _| match s.state() {
                 tcp::State::Listen | tcp::State::SynSent | tcp::State::SynReceived => {
                     s.register_send_waker(cx.waker());
@@ -129,52 +181,120 @@ impl<'a> TcpSocket<'a> {
         .await
     }
 
+    /// Read data from the socket.
+    ///
+    /// Returns how many bytes were read, or an error. If no data is available, it waits
+    /// until there is at least one byte available.
+    pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
+        self.io.read(buf).await
+    }
+
+    /// Write data to the socket.
+    ///
+    /// Returns how many bytes were written, or an error. If the socket is not ready to
+    /// accept data, it waits until it is.
+    pub async fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
+        self.io.write(buf).await
+    }
+
+    /// Flushes the written data to the socket.
+    ///
+    /// This waits until all data has been sent, and ACKed by the remote host. For a connection
+    /// closed with [`abort()`](TcpSocket::abort) it will wait for the TCP RST packet to be sent.
+    pub async fn flush(&mut self) -> Result<(), Error> {
+        self.io.flush().await
+    }
+
+    /// Set the timeout for the socket.
+    ///
+    /// If the timeout is set, the socket will be closed if no data is received for the
+    /// specified duration.
     pub fn set_timeout(&mut self, duration: Option<Duration>) {
-        unsafe { self.io.with_mut(|s, _| s.set_timeout(duration)) }
+        self.io
+            .with_mut(|s, _| s.set_timeout(duration.map(duration_to_smoltcp)))
     }
 
+    /// Set the keep-alive interval for the socket.
+    ///
+    /// If the keep-alive interval is set, the socket will send keep-alive packets after
+    /// the specified duration of inactivity.
+    ///
+    /// If not set, the socket will not send keep-alive packets.
     pub fn set_keep_alive(&mut self, interval: Option<Duration>) {
-        unsafe { self.io.with_mut(|s, _| s.set_keep_alive(interval)) }
+        self.io
+            .with_mut(|s, _| s.set_keep_alive(interval.map(duration_to_smoltcp)))
     }
 
+    /// Set the hop limit field in the IP header of sent packets.
     pub fn set_hop_limit(&mut self, hop_limit: Option<u8>) {
-        unsafe { self.io.with_mut(|s, _| s.set_hop_limit(hop_limit)) }
+        self.io.with_mut(|s, _| s.set_hop_limit(hop_limit))
     }
 
+    /// Get the local endpoint of the socket.
+    ///
+    /// Returns `None` if the socket is not bound (listening) or not connected.
     pub fn local_endpoint(&self) -> Option<IpEndpoint> {
-        unsafe { self.io.with(|s, _| s.local_endpoint()) }
+        self.io.with(|s, _| s.local_endpoint())
     }
 
+    /// Get the remote endpoint of the socket.
+    ///
+    /// Returns `None` if the socket is not connected.
     pub fn remote_endpoint(&self) -> Option<IpEndpoint> {
-        unsafe { self.io.with(|s, _| s.remote_endpoint()) }
+        self.io.with(|s, _| s.remote_endpoint())
     }
 
-    pub fn state(&self) -> tcp::State {
-        unsafe { self.io.with(|s, _| s.state()) }
+    /// Get the state of the socket.
+    pub fn state(&self) -> State {
+        self.io.with(|s, _| s.state())
     }
 
+    /// Close the write half of the socket.
+    ///
+    /// This closes only the write half of the socket. The read half side remains open, the
+    /// socket can still receive data.
+    ///
+    /// Data that has been written to the socket and not yet sent (or not yet ACKed) will still
+    /// still sent. The last segment of the pending to send data is sent with the FIN flag set.
     pub fn close(&mut self) {
-        unsafe { self.io.with_mut(|s, _| s.close()) }
-    }
-
+        self.io.with_mut(|s, _| s.close())
+    }
+
+    /// Forcibly close the socket.
+    ///
+    /// This instantly closes both the read and write halves of the socket. Any pending data
+    /// that has not been sent will be lost.
+    ///
+    /// Note that the TCP RST packet is not sent immediately - if the `TcpSocket` is dropped too soon
+    /// the remote host may not know the connection has been closed.
+    /// `abort()` callers should wait for a [`flush()`](TcpSocket::flush) call to complete before
+    /// dropping or reusing the socket.
     pub fn abort(&mut self) {
-        unsafe { self.io.with_mut(|s, _| s.abort()) }
+        self.io.with_mut(|s, _| s.abort())
     }
 
+    /// Get whether the socket is ready to send data, i.e. whether there is space in the send buffer.
     pub fn may_send(&self) -> bool {
-        unsafe { self.io.with(|s, _| s.may_send()) }
+        self.io.with(|s, _| s.may_send())
     }
 
+    /// return whether the recieve half of the full-duplex connection is open.
+    /// This function returns true if it’s possible to receive data from the remote endpoint.
+    /// It will return true while there is data in the receive buffer, and if there isn’t,
+    /// as long as the remote endpoint has not closed the connection.
     pub fn may_recv(&self) -> bool {
-        unsafe { self.io.with(|s, _| s.may_recv()) }
+        self.io.with(|s, _| s.may_recv())
+    }
+
+    /// Get whether the socket is ready to receive data, i.e. whether there is some pending data in the receive buffer.
+    pub fn can_recv(&self) -> bool {
+        self.io.with(|s, _| s.can_recv())
     }
 }
 
 impl<'a> Drop for TcpSocket<'a> {
     fn drop(&mut self) {
-        // safety: not accessed reentrantly.
-        let s = unsafe { &mut *self.io.stack.get() };
-        s.sockets.remove(self.io.handle);
+        self.io.stack.borrow_mut().sockets.remove(self.io.handle);
     }
 }
 
@@ -182,21 +302,19 @@ impl<'a> Drop for TcpSocket<'a> {
 
 #[derive(Copy, Clone)]
 struct TcpIo<'a> {
-    stack: &'a UnsafeCell<SocketStack>,
+    stack: &'a RefCell<SocketStack>,
     handle: SocketHandle,
 }
 
 impl<'d> TcpIo<'d> {
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with<R>(&self, f: impl FnOnce(&tcp::Socket, &Interface) -> R) -> R {
-        let s = &*self.stack.get();
+    fn with<R>(&self, f: impl FnOnce(&tcp::Socket, &Interface) -> R) -> R {
+        let s = &*self.stack.borrow();
         let socket = s.sockets.get::<tcp::Socket>(self.handle);
         f(socket, &s.iface)
     }
 
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with_mut<R>(&mut self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
-        let s = &mut *self.stack.get();
+    fn with_mut<R>(&mut self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
+        let s = &mut *self.stack.borrow_mut();
         let socket = s.sockets.get_mut::<tcp::Socket>(self.handle);
         let res = f(socket, &mut s.iface);
         s.waker.wake();
@@ -204,7 +322,7 @@ impl<'d> TcpIo<'d> {
     }
 
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
-        poll_fn(move |cx| unsafe {
+        poll_fn(move |cx| {
             // CAUTION: smoltcp semantics around EOF are different to what you'd expect
             // from posix-like IO, so we have to tweak things here.
             self.with_mut(|s, _| match s.recv_slice(buf) {
@@ -225,7 +343,7 @@ impl<'d> TcpIo<'d> {
     }
 
     async fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
-        poll_fn(move |cx| unsafe {
+        poll_fn(move |cx| {
             self.with_mut(|s, _| match s.send_slice(buf) {
                 // Not ready to send (no space in the tx buffer)
                 Ok(0) => {
@@ -242,95 +360,89 @@ impl<'d> TcpIo<'d> {
     }
 
     async fn flush(&mut self) -> Result<(), Error> {
-        poll_fn(move |_| {
-            Poll::Ready(Ok(())) // TODO: Is there a better implementation for this?
+        poll_fn(move |cx| {
+            self.with_mut(|s, _| {
+                let waiting_close = s.state() == tcp::State::Closed && s.remote_endpoint().is_some();
+                // If there are outstanding send operations, register for wake up and wait
+                // smoltcp issues wake-ups when octets are dequeued from the send buffer
+                if s.send_queue() > 0 || waiting_close {
+                    s.register_send_waker(cx.waker());
+                    Poll::Pending
+                // No outstanding sends, socket is flushed
+                } else {
+                    Poll::Ready(Ok(()))
+                }
+            })
         })
         .await
     }
 }
 
-impl embedded_io::Error for ConnectError {
-    fn kind(&self) -> embedded_io::ErrorKind {
-        embedded_io::ErrorKind::Other
-    }
-}
+#[cfg(feature = "nightly")]
+mod embedded_io_impls {
+    use super::*;
 
-impl embedded_io::Error for Error {
-    fn kind(&self) -> embedded_io::ErrorKind {
-        embedded_io::ErrorKind::Other
+    impl embedded_io::Error for ConnectError {
+        fn kind(&self) -> embedded_io::ErrorKind {
+            embedded_io::ErrorKind::Other
+        }
     }
-}
-
-impl<'d> embedded_io::Io for TcpSocket<'d> {
-    type Error = Error;
-}
 
-impl<'d> embedded_io::asynch::Read for TcpSocket<'d> {
-    type ReadFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
-
-    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::ReadFuture<'a> {
-        self.io.read(buf)
+    impl embedded_io::Error for Error {
+        fn kind(&self) -> embedded_io::ErrorKind {
+            embedded_io::ErrorKind::Other
+        }
     }
-}
 
-impl<'d> embedded_io::asynch::Write for TcpSocket<'d> {
-    type WriteFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
-
-    fn write<'a>(&'a mut self, buf: &'a [u8]) -> Self::WriteFuture<'a> {
-        self.io.write(buf)
+    impl<'d> embedded_io::Io for TcpSocket<'d> {
+        type Error = Error;
     }
 
-    type FlushFuture<'a> = impl Future<Output = Result<(), Self::Error>>
-    where
-        Self: 'a;
-
-    fn flush<'a>(&'a mut self) -> Self::FlushFuture<'a> {
-        self.io.flush()
+    impl<'d> embedded_io::asynch::Read for TcpSocket<'d> {
+        async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
+            self.io.read(buf).await
+        }
     }
-}
-
-impl<'d> embedded_io::Io for TcpReader<'d> {
-    type Error = Error;
-}
 
-impl<'d> embedded_io::asynch::Read for TcpReader<'d> {
-    type ReadFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
+    impl<'d> embedded_io::asynch::Write for TcpSocket<'d> {
+        async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
+            self.io.write(buf).await
+        }
 
-    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::ReadFuture<'a> {
-        self.io.read(buf)
+        async fn flush(&mut self) -> Result<(), Self::Error> {
+            self.io.flush().await
+        }
     }
-}
 
-impl<'d> embedded_io::Io for TcpWriter<'d> {
-    type Error = Error;
-}
+    impl<'d> embedded_io::Io for TcpReader<'d> {
+        type Error = Error;
+    }
 
-impl<'d> embedded_io::asynch::Write for TcpWriter<'d> {
-    type WriteFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
+    impl<'d> embedded_io::asynch::Read for TcpReader<'d> {
+        async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
+            self.io.read(buf).await
+        }
+    }
 
-    fn write<'a>(&'a mut self, buf: &'a [u8]) -> Self::WriteFuture<'a> {
-        self.io.write(buf)
+    impl<'d> embedded_io::Io for TcpWriter<'d> {
+        type Error = Error;
     }
 
-    type FlushFuture<'a> = impl Future<Output = Result<(), Self::Error>>
-    where
-        Self: 'a;
+    impl<'d> embedded_io::asynch::Write for TcpWriter<'d> {
+        async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
+            self.io.write(buf).await
+        }
 
-    fn flush<'a>(&'a mut self) -> Self::FlushFuture<'a> {
-        self.io.flush()
+        async fn flush(&mut self) -> Result<(), Self::Error> {
+            self.io.flush().await
+        }
     }
 }
 
-#[cfg(feature = "unstable-traits")]
+/// TCP client compatible with `embedded-nal-async` traits.
+#[cfg(all(feature = "unstable-traits", feature = "nightly"))]
 pub mod client {
+    use core::cell::UnsafeCell;
     use core::mem::MaybeUninit;
     use core::ptr::NonNull;
 
@@ -339,49 +451,56 @@ pub mod client {
 
     use super::*;
 
-    /// TCP client capable of creating up to N multiple connections with tx and rx buffers according to TX_SZ and RX_SZ.
-    pub struct TcpClient<'d, D: Device, const N: usize, const TX_SZ: usize = 1024, const RX_SZ: usize = 1024> {
+    /// TCP client connection pool compatible with `embedded-nal-async` traits.
+    ///
+    /// The pool is capable of managing up to N concurrent connections with tx and rx buffers according to TX_SZ and RX_SZ.
+    pub struct TcpClient<'d, D: Driver, const N: usize, const TX_SZ: usize = 1024, const RX_SZ: usize = 1024> {
         stack: &'d Stack<D>,
         state: &'d TcpClientState<N, TX_SZ, RX_SZ>,
     }
 
-    impl<'d, D: Device, const N: usize, const TX_SZ: usize, const RX_SZ: usize> TcpClient<'d, D, N, TX_SZ, RX_SZ> {
-        /// Create a new TcpClient
+    impl<'d, D: Driver, const N: usize, const TX_SZ: usize, const RX_SZ: usize> TcpClient<'d, D, N, TX_SZ, RX_SZ> {
+        /// Create a new `TcpClient`.
         pub fn new(stack: &'d Stack<D>, state: &'d TcpClientState<N, TX_SZ, RX_SZ>) -> Self {
             Self { stack, state }
         }
     }
 
-    impl<'d, D: Device, const N: usize, const TX_SZ: usize, const RX_SZ: usize> embedded_nal_async::TcpConnect
+    impl<'d, D: Driver, const N: usize, const TX_SZ: usize, const RX_SZ: usize> embedded_nal_async::TcpConnect
         for TcpClient<'d, D, N, TX_SZ, RX_SZ>
     {
         type Error = Error;
         type Connection<'m> = TcpConnection<'m, N, TX_SZ, RX_SZ> where Self: 'm;
-        type ConnectFuture<'m> = impl Future<Output = Result<Self::Connection<'m>, Self::Error>> + 'm
-    where
-        Self: 'm;
-
-        fn connect<'m>(&'m self, remote: embedded_nal_async::SocketAddr) -> Self::ConnectFuture<'m> {
-            async move {
-                let addr: crate::IpAddress = match remote.ip() {
-                    IpAddr::V4(addr) => crate::IpAddress::Ipv4(crate::Ipv4Address::from_bytes(&addr.octets())),
-                    #[cfg(feature = "proto-ipv6")]
-                    IpAddr::V6(addr) => crate::IpAddress::Ipv6(crate::Ipv6Address::from_bytes(&addr.octets())),
-                    #[cfg(not(feature = "proto-ipv6"))]
-                    IpAddr::V6(_) => panic!("ipv6 support not enabled"),
-                };
-                let remote_endpoint = (addr, remote.port());
-                let mut socket = TcpConnection::new(&self.stack, self.state)?;
-                socket
-                    .socket
-                    .connect(remote_endpoint)
-                    .await
-                    .map_err(|_| Error::ConnectionReset)?;
-                Ok(socket)
-            }
+
+        async fn connect<'a>(
+            &'a self,
+            remote: embedded_nal_async::SocketAddr,
+        ) -> Result<Self::Connection<'a>, Self::Error>
+        where
+            Self: 'a,
+        {
+            let addr: crate::IpAddress = match remote.ip() {
+                #[cfg(feature = "proto-ipv4")]
+                IpAddr::V4(addr) => crate::IpAddress::Ipv4(crate::Ipv4Address::from_bytes(&addr.octets())),
+                #[cfg(not(feature = "proto-ipv4"))]
+                IpAddr::V4(_) => panic!("ipv4 support not enabled"),
+                #[cfg(feature = "proto-ipv6")]
+                IpAddr::V6(addr) => crate::IpAddress::Ipv6(crate::Ipv6Address::from_bytes(&addr.octets())),
+                #[cfg(not(feature = "proto-ipv6"))]
+                IpAddr::V6(_) => panic!("ipv6 support not enabled"),
+            };
+            let remote_endpoint = (addr, remote.port());
+            let mut socket = TcpConnection::new(&self.stack, self.state)?;
+            socket
+                .socket
+                .connect(remote_endpoint)
+                .await
+                .map_err(|_| Error::ConnectionReset)?;
+            Ok(socket)
         }
     }
 
+    /// Opened TCP connection in a [`TcpClient`].
     pub struct TcpConnection<'d, const N: usize, const TX_SZ: usize, const RX_SZ: usize> {
         socket: TcpSocket<'d>,
         state: &'d TcpClientState<N, TX_SZ, RX_SZ>,
@@ -389,10 +508,10 @@ pub mod client {
     }
 
     impl<'d, const N: usize, const TX_SZ: usize, const RX_SZ: usize> TcpConnection<'d, N, TX_SZ, RX_SZ> {
-        fn new<D: Device>(stack: &'d Stack<D>, state: &'d TcpClientState<N, TX_SZ, RX_SZ>) -> Result<Self, Error> {
+        fn new<D: Driver>(stack: &'d Stack<D>, state: &'d TcpClientState<N, TX_SZ, RX_SZ>) -> Result<Self, Error> {
             let mut bufs = state.pool.alloc().ok_or(Error::ConnectionReset)?;
             Ok(Self {
-                socket: unsafe { TcpSocket::new(stack, &mut bufs.as_mut().0, &mut bufs.as_mut().1) },
+                socket: unsafe { TcpSocket::new(stack, &mut bufs.as_mut().1, &mut bufs.as_mut().0) },
                 state,
                 bufs,
             })
@@ -417,32 +536,20 @@ pub mod client {
     impl<'d, const N: usize, const TX_SZ: usize, const RX_SZ: usize> embedded_io::asynch::Read
         for TcpConnection<'d, N, TX_SZ, RX_SZ>
     {
-        type ReadFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
-
-        fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::ReadFuture<'a> {
-            self.socket.read(buf)
+        async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error> {
+            self.socket.read(buf).await
         }
     }
 
     impl<'d, const N: usize, const TX_SZ: usize, const RX_SZ: usize> embedded_io::asynch::Write
         for TcpConnection<'d, N, TX_SZ, RX_SZ>
     {
-        type WriteFuture<'a> = impl Future<Output = Result<usize, Self::Error>>
-    where
-        Self: 'a;
-
-        fn write<'a>(&'a mut self, buf: &'a [u8]) -> Self::WriteFuture<'a> {
-            self.socket.write(buf)
+        async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
+            self.socket.write(buf).await
         }
 
-        type FlushFuture<'a> = impl Future<Output = Result<(), Self::Error>>
-    where
-        Self: 'a;
-
-        fn flush<'a>(&'a mut self) -> Self::FlushFuture<'a> {
-            self.socket.flush()
+        async fn flush(&mut self) -> Result<(), Self::Error> {
+            self.socket.flush().await
         }
     }
 
@@ -452,6 +559,7 @@ pub mod client {
     }
 
     impl<const N: usize, const TX_SZ: usize, const RX_SZ: usize> TcpClientState<N, TX_SZ, RX_SZ> {
+        /// Create a new `TcpClientState`.
         pub const fn new() -> Self {
             Self { pool: Pool::new() }
         }
diff --git a/embassy-net/src/time.rs b/embassy-net/src/time.rs
new file mode 100644
index 000000000..b98d40fdc
--- /dev/null
+++ b/embassy-net/src/time.rs
@@ -0,0 +1,20 @@
+#![allow(unused)]
+
+use embassy_time::{Duration, Instant};
+use smoltcp::time::{Duration as SmolDuration, Instant as SmolInstant};
+
+pub(crate) fn instant_to_smoltcp(instant: Instant) -> SmolInstant {
+    SmolInstant::from_micros(instant.as_micros() as i64)
+}
+
+pub(crate) fn instant_from_smoltcp(instant: SmolInstant) -> Instant {
+    Instant::from_micros(instant.total_micros() as u64)
+}
+
+pub(crate) fn duration_to_smoltcp(duration: Duration) -> SmolDuration {
+    SmolDuration::from_micros(duration.as_micros())
+}
+
+pub(crate) fn duration_from_smoltcp(duration: SmolDuration) -> Duration {
+    Duration::from_micros(duration.total_micros())
+}
diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index 78b09a492..0d97b6db1 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -1,15 +1,19 @@
-use core::cell::UnsafeCell;
+//! UDP sockets.
+
+use core::cell::RefCell;
+use core::future::poll_fn;
 use core::mem;
-use core::task::Poll;
+use core::task::{Context, Poll};
 
-use futures::future::poll_fn;
+use embassy_net_driver::Driver;
 use smoltcp::iface::{Interface, SocketHandle};
-use smoltcp::socket::udp::{self, PacketMetadata};
+use smoltcp::socket::udp;
+pub use smoltcp::socket::udp::PacketMetadata;
 use smoltcp::wire::{IpEndpoint, IpListenEndpoint};
 
-use super::stack::SocketStack;
-use crate::{Device, Stack};
+use crate::{SocketStack, Stack};
 
+/// Error returned by [`UdpSocket::bind`].
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum BindError {
@@ -19,6 +23,7 @@ pub enum BindError {
     NoRoute,
 }
 
+/// Error returned by [`UdpSocket::recv_from`] and [`UdpSocket::send_to`].
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum Error {
@@ -26,21 +31,22 @@ pub enum Error {
     NoRoute,
 }
 
+/// An UDP socket.
 pub struct UdpSocket<'a> {
-    stack: &'a UnsafeCell<SocketStack>,
+    stack: &'a RefCell<SocketStack>,
     handle: SocketHandle,
 }
 
 impl<'a> UdpSocket<'a> {
-    pub fn new<D: Device>(
+    /// Create a new UDP socket using the provided stack and buffers.
+    pub fn new<D: Driver>(
         stack: &'a Stack<D>,
         rx_meta: &'a mut [PacketMetadata],
         rx_buffer: &'a mut [u8],
         tx_meta: &'a mut [PacketMetadata],
         tx_buffer: &'a mut [u8],
     ) -> Self {
-        // safety: not accessed reentrantly.
-        let s = unsafe { &mut *stack.socket.get() };
+        let s = &mut *stack.socket.borrow_mut();
 
         let rx_meta: &'static mut [PacketMetadata] = unsafe { mem::transmute(rx_meta) };
         let rx_buffer: &'static mut [u8] = unsafe { mem::transmute(rx_buffer) };
@@ -57,101 +63,131 @@ impl<'a> UdpSocket<'a> {
         }
     }
 
+    /// Bind the socket to a local endpoint.
     pub fn bind<T>(&mut self, endpoint: T) -> Result<(), BindError>
     where
         T: Into<IpListenEndpoint>,
     {
         let mut endpoint = endpoint.into();
 
-        // safety: not accessed reentrantly.
         if endpoint.port == 0 {
             // If user didn't specify port allocate a dynamic port.
-            endpoint.port = unsafe { &mut *self.stack.get() }.get_local_port();
+            endpoint.port = self.stack.borrow_mut().get_local_port();
         }
 
-        // safety: not accessed reentrantly.
-        match unsafe { self.with_mut(|s, _| s.bind(endpoint)) } {
+        match self.with_mut(|s, _| s.bind(endpoint)) {
             Ok(()) => Ok(()),
             Err(udp::BindError::InvalidState) => Err(BindError::InvalidState),
             Err(udp::BindError::Unaddressable) => Err(BindError::NoRoute),
         }
     }
 
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with<R>(&self, f: impl FnOnce(&udp::Socket, &Interface) -> R) -> R {
-        let s = &*self.stack.get();
+    fn with<R>(&self, f: impl FnOnce(&udp::Socket, &Interface) -> R) -> R {
+        let s = &*self.stack.borrow();
         let socket = s.sockets.get::<udp::Socket>(self.handle);
         f(socket, &s.iface)
     }
 
-    /// SAFETY: must not call reentrantly.
-    unsafe fn with_mut<R>(&self, f: impl FnOnce(&mut udp::Socket, &mut Interface) -> R) -> R {
-        let s = &mut *self.stack.get();
+    fn with_mut<R>(&self, f: impl FnOnce(&mut udp::Socket, &mut Interface) -> R) -> R {
+        let s = &mut *self.stack.borrow_mut();
         let socket = s.sockets.get_mut::<udp::Socket>(self.handle);
         let res = f(socket, &mut s.iface);
         s.waker.wake();
         res
     }
 
+    /// Receive a datagram.
+    ///
+    /// This method will wait until a datagram is received.
+    ///
+    /// Returns the number of bytes received and the remote endpoint.
     pub async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, IpEndpoint), Error> {
-        poll_fn(move |cx| unsafe {
-            self.with_mut(|s, _| match s.recv_slice(buf) {
-                Ok(x) => Poll::Ready(Ok(x)),
-                // No data ready
-                Err(udp::RecvError::Exhausted) => {
-                    //s.register_recv_waker(cx.waker());
-                    cx.waker().wake_by_ref();
-                    Poll::Pending
-                }
-            })
+        poll_fn(move |cx| self.poll_recv_from(buf, cx)).await
+    }
+
+    /// Receive a datagram.
+    ///
+    /// When no datagram is available, this method will return `Poll::Pending` and
+    /// register the current task to be notified when a datagram is received.
+    ///
+    /// When a datagram is received, this method will return `Poll::Ready` with the
+    /// number of bytes received and the remote endpoint.
+    pub fn poll_recv_from(&self, buf: &mut [u8], cx: &mut Context<'_>) -> Poll<Result<(usize, IpEndpoint), Error>> {
+        self.with_mut(|s, _| match s.recv_slice(buf) {
+            Ok((n, meta)) => Poll::Ready(Ok((n, meta.endpoint))),
+            // No data ready
+            Err(udp::RecvError::Exhausted) => {
+                s.register_recv_waker(cx.waker());
+                Poll::Pending
+            }
         })
-        .await
     }
 
+    /// Send a datagram to the specified remote endpoint.
+    ///
+    /// This method will wait until the datagram has been sent.
+    ///
+    /// When the remote endpoint is not reachable, this method will return `Err(Error::NoRoute)`
     pub async fn send_to<T>(&self, buf: &[u8], remote_endpoint: T) -> Result<(), Error>
     where
         T: Into<IpEndpoint>,
     {
-        let remote_endpoint = remote_endpoint.into();
-        poll_fn(move |cx| unsafe {
-            self.with_mut(|s, _| match s.send_slice(buf, remote_endpoint) {
-                // Entire datagram has been sent
-                Ok(()) => Poll::Ready(Ok(())),
-                Err(udp::SendError::BufferFull) => {
-                    s.register_send_waker(cx.waker());
-                    Poll::Pending
-                }
-                Err(udp::SendError::Unaddressable) => Poll::Ready(Err(Error::NoRoute)),
-            })
+        let remote_endpoint: IpEndpoint = remote_endpoint.into();
+        poll_fn(move |cx| self.poll_send_to(buf, remote_endpoint, cx)).await
+    }
+
+    /// Send a datagram to the specified remote endpoint.
+    ///
+    /// When the datagram has been sent, this method will return `Poll::Ready(Ok())`.
+    ///
+    /// When the socket's send buffer is full, this method will return `Poll::Pending`
+    /// and register the current task to be notified when the buffer has space available.
+    ///
+    /// When the remote endpoint is not reachable, this method will return `Poll::Ready(Err(Error::NoRoute))`.
+    pub fn poll_send_to<T>(&self, buf: &[u8], remote_endpoint: T, cx: &mut Context<'_>) -> Poll<Result<(), Error>>
+    where
+        T: Into<IpEndpoint>,
+    {
+        self.with_mut(|s, _| match s.send_slice(buf, remote_endpoint) {
+            // Entire datagram has been sent
+            Ok(()) => Poll::Ready(Ok(())),
+            Err(udp::SendError::BufferFull) => {
+                s.register_send_waker(cx.waker());
+                Poll::Pending
+            }
+            Err(udp::SendError::Unaddressable) => Poll::Ready(Err(Error::NoRoute)),
         })
-        .await
     }
 
+    /// Returns the local endpoint of the socket.
     pub fn endpoint(&self) -> IpListenEndpoint {
-        unsafe { self.with(|s, _| s.endpoint()) }
+        self.with(|s, _| s.endpoint())
     }
 
+    /// Returns whether the socket is open.
+
     pub fn is_open(&self) -> bool {
-        unsafe { self.with(|s, _| s.is_open()) }
+        self.with(|s, _| s.is_open())
     }
 
+    /// Close the socket.
     pub fn close(&mut self) {
-        unsafe { self.with_mut(|s, _| s.close()) }
+        self.with_mut(|s, _| s.close())
     }
 
+    /// Returns whether the socket is ready to send data, i.e. it has enough buffer space to hold a packet.
     pub fn may_send(&self) -> bool {
-        unsafe { self.with(|s, _| s.can_send()) }
+        self.with(|s, _| s.can_send())
     }
 
+    /// Returns whether the socket is ready to receive data, i.e. it has received a packet that's now in the buffer.
     pub fn may_recv(&self) -> bool {
-        unsafe { self.with(|s, _| s.can_recv()) }
+        self.with(|s, _| s.can_recv())
     }
 }
 
 impl Drop for UdpSocket<'_> {
     fn drop(&mut self) {
-        // safety: not accessed reentrantly.
-        let s = unsafe { &mut *self.stack.get() };
-        s.sockets.remove(self.handle);
+        self.stack.borrow_mut().sockets.remove(self.handle);
     }
 }