<commit_before>// Copyright (C) 2019 Alibaba Cloud Computing. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Common data structures for listener and endpoint.

cfg_if::cfg_if! {
    if #[cfg(unix)] {
        pub mod socket;
        #[cfg(feature = "vfio-device")]
        pub mod vfio;
        mod unix;
    } else if #[cfg(windows)] {
        mod tube;
        pub use tube::TubeEndpoint;
        mod windows;
    }
}

use base::RawDescriptor;
use std::fs::File;
use std::io::{IoSlice, IoSliceMut};
use std::mem;
use std::path::Path;

use data_model::DataInit;

use super::message::*;
use super::{Error, Result};
use crate::connection::Req;

/// Listener for accepting connections.
pub trait Listener: Sized {
    /// Type of an object created when a connection is accepted.
    type Connection;
    /// Type of endpoint created when a connection is accepted.
    type Endpoint;

    /// Accept an incoming connection.
    fn accept(&mut self) -> Result<Option<Self::Endpoint>>;

    /// Change blocking status on the listener.
    fn set_nonblocking(&self, block: bool) -> Result<()>;
}

/// Abstracts a vhost-user connection and related operations.
pub trait Endpoint<R: Req>: Send {
    /// Create a new stream by connecting to server at `str`.
    fn connect<P: AsRef<Path>>(path: P) -> Result<Self>
    where
        Self: Sized;

    /// Sends bytes from scatter-gather vectors with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    fn send_iovec(&mut self, iovs: &[IoSlice], fds: Option<&[RawDescriptor]>) -> Result<usize>;

    /// Reads bytes into the given scatter/gather vectors with optional attached file.
    ///
    /// # Arguements
    /// * `bufs` - A slice of buffers to store received data.
    /// * `allow_fd` - Indicates whether we can receive FDs.
    ///
    /// # Return:
    /// * - (number of bytes received, [received files]) on success.
    /// * - `Error::Disconnect` if the client closed.
    fn recv_into_bufs(
        &mut self,
        bufs: &mut [IoSliceMut],
        allow_fd: bool,
    ) -> Result<(usize, Option<Vec<File>>)>;

    /// Constructs the slave request endpoint for self.
    ///
    /// # Arguments
    /// * `files` - Files from which to create the endpoint
    fn create_slave_request_endpoint(
        &mut self,
        files: Option<Vec<File>>,
    ) -> Result<Box<dyn Endpoint<SlaveReq>>>;
}

// Advance the internal cursor of the slices.
// This is same with a nightly API `IoSlice::advance_slices` but for `&[u8]`.
fn advance_slices(bufs: &mut &mut [&[u8]], mut count: usize) {
    use std::mem::take;

    let mut idx = 0;
    for b in bufs.iter() {
        if count < b.len() {
            break;
        }
        count -= b.len();
        idx += 1;
    }
    *bufs = &mut take(bufs)[idx..];
    if !bufs.is_empty() {
        bufs[0] = &bufs[0][count..];
    }
}

// Advance the internal cursor of the slices.
// This is same with a nightly API `IoSliceMut::advance_slices` but for `&mut [u8]`.
fn advance_slices_mut(bufs: &mut &mut [&mut [u8]], mut count: usize) {
    use std::mem::take;

    let mut idx = 0;
    for b in bufs.iter() {
        if count < b.len() {
            break;
        }
        count -= b.len();
        idx += 1;
    }
    *bufs = &mut take(bufs)[idx..];
    if !bufs.is_empty() {
        let slice = take(&mut bufs[0]);
        let (_, remaining) = slice.split_at_mut(count);
        bufs[0] = remaining;
    }
}

