feat(driver): add recv_from_managed operation support (#709)

* feat(driver): add recv_from_managed operation support

* refactor: use mop to define the RecvFromManaged

Author: Sherlock-Holo

compio-driver/src/op.rs

@@ -25,7 +25,7 @@ pub use crate::sys::op::{
     WriteVectored, WriteVectoredAt,
 };
 #[cfg(io_uring)]
-pub use crate::sys::op::{ReadManaged, ReadManagedAt, RecvManaged};
+pub use crate::sys::op::{ReadManaged, ReadManagedAt, RecvFromManaged, RecvManaged};
 use crate::{Extra, OwnedFd, SharedFd, TakeBuffer, sys::aio::*};
 
 /// Trait to update the buffer length inside the [`BufResult`].
@@ -471,9 +471,10 @@ pub(crate) mod managed {
 
     use compio_buf::IntoInner;
     use pin_project_lite::pin_project;
+    use socket2::SockAddr;
 
-    use super::{Read, ReadAt, Recv};
-    use crate::{BorrowedBuffer, BufferPool, OwnedBuffer, TakeBuffer};
+    use super::{Read, ReadAt, Recv, RecvFrom};
+    use crate::{AsFd, BorrowedBuffer, BufferPool, OwnedBuffer, TakeBuffer};
 
     pin_project! {
         /// Read a file at specified position into managed buffer.
@@ -600,6 +601,48 @@ pub(crate) mod managed {
             Ok(res)
         }
     }
+
+    pin_project! {
+        /// Receive data and source address into managed buffer.
+        pub struct RecvFromManaged<S: AsFd> {
+            #[pin]
+            pub(crate) op: RecvFrom<OwnedBuffer, S>,
+        }
+    }
+
+    impl<S: AsFd> RecvFromManaged<S> {
+        /// Create [`RecvFromManaged`].
+        pub fn new(fd: S, pool: &BufferPool, len: usize, flags: i32) -> io::Result<Self> {
+            #[cfg(fusion)]
+            let pool = pool.as_poll();
+            Ok(Self {
+                op: RecvFrom::new(fd, pool.get_buffer(len)?, flags),
+            })
+        }
+    }
+
+    impl<S: AsFd> TakeBuffer for RecvFromManaged<S> {
+        type Buffer<'a> = (BorrowedBuffer<'a>, SockAddr);
+        type BufferPool = BufferPool;
+
+        fn take_buffer(
+            self,
+            buffer_pool: &Self::BufferPool,
+            result: io::Result<usize>,
+            _: u16,
+        ) -> io::Result<Self::Buffer<'_>> {
+            let result = result?;
+            #[cfg(fusion)]
+            let buffer_pool = buffer_pool.as_poll();
+            let (slice, addr_buffer, addr_size) = self.op.into_inner();
+            let addr = unsafe { SockAddr::new(addr_buffer, addr_size) };
+            // SAFETY: result is valid
+            let res = unsafe { buffer_pool.create_proxy(slice, result) };
+            #[cfg(fusion)]
+            let res = BorrowedBuffer::new_poll(res);
+            Ok((res, addr))
+        }
+    }
 }
 
 #[cfg(not(io_uring))]

compio-driver/src/sys/fusion/op.rs

