From b0142d0464593941fe7a97b741b21ab4af2ec4a6 Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Wed, 21 Oct 2026 06:01:27 +0300 Subject: [PATCH] Implement socket.recvmsg_into on Unix (#9268) Assisted-by: Codex:GPT-6 --- crates/host_env/src/socket.rs | 92 +++++++++++++++++++++------ crates/stdlib/src/socket.rs | 81 ++++++++++++++++++++++- extra_tests/snippets/stdlib_socket.py | 19 ++++++ 3 files changed, 170 insertions(+), 22 deletions(-) diff --git a/crates/host_env/src/socket.rs b/crates/host_env/src/socket.rs index 12ca8d0fb..4e6ba75f6 100644 --- a/crates/host_env/src/socket.rs +++ b/crates/host_env/src/socket.rs @@ -877,6 +877,15 @@ pub struct RecvMsgResult { pub address: Option, } +#[cfg(all(unix, not(target_os = "redox")))] +#[derive(Debug)] +pub struct RecvMsgIntoResult { + pub bytes_read: usize, + pub ancdata: Vec, + pub msg_flags: i32, + pub address: Option, +} + #[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> = 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>(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]], + ancbufsize: usize, + flags: i32, +) -> io::Result { + use core::mem::MaybeUninit; + + let as_iovec = |buf: &mut [MaybeUninit]| 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::>(); + &mut multiple + }; + // CMSG_* requires aligned headers; initialize padding before exposing bytes. + let mut anc_buf = vec![ + MaybeUninit::::zeroed(); + ancbufsize.div_ceil(core::mem::size_of::()) + ]; 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::() 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>(data_buf) - }; let control = unsafe { - core::slice::from_raw_parts(anc_buf.as_ptr().cast::(), msg.msg_controllen as usize) + core::slice::from_raw_parts( + anc_buf.as_ptr().cast::(), + (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, diff --git a/crates/stdlib/src/socket.rs b/crates/stdlib/src/socket.rs index f5aa1c0f5..427b5d7a2 100644 --- a/crates/stdlib/src/socket.rs +++ b/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, + buffers: PyObjectRef, + ancbufsize: OptionalArg, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + 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::>>()?; + 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, 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()); } diff --git a/extra_tests/snippets/stdlib_socket.py b/extra_tests/snippets/stdlib_socket.py index d81d57157..4a1fcd710 100644 --- a/extra_tests/snippets/stdlib_socket.py +++ b/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()