// Copyright 2017 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-BSD-3-Clause file.
// SPDX-License-Identifier: BSD-3-Clause

/* Copied from the crosvm Project, commit 186eb8b */

//! Wrapper for sending and receiving messages with file descriptors on sockets that accept
//! control messages (e.g. Unix domain sockets).

use std::fs::File;
use std::mem;
use std::mem::size_of;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::os::unix::net::{UnixDatagram, UnixStream};
use std::ptr::null_mut;

use crate::errno::{Error, Result};
use crate::fam::{FamStruct, FamStructWrapper};
use libc::{
    c_uint, c_void, cmsghdr, iovec, msghdr, recvmsg, sendmsg, CMSG_LEN, CMSG_SPACE, MSG_NOSIGNAL,
    SCM_RIGHTS, SOL_SOCKET,
};

#[cfg(not(target_env = "musl"))]
fn new_msghdr(iovecs: &mut [iovec]) -> msghdr {
    msghdr {
        msg_name: null_mut(),
        msg_namelen: 0,
        msg_iov: iovecs.as_mut_ptr(),
        #[cfg(any(target_os = "linux", target_os = "android"))]
        msg_iovlen: iovecs.len(),
        #[cfg(not(any(target_os = "linux", target_os = "android")))]
        msg_iovlen: iovecs
            .len()
            .try_into()
            .expect("iovecs.len() exceeds i32 range"),
        msg_control: null_mut(),
        msg_controllen: 0,
        msg_flags: 0,
    }
}

#[cfg(target_env = "musl")]
fn new_msghdr(iovecs: &mut [iovec]) -> msghdr {
    assert!(iovecs.len() <= (std::i32::MAX as usize));
    let mut msg: msghdr = unsafe { std::mem::zeroed() };
    msg.msg_name = null_mut();
    msg.msg_iov = iovecs.as_mut_ptr();
    msg.msg_iovlen = iovecs.len() as i32;
    msg.msg_control = null_mut();
    msg
}

#[cfg(all(
    not(target_env = "musl"),
    any(target_os = "linux", target_os = "android")
))]
fn set_msg_controllen(msg: &mut msghdr, cmsg_capacity: c_uint) {
    msg.msg_controllen = cmsg_capacity as libc::size_t;
}

#[cfg(any(
    target_env = "musl",
    not(any(target_os = "linux", target_os = "android"))
))]
fn set_msg_controllen(msg: &mut msghdr, cmsg_capacity: c_uint) {
    msg.msg_controllen = cmsg_capacity;
}

#[repr(transparent)]
struct CmsgHdr(cmsghdr);

impl Default for CmsgHdr {
    fn default() -> Self {
        // SAFETY: all-zero is valid for cmsghdr on all architectures/targets.
        Self(unsafe { mem::zeroed() })
    }
}

// SAFETY: `CmsgHdr` is a POD with a FAM struct at the end that is also a POD.
// We do not have the FAM represented by a zero-width field as would be customary if
// generated by bindgen, but this makes no difference. The length implementation correctly
// converts between the size in bytes of the cmsghdr structure + the payload, to the number
// of fds in the payload.
unsafe impl FamStruct for CmsgHdr {
    type Entry = RawFd;

    // on gnu, size_t is usize, while on musl it is u32
    #[allow(clippy::unnecessary_cast)]
    fn len(&self) -> usize {
        (self.0.cmsg_len as usize - size_of::<cmsghdr>()) / size_of::<RawFd>()
    }

    unsafe fn set_len(&mut self, len: usize) {
        // Different targets define cmsg_len of different types, but its always
        // at least 32 bit wide. The safety invariant on this function ensures us
        // that we're not being fed a value that exceeds Self::max_len(), so we
        // are good to cast indiscriminately.
        self.0.cmsg_len = CMSG_LEN((len * size_of::<RawFd>()) as _) as _;
    }

    fn max_len() -> usize {
        (u32::MAX as usize - size_of::<cmsghdr>()) / size_of::<RawFd>()
    }

    fn as_slice(&self) -> &[RawFd] {
        // SAFETY: By the invariants of the trait impl and set_len, we have a payload of self.len()
        // fds after the header itself
        unsafe { std::slice::from_raw_parts((&self.0 as *const cmsghdr).add(1).cast(), self.len()) }
    }

