From 002c94f533e8d27ad3d1ba285009f642328d6479 Mon Sep 17 00:00:00 2001
From: FrankReh <FrankReh@users.noreply.github.com>
Date: Sat, 3 Dec 2022 17:03:54 -0500
Subject: [PATCH] net: UDP socket gets local_addr, read_fixed and write_fixed
 (#189)

The type UdpSocket gets methods local_addr, read_fixed and write_fixed.

local_addr is useful when binding the socket to 127.0.0.1:0 to be able
to report what port the OS assigned.

The read_fixed and write_fixed are useful, like read and write, when the
UDP socket has been connected to an address.
---
 src/net/udp.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/src/net/udp.rs b/src/net/udp.rs
index 2df8e150..47d4dee6 100644
--- a/src/net/udp.rs
+++ b/src/net/udp.rs
@@ -1,4 +1,5 @@
 use crate::{
+    buf::fixed::FixedBuf,
     buf::{BoundedBuf, BoundedBufMut},
     io::{SharedFd, Socket},
 };
@@ -96,6 +97,34 @@ impl UdpSocket {
         Ok(UdpSocket { inner: socket })
     }
 
+    /// Returns the local address that this UDP socket is bound to.
+    ///
+    /// This can be useful, for example, when binding to port 0 to
+    /// figure out which port was actually bound.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+    /// use tokio_uring::net::UdpSocket;
+    ///
+    /// tokio_uring::start(async {
+    ///     let socket = UdpSocket::bind("127.0.0.1:8080".parse().unwrap()).await.unwrap();
+    ///     let addr = socket.local_addr().expect("Couldn't get local address");
+    ///     assert_eq!(addr, SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080)));
+    /// });
+    /// ```
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        let fd = self.inner.as_raw_fd();
+        // SAFETY: Our fd is the handle the kernel has given us for a UdpSocket.
+        // Create a std::net::UdpSocket long enough to call its local_addr method
+        // and then forget it so the socket is not closed here.
+        let s = unsafe { std::net::UdpSocket::from_raw_fd(fd) };
+        let local_addr = s.local_addr();
+        std::mem::forget(s);
+        local_addr
+    }
+
     /// Creates new `UdpSocket` from a previously bound `std::net::UdpSocket`.
     ///
     /// This function is intended to be used to wrap a UDP socket from the
@@ -206,12 +235,48 @@ impl UdpSocket {
         self.inner.read(buf).await
     }
 
+    /// Like [`read`], but using a pre-mapped buffer
+    /// registered with [`FixedBufRegistry`].
+    ///
+    /// [`read`]: Self::read
+    /// [`FixedBufRegistry`]: crate::buf::fixed::FixedBufRegistry
+    ///
+    /// # Errors
+    ///
+    /// In addition to errors that can be reported by `read`,
+    /// this operation fails if the buffer is not registered in the
+    /// current `tokio-uring` runtime.
+    pub async fn read_fixed<T>(&self, buf: T) -> crate::BufResult<usize, T>
+    where
+        T: BoundedBufMut<BufMut = FixedBuf>,
+    {
+        self.inner.read_fixed(buf).await
+    }
+
     /// Write some data to the socket from the buffer, returning the original buffer and
     /// quantity of data written.
     pub async fn write<T: BoundedBuf>(&self, buf: T) -> crate::BufResult<usize, T> {
         self.inner.write(buf).await
     }
 
+    /// Like [`write`], but using a pre-mapped buffer
+    /// registered with [`FixedBufRegistry`].
+    ///
+    /// [`write`]: Self::write
+    /// [`FixedBufRegistry`]: crate::buf::fixed::FixedBufRegistry
+    ///
+    /// # Errors
+    ///
+    /// In addition to errors that can be reported by `write`,
+    /// this operation fails if the buffer is not registered in the
+    /// current `tokio-uring` runtime.
+    pub async fn write_fixed<T>(&self, buf: T) -> crate::BufResult<usize, T>
+    where
+        T: BoundedBuf<Buf = FixedBuf>,
+    {
+        self.inner.write_fixed(buf).await
+    }
+
     /// Shuts down the read, write, or both halves of this connection.
     ///
     /// This function will cause all pending and future I/O on the specified portions to return