2 files changed, 112 insertions, 3 deletions

diff --git a/embassy-net/src/tcp.rs b/embassy-net/src/tcp.rs
index 043062e06..a00ced8f4 100644
--- a/embassy-net/src/tcp.rs
+++ b/embassy-net/src/tcp.rs
@@ -10,7 +10,7 @@
 
 use core::future::poll_fn;
 use core::mem;
-use core::task::Poll;
+use core::task::{Context, Poll};
 
 use embassy_time::Duration;
 use smoltcp::iface::{Interface, SocketHandle};
@@ -274,6 +274,16 @@ impl<'a> TcpSocket<'a> {
         .await
     }
 
+    /// Wait until the socket becomes readable.
+    ///
+    /// A socket becomes readable when the receive half of the full-duplex connection is open
+    /// (see [may_recv](#method.may_recv)), and there is some pending data in the receive buffer.
+    ///
+    /// This is the equivalent of [read](#method.read), without buffering any data.
+    pub async fn wait_read_ready(&self) {
+        poll_fn(move |cx| self.io.poll_read_ready(cx)).await
+    }
+
     /// Read data from the socket.
     ///
     /// Returns how many bytes were read, or an error. If no data is available, it waits
@@ -285,6 +295,16 @@ impl<'a> TcpSocket<'a> {
         self.io.read(buf).await
     }
 
+    /// Wait until the socket becomes writable.
+    ///
+    /// A socket becomes writable when the transmit half of the full-duplex connection is open
+    /// (see [may_send](#method.may_send)), and the transmit buffer is not full.
+    ///
+    /// This is the equivalent of [write](#method.write), without sending any data.
+    pub async fn wait_write_ready(&self) {
+        poll_fn(move |cx| self.io.poll_write_ready(cx)).await
+    }
+
     /// Write data to the socket.
     ///
     /// Returns how many bytes were written, or an error. If the socket is not ready to
@@ -376,11 +396,25 @@ impl<'a> TcpSocket<'a> {
         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.
+    /// Return whether the transmit half of the full-duplex connection is open.
+    ///
+    /// This function returns true if it's possible to send data and have it arrive
+    /// to the remote endpoint. However, it does not make any guarantees about the state
+    /// of the transmit buffer, and even if it returns true, [write](#method.write) may
+    /// not be able to enqueue any octets.
+    ///
+    /// In terms of the TCP state machine, the socket must be in the `ESTABLISHED` or
+    /// `CLOSE-WAIT` state.
     pub fn may_send(&self) -> bool {
         self.io.with(|s, _| s.may_send())
     }
 
+    /// Check whether the transmit half of the full-duplex connection is open
+    /// (see [may_send](#method.may_send)), and the transmit buffer is not full.
+    pub fn can_send(&self) -> bool {
+        self.io.with(|s, _| s.can_send())
+    }
+
     /// return whether the receive 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,
@@ -427,7 +461,7 @@ impl<'d> TcpIo<'d> {
         })
     }
 
-    fn with_mut<R>(&mut self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
+    fn with_mut<R>(&self, f: impl FnOnce(&mut tcp::Socket, &mut Interface) -> R) -> R {
         self.stack.with_mut(|i| {
             let socket = i.sockets.get_mut::<tcp::Socket>(self.handle);
             let res = f(socket, &mut i.iface);
@@ -436,6 +470,17 @@ impl<'d> TcpIo<'d> {
         })
     }
 
+    fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_recv() {
+                Poll::Ready(())
+            } else {
+                s.register_recv_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
         poll_fn(move |cx| {
             // CAUTION: smoltcp semantics around EOF are different to what you'd expect
@@ -464,6 +509,17 @@ impl<'d> TcpIo<'d> {
         .await
     }
 
+    fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_send() {
+                Poll::Ready(())
+            } else {
+                s.register_send_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     async fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
         poll_fn(move |cx| {
             self.with_mut(|s, _| match s.send_slice(buf) {
diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index ac650364e..76602edc2 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -103,6 +103,32 @@ impl<'a> UdpSocket<'a> {
         })
     }
 
+    /// Wait until the socket becomes readable.
+    ///
+    /// A socket is readable when a packet has been received, or when there are queued packets in
+    /// the buffer.
+    pub async fn wait_recv_ready(&self) {
+        poll_fn(move |cx| self.poll_recv_ready(cx)).await
+    }
+
+    /// Wait until a datagram can be read.
+    ///
+    /// When no datagram is readable, 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`.
+    pub fn poll_recv_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_recv() {
+                Poll::Ready(())
+            } else {
+                // socket buffer is empty wait until at least one byte has arrived
+                s.register_recv_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     /// Receive a datagram.
     ///
     /// This method will wait until a datagram is received.
@@ -164,6 +190,33 @@ impl<'a> UdpSocket<'a> {
         .await
     }
 
+    /// Wait until the socket becomes writable.
+    ///
+    /// A socket becomes writable when there is space in the buffer, from initial memory or after
+    /// dispatching datagrams on a full buffer.
+    pub async fn wait_send_ready(&self) {
+        poll_fn(move |cx| self.poll_send_ready(cx)).await
+    }
+
+    /// Wait until a datagram can be sent.
+    ///
+    /// When no datagram can be sent (i.e. the buffer is full), this method will return
+    /// `Poll::Pending` and register the current task to be notified when
+    /// space is freed in the buffer after a datagram has been dispatched.
+    ///
+    /// When a datagram can be sent, this method will return `Poll::Ready`.
+    pub fn poll_send_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
+        self.with_mut(|s, _| {
+            if s.can_send() {
+                Poll::Ready(())
+            } else {
+                // socket buffer is full wait until a datagram has been dispatched
+                s.register_send_waker(cx.waker());
+                Poll::Pending
+            }
+        })
+    }
+
     /// Send a datagram to the specified remote endpoint.
     ///
     /// This method will wait until the datagram has been sent.