From b81c14f442839de61f2217cf88e09d390d465f3d Mon Sep 17 00:00:00 2001
From: Roy Buitenhuis <roy.buitenhuis@technolution.nl>
Date: Wed, 12 Jul 2023 11:32:02 +0200
Subject: [PATCH] Add polling fn's for send_to and recv_from in UdpSocket.

---
 embassy-net/src/udp.rs | 27 +++++++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index bd7a32f7c..16f48bcb0 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -3,7 +3,7 @@
 use core::cell::RefCell;
 use core::future::poll_fn;
 use core::mem;
-use core::task::Poll;
+use core::task::{Poll, Context};
 
 use embassy_net_driver::Driver;
 use smoltcp::iface::{Interface, SocketHandle};
@@ -105,7 +105,14 @@ impl<'a> UdpSocket<'a> {
         poll_fn(move |cx| self.poll_recv_from(buf, cx)).await
     }
 
-    pub fn poll_recv_from(&self, buf: &mut[u8], cx: Context) -> Poll<Result<(usize, IpEndpoint), Error>> {
+    /// 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
@@ -117,18 +124,30 @@ impl<'a> UdpSocket<'a> {
     }
 
     /// 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: IpEndpoint = remote_endpoint.into();
         poll_fn(move |cx| self.poll_send_to(buf, remote_endpoint, cx)).await
     }
 
-    pub fn poll_send_to<T>(&self, buf: &[u8], remote_endpoint: T, cx: Context) -> Poll<Result<(), Error>>
+    /// 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>,
     {
-        let remote_endpoint = remote_endpoint.into();
         self.with_mut(|s, _| match s.send_slice(buf, remote_endpoint) {
             // Entire datagram has been sent
             Ok(()) => Poll::Ready(Ok(())),