@@ -115,6 +115,9 @@ op!(<S: AsFd> PathStat(dirfd: S, path: CString, follow_symlink: bool));
 
 macro_rules! mop {
     (<$($ty:ident: $trait:ident),* $(,)?> $name:ident( $($arg:ident: $arg_t:ty),* $(,)? ) with $pool:ident) => {
+        mop!{ < $($ty: $trait),* > $name ( $( $arg: $arg_t ),* ) with $pool, buffer: crate::BorrowedBuffer<'a> }
+    };
+    (<$($ty:ident: $trait:ident),* $(,)?> $name:ident( $($arg:ident: $arg_t:ty),* $(,)? ) with $pool:ident, buffer: $buffer:ty) => {
         ::paste::paste!{
             enum [< $name Inner >] <$($ty: $trait),*> {
                 Poll(crate::op::managed::$name<$($ty),*>),
@@ -159,7 +162,7 @@ macro_rules! mop {
 
             impl<$($ty: $trait),*> crate::TakeBuffer for $name<$($ty),*> {
                 type BufferPool = crate::BufferPool;
-                type Buffer<'a> = crate::BorrowedBuffer<'a>;
+                type Buffer<'a> = $buffer;
 
                 fn take_buffer(
                     self,
@@ -206,3 +209,4 @@ macro_rules! mop {
 mop!(<S: AsFd> ReadManagedAt(fd: S, offset: u64, pool: &BufferPool, len: usize) with pool);
 mop!(<S: AsFd> ReadManaged(fd: S, pool: &BufferPool, len: usize) with pool);
 mop!(<S: AsFd> RecvManaged(fd: S, pool: &BufferPool, len: usize, flags: i32) with pool);
+mop!(<S: AsFd> RecvFromManaged(fd: S, pool: &BufferPool, len: usize, flags: i32) with pool, buffer: (crate::BorrowedBuffer<'a>, SockAddr));

compio-driver/src/sys/iocp/op.rs

@@ -578,6 +578,16 @@ unsafe impl<S: AsFd> OpCode for RecvManaged<S> {
     }
 }
 
+unsafe impl<S: AsFd> OpCode for RecvFromManaged<S> {
+    unsafe fn operate(self: Pin<&mut Self>, optr: *mut OVERLAPPED) -> Poll<io::Result<usize>> {
+        unsafe { self.project().op.operate(optr) }
+    }
+
+    fn cancel(self: Pin<&mut Self>, optr: *mut OVERLAPPED) -> io::Result<()> {
+        self.project().op.cancel(optr)
+    }
+}
+
 pin_project! {
     /// Receive data from remote into vectored buffer.
     pub struct RecvVectored<T: IoVectoredBufMut, S> {

compio-driver/src/sys/iour/op.rs

@@ -917,6 +917,8 @@ mod buf_ring {
     };
 
     use io_uring::{opcode, squeue::Flags, types::Fd};
+    use pin_project_lite::pin_project;
+    use socket2::{SockAddr, SockAddrStorage, socklen_t};
 
     use super::OpCode;
     use crate::{BorrowedBuffer, BufferPool, OpEntry, TakeBuffer};
@@ -1096,6 +1098,83 @@ mod buf_ring {
             res
         }
     }
+
+    pin_project! {
+        /// Receive data and source address into managed buffer.
+        pub struct RecvFromManaged<S> {
+            fd: S,
+            buffer_group: u16,
+            flags: i32,
+            addr: SockAddrStorage,
+            addr_len: socklen_t,
+            iovec: libc::iovec,
+            msg: libc::msghdr,
+            _p: PhantomPinned,
+        }
+    }
+
+    impl<S> RecvFromManaged<S> {
+        /// Create [`RecvFromManaged`].
+        pub fn new(fd: S, buffer_pool: &BufferPool, len: usize, flags: i32) -> io::Result<Self> {
+            #[cfg(fusion)]
+            let buffer_pool = buffer_pool.as_io_uring();
+            let len: u32 = len.try_into().map_err(|_| {
+                io::Error::new(io::ErrorKind::InvalidInput, "required length too long")
+            })?;
+            let addr = SockAddrStorage::zeroed();
+            Ok(Self {
+                fd,
+                buffer_group: buffer_pool.buffer_group(),
+                flags,
+                addr_len: addr.size_of() as _,
+                addr,
+                iovec: libc::iovec {
+                    iov_base: ptr::null_mut(),
+                    iov_len: len as _,
+                },
+                msg: unsafe { std::mem::zeroed() },
+                _p: PhantomPinned,
+            })
+        }
+    }
+
+    unsafe impl<S: AsFd> OpCode for RecvFromManaged<S> {
+        fn create_entry(self: Pin<&mut Self>) -> OpEntry {
+            let this = self.project();
+            this.msg.msg_name = this.addr as *mut _ as _;
+            this.msg.msg_namelen = *this.addr_len;
+            this.msg.msg_iov = this.iovec as *const _ as *mut _;
+            this.msg.msg_iovlen = 1;
+            opcode::RecvMsg::new(Fd(this.fd.as_fd().as_raw_fd()), this.msg)
+                .flags(*this.flags as _)
+                .buf_group(*this.buffer_group)
+                .build()
+                .flags(Flags::BUFFER_SELECT)
+                .into()
+        }
+    }
+
+    impl<S> TakeBuffer for RecvFromManaged<S> {
+        type Buffer<'a> = (BorrowedBuffer<'a>, SockAddr);
+        type BufferPool = BufferPool;
+
+        fn take_buffer(
+            self,
+            buffer_pool: &Self::BufferPool,
+            result: io::Result<usize>,
+            buffer_id: u16,
+        ) -> io::Result<Self::Buffer<'_>> {
+            #[cfg(fusion)]
+            let buffer_pool = buffer_pool.as_io_uring();
+            let result = result.inspect_err(|_| buffer_pool.reuse_buffer(buffer_id))?;
+            let addr = unsafe { SockAddr::new(self.addr, self.addr_len) };
+            // SAFETY: result is valid
+            let buffer = unsafe { buffer_pool.get_buffer(buffer_id, result) }?;
+            #[cfg(fusion)]
+            let buffer = BorrowedBuffer::new_io_uring(buffer);
+            Ok((buffer, addr))
+        }
+    }
 }
 
-pub use buf_ring::{ReadManaged, ReadManagedAt, RecvManaged};
+pub use buf_ring::{ReadManaged, ReadManagedAt, RecvFromManaged, RecvManaged};

compio-driver/src/sys/poll/op.rs

@@ -858,6 +858,20 @@ unsafe impl<S: AsFd> OpCode for crate::op::managed::RecvManaged<S> {
     }
 }
 
+unsafe impl<S: AsFd> OpCode for crate::op::managed::RecvFromManaged<S> {
+    fn pre_submit(self: Pin<&mut Self>) -> io::Result<Decision> {
+        self.project().op.pre_submit()
+    }
+
+    fn op_type(self: Pin<&mut Self>) -> Option<OpType> {
+        self.project().op.op_type()
+    }
+
+    fn operate(self: Pin<&mut Self>) -> Poll<io::Result<usize>> {
+        self.project().op.operate()
+    }
+}
+
 impl<T: IoVectoredBufMut, S: AsFd> RecvVectored<T, S> {
     unsafe fn call(self: Pin<&mut Self>) -> libc::ssize_t {
         let this = self.project();

compio-driver/src/sys/stub/op.rs

@@ -256,3 +256,5 @@ impl<S: AsFd> OpCode for crate::op::managed::ReadManagedAt<S> {}
 impl<S: AsFd> OpCode for crate::op::managed::ReadManaged<S> {}
 
 impl<S: AsFd> OpCode for crate::op::managed::RecvManaged<S> {}
+
+impl<S: AsFd> OpCode for crate::op::managed::RecvFromManaged<S> {}