embassy_net/
raw.rs

1//! Raw sockets.
2
3use core::future::{poll_fn, Future};
4use core::mem;
5use core::task::{Context, Poll};
6
7use embassy_net_driver::Driver;
8use smoltcp::iface::{Interface, SocketHandle};
9use smoltcp::socket::raw;
10pub use smoltcp::socket::raw::PacketMetadata;
11pub use smoltcp::wire::{IpProtocol, IpVersion};
12
13use crate::Stack;
14
15/// Error returned by [`RawSocket::recv`] and [`RawSocket::send`].
16#[derive(PartialEq, Eq, Clone, Copy, Debug)]
17#[cfg_attr(feature = "defmt", derive(defmt::Format))]
18pub enum RecvError {
19    /// Provided buffer was smaller than the received packet.
20    Truncated,
21}
22
23/// An Raw socket.
24pub struct RawSocket<'a> {
25    stack: Stack<'a>,
26    handle: SocketHandle,
27}
28
29impl<'a> RawSocket<'a> {
30    /// Create a new Raw socket using the provided stack and buffers.
31    pub fn new<D: Driver>(
32        stack: Stack<'a>,
33        ip_version: IpVersion,
34        ip_protocol: IpProtocol,
35        rx_meta: &'a mut [PacketMetadata],
36        rx_buffer: &'a mut [u8],
37        tx_meta: &'a mut [PacketMetadata],
38        tx_buffer: &'a mut [u8],
39    ) -> Self {
40        let handle = stack.with_mut(|i| {
41            let rx_meta: &'static mut [PacketMetadata] = unsafe { mem::transmute(rx_meta) };
42            let rx_buffer: &'static mut [u8] = unsafe { mem::transmute(rx_buffer) };
43            let tx_meta: &'static mut [PacketMetadata] = unsafe { mem::transmute(tx_meta) };
44            let tx_buffer: &'static mut [u8] = unsafe { mem::transmute(tx_buffer) };
45            i.sockets.add(raw::Socket::new(
46                ip_version,
47                ip_protocol,
48                raw::PacketBuffer::new(rx_meta, rx_buffer),
49                raw::PacketBuffer::new(tx_meta, tx_buffer),
50            ))
51        });
52
53        Self { stack, handle }
54    }
55
56    fn with_mut<R>(&self, f: impl FnOnce(&mut raw::Socket, &mut Interface) -> R) -> R {
57        self.stack.with_mut(|i| {
58            let socket = i.sockets.get_mut::<raw::Socket>(self.handle);
59            let res = f(socket, &mut i.iface);
60            i.waker.wake();
61            res
62        })
63    }
64
65    /// Wait until the socket becomes readable.
66    ///
67    /// A socket is readable when a packet has been received, or when there are queued packets in
68    /// the buffer.
69    pub fn wait_recv_ready(&self) -> impl Future<Output = ()> + '_ {
70        poll_fn(move |cx| self.poll_recv_ready(cx))
71    }
72
73    /// Receive a datagram.
74    ///
75    /// This method will wait until a datagram is received.
76    pub async fn recv(&self, buf: &mut [u8]) -> Result<usize, RecvError> {
77        poll_fn(move |cx| self.poll_recv(buf, cx)).await
78    }
79
80    /// Wait until a datagram can be read.
81    ///
82    /// When no datagram is readable, this method will return `Poll::Pending` and
83    /// register the current task to be notified when a datagram is received.
84    ///
85    /// When a datagram is received, this method will return `Poll::Ready`.
86    pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
87        self.with_mut(|s, _| {
88            if s.can_recv() {
89                Poll::Ready(())
90            } else {
91                // socket buffer is empty wait until at least one byte has arrived
92                s.register_recv_waker(cx.waker());
93                Poll::Pending
94            }
95        })
96    }
97
98    /// Receive a datagram.
99    ///
100    /// When no datagram is available, this method will return `Poll::Pending` and
101    /// register the current task to be notified when a datagram is received.
102    pub fn poll_recv(&self, buf: &mut [u8], cx: &mut Context<'_>) -> Poll<Result<usize, RecvError>> {
103        self.with_mut(|s, _| match s.recv_slice(buf) {
104            Ok(n) => Poll::Ready(Ok(n)),
105            // No data ready
106            Err(raw::RecvError::Truncated) => Poll::Ready(Err(RecvError::Truncated)),
107            Err(raw::RecvError::Exhausted) => {
108                s.register_recv_waker(cx.waker());
109                Poll::Pending
110            }
111        })
112    }
113
114    /// Wait until the socket becomes writable.
115    ///
116    /// A socket becomes writable when there is space in the buffer, from initial memory or after
117    /// dispatching datagrams on a full buffer.
118    pub fn wait_send_ready(&self) -> impl Future<Output = ()> + '_ {
119        poll_fn(move |cx| self.poll_send_ready(cx))
120    }
121
122    /// Wait until a datagram can be sent.
123    ///
124    /// When no datagram can be sent (i.e. the buffer is full), this method will return
125    /// `Poll::Pending` and register the current task to be notified when
126    /// space is freed in the buffer after a datagram has been dispatched.
127    ///
128    /// When a datagram can be sent, this method will return `Poll::Ready`.
129    pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
130        self.with_mut(|s, _| {
131            if s.can_send() {
132                Poll::Ready(())
133            } else {
134                // socket buffer is full wait until a datagram has been dispatched
135                s.register_send_waker(cx.waker());
136                Poll::Pending
137            }
138        })
139    }
140
141    /// Send a datagram.
142    ///
143    /// This method will wait until the datagram has been sent.`
144    pub fn send<'s>(&'s self, buf: &'s [u8]) -> impl Future<Output = ()> + 's {
145        poll_fn(|cx| self.poll_send(buf, cx))
146    }
147
148    /// Send a datagram.
149    ///
150    /// When the datagram has been sent, this method will return `Poll::Ready(Ok())`.
151    ///
152    /// When the socket's send buffer is full, this method will return `Poll::Pending`
153    /// and register the current task to be notified when the buffer has space available.
154    pub fn poll_send(&self, buf: &[u8], cx: &mut Context<'_>) -> Poll<()> {
155        self.with_mut(|s, _| match s.send_slice(buf) {
156            // Entire datagram has been sent
157            Ok(()) => Poll::Ready(()),
158            Err(raw::SendError::BufferFull) => {
159                s.register_send_waker(cx.waker());
160                Poll::Pending
161            }
162        })
163    }
164
165    /// Flush the socket.
166    ///
167    /// This method will wait until the socket is flushed.
168    pub fn flush(&mut self) -> impl Future<Output = ()> + '_ {
169        poll_fn(|cx| {
170            self.with_mut(|s, _| {
171                if s.send_queue() == 0 {
172                    Poll::Ready(())
173                } else {
174                    s.register_send_waker(cx.waker());
175                    Poll::Pending
176                }
177            })
178        })
179    }
180}
181
182impl Drop for RawSocket<'_> {
183    fn drop(&mut self) {
184        self.stack.with_mut(|i| i.sockets.remove(self.handle));
185    }
186}
187
188fn _assert_covariant<'a, 'b: 'a>(x: RawSocket<'b>) -> RawSocket<'a> {
189    x
190}