Add async socket layer

dgram/Cargo.toml

@@ -3,11 +3,8 @@ name = "dgram"
 version = "0.1.0"
 edition = "2021"
 
-[features]
-async = []
-
 [dependencies]
 libc = "0.2.76"
 nix = "0.26.2"
 smallvec = { version = "1.10", features = ["union"] }
-
+tokio = { version = "1.29", features = ["full", "test-util"] }

dgram/src/async_socket.rs

@@ -0,0 +1,94 @@
+use crate::syscalls::RecvData;
+use std::io::ErrorKind;
+use std::io::Result;
+use std::net::SocketAddr;
+use std::os::fd::AsRawFd;
+use std::time::Instant;
+use tokio::net::UdpSocket;
+
+#[cfg(target_os = "linux")]
+mod linux_imports {
+    pub(super) use crate::syscalls::recv_msg;
+    pub(super) use crate::syscalls::send_msg;
+    pub(super) use nix::sys::socket::MsgFlags;
+    pub(super) use nix::sys::socket::SockaddrStorage;
+    pub(super) use tokio::io::Interest;
+}
+
+#[cfg(target_os = "linux")]
+use self::linux_imports::*;
+
+#[cfg(target_os = "linux")]
+pub async fn send_to(
+    socket: &UdpSocket, send_buf: &[u8], segment_size: usize, num_pkts: usize,
+    tx_time: Option<Instant>, client_addr: &SocketAddr,
+) -> Result<usize> {
+    loop {
+        // Important to use try_io so that Tokio can clear the socket's readiness
+        // flag
+        let res = socket.try_io(Interest::WRITABLE, || {
+            let fd = socket.as_raw_fd();
+            send_msg(
+                fd,
+                send_buf,
+                segment_size,
+                num_pkts,
+                tx_time,
+                &SockaddrStorage::from(*client_addr),
+            )
+            .map_err(Into::into)
+        });
+
+        match res {
+            Err(e) if e.kind() == ErrorKind::WouldBlock =>
+                socket.writable().await?,
+            res => return res,
+        }
+    }
+}
+
+#[cfg(target_os = "linux")]
+pub async fn recv_from(
+    socket: &UdpSocket, read_buf: &mut [u8], cmsg_space: &mut Vec<u8>,
+    msg_flags: MsgFlags,
+) -> Result<RecvData> {
+    loop {
+        // Important to use try_io so that Tokio can clear the socket's readiness
+        // flag
+        let res = socket.try_io(Interest::READABLE, || {
+            let fd = socket.as_raw_fd();
+            recv_msg(fd, read_buf, cmsg_space, Some(msg_flags))
+                .map_err(Into::into)
+        });
+
+        match res {
+            Err(e) if e.kind() == ErrorKind::WouldBlock =>
+                socket.readable().await?,
+            _ => return res,
+        }
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+pub async fn send_to(
+    socket: &tokio::net::UdpSocket, client_addr: SocketAddr, send_buf: &[u8],
+    _segment_size: usize, _num_pkts: usize, _tx_time: Option<Instant>,
+) -> io::Result<usize> {
+    socket.send_to(send_buf, client_addr).await
+}
+
+// Signature changes because we can't use MessageFlags outside of a *NIX context
+#[cfg(not(target_os = "linux"))]
+pub async fn recv_from(
+    socket: &UdpSocket, read_buf: &mut [u8], _cmsg_space: &mut Vec<u8>,
+) -> Result<RecvData> {
+    let recv = socket.recv(read_buf).await?;
+
+    Ok(RecvData {
+        bytes: recv,
+        peer_addr: None,
+        cmsgs: vec![],
+        gro: None,
+        rx_time: None,
+    })
+}

dgram/src/syscalls.rs

@@ -313,27 +313,4 @@ mod tests {
 
         Ok(())
     }
-
-    #[test]
-    fn recv_from_control_message() -> Result<()> {
-        let (send, recv) = new_sockets()?;
-        let addr = getsockname::<SockaddrStorage>(recv).unwrap();
-
-        let send_buf = b"jets";
-        let iov = [IoSlice::new(send_buf)];
-        sendmsg(send, &iov, &vec![], MsgFlags::empty(), Some(&addr))?;
-
-        let mut cmsg_space = cmsg_space!(TimeVal);
-        let mut read_buf = [0; 4];
-
-        let recv_data = recv_msg(recv, &mut read_buf, &mut cmsg_space, None)?;
-
-        println!("cmsgs: {:?}", recv_data.cmsgs);
-
-        // TODO: test cmsgs get transferred to cmsg_space
-        assert_eq!(recv_data.bytes, 4);
-        assert_eq!(&read_buf, b"jets");
-
-        Ok(())
-    }
 }