From 49a9dc6743a8d90c46a51a42706943acf39a5d85 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Fri, 31 Dec 2021 20:47:34 +0900
Subject: [PATCH] net: add buffer size methods to UdpSocket (#4363)

This adds the following methods:

- UdpSocket::set_send_buffer_size
- UdpSocket::send_buffer_size
- UdpSocket::set_recv_buffer_size
- UdpSocket::recv_buffer_size
---
 tokio/src/net/tcp/stream.rs |  6 ++--
 tokio/src/net/udp.rs        | 72 +++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/tokio/src/net/tcp/stream.rs b/tokio/src/net/tcp/stream.rs
index b47d97d5564..abfc3c6612b 100644
--- a/tokio/src/net/tcp/stream.rs
+++ b/tokio/src/net/tcp/stream.rs
@@ -1090,7 +1090,7 @@ impl TcpStream {
     /// # }
     /// ```
     pub fn linger(&self) -> io::Result<Option<Duration>> {
-        let socket = self.to_socket();
+        let socket = self.as_socket();
         socket.linger()
     }
 
@@ -1116,11 +1116,11 @@ impl TcpStream {
     /// # }
     /// ```
     pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
-        let socket = self.to_socket();
+        let socket = self.as_socket();
         socket.set_linger(dur)
     }
 
-    fn to_socket(&self) -> socket2::SockRef<'_> {
+    fn as_socket(&self) -> socket2::SockRef<'_> {
         socket2::SockRef::from(self)
     }
 
diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs
index 7dc72af35bd..24de488aaa7 100644
--- a/tokio/src/net/udp.rs
+++ b/tokio/src/net/udp.rs
@@ -253,6 +253,78 @@ impl UdpSocket {
         }
     }
 
+    /// Sets the size of the UDP send buffer on this socket.
+    ///
+    /// On most operating systems, this sets the `SO_SNDBUF` socket option.
+    pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> {
+        self.as_socket().set_send_buffer_size(size as usize)
+    }
+
+    /// Returns the size of the UDP send buffer for this socket.
+    ///
+    /// On most operating systems, this is the value of the `SO_SNDBUF` socket
+    /// option.
+    ///
+    /// Note that if [`set_send_buffer_size`] has been called on this socket
+    /// previously, the value returned by this function may not be the same as
+    /// the argument provided to `set_send_buffer_size`. This is for the
+    /// following reasons:
+    ///
+    /// * Most operating systems have minimum and maximum allowed sizes for the
+    ///   send buffer, and will clamp the provided value if it is below the
+    ///   minimum or above the maximum. The minimum and maximum buffer sizes are
+    ///   OS-dependent.
+    /// * Linux will double the buffer size to account for internal bookkeeping
+    ///   data, and returns the doubled value from `getsockopt(2)`. As per `man
+    ///   7 socket`:
+    ///   > Sets or gets the maximum socket send buffer in bytes. The
+    ///   > kernel doubles this value (to allow space for bookkeeping
+    ///   > overhead) when it is set using `setsockopt(2)`, and this doubled
+    ///   > value is returned by `getsockopt(2)`.
+    ///
+    /// [`set_send_buffer_size`]: Self::set_send_buffer_size
+    pub fn send_buffer_size(&self) -> io::Result<u32> {
+        self.as_socket().send_buffer_size().map(|n| n as u32)
+    }
+
+    /// Sets the size of the UDP receive buffer on this socket.
+    ///
+    /// On most operating systems, this sets the `SO_RCVBUF` socket option.
+    pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> {
+        self.as_socket().set_recv_buffer_size(size as usize)
+    }
+
+    /// Returns the size of the UDP receive buffer for this socket.
+    ///
+    /// On most operating systems, this is the value of the `SO_RCVBUF` socket
+    /// option.
+    ///
+    /// Note that if [`set_recv_buffer_size`] has been called on this socket
+    /// previously, the value returned by this function may not be the same as
+    /// the argument provided to `set_send_buffer_size`. This is for the
+    /// following reasons:
+    ///
+    /// * Most operating systems have minimum and maximum allowed sizes for the
+    ///   receive buffer, and will clamp the provided value if it is below the
+    ///   minimum or above the maximum. The minimum and maximum buffer sizes are
+    ///   OS-dependent.
+    /// * Linux will double the buffer size to account for internal bookkeeping
+    ///   data, and returns the doubled value from `getsockopt(2)`. As per `man
+    ///   7 socket`:
+    ///   > Sets or gets the maximum socket send buffer in bytes. The
+    ///   > kernel doubles this value (to allow space for bookkeeping
+    ///   > overhead) when it is set using `setsockopt(2)`, and this doubled
+    ///   > value is returned by `getsockopt(2)`.
+    ///
+    /// [`set_recv_buffer_size`]: Self::set_recv_buffer_size
+    pub fn recv_buffer_size(&self) -> io::Result<u32> {
+        self.as_socket().recv_buffer_size().map(|n| n as u32)
+    }
+
+    fn as_socket(&self) -> socket2::SockRef<'_> {
+        socket2::SockRef::from(self)
+    }
+
     /// Returns the local address that this socket is bound to.
     ///
     /// # Example