    fn as_mut_slice(&mut self) -> &mut [RawFd] {
        // SAFETY: By the invariants of the trait impl and set_len, we have a payload of self.len()
        // fds after the header itself
        unsafe {
            std::slice::from_raw_parts_mut((&mut self.0 as *mut cmsghdr).add(1).cast(), self.len())
        }
    }
}

type CmsgBuffer = FamStructWrapper<CmsgHdr>;

fn raw_sendmsg<D: IntoIovec>(fd: RawFd, out_data: &[D], out_fds: &[RawFd]) -> Result<usize> {
    let mut cmsg_buffer =
        CmsgBuffer::from_entries(out_fds).map_err(|_| Error::new(libc::ENOMEM))?;

    let mut iovecs = Vec::with_capacity(out_data.len());
    for data in out_data {
        iovecs.push(iovec {
            iov_base: data.as_ptr() as *mut c_void,
            iov_len: data.size(),
        });
    }

    let mut msg = new_msghdr(&mut iovecs);

    if !out_fds.is_empty() {
        // SAFETY: We do not touch the cmsg_len field.
        unsafe {
            let hdr = cmsg_buffer.as_mut_fam_struct();
            hdr.0.cmsg_level = SOL_SOCKET;
            hdr.0.cmsg_type = SCM_RIGHTS;
        }

        msg.msg_control = cmsg_buffer.as_mut_fam_struct_ptr() as *mut c_void;
        // SAFETY: CMSG_SPACE has no invariants to uphold
        unsafe {
            set_msg_controllen(&mut msg, CMSG_SPACE(size_of_val(out_fds) as _));
        }
    }

    // SAFETY: Safe because the msghdr was properly constructed from valid (or null) pointers of
    // the indicated length and we check the return value.
    let write_count = unsafe { sendmsg(fd, &msg, MSG_NOSIGNAL) };

    if write_count == -1 {
        Err(Error::last())
    } else {
        Ok(write_count as usize)
    }
}

#[allow(clippy::unnecessary_cast)]
unsafe fn raw_recvmsg(
    fd: RawFd,
    iovecs: &mut [iovec],
    in_fds: &mut [RawFd],
) -> Result<(usize, usize)> {
    let mut cmsg_buffer = CmsgBuffer::new(in_fds.len()).map_err(|_| Error::new(libc::ENOMEM))?;
    let mut msg = new_msghdr(iovecs);
    // Due to alignment constraints, this might be > in_fds.len()!
    let cmsg_capacity = cmsg_buffer.as_fam_struct_ref().len();

    if !in_fds.is_empty() {
        // MSG control len is size_of(cmsghdr) + size_of(RawFd) * in_fds.len().
        msg.msg_control = cmsg_buffer.as_mut_fam_struct_ptr() as *mut c_void;
        // SAFETY: CMSG_SPACE has no invariants to uphold
        unsafe {
            set_msg_controllen(&mut msg, CMSG_SPACE(size_of_val(in_fds) as _));
        }
    }

    // Safe because the msghdr was properly constructed from valid (or null) pointers of the
    // indicated length and we check the return value.
    // TODO: Should we handle MSG_TRUNC in a specific way?
    let total_read = recvmsg(fd, &mut msg, 0);
    if total_read == -1 {
        return Err(Error::last());
    }

    let mut copied_fds_count = 0;
    // If the control data was truncated, then this might be a sign of incorrect communication
    // protocol. If MSG_CTRUNC was set we must close the fds from the control data.
    let mut teardown_control_data = msg.msg_flags & libc::MSG_CTRUNC != 0;

    if !msg.msg_control.is_null() {
        let cmsg = cmsg_buffer.as_mut_fam_struct();

        if cmsg.0.cmsg_level == SOL_SOCKET && cmsg.0.cmsg_type == SCM_RIGHTS {
            // SAFETY: On some OSes (MacOS), when CMSG_TRUNC is set, the kernel sets
            // the cmsg->len field to the untruncated length of the message sent,
            // even if the buffer is actually smaller. Compensate for this by reducing
            // the length to be at most the size of the memory we have allocated.
            unsafe { cmsg.set_len(cmsg.len().min(cmsg_capacity)) }
            let fds = cmsg.as_slice();
            // It could be that while constructing the cmsg structures, alignment constraints made
            // our allocation big enough that it fits more than the in_fds.len() file descriptors
            // we intended to receive. Treat this the same way we would treat a truncated message,
            // because there is no way for us to communicate these extra FDs back to the caller.
            teardown_control_data |= fds.len() > in_fds.len();
            if teardown_control_data {
                for fd in fds {
                    libc::close(*fd);
                }
            } else {
                in_fds[..fds.len()].copy_from_slice(fds);

                copied_fds_count = fds.len();
            }
        }
    }

    if teardown_control_data {
        return Err(Error::new(libc::ENOBUFS));
    }

    Ok((total_read as usize, copied_fds_count))
}

