From 6ddd0e978f5a4018a38b00f206f931ce54615d6f Mon Sep 17 00:00:00 2001
From: Icelk <main@icelk.dev>
Date: Fri, 21 Jul 2023 01:13:43 +0200
Subject: [PATCH] Add `msg_control` access for `UdpSocket::recvmsg`.

---
 src/io/recvmsg.rs | 25 +++++++++++++++++++------
 src/io/socket.rs  |  7 ++++---
 src/net/udp.rs    | 10 ++++++----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/io/recvmsg.rs b/src/io/recvmsg.rs
index 3cae2e50..d95c55d1 100644
--- a/src/io/recvmsg.rs
+++ b/src/io/recvmsg.rs
@@ -7,18 +7,23 @@ use std::{
     {boxed::Box, io, net::SocketAddr},
 };
 
-pub(crate) struct RecvMsg<T> {
+pub(crate) struct RecvMsg<T, U = Vec<u8>> {
     #[allow(dead_code)]
     fd: SharedFd,
     pub(crate) buf: Vec<T>,
     #[allow(dead_code)]
     io_slices: Vec<IoSliceMut<'static>>,
     pub(crate) socket_addr: Box<SockAddr>,
+    pub(crate) msg_control: Option<U>,
     pub(crate) msghdr: Box<libc::msghdr>,
 }
 
-impl<T: BoundedBufMut> Op<RecvMsg<T>> {
-    pub(crate) fn recvmsg(fd: &SharedFd, mut bufs: Vec<T>) -> io::Result<Op<RecvMsg<T>>> {
+impl<T: BoundedBufMut, U: BoundedBufMut> Op<RecvMsg<T, U>> {
+    pub(crate) fn recvmsg(
+        fd: &SharedFd,
+        mut bufs: Vec<T>,
+        mut msg_control: Option<U>,
+    ) -> io::Result<Op<RecvMsg<T, U>>> {
         use io_uring::{opcode, types};
 
         let mut io_slices = Vec::with_capacity(bufs.len());
@@ -35,6 +40,10 @@ impl<T: BoundedBufMut> Op<RecvMsg<T>> {
         msghdr.msg_iovlen = io_slices.len() as _;
         msghdr.msg_name = socket_addr.as_ptr() as *mut libc::c_void;
         msghdr.msg_namelen = socket_addr.len();
+        if let Some(msg_control) = &mut msg_control {
+            msghdr.msg_control = msg_control.stable_mut_ptr().cast();
+            msghdr.msg_controllen = msg_control.bytes_total();
+        }
 
         CONTEXT.with(|x| {
             x.handle().expect("Not in a runtime context").submit_op(
@@ -43,6 +52,7 @@ impl<T: BoundedBufMut> Op<RecvMsg<T>> {
                     buf: bufs,
                     io_slices,
                     socket_addr,
+                    msg_control,
                     msghdr,
                 },
                 |recv_from| {
@@ -57,11 +67,12 @@ impl<T: BoundedBufMut> Op<RecvMsg<T>> {
     }
 }
 
-impl<T> Completable for RecvMsg<T>
+impl<T, U> Completable for RecvMsg<T, U>
 where
     T: BoundedBufMut,
+    U: BoundedBufMut,
 {
-    type Output = BufResult<(usize, SocketAddr), Vec<T>>;
+    type Output = BufResult<(usize, SocketAddr, Option<U>), Vec<T>>;
 
     fn complete(self, cqe: CqeResult) -> Self::Output {
         // Convert the operation result to `usize`
@@ -71,6 +82,8 @@ where
 
         let socket_addr = (*self.socket_addr).as_socket();
 
+        let msg_control = self.msg_control;
+
         let res = res.map(|n| {
             let socket_addr: SocketAddr = socket_addr.unwrap();
 
@@ -89,7 +102,7 @@ where
                     break;
                 }
             }
-            (n, socket_addr)
+            (n, socket_addr, msg_control)
         });
 
         (res, bufs)
diff --git a/src/io/socket.rs b/src/io/socket.rs
index dda1bb36..1a95baf1 100644
--- a/src/io/socket.rs
+++ b/src/io/socket.rs
@@ -189,11 +189,12 @@ impl Socket {
         op.await
     }
 
-    pub(crate) async fn recvmsg<T: BoundedBufMut>(
+    pub(crate) async fn recvmsg<T: BoundedBufMut, U: BoundedBufMut>(
         &self,
         buf: Vec<T>,
-    ) -> crate::BufResult<(usize, SocketAddr), Vec<T>> {
-        let op = Op::recvmsg(&self.fd, buf).unwrap();
+        msg_control: Option<U>,
+    ) -> crate::BufResult<(usize, SocketAddr, Option<U>), Vec<T>> {
+        let op = Op::recvmsg(&self.fd, buf, msg_control).unwrap();
         op.await
     }
 
diff --git a/src/net/udp.rs b/src/net/udp.rs
index cb0cef66..c224b319 100644
--- a/src/net/udp.rs
+++ b/src/net/udp.rs
@@ -305,12 +305,14 @@ impl UdpSocket {
 
     /// Receives a single datagram message on the socket, into multiple buffers
     ///
-    /// On success, returns the number of bytes read and the origin.
-    pub async fn recvmsg<T: BoundedBufMut>(
+    /// On success, returns the number of bytes read, the origin, and the msg_control, as modified
+    /// by the kernel.
+    pub async fn recvmsg<T: BoundedBufMut, U: BoundedBufMut>(
         &self,
         buf: Vec<T>,
-    ) -> crate::BufResult<(usize, SocketAddr), Vec<T>> {
-        self.inner.recvmsg(buf).await
+        msg_control: Option<U>,
+    ) -> crate::BufResult<(usize, SocketAddr, Option<U>), Vec<T>> {
+        self.inner.recvmsg(buf, msg_control).await
     }
 
     /// Reads a packet of data from the socket into the buffer.