/// Abstracts VVU message parsing, sending and receiving.
pub trait EndpointExt<R: Req>: Endpoint<R> {
    /// Sends all bytes from scatter-gather vectors with optional attached file descriptors. Will
    /// loop until all data has been transfered.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    ///
    /// # TODO
    /// This function takes a slice of `&[u8]` instead of `IoSlice` because the internal
    /// cursor needs to be moved by `advance_slices()`.
    /// Once `IoSlice::advance_slices()` becomes stable, this should be updated.
    /// <https://github.com/rust-lang/rust/issues/62726>.
    fn send_iovec_all(
        &mut self,
        mut iovs: &mut [&[u8]],
        mut fds: Option<&[RawDescriptor]>,
    ) -> Result<usize> {
        // Guarantee that `iovs` becomes empty if it doesn't contain any data.
        advance_slices(&mut iovs, 0);

        let mut data_sent = 0;
        while !iovs.is_empty() {
            let iovec: Vec<_> = iovs.iter_mut().map(|i| IoSlice::new(i)).collect();
            match self.send_iovec(&iovec, fds) {
                Ok(0) => {
                    break;
                }
                Ok(n) => {
                    data_sent += n;
                    fds = None;
                    advance_slices(&mut iovs, n);
                }
                Err(e) => match e {
                    Error::SocketRetry(_) => {}
                    _ => return Err(e),
                },
            }
        }
        Ok(data_sent)
    }

    /// Sends bytes from a slice with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    #[cfg(test)]
    fn send_slice(&mut self, data: IoSlice, fds: Option<&[RawDescriptor]>) -> Result<usize> {
        self.send_iovec(&[data], fds)
    }