/// Trait for file descriptors can send and receive socket control messages via `sendmsg` and
/// `recvmsg`.
///
/// # Examples
///
/// ```
/// use std::os::fd::{AsRawFd, FromRawFd};
/// use std::os::unix::net::UnixDatagram;
///
/// use libc::{c_void, iovec};
/// use vmm_sys_util::event::{new_event_consumer_and_notifier, EventFlag, EventNotifier};
/// use vmm_sys_util::sock_ctrl_msg::ScmSocket;
///
/// let (s1, s2) = UnixDatagram::pair().expect("failed to create socket pair");
/// let (consumer, fd_to_send) = new_event_consumer_and_notifier(EventFlag::empty())
///     .expect("Failed to create notifier and consumer");
///
/// let write_count = s1
///     .send_with_fds(&[[237].as_ref()], &[fd_to_send.as_raw_fd()])
///     .expect("failed to send fd");
///
/// let mut files = [0; 2];
/// let mut buf = [0u8];
/// let mut iovecs = [iovec {
///     iov_base: buf.as_mut_ptr() as *mut c_void,
///     iov_len: buf.len(),
/// }];
/// let (read_count, file_count) = unsafe {
///     s2.recv_with_fds(&mut iovecs[..], &mut files)
///         .expect("failed to recv fd")
/// };
///
/// let mut notifier = unsafe { EventNotifier::from_raw_fd(files[0]) };
/// notifier.notify().unwrap();
/// assert!(consumer.consume().is_ok());
/// ```
pub trait ScmSocket {
    /// Gets the file descriptor of this socket.
    fn socket_fd(&self) -> RawFd;

    /// Sends the given data and file descriptor over the socket.
    ///
    /// On success, returns the number of bytes sent.
    ///
    /// # Arguments
    ///
    /// * `buf` - A buffer of data to send on the `socket`.
    /// * `fd` - A file descriptors to be sent.
    fn send_with_fd<D: IntoIovec>(&self, buf: D, fd: RawFd) -> Result<usize> {
        self.send_with_fds(&[buf], &[fd])
    }

    /// Sends the given data and file descriptors over the socket.
    ///
    /// On success, returns the number of bytes sent.
    ///
    /// # Arguments
    ///
    /// * `bufs` - A list of data buffer to send on the `socket`.
    /// * `fds` - A list of file descriptors to be sent.
    fn send_with_fds<D: IntoIovec>(&self, bufs: &[D], fds: &[RawFd]) -> Result<usize> {
        raw_sendmsg(self.socket_fd(), bufs, fds)
    }

    /// Receives data and potentially a file descriptor from the socket.
    ///
    /// On success, returns the number of bytes and an optional file descriptor.
    ///
    /// # Arguments
    ///
    /// * `buf` - A buffer to receive data from the socket.
    fn recv_with_fd(&self, buf: &mut [u8]) -> Result<(usize, Option<File>)> {
        let mut fd = [0];
        let mut iovecs = [iovec {
            iov_base: buf.as_mut_ptr() as *mut c_void,
            iov_len: buf.len(),
        }];

        // SAFETY: Safe because we have mutably borrowed buf and it's safe to write arbitrary data
        // to a slice.
        let (read_count, fd_count) = unsafe { self.recv_with_fds(&mut iovecs[..], &mut fd)? };
        let file = if fd_count == 0 {
            None
        } else {
            // SAFETY: Safe because the first fd from recv_with_fds is owned by us and valid
            // because this branch was taken.
            Some(unsafe { File::from_raw_fd(fd[0]) })
        };
        Ok((read_count, file))
    }

