From dad57e8d1cc0045f53585efe6398b19d391e72af Mon Sep 17 00:00:00 2001
From: Fan Jiang <ProfFan@users.noreply.github.com>
Date: Sat, 19 Oct 2024 14:52:37 -0400
Subject: [PATCH] Add `send_queue`/`recv_queue` to UDP/ICMP/Raw Socket

---
 src/socket/icmp.rs           | 10 ++++++++++
 src/socket/raw.rs            | 10 ++++++++++
 src/socket/udp.rs            | 16 ++++++++++++++++
 src/storage/packet_buffer.rs |  5 +++++
 4 files changed, 41 insertions(+)

diff --git a/src/socket/icmp.rs b/src/socket/icmp.rs
index 2a9258e49..dee3416a3 100644
--- a/src/socket/icmp.rs
+++ b/src/socket/icmp.rs
@@ -412,6 +412,16 @@ impl<'a> Socket<'a> {
         Ok((length, endpoint))
     }
 
+    /// Return the amount of octets queued in the transmit buffer.
+    pub fn send_queue(&self) -> usize {
+        self.tx_buffer.payload_bytes_count()
+    }
+
+    /// Return the amount of octets queued in the receive buffer.
+    pub fn recv_queue(&self) -> usize {
+        self.rx_buffer.payload_bytes_count()
+    }
+
     /// Fitler determining whether the socket accepts a given ICMPv4 packet.
     /// Accepted packets are enqueued into the socket's receive buffer.
     #[cfg(feature = "proto-ipv4")]
diff --git a/src/socket/raw.rs b/src/socket/raw.rs
index 7df7830e3..c79f99b95 100644
--- a/src/socket/raw.rs
+++ b/src/socket/raw.rs
@@ -327,6 +327,16 @@ impl<'a> Socket<'a> {
         Ok(length)
     }
 
+    /// Return the amount of octets queued in the transmit buffer.
+    pub fn send_queue(&self) -> usize {
+        self.tx_buffer.payload_bytes_count()
+    }
+
+    /// Return the amount of octets queued in the receive buffer.
+    pub fn recv_queue(&self) -> usize {
+        self.rx_buffer.payload_bytes_count()
+    }
+
     pub(crate) fn accepts(&self, ip_repr: &IpRepr) -> bool {
         if ip_repr.version() != self.ip_version {
             return false;
diff --git a/src/socket/udp.rs b/src/socket/udp.rs
index 92ccfe657..1e09116f3 100644
--- a/src/socket/udp.rs
+++ b/src/socket/udp.rs
@@ -459,6 +459,22 @@ impl<'a> Socket<'a> {
         Ok((length, endpoint))
     }
 
+    /// Return the amount of octets queued in the transmit buffer.
+    ///
+    /// Note that the Berkeley sockets interface does not have an equivalent of this API.
+    pub fn send_queue(&self) -> usize {
+        self.tx_buffer.payload_bytes_count()
+    }
+
+    /// Return the amount of octets queued in the receive buffer. This value can be larger than
+    /// the slice read by the next `recv` or `peek` call because it includes all queued octets,
+    /// and not only the octets that may be returned as a contiguous slice.
+    ///
+    /// Note that the Berkeley sockets interface does not have an equivalent of this API.
+    pub fn recv_queue(&self) -> usize {
+        self.rx_buffer.payload_bytes_count()
+    }
+
     pub(crate) fn accepts(&self, cx: &mut Context, ip_repr: &IpRepr, repr: &UdpRepr) -> bool {
         if self.endpoint.port != repr.dst_port {
             return false;
diff --git a/src/storage/packet_buffer.rs b/src/storage/packet_buffer.rs
index 28119fa10..f9121e611 100644
--- a/src/storage/packet_buffer.rs
+++ b/src/storage/packet_buffer.rs
@@ -239,6 +239,11 @@ impl<'a, H> PacketBuffer<'a, H> {
         self.payload_ring.capacity()
     }
 
+    /// Return the current number of bytes in the payload ring buffer.
+    pub fn payload_bytes_count(&self) -> usize {
+        self.payload_ring.len()
+    }
+
     /// Reset the packet buffer and clear any staged.
     #[allow(unused)]
     pub(crate) fn reset(&mut self) {