PR suggestion: Set the length later instead of using `from_raw`

2 files changed, 114 insertions, 17 deletions

diff --git a/src/sys/socket/addr.rs b/src/sys/socket/addr.rs
index 8c5a56fe..12cdc7a4 100644
--- a/src/sys/socket/addr.rs
+++ b/src/sys/socket/addr.rs
@@ -763,6 +763,22 @@ impl SockaddrLike for UnixAddr {
     {
         mem::size_of::<libc::sockaddr_un>() as libc::socklen_t
     }
+
+    unsafe fn set_length(&mut self, new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+        cfg_if! {
+            if #[cfg(any(target_os = "android",
+                         target_os = "fuchsia",
+                         target_os = "illumos",
+                         target_os = "linux",
+                         target_os = "redox",
+                ))] {
+                self.sun_len = new_length as u8;
+            } else {
+                self.sun.sun_len = new_length as u8;
+            }
+        };
+        Ok(())
+    }
 }
 
 impl AsRef<libc::sockaddr_un> for UnixAddr {
@@ -912,8 +928,30 @@ pub trait SockaddrLike: private::SockaddrLikePriv {
     {
         mem::size_of::<Self>() as libc::socklen_t
     }
+
+    /// Set the length of this socket address
+    ///
+    /// This method may only be called on socket addresses whose lenghts are dynamic, and it
+    /// returns an error if called on a type whose length is static.
+    ///
+    /// # Safety
+    ///
+    /// `new_length` must be a valid length for this type of address. Specifically, reads of that
+    /// length from `self` must be valid.
+    unsafe fn set_length(&mut self, new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic>;
 }
 
+/// The error returned by [`SockaddrLike::set_length`] on an address whose length is statically
+/// fixed.
+#[derive(Copy, Clone, Debug)]
+pub struct SocketAddressLengthNotDynamic;
+impl fmt::Display for SocketAddressLengthNotDynamic {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("Attempted to set length on socket whose length is statically fixed")
+    }
+}
+impl std::error::Error for SocketAddressLengthNotDynamic {}
+
 impl private::SockaddrLikePriv for () {
     fn as_mut_ptr(&mut self) -> *mut libc::sockaddr {
         ptr::null_mut()
@@ -946,6 +984,10 @@ impl SockaddrLike for () {
     fn len(&self) -> libc::socklen_t {
         0
     }
+
+    unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+        Err(SocketAddressLengthNotDynamic)
+    }
 }
 
 /// An IPv4 socket address
@@ -1015,6 +1057,10 @@ impl SockaddrLike for SockaddrIn {
         }
         Some(Self(ptr::read_unaligned(addr as *const _)))
     }
+
+    unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+        Err(SocketAddressLengthNotDynamic)
+    }
 }
 
 #[cfg(feature = "net")]
@@ -1134,6 +1180,10 @@ impl SockaddrLike for SockaddrIn6 {
         }
         Some(Self(ptr::read_unaligned(addr as *const _)))
     }
+
+    unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+        Err(SocketAddressLengthNotDynamic)
+    }
 }
 
 #[cfg(feature = "net")]
@@ -1361,6 +1411,27 @@ impl SockaddrLike for SockaddrStorage {
             None => mem::size_of_val(self) as libc::socklen_t,
         }
     }