    /// Receives data and file descriptors from the socket.
    ///
    /// On success, returns the number of bytes and file descriptors received as a tuple
    /// `(bytes count, files count)`.
    ///
    /// # Arguments
    ///
    /// * `iovecs` - A list of iovec to receive data from the socket.
    /// * `fds` - A slice of `RawFd`s to put the received file descriptors into. On success, the
    ///   number of valid file descriptors is indicated by the second element of the
    ///   returned tuple. The caller owns these file descriptors, but they will not be
    ///   closed on drop like a `File`-like type would be. It is recommended that each valid
    ///   file descriptor gets wrapped in a drop type that closes it after this returns.
    ///
    /// # Safety
    ///
    /// It is the callers responsibility to ensure it is safe for arbitrary data to be
    /// written to the iovec pointers.
    unsafe fn recv_with_fds(
        &self,
        iovecs: &mut [iovec],
        fds: &mut [RawFd],
    ) -> Result<(usize, usize)> {
        raw_recvmsg(self.socket_fd(), iovecs, fds)
    }
}

impl ScmSocket for UnixDatagram {
    fn socket_fd(&self) -> RawFd {
        self.as_raw_fd()
    }
}

impl ScmSocket for UnixStream {
    fn socket_fd(&self) -> RawFd {
        self.as_raw_fd()
    }
}

/// Trait for types that can be converted into an `iovec` that can be referenced by a syscall for
/// the lifetime of this object.
///
/// # Safety
///
/// This is marked unsafe because the implementation must ensure that the returned pointer and size
/// is valid and that the lifetime of the returned pointer is at least that of the trait object.
pub unsafe trait IntoIovec {
    /// Gets the base pointer of this `iovec`.
    fn as_ptr(&self) -> *const c_void;

    /// Gets the size in bytes of this `iovec`.
    fn size(&self) -> usize;
}

// SAFETY: Safe because this slice can not have another mutable reference and it's pointer and
// size are guaranteed to be valid.
unsafe impl IntoIovec for &[u8] {
    // Clippy false positive: https://github.com/rust-lang/rust-clippy/issues/3480
    #[allow(clippy::useless_asref)]
    fn as_ptr(&self) -> *const c_void {
        self.as_ref().as_ptr() as *const c_void
    }

