From b4ee17fb4f1a77c26fc08e9fe9e0d343d1c059b4 Mon Sep 17 00:00:00 2001
From: Fan Jiang <ProfFan@users.noreply.github.com>
Date: Sun, 20 Oct 2024 12:29:38 -0400
Subject: [PATCH] net: Add flush for UDP and Raw sockets.

---
 embassy-net/Cargo.toml |  2 +-
 embassy-net/src/raw.rs | 17 +++++++++++++++++
 embassy-net/src/udp.rs | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/embassy-net/Cargo.toml b/embassy-net/Cargo.toml
index a33c693fc..8a29aca79 100644
--- a/embassy-net/Cargo.toml
+++ b/embassy-net/Cargo.toml
@@ -68,7 +68,7 @@ multicast = ["smoltcp/multicast"]
 defmt = { version = "0.3", optional = true }
 log = { version = "0.4.14", optional = true }
 
-smoltcp = { git="https://github.com/smoltcp-rs/smoltcp", rev="b65e1b64dc9b66fa984a2ad34e90685cb0b606de", default-features = false, features = [
+smoltcp = { git="https://github.com/smoltcp-rs/smoltcp", rev="fe0b4d102253465850cd1cf39cd33d4721a4a8d5", default-features = false, features = [
   "socket",
   "async",
 ] }
diff --git a/embassy-net/src/raw.rs b/embassy-net/src/raw.rs
index 1098dc208..ace325a46 100644
--- a/embassy-net/src/raw.rs
+++ b/embassy-net/src/raw.rs
@@ -108,6 +108,23 @@ impl<'a> RawSocket<'a> {
             }
         })
     }
+
+    /// Flush the socket.
+    ///
+    /// This method will wait until the socket is flushed.
+    pub async fn flush(&mut self) {
+        poll_fn(move |cx| {
+            self.with_mut(|s, _| {
+                if s.send_queue() == 0 {
+                    Poll::Ready(())
+                } else {
+                    s.register_send_waker(cx.waker());
+                    Poll::Pending
+                }
+            })
+        })
+        .await
+    }
 }
 
 impl Drop for RawSocket<'_> {
diff --git a/embassy-net/src/udp.rs b/embassy-net/src/udp.rs
index 3eb6e2f83..ac650364e 100644
--- a/embassy-net/src/udp.rs
+++ b/embassy-net/src/udp.rs
@@ -241,6 +241,23 @@ impl<'a> UdpSocket<'a> {
         .await
     }
 
+    /// Flush the socket.
+    ///
+    /// This method will wait until the socket is flushed.
+    pub async fn flush(&mut self) {
+        poll_fn(move |cx| {
+            self.with_mut(|s, _| {
+                if s.send_queue() == 0 {
+                    Poll::Ready(())
+                } else {
+                    s.register_send_waker(cx.waker());
+                    Poll::Pending
+                }
+            })
+        })
+        .await
+    }
+
     /// Returns the local endpoint of the socket.
     pub fn endpoint(&self) -> IpListenEndpoint {
         self.with(|s, _| s.endpoint())