+
+    unsafe fn set_length(&mut self, new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+        match self.as_unix_addr_mut() {
+            Some(addr) => {
+                cfg_if! {
+                    if #[cfg(any(target_os = "android",
+                                 target_os = "fuchsia",
+                                 target_os = "illumos",
+                                 target_os = "linux",
+                                 target_os = "redox",
+                        ))] {
+                        addr.sun_len = new_length as u8;
+                    } else {
+                        addr.sun.sun_len = new_length as u8;
+                    }
+                }
+                Ok(())
+            },
+            None => Err(SocketAddressLengthNotDynamic),
+        }
+    }
 }
 
 macro_rules! accessors {
@@ -1754,6 +1825,10 @@ pub mod netlink {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_nl> for NetlinkAddr {
@@ -1803,6 +1878,10 @@ pub mod alg {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_alg> for AlgAddr {
@@ -1902,7 +1981,7 @@ pub mod sys_control {
     use std::{fmt, mem, ptr};
     use std::os::unix::io::RawFd;
     use crate::{Errno, Result};
-    use super::{private, SockaddrLike};
+    use super::{private, SockaddrLike, SocketAddressLengthNotDynamic};
 
     // FIXME: Move type into `libc`
     #[repr(C)]
@@ -1943,6 +2022,10 @@ pub mod sys_control {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_ctl> for SysControlAddr {
@@ -2007,7 +2090,7 @@ pub mod sys_control {
 mod datalink {
     feature! {
     #![feature = "net"]
-    use super::{fmt, mem, private, ptr, SockaddrLike};
+    use super::{fmt, mem, private, ptr, SockaddrLike, SocketAddressLengthNotDynamic};
 
     /// Hardware Address
     #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
@@ -2085,6 +2168,10 @@ mod datalink {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_ll> for LinkAddr {
@@ -2110,7 +2197,7 @@ mod datalink {
 mod datalink {
     feature! {
     #![feature = "net"]
-    use super::{fmt, mem, private, ptr, SockaddrLike};
+    use super::{fmt, mem, private, ptr, SockaddrLike, SocketAddressLengthNotDynamic};
 
     /// Hardware Address
     #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
@@ -2209,6 +2296,10 @@ mod datalink {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_dl> for LinkAddr {
@@ -2257,6 +2348,10 @@ pub mod vsock {
             }
             Some(Self(ptr::read_unaligned(addr as *const _)))
         }
+
+        unsafe fn set_length(&mut self, _new_length: usize) -> std::result::Result<(), SocketAddressLengthNotDynamic> {
+            Err(SocketAddressLengthNotDynamic)
+        }
     }
 
     impl AsRef<libc::sockaddr_vm> for VsockAddr {
diff --git a/src/sys/socket/mod.rs b/src/sys/socket/mod.rs
index 1bbfa541..78b7e8b0 100644
--- a/src/sys/socket/mod.rs
+++ b/src/sys/socket/mod.rs
@@ -1627,7 +1627,7 @@ impl<S> MultiHeaders<S> {
                     None => (std::ptr::null(), 0),
                 };
                 let msg_hdr = unsafe {
-                    pack_mhdr_to_receive(std::ptr::null(), 0, ptr, cap, <S as addr::private::SockaddrLikePriv>::as_mut_ptr(address.assume_init_mut()).cast())
+                    pack_mhdr_to_receive(std::ptr::null(), 0, ptr, cap, address.assume_init_mut())
                 };
                 libc::mmsghdr {
                     msg_hdr,
@@ -1763,7 +1763,7 @@ where
                 mmsghdr.msg_hdr,
                 mmsghdr.msg_len as isize,
                 self.rmm.msg_controllen,
-                Some(address),
+                address,
             )
         })
     }
@@ -1916,7 +1916,7 @@ unsafe fn read_mhdr<'a, 'i, S>(
     mhdr: msghdr,
     r: isize,
     msg_controllen: usize,
-    address: Option<S>,
+    mut address: S,
 ) -> RecvMsg<'a, 'i, S>
     where S: SockaddrLike
 {
@@ -1932,10 +1932,15 @@ unsafe fn read_mhdr<'a, 'i, S>(
         }.as_ref()
     };
 
+    // Ignore errors if this socket address has statically-known length
+    //
+    // This is to ensure that unix socket addresses have their length set appropriately.
+    let _ = unsafe { address.set_length(mhdr.msg_namelen as usize) };
+
     RecvMsg {
         bytes: r as usize,
         cmsghdr,
-        address,
+        address: Some(address),
         flags: MsgFlags::from_bits_truncate(mhdr.msg_flags),
         mhdr,
         iobufs: std::marker::PhantomData,
@@ -1953,19 +1958,19 @@ unsafe fn read_mhdr<'a, 'i, S>(
 /// headers are not used
 ///
 /// Buffers must remain valid for the whole lifetime of msghdr
-unsafe fn pack_mhdr_to_receive(
+unsafe fn pack_mhdr_to_receive<S>(
     iov_buffer: *const IoSliceMut,
     iov_buffer_len: usize,
     cmsg_buffer: *const u8,
     cmsg_capacity: usize,
-    address: *mut libc::sockaddr_storage,
-) -> msghdr {
+    address: *mut S,
+) -> msghdr where S: SockaddrLike {
     // 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::<libc::sockaddr_storage>() as u32;
+    (*p).msg_namelen = S::size();
     (*p).msg_iov = iov_buffer as *mut iovec;
     (*p).msg_iovlen = iov_buffer_len as _;
     (*p).msg_control = cmsg_buffer as *mut c_void;
@@ -2047,23 +2052,20 @@ pub fn recvmsg<'a, 'outer, 'inner, S>(fd: RawFd, iov: &'outer mut [IoSliceMut<'i
     where S: SockaddrLike + 'a,
     'inner: 'outer
 {
-    let mut address: libc::sockaddr_storage = unsafe { mem::MaybeUninit::zeroed().assume_init() };
-    let address_ptr: *mut libc::sockaddr_storage = &mut address as *mut libc::sockaddr_storage;
+    let mut address: mem::MaybeUninit<S> = mem::MaybeUninit::zeroed();
 
     let (msg_control, msg_controllen) = cmsg_buffer.as_mut()
         .map(|v| (v.as_mut_ptr(), v.capacity()))
         .unwrap_or((ptr::null_mut(), 0));
     let mut mhdr = unsafe {
-        pack_mhdr_to_receive(iov.as_ref().as_ptr(), iov.len(), msg_control, msg_controllen, address_ptr)
+        pack_mhdr_to_receive(iov.as_ref().as_ptr(), iov.len(), msg_control, msg_controllen, address.as_mut_ptr())
     };
 
     let ret = unsafe { libc::recvmsg(fd, &mut mhdr, flags.bits()) };
 
     let r = Errno::result(ret)?;
 
-    let address = unsafe { S::from_raw(address_ptr.cast::<libc::sockaddr>(), Some(mhdr.msg_namelen)) };
-
-    Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address) })
+    Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address.assume_init()) })
 }
 }