feat(net): basic recv_msg/send_msg on UdpSocket

Author: AsakuraMizu

compio-net/src/socket.rs

@@ -6,8 +6,9 @@ use compio_driver::op::CreateSocket;
 use compio_driver::{
     impl_raw_fd,
     op::{
-        Accept, BufResultExt, CloseSocket, Connect, Recv, RecvFrom, RecvFromVectored,
-        RecvResultExt, RecvVectored, Send, SendTo, SendToVectored, SendVectored, ShutdownSocket,
+        Accept, BufResultExt, CloseSocket, Connect, MsgBuf, Recv, RecvFrom, RecvFromVectored,
+        RecvMsg, RecvMsgVectored, RecvResultExt, RecvVectored, Send, SendMsg, SendMsgVectored,
+        SendTo, SendToVectored, SendVectored, ShutdownSocket,
     },
     ToSharedFd,
 };
@@ -256,6 +257,36 @@ impl Socket {
             .map_advanced()
     }
 
+    pub async fn recv_msg<T: IoBufMut, C: IoBufMut>(
+        &self,
+        buffer: T,
+        control: C,
+    ) -> BufResult<(usize, SockAddr), (T, C)> {
+        let fd = self.to_shared_fd();
+        let op = RecvMsg::new(fd, MsgBuf::new(buffer, control));
+        compio_runtime::submit(op)
+            .await
+            .into_inner()
+            .map_addr()
+            .map_advanced()
+            .map_buffer(MsgBuf::into_tuple)
+    }
+
+    pub async fn recv_msg_vectored<T: IoVectoredBufMut, C: IoBufMut>(
+        &self,
+        buffer: T,
+        control: C,
+    ) -> BufResult<(usize, SockAddr), (T, C)> {
+        let fd = self.to_shared_fd();
+        let op = RecvMsgVectored::new(fd, MsgBuf::new(buffer, control));
+        compio_runtime::submit(op)
+            .await
+            .into_inner()
+            .map_addr()
+            .map_advanced()
+            .map_buffer(MsgBuf::into_tuple)
+    }
+
     pub async fn send_to<T: IoBuf>(&self, buffer: T, addr: &SockAddr) -> BufResult<usize, T> {
         let fd = self.to_shared_fd();
         let op = SendTo::new(fd, buffer, addr.clone());
@@ -271,6 +302,34 @@ impl Socket {
         let op = SendToVectored::new(fd, buffer, addr.clone());
         compio_runtime::submit(op).await.into_inner()
     }
+
+    pub async fn send_msg<T: IoBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: &SockAddr,
+    ) -> BufResult<usize, (T, C)> {
+        let fd = self.to_shared_fd();
+        let op = SendMsg::new(fd, MsgBuf::new(buffer, control), addr.clone());
+        compio_runtime::submit(op)
+            .await
+            .into_inner()
+            .map_buffer(MsgBuf::into_tuple)
+    }
+
+    pub async fn send_msg_vectored<T: IoVectoredBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: &SockAddr,
+    ) -> BufResult<usize, (T, C)> {
+        let fd = self.to_shared_fd();
+        let op = SendMsgVectored::new(fd, MsgBuf::new(buffer, control), addr.clone());
+        compio_runtime::submit(op)
+            .await
+            .into_inner()
+            .map_buffer(MsgBuf::into_tuple)
+    }
 }
 
 impl_raw_fd!(Socket, Socket2, socket, socket);

compio-net/src/udp.rs

@@ -222,6 +222,32 @@ impl UdpSocket {
             .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
     }
 
+    /// Receives a single datagram message and ancillary data on the socket. On
+    /// success, returns the number of bytes received and the origin.
+    pub async fn recv_msg<T: IoBufMut, C: IoBufMut>(
+        &self,
+        buffer: T,
+        control: C,
+    ) -> BufResult<(usize, SocketAddr), (T, C)> {
+        self.inner
+            .recv_msg(buffer, control)
+            .await
+            .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
+    }
+
+    /// Receives a single datagram message and ancillary data on the socket. On
+    /// success, returns the number of bytes received and the origin.
+    pub async fn recv_msg_vectored<T: IoVectoredBufMut, C: IoBufMut>(
+        &self,
+        buffer: T,
+        control: C,
+    ) -> BufResult<(usize, SocketAddr), (T, C)> {
+        self.inner
+            .recv_msg_vectored(buffer, control)
+            .await
+            .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
+    }
+
     /// Sends data on the socket to the given address. On success, returns the
     /// number of bytes sent.
     pub async fn send_to<T: IoBuf>(
@@ -249,6 +275,46 @@ impl UdpSocket {
         })
         .await
     }
+
+    /// Sends data on the socket to the given address accompanied by ancillary
+    /// data. On success, returns the number of bytes sent.
+    pub async fn send_msg<T: IoBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: impl ToSocketAddrsAsync,
+    ) -> BufResult<usize, (T, C)> {
+        super::first_addr_buf(
+            addr,
+            (buffer, control),
+            |addr, (buffer, control)| async move {
+                self.inner
+                    .send_msg(buffer, control, &SockAddr::from(addr))
+                    .await
+            },
+        )
+        .await
+    }
+
+    /// Sends data on the socket to the given address accompanied by ancillary
+    /// data. On success, returns the number of bytes sent.
+    pub async fn send_msg_vectored<T: IoVectoredBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: impl ToSocketAddrsAsync,
+    ) -> BufResult<usize, (T, C)> {
+        super::first_addr_buf(
+            addr,
+            (buffer, control),
+            |addr, (buffer, control)| async move {
+                self.inner
+                    .send_msg_vectored(buffer, control, &SockAddr::from(addr))
+                    .await
+            },
+        )
+        .await
+    }
 }
 
 impl_raw_fd!(UdpSocket, socket2::Socket, inner, socket);