Skip to content

Commit

Permalink
Reduce unsafety
Browse files Browse the repository at this point in the history
  • Loading branch information
Gleb Pomykalov committed Apr 22, 2020
1 parent 101a405 commit cbeb0d2
Showing 1 changed file with 74 additions and 73 deletions.
147 changes: 74 additions & 73 deletions src/sys/socket/mod.rs
View file
Expand Up @@ -774,17 +774,13 @@ impl<'a> ControlMessage<'a> {
pub fn sendmsg(fd: RawFd, iov: &[IoVec<&[u8]>], cmsgs: &[ControlMessage],
flags: MsgFlags, addr: Option<&SockAddr>) -> Result<usize>
{
let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();

let capacity = cmsgs.iter().map(|c| c.space()).sum();

// First size the buffer needed to hold the cmsgs. It must be zeroed,
// because subsequent code will not clear the padding bytes.
let mut cmsg_buffer = vec![0u8; capacity];

pack_mhdr_to_send(mhdr.as_mut_ptr(), &mut cmsg_buffer[..], &iov, &cmsgs, addr);

let mhdr = unsafe { mhdr.assume_init() };
let mhdr = pack_mhdr_to_send(&mut cmsg_buffer[..], &iov, &cmsgs, addr);

let ret = unsafe { libc::sendmsg(fd, &mhdr, flags.bits()) };

Expand Down Expand Up @@ -836,8 +832,7 @@ pub struct SendMmsgData<'a, I, C>
))]
pub fn sendmmsg<'a, I, C>(
fd: RawFd,
data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>,
IntoIter=impl ExactSizeIterator + Iterator<Item=&'a SendMmsgData<'a, I, C>>>,
data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>>,
flags: MsgFlags
) -> Result<Vec<usize>>
where
Expand All @@ -846,31 +841,29 @@ pub fn sendmmsg<'a, I, C>(
{
let iter = data.into_iter();

let num_messages = iter.len();
let size_hint = iter.size_hint();
let reserve_items = size_hint.1.unwrap_or(size_hint.0);

let mut output = Vec::<libc::mmsghdr>::with_capacity(num_messages);
unsafe {
output.set_len(num_messages);
}
let mut output = Vec::<libc::mmsghdr>::with_capacity(reserve_items);

let mut cmsgs_buffer = vec![0u8; 0];

iter.enumerate().for_each(|(i, d)| {
let element = &mut output[i];

for d in iter {
let cmsgs_start = cmsgs_buffer.len();
let cmsgs_required_capacity: usize = d.cmsgs.as_ref().iter().map(|c| c.space()).sum();
let cmsgs_buffer_need_capacity = cmsgs_start + cmsgs_required_capacity;
cmsgs_buffer.resize(cmsgs_buffer_need_capacity, 0);

pack_mhdr_to_send(
&mut element.msg_hdr,
&mut cmsgs_buffer[cmsgs_start..],
&d.iov,
&d.cmsgs,
d.addr.as_ref()
);
});
output.push(libc::mmsghdr {
msg_hdr: pack_mhdr_to_send(
&mut cmsgs_buffer[cmsgs_start..],
&d.iov,
&d.cmsgs,
d.addr.as_ref()
),
msg_len: 0,
});
};

let ret = unsafe { libc::sendmmsg(fd, output.as_mut_ptr(), output.len() as _, flags.bits() as _) };

Expand Down Expand Up @@ -947,25 +940,27 @@ pub fn recvmmsg<'a, I>(
let num_messages = iter.len();

let mut output: Vec<libc::mmsghdr> = Vec::with_capacity(num_messages);
let mut address: Vec<sockaddr_storage> = Vec::with_capacity(num_messages);

unsafe {
output.set_len(num_messages);
address.set_len(num_messages);
}
// Addresses should be pre-allocated and never change the address during building
// of the input data for `recvmmsg`
let mut addresses: Vec<sockaddr_storage> = vec![unsafe { mem::zeroed() }; num_messages];

let results: Vec<_> = iter.enumerate().map(|(i, d)| {
let element = &mut output[i];

let msg_controllen = unsafe {
let (msg_controllen, mhdr) = unsafe {
pack_mhdr_to_receive(
&mut element.msg_hdr,
d.iov.as_ref(),
&mut d.cmsg_buffer,
&mut address[i]
&mut addresses[i],
)
};

output.push(
libc::mmsghdr {
msg_hdr: mhdr,
msg_len: 0,
}
);

(msg_controllen as usize, &mut d.cmsg_buffer)
}).collect();

Expand All @@ -981,15 +976,15 @@ pub fn recvmmsg<'a, I>(

Ok(output
.into_iter()
.zip(address.into_iter())
.zip(addresses.into_iter())
.zip(results.into_iter())
.map(|((mmsghdr, mut address), (msg_controllen, cmsg_buffer))| {
.map(|((mmsghdr, address), (msg_controllen, cmsg_buffer))| {
unsafe {
read_mhdr(
mmsghdr.msg_hdr,
r as isize,
msg_controllen,
&mut address,
address,
cmsg_buffer
)
}
Expand All @@ -1001,7 +996,7 @@ unsafe fn read_mhdr<'a, 'b>(
mhdr: msghdr,
r: isize,
msg_controllen: usize,
address: *mut sockaddr_storage,
address: sockaddr_storage,
cmsg_buffer: &'a mut Option<&'b mut Vec<u8>>
) -> RecvMsg<'b> {
let cmsghdr = {
Expand All @@ -1020,7 +1015,7 @@ unsafe fn read_mhdr<'a, 'b>(
};

let address = sockaddr_storage_to_addr(
&*address ,
&address ,
mhdr.msg_namelen as usize
).ok();

Expand All @@ -1034,42 +1029,46 @@ unsafe fn read_mhdr<'a, 'b>(
}

unsafe fn pack_mhdr_to_receive<'a, I>(
out: *mut msghdr,
iov: I,
cmsg_buffer: &mut Option<&mut Vec<u8>>,
address: *mut sockaddr_storage,
) -> usize
) -> (usize, msghdr)
where
I: AsRef<[IoVec<&'a mut [u8]>]> + 'a,
{
let (msg_control, msg_controllen) = cmsg_buffer.as_mut()
.map(|v| (v.as_mut_ptr(), v.capacity()))
.unwrap_or((ptr::null_mut(), 0));

(*out).msg_name = address as *mut c_void;
(*out).msg_namelen = mem::size_of::<sockaddr_storage>() as socklen_t;
(*out).msg_iov = iov.as_ref().as_ptr() as *mut iovec;
(*out).msg_iovlen = iov.as_ref().len() as _;
(*out).msg_control = msg_control as *mut c_void;
(*out).msg_controllen = msg_controllen as _;
(*out).msg_flags = 0;
let mhdr = {
// Musl's msghdr has private fields, so this is the only way to
// initialize it.
let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();
let p = mhdr.as_mut_ptr();
(*p).msg_name = address as *mut c_void;
(*p).msg_namelen = mem::size_of::<sockaddr_storage>() as socklen_t;
(*p).msg_iov = iov.as_ref().as_ptr() as *mut iovec;
(*p).msg_iovlen = iov.as_ref().len() as _;
(*p).msg_control = msg_control as *mut c_void;
(*p).msg_controllen = msg_controllen as _;
(*p).msg_flags = 0;
mhdr.assume_init()
};

msg_controllen
(msg_controllen, mhdr)
}


fn pack_mhdr_to_send<'a, I, C>(
out: *mut msghdr,
cmsg_buffer: &mut [u8],
iov: I,
cmsgs: C,
addr: Option<&SockAddr>
)
) -> msghdr
where
I: AsRef<[IoVec<&'a [u8]>]>,
C: AsRef<[ControlMessage<'a>]>
{
let cmsg_capacity = cmsg_buffer.len();
let capacity = cmsg_buffer.len();

// Next encode the sending address, if provided
let (name, namelen) = match addr {
Expand All @@ -1081,38 +1080,43 @@ fn pack_mhdr_to_send<'a, I, C>(
};

// The message header must be initialized before the individual cmsgs.
let cmsg_ptr = if cmsg_capacity > 0 {
let cmsg_ptr = if capacity > 0 {
cmsg_buffer.as_ptr() as *mut c_void
} else {
ptr::null_mut()
};

// Musl's msghdr has private fields, so this is the only way to
// initialize it.
unsafe {
(*out).msg_name = name as *mut _;
(*out).msg_namelen = namelen;
let mhdr = unsafe {
// Musl's msghdr has private fields, so this is the only way to
// initialize it.
let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();
let p = mhdr.as_mut_ptr();
(*p).msg_name = name as *mut _;
(*p).msg_namelen = namelen;
// transmute iov into a mutable pointer. sendmsg doesn't really mutate
// the buffer, but the standard says that it takes a mutable pointer
(*out).msg_iov = iov.as_ref().as_ptr() as *mut _;
(*out).msg_iovlen = iov.as_ref().len() as _;
(*out).msg_control = cmsg_ptr;
(*out).msg_controllen = cmsg_capacity as _;
(*out).msg_flags = 0;
}
(*p).msg_iov = iov.as_ref().as_ptr() as *mut _;
(*p).msg_iovlen = iov.as_ref().len() as _;
(*p).msg_control = cmsg_ptr;
(*p).msg_controllen = capacity as _;
(*p).msg_flags = 0;
mhdr.assume_init()
};

// Encode each cmsg. This must happen after initializing the header because
// CMSG_NEXT_HDR and friends read the msg_control and msg_controllen fields.
// CMSG_FIRSTHDR is always safe
let mut pmhdr: *mut cmsghdr = unsafe { CMSG_FIRSTHDR(out) };
let mut pmhdr: *mut cmsghdr = unsafe { CMSG_FIRSTHDR(&mhdr as *const msghdr) };
for cmsg in cmsgs.as_ref() {
assert_ne!(pmhdr, ptr::null_mut());
// Safe because we know that pmhdr is valid, and we initialized it with
// sufficient space
unsafe { cmsg.encode_into(pmhdr); }
unsafe { cmsg.encode_into(pmhdr) };
// Safe because mhdr is valid
pmhdr = unsafe { CMSG_NXTHDR(out, pmhdr) };
pmhdr = unsafe { CMSG_NXTHDR(&mhdr as *const msghdr, pmhdr) };
}

mhdr
}

/// Receive message in scatter-gather vectors from a socket, and
Expand All @@ -1133,20 +1137,17 @@ pub fn recvmsg<'a>(fd: RawFd, iov: &[IoVec<&mut [u8]>],
mut cmsg_buffer: Option<&'a mut Vec<u8>>,
flags: MsgFlags) -> Result<RecvMsg<'a>>
{
let mut out = mem::MaybeUninit::<msghdr>::zeroed();
let mut address = mem::MaybeUninit::uninit();

let msg_controllen = unsafe {
pack_mhdr_to_receive(out.as_mut_ptr(), &iov, &mut cmsg_buffer, address.as_mut_ptr())
let (msg_controllen, mut mhdr) = unsafe {
pack_mhdr_to_receive(&iov, &mut cmsg_buffer, address.as_mut_ptr())
};

let mut mhdr = unsafe { out.assume_init() };

let ret = unsafe { libc::recvmsg(fd, &mut mhdr, flags.bits()) };

let r = Errno::result(ret)?;

Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address.as_mut_ptr(), &mut cmsg_buffer) })
Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address.assume_init(), &mut cmsg_buffer) })
}


Expand Down

0 comments on commit cbeb0d2

Please sign in to comment.