Implement socket.recvmsg_into on Unix (#9268)

Assisted-by: Codex:GPT-6

crates/host_env/src/socket.rs

@@ -877,6 +877,15 @@ pub struct RecvMsgResult {
     pub address: Option<RawSocketAddress>,
 }
 
+#[cfg(all(unix, not(target_os = "redox")))]
+#[derive(Debug)]
+pub struct RecvMsgIntoResult {
+    pub bytes_read: usize,
+    pub ancdata: Vec<AncillaryMessage>,
+    pub msg_flags: i32,
+    pub address: Option<RawSocketAddress>,
+}
+
 #[cfg(all(unix, not(target_os = "redox")))]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum AncillaryPackError {
@@ -1001,33 +1010,79 @@ pub fn recvmsg(
         MaybeUninit::uninit()
     };
     let mut data_buf = vec![initial; bufsize];
-    let mut anc_buf: Vec<MaybeUninit<u8>> = vec![MaybeUninit::uninit(); ancbufsize];
+    let result = recvmsg_into(fd, &mut [&mut data_buf], ancbufsize, flags)?;
+    let data = unsafe {
+        // Datagram MSG_TRUNC may report more bytes than the supplied storage.
+        data_buf.set_len(result.bytes_read.min(bufsize));
+        core::mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(data_buf)
+    };
+    Ok(RecvMsgResult {
+        data,
+        ancdata: result.ancdata,
+        msg_flags: result.msg_flags,
+        address: result.address,
+    })
+}
+
+#[cfg(all(unix, not(target_os = "redox")))]
+pub fn recvmsg_into(
+    fd: BorrowedFd<'_>,
+    buffers: &mut [&mut [core::mem::MaybeUninit<u8>]],
+    ancbufsize: usize,
+    flags: i32,
+) -> io::Result<RecvMsgIntoResult> {
+    use core::mem::MaybeUninit;
+
+    let as_iovec = |buf: &mut [MaybeUninit<u8>]| libc::iovec {
+        iov_base: buf.as_mut_ptr().cast(),
+        iov_len: buf.len(),
+    };
+    let mut single;
+    let mut multiple;
+    let iov: &mut [libc::iovec] = if let [buf] = buffers {
+        single = [as_iovec(buf)];
+        &mut single
+    } else {
+        multiple = buffers
+            .iter_mut()
+            .map(|buf| as_iovec(buf))
+            .collect::<Vec<_>>();
+        &mut multiple
+    };
+    // CMSG_* requires aligned headers; initialize padding before exposing bytes.
+    let mut anc_buf = vec![
+        MaybeUninit::<libc::cmsghdr>::zeroed();
+        ancbufsize.div_ceil(core::mem::size_of::<libc::cmsghdr>())
+    ];
     let mut addr_storage: libc::sockaddr_storage = unsafe { core::mem::zeroed() };
-
-    let mut iov = [libc::iovec {
-        iov_base: data_buf.as_mut_ptr().cast(),
-        iov_len: bufsize,
-    }];
-
     let mut msg: libc::msghdr = unsafe { core::mem::zeroed() };
     msg.msg_name = (&mut addr_storage as *mut libc::sockaddr_storage).cast();
     msg.msg_namelen = core::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
     msg.msg_iov = iov.as_mut_ptr();
-    msg.msg_iovlen = 1;
+    // These msghdr fields are narrower on some Unix targets.
+    #[allow(clippy::useless_conversion)]
+    let iovlen = iov
+        .len()
+        .try_into()
+        .map_err(|_| io::Error::from(io::ErrorKind::InvalidInput))?;
+    msg.msg_iovlen = iovlen;
     if ancbufsize > 0 {
         msg.msg_control = anc_buf.as_mut_ptr().cast();
-        msg.msg_controllen = ancbufsize as _;
+        #[allow(clippy::useless_conversion)]
+        let controllen = ancbufsize
+            .try_into()
+            .map_err(|_| io::Error::from(io::ErrorKind::InvalidInput))?;
+        msg.msg_controllen = controllen;
     }
 
+    // SAFETY: all iovecs refer to disjoint writable slices; the control and address
+    // buffers are aligned and live until the syscall returns.
     let ret = unsafe { libc::recvmsg(fd.as_raw_fd(), &mut msg, flags) }.check_libc_neg()?;
-
-    let data = unsafe {
-        // Datagram MSG_TRUNC may report more bytes than the supplied storage.
-        data_buf.set_len((ret as usize).min(bufsize));
-        core::mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(data_buf)
-    };
     let control = unsafe {
-        core::slice::from_raw_parts(anc_buf.as_ptr().cast::<u8>(), msg.msg_controllen as usize)
+        core::slice::from_raw_parts(
+            anc_buf.as_ptr().cast::<u8>(),
+            (msg.msg_controllen as usize).min(ancbufsize),
+        )
     };
     let ancdata = parse_ancillary_messages(control);
     let address = if msg.msg_namelen > 0 {
@@ -1041,9 +1096,8 @@ pub fn recvmsg(
     } else {
         None
     };
-
-    Ok(RecvMsgResult {
-        data,
+    Ok(RecvMsgIntoResult {
+        bytes_read: ret as usize,
         ancdata,
         msg_flags: msg.msg_flags,
         address,

crates/stdlib/src/socket.rs

@@ -2468,7 +2468,7 @@ mod _socket {
             msg.data.shrink_to_fit();
 
             // Build ancdata list
-            let ancdata = Self::parse_ancillary_data(&msg.ancdata, vm);
+            let ancdata = Self::parse_ancillary_data(msg.ancdata, vm);
 
             // Build address tuple
             let address = if let Some(address) = msg.address {
@@ -2488,10 +2488,85 @@ mod _socket {
             ]))
         }
 
+        #[cfg(all(unix, not(target_os = "redox")))]
+        #[pymethod]
+        fn recvmsg_into(
+            zelf: &Py<Self>,
+            buffers: PyObjectRef,
+            ancbufsize: OptionalArg<isize>,
+            flags: OptionalArg<i32>,
+            vm: &VirtualMachine,
+        ) -> PyResult<PyTupleRef> {
+            let buffer_items = vm
+                .extract_elements_with(&buffers, Ok)
+                .map_err(|_| vm.new_type_error("recvmsg_into() argument 1 must be an iterable"))?;
+            let buffers = buffer_items
+                .iter()
+                .map(|obj| ArgMemoryBuffer::try_from_object(vm, obj.clone()))
+                .collect::<PyResult<Vec<_>>>()?;
+            let ancbufsize = ancbufsize.unwrap_or(0);
+            if !(0..=i32::MAX as isize).contains(&ancbufsize) {
+                return Err(vm.new_value_error("invalid ancillary data buffer length"));
+            }
+            let flags = flags.unwrap_or(0);
+            let total = buffers
+                .iter()
+                .try_fold(0usize, |sum, buf| sum.checked_add(buf.len()))
+                .filter(|&total| isize::try_from(total).is_ok())
+                .ok_or_else(|| vm.new_overflow_error("recvmsg_into buffers are too large"))?;
+            let mut scratch = alloc_recv_scratch(total, flags, vm)?;
+            let msg = zelf
+                .sock_op(vm, SockWaitKind::Read, || {
+                    let mut remaining = &mut scratch.spare_capacity_mut()[..total];
+                    let mut slices: Vec<_> = buffers
+                        .iter()
+                        .map(|buf| {
+                            let (part, rest) =
+                                core::mem::take(&mut remaining).split_at_mut(buf.len());
+                            remaining = rest;
+                            part
+                        })
+                        .collect();
+                    let sock = zelf.sock_snapshot()?;
+                    let fd = unsafe { std::os::fd::BorrowedFd::borrow_raw(sock_fileno(&sock)) };
+                    host_socket::recvmsg_into(fd, &mut slices, ancbufsize as usize, flags)
+                })
+                .map_err(|e| e.into_pyexception(vm))?;
+            // MSG_TRUNC can report a larger datagram or consume TCP data without
+            // writing. alloc_recv_scratch initializes that case before the syscall.
+            unsafe { scratch.set_len(msg.bytes_read.min(total)) };
+            if !zelf.recv_discards_data(flags) {
+                let mut remaining = scratch.as_slice();
+                for buffer in &buffers {
+                    let written = remaining.len().min(buffer.len());
+                    if written == 0 {
+                        continue;
+                    }
+                    buffer.borrow_buf_mut()[..written].copy_from_slice(&remaining[..written]);
+                    remaining = &remaining[written..];
+                }
+            }
+            let ancdata = Self::parse_ancillary_data(msg.ancdata, vm);
+            let address = if let Some(address) = msg.address {
+                let storage: host_socket::raw::SockAddrStorage =
+                    unsafe { core::mem::transmute(address.storage) };
+                let addr = unsafe { host_socket::raw::SockAddr::new(storage, address.len as _) };
+                get_addr_tuple(&addr, zelf.proto.load(Ordering::Relaxed), vm)
+            } else {
+                vm.ctx.none()
+            };
+            Ok(vm.ctx.new_tuple(vec![
+                vm.ctx.new_int(msg.bytes_read).into(),
+                ancdata,
+                vm.ctx.new_int(msg.msg_flags).into(),
+                address,
+            ]))
+        }
+
         /// Parse ancillary data from a received message header
         #[cfg(all(unix, not(target_os = "redox")))]
         fn parse_ancillary_data(
-            control: &[host_socket::AncillaryMessage],
+            control: Vec<host_socket::AncillaryMessage>,
             vm: &VirtualMachine,
         ) -> PyObjectRef {
             let mut result = Vec::new();
@@ -2499,7 +2574,7 @@ mod _socket {
                 let tuple = vm.ctx.new_tuple(vec![
                     vm.ctx.new_int(cmsg.level).into(),
                     vm.ctx.new_int(cmsg.kind).into(),
-                    vm.ctx.new_bytes(cmsg.data.clone()).into(),
+                    vm.ctx.new_bytes(cmsg.data).into(),
                 ]);
                 result.push(tuple.into());
             }

extra_tests/snippets/stdlib_socket.py

@@ -296,3 +296,22 @@ if hasattr(socket, "NETLINK_USERSOCK"):
         assert receiver.recvfrom(32) == (b"message", sender.getsockname())
         sender.connect(receiver.getsockname())
         assert sender.getpeername() == receiver.getsockname()
+
+
+# Exhaust the iterable before exporting its writable buffers.
+if hasattr(socket.socket, "recvmsg_into"):
+    first, second = bytearray(2), bytearray(2)
+
+    def receive_buffers():
+        yield first
+        first.extend(b"xx")
+        yield second
+
+    reader, writer = socket.socketpair()
+    try:
+        writer.sendall(b"abcdef")
+        assert reader.recvmsg_into(receive_buffers())[0] == 6
+        assert (first, second) == (b"abcd", b"ef")
+    finally:
+        reader.close()
+        writer.close()