    fn size(&self) -> usize {
        self.len()
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::undocumented_unsafe_blocks)]
    use super::*;
    use std::io::{pipe, Read};

    use std::io::Write;
    use std::os::unix::net::UnixDatagram;
    use std::slice::from_raw_parts;

    #[test]
    fn send_recv_no_fd() {
        let (s1, s2) = UnixDatagram::pair().expect("failed to create socket pair");

        let write_count = s1
            .send_with_fds(&[[1u8, 1, 2].as_ref(), [21u8, 34, 55].as_ref()], &[])
            .expect("failed to send data");

        assert_eq!(write_count, 6);

        let mut buf = [0u8; 6];
        let mut files = [0; 1];
        let mut iovecs = [iovec {
            iov_base: buf.as_mut_ptr() as *mut c_void,
            iov_len: buf.len(),
        }];
        let (read_count, file_count) = unsafe {
            s2.recv_with_fds(&mut iovecs[..], &mut files)
                .expect("failed to recv data")
        };

        assert_eq!(read_count, 6);
        assert_eq!(file_count, 0);
        assert_eq!(buf, [1, 1, 2, 21, 34, 55]);
    }

    #[test]
    fn send_recv_only_fd() {
        let (s1, s2) = UnixDatagram::pair().expect("failed to create socket pair");

        let (mut evt_consumer, evt_notifier) = pipe().expect("failed to create pipe");
        let write_count = s1
            .send_with_fd([].as_ref(), evt_notifier.as_raw_fd())
            .expect("failed to send fd");

        assert_eq!(write_count, 0);

        let (read_count, file_opt) = s2.recv_with_fd(&mut []).expect("failed to recv fd");

        let mut file = file_opt.unwrap();

        assert_eq!(read_count, 0);
        assert!(file.as_raw_fd() >= 0);
        assert_ne!(file.as_raw_fd(), s1.as_raw_fd());
        assert_ne!(file.as_raw_fd(), s2.as_raw_fd());
        assert_ne!(file.as_raw_fd(), evt_notifier.as_raw_fd());

        file.write_all(unsafe { from_raw_parts(&1203u64 as *const u64 as *const u8, 8) })
            .expect("failed to write to sent fd");

        let mut buf = [0u8; std::mem::size_of::<u64>()];
        evt_consumer
            .read_exact(buf.as_mut_slice())
            .expect("Failed to read from PipeReader");
        assert_eq!(u64::from_ne_bytes(buf), 1203);
    }

    #[test]
    fn send_recv_with_fd() {
        let (s1, s2) = UnixDatagram::pair().expect("failed to create socket pair");

        let (mut evt_consumer, evt_notifier) = pipe().expect("failed to create pipe");
        let write_count = s1
            .send_with_fds(&[[237].as_ref()], &[evt_notifier.as_raw_fd()])
            .expect("failed to send fd");

        assert_eq!(write_count, 1);

        let mut files = [0; 2];
        let mut buf = [0u8];
        let mut iovecs = [iovec {
            iov_base: buf.as_mut_ptr() as *mut c_void,
            iov_len: buf.len(),
        }];
        let (read_count, file_count) = unsafe {
            s2.recv_with_fds(&mut iovecs[..], &mut files)
                .expect("failed to recv fd")
        };

        assert_eq!(read_count, 1);
        assert_eq!(buf[0], 237);
        assert_eq!(file_count, 1);
        assert!(files[0] >= 0);
        assert_ne!(files[0], s1.as_raw_fd());
        assert_ne!(files[0], s2.as_raw_fd());
        assert_ne!(files[0], evt_notifier.as_raw_fd());

        let mut file = unsafe { File::from_raw_fd(files[0]) };

        file.write_all(unsafe { from_raw_parts(&1203u64 as *const u64 as *const u8, 8) })
            .expect("failed to write to sent fd");

        let mut buf = [0u8; std::mem::size_of::<u64>()];
        evt_consumer
            .read_exact(buf.as_mut_slice())
            .expect("Failed to read from PipeReader");
        assert_eq!(u64::from_ne_bytes(buf), 1203);
    }

    #[test]
    // Exercise the code paths that activate the issue of receiving the all the ancillary data,
    // but missing to provide enough buffer space to store it.
    fn send_more_recv_less() {
        // macos does not set MSG_CTRUNC if we pass a zero-size buffer for control data, even
        // if the sender does provide control data, while linux does
        #[cfg(any(target_os = "linux", target_os = "android"))]
        let start = 0;
        #[cfg(not(any(target_os = "linux", target_os = "android")))]
        let start = 1;

        for too_small in start..3 {
            let (s1, s2) = UnixDatagram::pair().expect("failed to create socket pair");

            let (_, evt_notifier1) = pipe().expect("failed to create pipe");
            let (_, evt_notifier2) = pipe().expect("failed to create pipe");
            let (_, evt_notifier3) = pipe().expect("failed to create pipe");
            let (_, evt_notifier4) = pipe().expect("failed to create pipe");
            let write_count = s1
                .send_with_fds(
                    &[[237].as_ref()],
                    &[
                        evt_notifier1.as_raw_fd(),
                        evt_notifier2.as_raw_fd(),
                        evt_notifier3.as_raw_fd(),
                        evt_notifier4.as_raw_fd(),
                    ],
                )
                .expect("failed to send fd");

            assert_eq!(write_count, 1);

            let mut files = vec![0; too_small];
            let mut buf = [0u8];
            let mut iovecs = [iovec {
                iov_base: buf.as_mut_ptr() as *mut c_void,
                iov_len: buf.len(),
            }];
            unsafe { s2.recv_with_fds(&mut iovecs[..], &mut files).unwrap_err() };
        }
    }
}
