diff --git a/crates/host_env/src/socket.rs b/crates/host_env/src/socket.rs index 4e6ba75f613..b16175b1ff3 100644 --- a/crates/host_env/src/socket.rs +++ b/crates/host_env/src/socket.rs @@ -921,42 +921,199 @@ pub fn checked_cmsg_space(len: usize) -> Option { } #[cfg(all(unix, not(target_os = "redox")))] -pub fn pack_ancillary_messages(cmsgs: &[(i32, i32, &[u8])]) -> Result, AncillaryPackError> { - use core::{mem, ptr}; - - if cmsgs.is_empty() { - return Ok(vec![]); - } +fn ancillary_capacity(cmsgs: &[(i32, i32, &[u8])]) -> Result { + cmsgs.iter().try_fold(0usize, |sum, (_, _, data)| { + let space = checked_cmsg_space(data.len()).ok_or(AncillaryPackError::ItemTooLarge)?; + sum.checked_add(space) + // POSIX permits msg_controllen to be a signed integer, and some + // targets expose a narrower field than usize. + .filter(|&len| i32::try_from(len).is_ok()) + .ok_or(AncillaryPackError::TooMuchData) + }) +} - let capacity = cmsgs - .iter() - .map(|(_, _, buf)| buf.len()) - .try_fold(0usize, |sum, len| { - let space = checked_cmsg_space(len).ok_or(AncillaryPackError::ItemTooLarge)?; - usize::checked_add(sum, space).ok_or(AncillaryPackError::TooMuchData) - })?; - - let mut cmsg_buffer = vec![0u8; capacity]; - let mut mhdr = unsafe { mem::zeroed::() }; - mhdr.msg_control = cmsg_buffer.as_mut_ptr().cast(); - mhdr.msg_controllen = capacity as _; - - let mut pmhdr: *mut libc::cmsghdr = unsafe { libc::CMSG_FIRSTHDR(&mhdr) }; - for (lvl, typ, data) in cmsgs { - if pmhdr.is_null() { - return Err(AncillaryPackError::UnexpectedNullHeader); +#[cfg(all(unix, not(target_os = "redox")))] +fn encode_ancillary_messages( + cmsgs: &[(i32, i32, &[u8])], + control: &mut [u8], +) -> Result<(), AncillaryPackError> { + let data_offset = checked_cmsg_len(0).ok_or(AncillaryPackError::ItemTooLarge)?; + let mut offset = 0usize; + for (level, kind, data) in cmsgs { + let length = checked_cmsg_len(data.len()).ok_or(AncillaryPackError::ItemTooLarge)?; + let space = checked_cmsg_space(data.len()).ok_or(AncillaryPackError::ItemTooLarge)?; + let end = offset + .checked_add(space) + .ok_or(AncillaryPackError::TooMuchData)?; + let slot = control + .get_mut(offset..end) + .ok_or(AncillaryPackError::TooMuchData)?; + let payload_end = data_offset + .checked_add(data.len()) + .ok_or(AncillaryPackError::ItemTooLarge)?; + if slot.len() < core::mem::size_of::() + || payload_end != length + || payload_end > slot.len() + { + return Err(AncillaryPackError::ItemTooLarge); } - let cmsg_len = checked_cmsg_len(data.len()).ok_or(AncillaryPackError::ItemTooLarge)?; + let header = slot.as_mut_ptr().cast::(); + // SAFETY: The slot contains a complete header. Field pointers may be + // unaligned, including on targets without padding between messages. unsafe { - (*pmhdr).cmsg_level = *lvl; - (*pmhdr).cmsg_type = *typ; - (*pmhdr).cmsg_len = cmsg_len as _; - ptr::copy_nonoverlapping(data.as_ptr(), libc::CMSG_DATA(pmhdr), data.len()); - pmhdr = libc::CMSG_NXTHDR(&mhdr, pmhdr); + core::ptr::addr_of_mut!((*header).cmsg_level).write_unaligned(*level); + core::ptr::addr_of_mut!((*header).cmsg_type).write_unaligned(*kind); + core::ptr::addr_of_mut!((*header).cmsg_len).write_unaligned(length as _); } + slot[data_offset..payload_end].copy_from_slice(data); + offset = end; } + Ok(()) +} - Ok(cmsg_buffer) +/// A control buffer with enough alignment for platform CMSG layouts. +#[cfg(all(unix, not(target_os = "redox")))] +pub struct AlignedAncillaryBuffer { + chunks: Vec, + len: usize, +} + +// Some supported platforms require 16-byte CMSG alignment, even when +// cmsghdr itself has a smaller Rust alignment. +#[cfg(all(unix, not(target_os = "redox")))] +#[repr(C, align(16))] +#[derive(Clone)] +struct AncillaryChunk { + _bytes: [u8; 16], +} + +#[cfg(all(unix, not(target_os = "redox")))] +impl AlignedAncillaryBuffer { + fn zeroed(len: usize) -> Option { + let count = len.div_ceil(core::mem::size_of::()); + count + .checked_mul(core::mem::size_of::()) + .filter(|&size| isize::try_from(size).is_ok())?; + Some(Self { + chunks: vec![AncillaryChunk { _bytes: [0; 16] }; count], + len, + }) + } + + pub fn as_bytes(&self) -> &[u8] { + // SAFETY: Every byte in each chunk is initialized. The logical length + // is at most the size of the allocation and the chunks outlive the slice. + unsafe { core::slice::from_raw_parts(self.chunks.as_ptr().cast(), self.len) } + } + + fn as_mut_bytes(&mut self) -> &mut [u8] { + // SAFETY: The chunks are fully initialized and exclusively borrowed; + // len fits in their allocation and every bit pattern is valid for u8. + unsafe { core::slice::from_raw_parts_mut(self.chunks.as_mut_ptr().cast(), self.len) } + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } +} + +#[cfg(all(unix, not(target_os = "redox")))] +pub fn pack_ancillary_messages_aligned( + cmsgs: &[(i32, i32, &[u8])], +) -> Result { + let len = ancillary_capacity(cmsgs)?; + let mut control = AlignedAncillaryBuffer::zeroed(len).ok_or(AncillaryPackError::TooMuchData)?; + encode_ancillary_messages(cmsgs, control.as_mut_bytes())?; + Ok(control) +} + +#[cfg(all(unix, not(target_os = "redox")))] +/// Encode control messages as bytes. The returned allocation does not promise +/// CMSG alignment; use `pack_ancillary_messages_aligned` for a syscall buffer. +pub fn pack_ancillary_messages(cmsgs: &[(i32, i32, &[u8])]) -> Result, AncillaryPackError> { + let len = ancillary_capacity(cmsgs)?; + let mut control = vec![0; len]; + encode_ancillary_messages(cmsgs, &mut control)?; + Ok(control) +} + +#[cfg(test)] +#[cfg(all(unix, not(target_os = "redox")))] +mod ancillary_pack_tests { + use super::{ + AlignedAncillaryBuffer, checked_cmsg_len, checked_cmsg_space, encode_ancillary_messages, + pack_ancillary_messages, pack_ancillary_messages_aligned, + }; + + #[test] + fn encodes_aligned_and_unaligned_storage_identically() { + let messages: &[(i32, i32, &[u8])] = &[(1, 2, b""), (3, 4, b"abc"), (5, 6, b"z")]; + let aligned = pack_ancillary_messages_aligned(messages).unwrap(); + let bytes = aligned.as_bytes(); + assert_eq!(bytes.as_ptr() as usize % 16, 0); + assert_eq!(pack_ancillary_messages(messages).unwrap().as_slice(), bytes); + + let mut offset = 0; + let data_offset = checked_cmsg_len(0).unwrap(); + for (level, kind, data) in messages { + let length = checked_cmsg_len(data.len()).unwrap(); + let space = checked_cmsg_space(data.len()).unwrap(); + assert!(offset + core::mem::size_of::() <= bytes.len()); + // SAFETY: Each slot contains a complete header, which may be + // unaligned on platforms that do not pad between messages. + let header = unsafe { + bytes + .as_ptr() + .add(offset) + .cast::() + .read_unaligned() + }; + assert_eq!((header.cmsg_level, header.cmsg_type), (*level, *kind)); + #[allow( + clippy::unnecessary_cast, + reason = "cmsg_len is u32 on some Unix targets" + )] + let header_len = header.cmsg_len as usize; + assert_eq!(header_len, length); + assert_eq!(&bytes[offset + data_offset..offset + length], *data); + assert!( + bytes[offset + length..offset + space] + .iter() + .all(|&b| b == 0) + ); + offset += space; + } + assert_eq!(offset, bytes.len()); + + let align = core::mem::align_of::(); + let mut storage = vec![0u8; bytes.len() + align]; + let shift = (0..align) + .find(|&n| !(storage.as_ptr() as usize + n).is_multiple_of(align)) + .unwrap_or(0); + let unaligned = &mut storage[shift..shift + bytes.len()]; + encode_ancillary_messages(messages, unaligned).unwrap(); + assert_eq!(&*unaligned, bytes); + } + + #[test] + fn empty_control_buffers_are_empty() { + assert!(pack_ancillary_messages(&[]).unwrap().is_empty()); + let aligned = pack_ancillary_messages_aligned(&[]).unwrap(); + assert!(aligned.is_empty()); + assert!(aligned.as_bytes().is_empty()); + } + + #[test] + fn receive_control_storage_is_aligned_and_initialized() { + for len in [0, 1, 16, 17] { + let mut buffer = AlignedAncillaryBuffer::zeroed(len).unwrap(); + let bytes = buffer.as_mut_bytes(); + assert_eq!(bytes.len(), len); + assert_eq!(bytes.as_ptr() as usize % 16, 0); + assert!(bytes.iter().all(|&byte| byte == 0)); + } + assert!(AlignedAncillaryBuffer::zeroed(usize::MAX).is_none()); + } } #[cfg(all(unix, not(target_os = "redox")))] @@ -1049,11 +1206,13 @@ pub fn recvmsg_into( .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::()) - ]; + #[allow(clippy::useless_conversion)] + let controllen = ancbufsize + .try_into() + .map_err(|_| io::Error::from(io::ErrorKind::InvalidInput))?; + // CMSG alignment can exceed cmsghdr alignment on some targets. + let mut anc_buf = AlignedAncillaryBuffer::zeroed(ancbufsize) + .ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?; 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(); @@ -1067,23 +1226,14 @@ pub fn recvmsg_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(); - #[allow(clippy::useless_conversion)] - let controllen = ancbufsize - .try_into() - .map_err(|_| io::Error::from(io::ErrorKind::InvalidInput))?; + msg.msg_control = anc_buf.as_mut_bytes().as_mut_ptr().cast(); 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 control = unsafe { - core::slice::from_raw_parts( - anc_buf.as_ptr().cast::(), - (msg.msg_controllen as usize).min(ancbufsize), - ) - }; + let control = &anc_buf.as_bytes()[..(msg.msg_controllen as usize).min(ancbufsize)]; let ancdata = parse_ancillary_messages(control); let address = if msg.msg_namelen > 0 { let storage = unsafe { diff --git a/crates/stdlib/src/socket.rs b/crates/stdlib/src/socket.rs index 1b79f6ded52..4a07ac0db6d 100644 --- a/crates/stdlib/src/socket.rs +++ b/crates/stdlib/src/socket.rs @@ -2390,7 +2390,7 @@ mod _socket { .collect::>>()?; control_buf = Self::pack_cmsgs_to_send(&cmsgs, vm)?; if !control_buf.is_empty() { - msg = msg.with_control(&control_buf); + msg = msg.with_control(control_buf.as_bytes()); } } @@ -2590,10 +2590,7 @@ mod _socket { fn pack_cmsgs_to_send( cmsgs: &[(i32, i32, ArgBytesLike)], vm: &VirtualMachine, - ) -> PyResult> { - if cmsgs.is_empty() { - return Ok(vec![]); - } + ) -> PyResult { let data = cmsgs .iter() .map(|(lvl, typ, buf)| { @@ -2606,7 +2603,8 @@ mod _socket { .map(|(lvl, typ, data)| (*lvl, *typ, data.as_slice())) .collect::>(); - host_socket::pack_ancillary_messages(&data_refs).map_err(|err| err.to_pyexception(vm)) + host_socket::pack_ancillary_messages_aligned(&data_refs) + .map_err(|err| err.to_pyexception(vm)) } #[pymethod]