msg_socket: implement MsgOnSocket for Vec and tuples

These container types are similar to arrays except tuples have
heterogeneous data types and Vec has a dynamic number of elements.

BUG=None
TEST=cargo test -p msg_socket

Change-Id: I2cbbaeb7f13b7700294ac50751530510098ba16d
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2168588
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-by: Dylan Reid <dgreid@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Tested-by: Zach Reizner <zachr@chromium.org>
Commit-Queue: Zach Reizner <zachr@chromium.org>

3 files changed, 431 insertions, 81 deletions

diff --git a/msg_socket/src/msg_on_socket.rs b/msg_socket/src/msg_on_socket.rs
index eeaf1bd..3e87102 100644
--- a/msg_socket/src/msg_on_socket.rs
+++ b/msg_socket/src/msg_on_socket.rs
@@ -2,17 +2,20 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+mod slice;
+mod tuple;
+
 use std::fmt::{self, Display};
 use std::fs::File;
 use std::mem::{size_of, transmute_copy, MaybeUninit};
 use std::net::{TcpListener, TcpStream, UdpSocket};
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
-use std::ptr::drop_in_place;
 use std::result;
 use std::sync::Arc;
 
 use data_model::*;
+use slice::{slice_read_helper, slice_write_helper};
 use sync::Mutex;
 use sys_util::{Error as SysError, EventFd};
 
@@ -283,20 +286,6 @@ rawfd_impl!(UdpSocket);
 rawfd_impl!(UnixListener);
 rawfd_impl!(UnixDatagram);
 