    /// Sends a header-only message with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - PartialMessage: received a partial message.
    /// * - backend specific errors
    fn send_header(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        let mut iovs = [hdr.as_slice()];
        let bytes = self.send_iovec_all(&mut iovs[..], fds)?;
        if bytes != mem::size_of::<VhostUserMsgHeader<R>>() {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Send a message with header and body. Optional file descriptors may be attached to
    /// the message.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - OversizedMsg: message size is too big.
    /// * - PartialMessage: received a partial message.
    /// * - backend specific errors
    fn send_message<T: Sized + DataInit>(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        body: &T,
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::OversizedMsg);
        }

        // We send the header and the body separately here. This is necessary on Windows. Otherwise
        // the recv side cannot read the header independently (the transport is message oriented).
        let mut bytes = self.send_iovec_all(&mut [hdr.as_slice()], fds)?;
        bytes += self.send_iovec_all(&mut [body.as_slice()], None)?;
        if bytes != mem::size_of::<VhostUserMsgHeader<R>>() + mem::size_of::<T>() {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Send a message with header, body and payload. Optional file descriptors
    /// may also be attached to the message.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - OversizedMsg: message size is too big.
    /// * - PartialMessage: received a partial message.
    /// * - IncorrectFds: wrong number of attached fds.
    /// * - backend specific errors
    fn send_message_with_payload<T: Sized + DataInit>(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        body: &T,
        payload: &[u8],
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        let len = payload.len();
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::OversizedMsg);
        }
        if len > MAX_MSG_SIZE - mem::size_of::<T>() {
            return Err(Error::OversizedMsg);
        }
        if let Some(fd_arr) = fds {
            if fd_arr.len() > MAX_ATTACHED_FD_ENTRIES {
                return Err(Error::IncorrectFds);
            }
        }

        let total = mem::size_of::<VhostUserMsgHeader<R>>() + mem::size_of::<T>() + len;

        // We send the header and the body separately here. This is necessary on Windows. Otherwise
        // the recv side cannot read the header independently (the transport is message oriented).
        let mut len = self.send_iovec_all(&mut [hdr.as_slice()], fds)?;
        len += self.send_iovec_all(&mut [body.as_slice(), payload], None)?;

        if len != total {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Reads `len` bytes at most.
    ///
    /// # Return:
    /// * - (number of bytes received, buf) on success
    fn recv_data(&mut self, len: usize) -> Result<Vec<u8>> {
        let mut buf = vec![0u8; len];
        let (data_len, _) =
            self.recv_into_bufs(&mut [IoSliceMut::new(&mut buf)], false /* allow_fd */)?;
        buf.truncate(data_len);
        Ok(buf)
    }

    /// Reads all bytes into the given scatter/gather vectors with optional attached files. Will
    /// loop until all data has been transferred.
    ///
    /// # Return:
    /// * - (number of bytes received, [received fds]) on success
    /// * - `Disconnect` - client is closed
    ///
    /// # TODO
    /// This function takes a slice of `&mut [u8]` instead of `IoSliceMut` because the internal
    /// cursor needs to be moved by `advance_slices_mut()`.
    /// Once `IoSliceMut::advance_slices()` becomes stable, this should be updated.
    /// <https://github.com/rust-lang/rust/issues/62726>.
    fn recv_into_bufs_all(
        &mut self,
        mut bufs: &mut [&mut [u8]],
    ) -> Result<(usize, Option<Vec<File>>)> {
        let buf_lens: Vec<usize> = bufs.iter().map(|b| b.len()).collect();
        let data_total: usize = buf_lens.iter().sum();
        let mut data_read = 0;
        let mut rfds = None;

        while (data_total - data_read) > 0 {
            let mut slices: Vec<IoSliceMut> = bufs.iter_mut().map(|b| IoSliceMut::new(b)).collect();
            let res = self.recv_into_bufs(&mut slices, true);
            match res {
                Ok((0, _)) => return Ok((data_read, rfds)),
                Ok((n, fds)) => {
                    if data_read == 0 {
                        rfds = fds;
                    }<commit_msg>vmm_vhost: eliminate unnecessary collect() into a Vec<commit_after>// Copyright (C) 2019 Alibaba Cloud Computing. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Common data structures for listener and endpoint.

cfg_if::cfg_if! {
    if #[cfg(unix)] {
        pub mod socket;
        #[cfg(feature = "vfio-device")]
        pub mod vfio;
        mod unix;
    } else if #[cfg(windows)] {
        mod tube;
        pub use tube::TubeEndpoint;
        mod windows;
    }
}

use base::RawDescriptor;
use std::fs::File;
use std::io::{IoSlice, IoSliceMut};
use std::mem;
use std::path::Path;

use data_model::DataInit;

use super::message::*;
use super::{Error, Result};
use crate::connection::Req;

/// Listener for accepting connections.
pub trait Listener: Sized {
    /// Type of an object created when a connection is accepted.
    type Connection;
    /// Type of endpoint created when a connection is accepted.
    type Endpoint;

    /// Accept an incoming connection.
    fn accept(&mut self) -> Result<Option<Self::Endpoint>>;

    /// Change blocking status on the listener.
    fn set_nonblocking(&self, block: bool) -> Result<()>;
}

/// Abstracts a vhost-user connection and related operations.
pub trait Endpoint<R: Req>: Send {
    /// Create a new stream by connecting to server at `str`.
    fn connect<P: AsRef<Path>>(path: P) -> Result<Self>
    where
        Self: Sized;

    /// Sends bytes from scatter-gather vectors with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    fn send_iovec(&mut self, iovs: &[IoSlice], fds: Option<&[RawDescriptor]>) -> Result<usize>;

    /// Reads bytes into the given scatter/gather vectors with optional attached file.
    ///
    /// # Arguements
    /// * `bufs` - A slice of buffers to store received data.
    /// * `allow_fd` - Indicates whether we can receive FDs.
    ///
    /// # Return:
    /// * - (number of bytes received, [received files]) on success.
    /// * - `Error::Disconnect` if the client closed.
    fn recv_into_bufs(
        &mut self,
        bufs: &mut [IoSliceMut],
        allow_fd: bool,
    ) -> Result<(usize, Option<Vec<File>>)>;

    /// Constructs the slave request endpoint for self.
    ///
    /// # Arguments
    /// * `files` - Files from which to create the endpoint
    fn create_slave_request_endpoint(
        &mut self,
        files: Option<Vec<File>>,
    ) -> Result<Box<dyn Endpoint<SlaveReq>>>;
}

// Advance the internal cursor of the slices.
// This is same with a nightly API `IoSlice::advance_slices` but for `&[u8]`.
fn advance_slices(bufs: &mut &mut [&[u8]], mut count: usize) {
    use std::mem::take;

    let mut idx = 0;
    for b in bufs.iter() {
        if count < b.len() {
            break;
        }
        count -= b.len();
        idx += 1;
    }
    *bufs = &mut take(bufs)[idx..];
    if !bufs.is_empty() {
        bufs[0] = &bufs[0][count..];
    }
}

// Advance the internal cursor of the slices.
// This is same with a nightly API `IoSliceMut::advance_slices` but for `&mut [u8]`.
fn advance_slices_mut(bufs: &mut &mut [&mut [u8]], mut count: usize) {
    use std::mem::take;

    let mut idx = 0;
    for b in bufs.iter() {
        if count < b.len() {
            break;
        }
        count -= b.len();
        idx += 1;
    }
    *bufs = &mut take(bufs)[idx..];
    if !bufs.is_empty() {
        let slice = take(&mut bufs[0]);
        let (_, remaining) = slice.split_at_mut(count);
        bufs[0] = remaining;
    }
}

/// Abstracts VVU message parsing, sending and receiving.
pub trait EndpointExt<R: Req>: Endpoint<R> {
    /// Sends all bytes from scatter-gather vectors with optional attached file descriptors. Will
    /// loop until all data has been transfered.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    ///
    /// # TODO
    /// This function takes a slice of `&[u8]` instead of `IoSlice` because the internal
    /// cursor needs to be moved by `advance_slices()`.
    /// Once `IoSlice::advance_slices()` becomes stable, this should be updated.
    /// <https://github.com/rust-lang/rust/issues/62726>.
    fn send_iovec_all(
        &mut self,
        mut iovs: &mut [&[u8]],
        mut fds: Option<&[RawDescriptor]>,
    ) -> Result<usize> {
        // Guarantee that `iovs` becomes empty if it doesn't contain any data.
        advance_slices(&mut iovs, 0);

        let mut data_sent = 0;
        while !iovs.is_empty() {
            let iovec: Vec<_> = iovs.iter_mut().map(|i| IoSlice::new(i)).collect();
            match self.send_iovec(&iovec, fds) {
                Ok(0) => {
                    break;
                }
                Ok(n) => {
                    data_sent += n;
                    fds = None;
                    advance_slices(&mut iovs, n);
                }
                Err(e) => match e {
                    Error::SocketRetry(_) => {}
                    _ => return Err(e),
                },
            }
        }
        Ok(data_sent)
    }

    /// Sends bytes from a slice with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    #[cfg(test)]
    fn send_slice(&mut self, data: IoSlice, fds: Option<&[RawDescriptor]>) -> Result<usize> {
        self.send_iovec(&[data], fds)
    }

    /// Sends a header-only message with optional attached file descriptors.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - PartialMessage: received a partial message.
    /// * - backend specific errors
    fn send_header(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        let mut iovs = [hdr.as_slice()];
        let bytes = self.send_iovec_all(&mut iovs[..], fds)?;
        if bytes != mem::size_of::<VhostUserMsgHeader<R>>() {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Send a message with header and body. Optional file descriptors may be attached to
    /// the message.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - OversizedMsg: message size is too big.
    /// * - PartialMessage: received a partial message.
    /// * - backend specific errors
    fn send_message<T: Sized + DataInit>(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        body: &T,
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::OversizedMsg);
        }

        // We send the header and the body separately here. This is necessary on Windows. Otherwise
        // the recv side cannot read the header independently (the transport is message oriented).
        let mut bytes = self.send_iovec_all(&mut [hdr.as_slice()], fds)?;
        bytes += self.send_iovec_all(&mut [body.as_slice()], None)?;
        if bytes != mem::size_of::<VhostUserMsgHeader<R>>() + mem::size_of::<T>() {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Send a message with header, body and payload. Optional file descriptors
    /// may also be attached to the message.
    ///
    /// # Return:
    /// * - number of bytes sent on success
    /// * - OversizedMsg: message size is too big.
    /// * - PartialMessage: received a partial message.
    /// * - IncorrectFds: wrong number of attached fds.
    /// * - backend specific errors
    fn send_message_with_payload<T: Sized + DataInit>(
        &mut self,
        hdr: &VhostUserMsgHeader<R>,
        body: &T,
        payload: &[u8],
        fds: Option<&[RawDescriptor]>,
    ) -> Result<()> {
        let len = payload.len();
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::OversizedMsg);
        }
        if len > MAX_MSG_SIZE - mem::size_of::<T>() {
            return Err(Error::OversizedMsg);
        }
        if let Some(fd_arr) = fds {
            if fd_arr.len() > MAX_ATTACHED_FD_ENTRIES {
                return Err(Error::IncorrectFds);
            }
        }

        let total = mem::size_of::<VhostUserMsgHeader<R>>() + mem::size_of::<T>() + len;

        // We send the header and the body separately here. This is necessary on Windows. Otherwise
        // the recv side cannot read the header independently (the transport is message oriented).
        let mut len = self.send_iovec_all(&mut [hdr.as_slice()], fds)?;
        len += self.send_iovec_all(&mut [body.as_slice(), payload], None)?;

        if len != total {
            return Err(Error::PartialMessage);
        }
        Ok(())
    }

    /// Reads `len` bytes at most.
    ///
    /// # Return:
    /// * - (number of bytes received, buf) on success
    fn recv_data(&mut self, len: usize) -> Result<Vec<u8>> {
        let mut buf = vec![0u8; len];
        let (data_len, _) =
            self.recv_into_bufs(&mut [IoSliceMut::new(&mut buf)], false /* allow_fd */)?;
        buf.truncate(data_len);
        Ok(buf)
    }

    /// Reads all bytes into the given scatter/gather vectors with optional attached files. Will
    /// loop until all data has been transferred.
    ///
    /// # Return:
    /// * - (number of bytes received, [received fds]) on success
    /// * - `Disconnect` - client is closed
    ///
    /// # TODO
    /// This function takes a slice of `&mut [u8]` instead of `IoSliceMut` because the internal
    /// cursor needs to be moved by `advance_slices_mut()`.
    /// Once `IoSliceMut::advance_slices()` becomes stable, this should be updated.
    /// <https://github.com/rust-lang/rust/issues/62726>.
    fn recv_into_bufs_all(
        &mut self,
        mut bufs: &mut [&mut [u8]],
    ) -> Result<(usize, Option<Vec<File>>)> {
        let data_total: usize = bufs.iter().map(|b| b.len()).sum();
        let mut data_read = 0;
        let mut rfds = None;

        while (data_total - data_read) > 0 {
            let mut slices: Vec<IoSliceMut> = bufs.iter_mut().map(|b| IoSliceMut::new(b)).collect();
            let res = self.recv_into_bufs(&mut slices, true);
            match res {
                Ok((0, _)) => return Ok((data_read, rfds)),
                Ok((n, fds)) => {
                    if data_read == 0 {
                        rfds = fds;
                    }
{'fraction_non_alphanumeric': 0.14401937046004842, 'fraction_numerical': 0.004164648910411622, 'mean_word_length': 2.015625, 'pattern_counts': {'":': 0, '<': 116, '<?xml version=': 0, '>': 162, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '117853', 'n_tokens_mistral': 6490, 'n_tokens_neox': 6396, 'n_words': 1890}