Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 72 additions & 18 deletions crates/host_env/src/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 mut addr_storage: libc::sockaddr_storage = unsafe { core::mem::zeroed() };
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,
})
}

let mut iov = [libc::iovec {
iov_base: data_buf.as_mut_ptr().cast(),
iov_len: bufsize,
}];
#[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 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 {
Expand All @@ -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,
Expand Down
81 changes: 78 additions & 3 deletions crates/stdlib/src/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -2488,18 +2488,93 @@ 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"))?;
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved
Comment thread
1ndahous3 marked this conversation as resolved.
Resolved
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();
for cmsg in control {
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());
}
Expand Down
19 changes: 19 additions & 0 deletions extra_tests/snippets/stdlib_socket.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,3 +296,22 @@ def send_after_mutation(signum, frame):
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()
Loading