-// Converts a slice into an array of fixed size inferred from by the return value. Panics if the
-// slice is too small, but will tolerate slices that are too large.
-fn slice_to_array<T, O>(s: &[T]) -> O
-where
-    T: Copy,
-    O: Default + AsMut<[T]>,
-{
-    let mut o = O::default();
-    let o_slice = o.as_mut();
-    let len = o_slice.len();
-    o_slice.copy_from_slice(&s[..len]);
-    o
-}
-
 // usize could be different sizes on different targets. We always use u64.
 impl MsgOnSocket for usize {
     fn msg_size(&self) -> usize {
@@ -381,6 +370,35 @@ le_impl!(Le16, u16);
 le_impl!(Le32, u32);
 le_impl!(Le64, u64);
 
+fn simple_read<T: MsgOnSocket>(buffer: &[u8], offset: &mut usize) -> MsgResult<T> {
+    assert!(!T::uses_fd());
+    // Safety for T::read_from_buffer depends on the given FDs being valid, but we pass no FDs.
+    let (v, _) = unsafe { T::read_from_buffer(&buffer[*offset..], &[])? };
+    *offset += v.msg_size();
+    Ok(v)
+}
+
+fn simple_write<T: MsgOnSocket>(v: T, buffer: &mut [u8], offset: &mut usize) -> MsgResult<()> {
+    assert!(!T::uses_fd());
+    v.write_to_buffer(&mut buffer[*offset..], &mut [])?;
+    *offset += v.msg_size();
+    Ok(())
+}
+
+// Converts a slice into an array of fixed size inferred from by the return value. Panics if the
+// slice is too small, but will tolerate slices that are too large.
+fn slice_to_array<T, O>(s: &[T]) -> O
+where
+    T: Copy,
+    O: Default + AsMut<[T]>,
+{
+    let mut o = O::default();
+    let o_slice = o.as_mut();
+    let len = o_slice.len();
+    o_slice.copy_from_slice(&s[..len]);
+    o
+}
+
 macro_rules! array_impls {
     ($N:expr, $t: ident $($ts:ident)*)
     => {
@@ -414,46 +432,7 @@ macro_rules! array_impls {
                 // themselves don't require initializing.
                 let mut msgs: [MaybeUninit<T>; $N] =  MaybeUninit::uninit().assume_init();
 
-                let mut offset = 0usize;
-                let mut fd_offset = 0usize;
-
-                // In case of an error, we need to keep track of how many elements got initialized.
-                // In order to perform the necessary drops, the below loop is executed in a closure
-                // to capture errors without returning.
-                let mut last_index = 0;
-                let res = (|| {
-                    for msg in &mut msgs[..] {
-                        let element_size = match T::fixed_size() {
-                            Some(s) => s,
-                            None => {
-                                let (element_size, _) = u64::read_from_buffer(&buffer[offset..], &[])?;
-                                offset += element_size.msg_size();
-                                element_size as usize
-                            }
-                        };
-                        let (m, fd_size) =
-                            T::read_from_buffer(&buffer[offset..], &fds[fd_offset..])?;
-                        *msg = MaybeUninit::new(m);
-                        offset += element_size;
-                        fd_offset += fd_size;
-                        last_index += 1;
-                    }
-                    Ok(())
-                })();
-
-                // Because `MaybeUninit` will not automatically call drops, we have to drop the
-                // partially initialized array manually in the case of an error.
-                if let Err(e) = res {
-                    for msg in &mut msgs[..last_index] {
-                        // The call to `as_mut_ptr()` turns the `MaybeUninit` element of the array
-                        // into a pointer, which can be used with `drop_in_place` to call the
-                        // destructor without moving the element, which is impossible. This is safe
-                        // because `last_index` prevents this loop from traversing into the
-                        // uninitialized parts of the array.
-                        drop_in_place(msg.as_mut_ptr());
-                    }
-                    return Err(e)
-                }
+                let fd_count = slice_read_helper(buffer, fds, &mut msgs)?;
 
                 // Also taken from the canonical example, we initialized every member of the array
                 // in the first loop of this function, so it is safe to `transmute_copy` the array
@@ -463,7 +442,7 @@ macro_rules! array_impls {
                 // Because this function operates on generic data, the type is "dependently-sized"
                 // and so the compiler will not check that the size of the input and output match.
                 // See this issue for details: https://github.com/rust-lang/rust/issues/61956
-                Ok((transmute_copy::<_, [T; $N]>(&msgs), fd_offset))
+                Ok((transmute_copy::<_, [T; $N]>(&msgs), fd_count))
             }
 
             fn write_to_buffer(
@@ -471,25 +450,7 @@ macro_rules! array_impls {
                 buffer: &mut [u8],
                 fds: &mut [RawFd],
                 ) -> MsgResult<usize> {
-                let mut offset = 0usize;
-                let mut fd_offset = 0usize;
-                for idx in 0..$N {
-                    let element_size = match T::fixed_size() {
-                        Some(s) => s,
-                        None => {
-                            let element_size = self[idx].msg_size() as u64;
-                            element_size.write_to_buffer(&mut buffer[offset..], &mut [])?;
-                            offset += element_size.msg_size();
-                            element_size as usize
-                        }
-                    };
-                    let fd_size = self[idx].write_to_buffer(&mut buffer[offset..],
-                                                            &mut fds[fd_offset..])?;
-                    offset += element_size;
-                    fd_offset += fd_size;
-                }
-
-                Ok(fd_offset)
+                slice_write_helper(self, buffer, fds)
             }
         }
         #[cfg(test)]
@@ -504,7 +465,7 @@ macro_rules! array_impls {
                 array.write_to_buffer(&mut buffer, &mut []).unwrap();
                 let read_array = unsafe { ArrayType::read_from_buffer(&buffer, &[]) }.unwrap().0;
 
-                assert_eq!(array, read_array);
+                assert!(array.iter().eq(read_array.iter()));
             }
 
             #[test]
@@ -515,7 +476,7 @@ macro_rules! array_impls {
                 array.write_to_buffer(&mut buffer, &mut []).unwrap();
                 let read_array = unsafe { ArrayType::read_from_buffer(&buffer, &[]) }.unwrap().0;
 
-                assert_eq!(array, read_array);
+                assert!(array.iter().eq(read_array.iter()));
             }
         }
         array_impls!(($N - 1), $($ts)*);
@@ -524,9 +485,9 @@ macro_rules! array_impls {
 }
 
 array_impls! {
-    32, tmp1 tmp2 tmp3 tmp4 tmp5 tmp6 tmp7 tmp8 tmp9 tmp10 tmp11 tmp12 tmp13 tmp14 tmp15 tmp16
+    64, tmp1 tmp2 tmp3 tmp4 tmp5 tmp6 tmp7 tmp8 tmp9 tmp10 tmp11 tmp12 tmp13 tmp14 tmp15 tmp16
         tmp17 tmp18 tmp19 tmp20 tmp21 tmp22 tmp23 tmp24 tmp25 tmp26 tmp27 tmp28 tmp29 tmp30 tmp31
-        tmp32
+        tmp32 tmp33 tmp34 tmp35 tmp36 tmp37 tmp38 tmp39 tmp40 tmp41 tmp42 tmp43 tmp44 tmp45 tmp46
+        tmp47 tmp48 tmp49 tmp50 tmp51 tmp52 tmp53 tmp54 tmp55 tmp56 tmp57 tmp58 tmp59 tmp60 tmp61
+        tmp62 tmp63 tmp64
 }
-
-// TODO(jkwang) Define MsgOnSocket for tuple?
diff --git a/msg_socket/src/msg_on_socket/slice.rs b/msg_socket/src/msg_on_socket/slice.rs
new file mode 100644
index 0000000..7b6ef28
--- /dev/null
+++ b/msg_socket/src/msg_on_socket/slice.rs
@@ -0,0 +1,184 @@
+// Copyright 2020 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::mem::{size_of, ManuallyDrop, MaybeUninit};
+use std::os::unix::io::RawFd;
+use std::ptr::drop_in_place;
+
+use crate::{MsgOnSocket, MsgResult};
+
+use super::{simple_read, simple_write};
+
+/// Helper used by the types that read a slice of homegenously typed data.
+///
+/// # Safety
+/// This function has the same safety requirements as `T::read_from_buffer`, with the additional
+/// requirements that the `msgs` are only used on success of this function
+pub unsafe fn slice_read_helper<T: MsgOnSocket>(
+    buffer: &[u8],
+    fds: &[RawFd],
+    msgs: &mut [MaybeUninit<T>],
+) -> MsgResult<usize> {
+    let mut offset = 0usize;
+    let mut fd_offset = 0usize;
+
+    // In case of an error, we need to keep track of how many elements got initialized.
+    // In order to perform the necessary drops, the below loop is executed in a closure
+    // to capture errors without returning.
+    let mut last_index = 0;
+    let res = (|| {
+        for msg in &mut msgs[..] {
+            let element_size = match T::fixed_size() {
+                Some(s) => s,
+                None => simple_read::<u64>(buffer, &mut offset)? as usize,
+            };
+            // Assuming the unsafe caller gave valid FDs, this call should be safe.
+            let (m, fd_size) = T::read_from_buffer(&buffer[offset..], &fds[fd_offset..])?;
+            *msg = MaybeUninit::new(m);
+            offset += element_size;
+            fd_offset += fd_size;
+            last_index += 1;
+        }
+        Ok(())
+    })();
+
+    // Because `MaybeUninit` will not automatically call drops, we have to drop the
+    // partially initialized array manually in the case of an error.
+    if let Err(e) = res {
+        for msg in &mut msgs[..last_index] {
+            // The call to `as_mut_ptr()` turns the `MaybeUninit` element of the array
+            // into a pointer, which can be used with `drop_in_place` to call the
+            // destructor without moving the element, which is impossible. This is safe
+            // because `last_index` prevents this loop from traversing into the
+            // uninitialized parts of the array.
+            drop_in_place(msg.as_mut_ptr());
+        }
+        return Err(e);
+    }
+
+    Ok(fd_offset)
+}
+
+/// Helper used by the types that write a slice of homegenously typed data.
+pub fn slice_write_helper<T: MsgOnSocket>(
+    msgs: &[T],
+    buffer: &mut [u8],
+    fds: &mut [RawFd],
+) -> MsgResult<usize> {
+    let mut offset = 0usize;
+    let mut fd_offset = 0usize;
+    for msg in msgs {
+        let element_size = match T::fixed_size() {
+            Some(s) => s,
+            None => {
+                let element_size = msg.msg_size();
+                simple_write(element_size as u64, buffer, &mut offset)?;
+                element_size as usize
+            }
+        };
+        let fd_size = msg.write_to_buffer(&mut buffer[offset..], &mut fds[fd_offset..])?;
+        offset += element_size;
+        fd_offset += fd_size;
+    }
+
+    Ok(fd_offset)
+}
+
+impl<T: MsgOnSocket> MsgOnSocket for Vec<T> {
+    fn uses_fd() -> bool {
+        T::uses_fd()
+    }
+
+    fn fixed_size() -> Option<usize> {
+        None
+    }
+
+    fn msg_size(&self) -> usize {
+        let vec_size = match T::fixed_size() {
+            Some(s) => s * self.len(),
+            None => self.iter().map(|i| i.msg_size() + size_of::<u64>()).sum(),
+        };
+        size_of::<u64>() + vec_size
+    }
+
+    fn fd_count(&self) -> usize {
+        if T::uses_fd() {
+            self.iter().map(|i| i.fd_count()).sum()
+        } else {
+            0
+        }
+    }
+
+    unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
+        let mut offset = 0;
+        let len = simple_read::<u64>(buffer, &mut offset)? as usize;
+        let mut msgs: Vec<MaybeUninit<T>> = Vec::with_capacity(len);
+        msgs.set_len(len);
+        let fd_count = slice_read_helper(&buffer[offset..], fds, &mut msgs)?;
+        let mut msgs = ManuallyDrop::new(msgs);
+        Ok((
+            Vec::from_raw_parts(msgs.as_mut_ptr() as *mut T, msgs.len(), msgs.capacity()),
+            fd_count,
+        ))
+    }
+
+    fn write_to_buffer(&self, buffer: &mut [u8], fds: &mut [RawFd]) -> MsgResult<usize> {
+        let mut offset = 0;
+        simple_write(self.len() as u64, buffer, &mut offset)?;
+        slice_write_helper(self, &mut buffer[offset..], fds)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn read_write_1_fixed() {
+        let vec = vec![1u32];
+        let mut buffer = vec![0; vec.msg_size()];
+        vec.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_vec = unsafe { <Vec<u32>>::read_from_buffer(&buffer, &[]) }
+            .unwrap()
+            .0;
+
+        assert_eq!(vec, read_vec);
+    }
+
+    #[test]
+    fn read_write_8_fixed() {
+        let vec = vec![1u16, 1, 3, 5, 8, 13, 21, 34];
+        let mut buffer = vec![0; vec.msg_size()];
+        vec.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_vec = unsafe { <Vec<u16>>::read_from_buffer(&buffer, &[]) }
+            .unwrap()
+            .0;
+        assert_eq!(vec, read_vec);
+    }
+
+    #[test]
+    fn read_write_1() {
+        let vec = vec![Some(1u64)];
+        let mut buffer = vec![0; vec.msg_size()];
+        println!("{:?}", vec.msg_size());
+        vec.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_vec = unsafe { <Vec<_>>::read_from_buffer(&buffer, &[]) }
+            .unwrap()
+            .0;
+
+        assert_eq!(vec, read_vec);
+    }
+
+    #[test]
+    fn read_write_4() {
+        let vec = vec![Some(12u16), Some(0), None, None];
+        let mut buffer = vec![0; vec.msg_size()];
+        vec.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_vec = unsafe { <Vec<_>>::read_from_buffer(&buffer, &[]) }
+            .unwrap()
+            .0;
+
+        assert_eq!(vec, read_vec);
+    }
+}
diff --git a/msg_socket/src/msg_on_socket/tuple.rs b/msg_socket/src/msg_on_socket/tuple.rs
new file mode 100644
index 0000000..f960ce5
--- /dev/null
+++ b/msg_socket/src/msg_on_socket/tuple.rs
@@ -0,0 +1,205 @@
+// Copyright 2020 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::mem::size_of;
+use std::os::unix::io::RawFd;
+
+use crate::{MsgOnSocket, MsgResult};
+
+use super::{simple_read, simple_write};
+
+// Returns the size of one part of a tuple.
+fn tuple_size_helper<T: MsgOnSocket>(v: &T) -> usize {
+    T::fixed_size().unwrap_or_else(|| v.msg_size() + size_of::<u64>())
+}
+
+unsafe fn tuple_read_helper<T: MsgOnSocket>(
+    buffer: &[u8],
+    fds: &[RawFd],
+    buffer_index: &mut usize,
+    fd_index: &mut usize,
+) -> MsgResult<T> {
+    let end = match T::fixed_size() {
+        Some(_) => buffer.len(),
+        None => {
+            let len = simple_read::<u64>(buffer, buffer_index)? as usize;
+            *buffer_index + len
+        }
+    };
+    let (v, fd_read) = T::read_from_buffer(&buffer[*buffer_index..end], &fds[*fd_index..])?;
+    *buffer_index += v.msg_size();
+    *fd_index += fd_read;
+    Ok(v)
+}
+
+fn tuple_write_helper<T: MsgOnSocket>(
+    v: &T,
+    buffer: &mut [u8],
+    fds: &mut [RawFd],
+    buffer_index: &mut usize,
+    fd_index: &mut usize,
+) -> MsgResult<()> {
+    let end = match T::fixed_size() {
+        Some(_) => buffer.len(),
+        None => {
+            let len = v.msg_size();
+            simple_write(len as u64, buffer, buffer_index)?;
+            *buffer_index + len
+        }
+    };
+    let fd_written = v.write_to_buffer(&mut buffer[*buffer_index..end], &mut fds[*fd_index..])?;
+    *buffer_index += v.msg_size();
+    *fd_index += fd_written;
+    Ok(())
+}
+
+macro_rules! tuple_impls {
+    () => {};
+    ($t: ident) => {
+        #[allow(unused_variables, non_snake_case)]
+        impl<$t: MsgOnSocket> MsgOnSocket for ($t,) {
+            fn uses_fd() -> bool {
+                $t::uses_fd()
+            }
+
+            fn fd_count(&self) -> usize {
+                self.0.fd_count()
+            }
+
+            fn fixed_size() -> Option<usize> {
+                $t::fixed_size()
+            }
+
+            fn msg_size(&self) -> usize {
+                self.0.msg_size()
+            }
+
+            unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
+                let (t, s) = $t::read_from_buffer(buffer, fds)?;
+                Ok(((t,), s))
+            }
+
+            fn write_to_buffer(
+                &self,
+                buffer: &mut [u8],
+                fds: &mut [RawFd],
+            ) -> MsgResult<usize> {
+                self.0.write_to_buffer(buffer, fds)
+            }
+        }
+    };
+    ($t: ident, $($ts:ident),*) => {
+        #[allow(unused_variables, non_snake_case)]
+        impl<$t: MsgOnSocket $(, $ts: MsgOnSocket)*> MsgOnSocket for ($t$(, $ts)*) {
+            fn uses_fd() -> bool {
+                $t::uses_fd() $(|| $ts::uses_fd())*
+            }
+
+            fn fd_count(&self) -> usize {
+                if Self::uses_fd() {
+                    return 0;
+                }
+                let ($t $(,$ts)*) = self;
+                $t.fd_count() $(+ $ts.fd_count())*
+            }
+
+            fn fixed_size() -> Option<usize> {
+                // Returns None if any element is not fixed size.
+                Some($t::fixed_size()? $(+ $ts::fixed_size()?)*)
+            }
+
+            fn msg_size(&self) -> usize {
+                if let Some(size) = Self::fixed_size() {
+                    return size
+                }
+
+                let ($t $(,$ts)*) = self;
+                tuple_size_helper($t) $(+ tuple_size_helper($ts))*
+            }
+
+            unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
+                let mut buffer_index = 0;
+                let mut fd_index = 0;
+                Ok((
+                        (
+                            tuple_read_helper(buffer, fds, &mut buffer_index, &mut fd_index)?,
+                            $({
+                                // Dummy let used to trigger the correct number of iterations.
+                                let $ts = ();
+                                tuple_read_helper(buffer, fds, &mut buffer_index, &mut fd_index)?
+                            },)*
+                        ),
+                        fd_index
+                ))
+            }
+
+            fn write_to_buffer(
+                &self,
+                buffer: &mut [u8],
+                fds: &mut [RawFd],
+            ) -> MsgResult<usize> {
+                let mut buffer_index = 0;
+                let mut fd_index = 0;
+                let ($t $(,$ts)*) = self;
+                tuple_write_helper($t, buffer, fds, &mut buffer_index, &mut fd_index)?;
+                $(
+                    tuple_write_helper($ts, buffer, fds, &mut buffer_index, &mut fd_index)?;
+                )*
+                Ok(fd_index)
+            }
+        }
+        tuple_impls!{ $($ts),* }
+    }
+}
+
+// Imlpement tuple for up to 8 elements.
+tuple_impls! { A, B, C, D, E, F, G, H }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn read_write_1_fixed() {
+        let tuple = (1,);
+        let mut buffer = vec![0; tuple.msg_size()];
+        tuple.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_tuple = unsafe { <(u32,)>::read_from_buffer(&buffer, &[]) }
+            .unwrap()
+            .0;
+
+        assert_eq!(tuple, read_tuple);
+    }
+
+    #[test]
+    fn read_write_8_fixed() {
+        let tuple = (1u32, 2u8, 3u16, 4u64, 5u32, 6u16, 7u8, 8u8);
+        let mut buffer = vec![0; tuple.msg_size()];
+        tuple.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_tuple = unsafe { <_>::read_from_buffer(&buffer, &[]) }.unwrap().0;
+
+        assert_eq!(tuple, read_tuple);
+    }
+
+    #[test]
+    fn read_write_1() {
+        let tuple = (Some(1u64),);
+        let mut buffer = vec![0; tuple.msg_size()];
+        tuple.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_tuple = unsafe { <_>::read_from_buffer(&buffer, &[]) }.unwrap().0;
+
+        assert_eq!(tuple, read_tuple);
+    }
+
+    #[test]
+    fn read_write_4() {
+        let tuple = (Some(12u16), Some(false), None::<u8>, None::<u64>);
+        let mut buffer = vec![0; tuple.msg_size()];
+        println!("{:?}", tuple.msg_size());
+        tuple.write_to_buffer(&mut buffer, &mut []).unwrap();
+        let read_tuple = unsafe { <_>::read_from_buffer(&buffer, &[]) }.unwrap().0;
+
+        assert_eq!(tuple, read_tuple);
+    }
+}