virtio: allow IoVecBufferMut to hold multiple DescriptorChain objects

Allow IoVecBufferMut objects to store multiple DescriptorChain objects,
so that we can describe guest memory meant to be used for receiving data
(for example memory used for network RX) as a single (sparse) memory
region.

This will allow us to always keep track all the available memory we have
for performing RX and use `readv` for copying memory from the TAP device
inside guest memory avoiding the extra copy. In the future, it will also
facilitate the implementation of mergeable buffers for the RX path of
the network device.

Co-authored-by: Babis Chalios <[email protected]>
Signed-off-by: Egor Lazarchuk <[email protected]>

Author: bchalios

src/vmm/src/devices/virtio/iovec.rs

@@ -9,6 +9,7 @@ use vm_memory::{
     GuestMemory, GuestMemoryError, ReadVolatile, VolatileMemoryError, VolatileSlice, WriteVolatile,
 };
 
+use super::iov_deque::{IovDeque, IovDequeError};
 use crate::devices::virtio::queue::DescriptorChain;
 use crate::vstate::memory::GuestMemoryMmap;
 
@@ -20,8 +21,12 @@ pub enum IoVecError {
     ReadOnlyDescriptor,
     /// Tried to create an `IoVec` or `IoVecMut` from a descriptor chain that was too large
     OverflowedDescriptor,
+    /// Tried to push to full IovDeque.
+    IovDequeOverflow,
     /// Guest memory error: {0}
     GuestMemory(#[from] GuestMemoryError),
+    /// Error with underlying `IovDeque`: {0}
+    IovDeque(#[from] IovDequeError),
 }
 
 /// This is essentially a wrapper of a `Vec<libc::iovec>` which can be passed to `libc::writev`.
@@ -205,16 +210,26 @@ impl IoVecBuffer {
     }
 }
 
+#[derive(Debug)]
+pub struct ParsedDescriptorChain {
+    pub head_index: u16,
+    pub length: u32,
+    pub nr_iovecs: u16,
+}
+
 /// This is essentially a wrapper of a `Vec<libc::iovec>` which can be passed to `libc::readv`.
 ///
 /// It describes a write-only buffer passed to us by the guest that is scattered across multiple
 /// memory regions. Additionally, this wrapper provides methods that allow reading arbitrary ranges
 /// of data from that buffer.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct IoVecBufferMut {
     // container of the memory regions included in this IO vector
-    vecs: Vec<iovec>,
+    vecs: IovDeque,
     // Total length of the IoVecBufferMut
+    // We use `u32` here because we use this type in devices which
+    // should not give us huge buffers. In any case this
+    // value will not overflow as we explicitly check for this case.
     len: u32,
 }
 
@@ -223,50 +238,113 @@ pub struct IoVecBufferMut {
 unsafe impl Send for IoVecBufferMut {}
 
 impl IoVecBufferMut {
-    /// Create an `IoVecBuffer` from a `DescriptorChain`
+    /// Append a `DescriptorChain` in this `IoVecBufferMut`
     ///
     /// # Safety
     ///
     /// The descriptor chain cannot be referencing the same memory location as another chain
-    pub unsafe fn load_descriptor_chain(
+    pub unsafe fn append_descriptor_chain(
         &mut self,
         mem: &GuestMemoryMmap,
         head: DescriptorChain,
-    ) -> Result<(), IoVecError> {
-        self.clear();
-
+    ) -> Result<ParsedDescriptorChain, IoVecError> {
+        let head_index = head.index;
         let mut next_descriptor = Some(head);
+        let mut length = 0u32;
+        let mut nr_iovecs = 0u16;
         while let Some(desc) = next_descriptor {
             if !desc.is_write_only() {
+                self.vecs.pop_back(nr_iovecs);
                 return Err(IoVecError::ReadOnlyDescriptor);
             }
 
             // We use get_slice instead of `get_host_address` here in order to have the whole
             // range of the descriptor chain checked, i.e. [addr, addr + len) is a valid memory
             // region in the GuestMemoryMmap.
-            let slice = mem.get_slice(desc.addr, desc.len as usize)?;
-
+            let slice = mem.get_slice(desc.addr, desc.len as usize).map_err(|err| {
+                self.vecs.pop_back(nr_iovecs);
+                err
+            })?;
             // We need to mark the area of guest memory that will be mutated through this
             // IoVecBufferMut as dirty ahead of time, as we loose access to all
             // vm-memory related information after converting down to iovecs.
             slice.bitmap().mark_dirty(0, desc.len as usize);
-
             let iov_base = slice.ptr_guard_mut().as_ptr().cast::<c_void>();
-            self.vecs.push(iovec {
+
+            if self.vecs.is_full() {
+                self.vecs.pop_back(nr_iovecs);
+                return Err(IoVecError::IovDequeOverflow);
+            }
+
+            self.vecs.push_back(iovec {
                 iov_base,
                 iov_len: desc.len as size_t,
             });
-            self.len = self
-                .len
+
+            nr_iovecs += 1;
+            length = length
                 .checked_add(desc.len)
-                .ok_or(IoVecError::OverflowedDescriptor)?;
+                .ok_or(IoVecError::OverflowedDescriptor)
+                .map_err(|err| {
+                    self.vecs.pop_back(nr_iovecs);
+                    err
+                })?;
 
             next_descriptor = desc.next_descriptor();
         }
 
+        self.len = self.len.checked_add(length).ok_or_else(|| {
+            self.vecs.pop_back(nr_iovecs);
+            IoVecError::OverflowedDescriptor
+        })?;
+
+        Ok(ParsedDescriptorChain {
+            head_index,
+            length,
+            nr_iovecs,
+        })
+    }
+
+    /// Create an empty `IoVecBufferMut`.
+    pub fn new() -> Result<Self, IovDequeError> {
+        let vecs = IovDeque::new()?;
+        Ok(Self { vecs, len: 0 })
+    }
+
+    /// Create an `IoVecBufferMut` from a `DescriptorChain`
+    ///
+    /// This will clear any previous `iovec` objects in the buffer and load the new
+    /// [`DescriptorChain`].
+    ///
+    /// # Safety
+    ///
+    /// The descriptor chain cannot be referencing the same memory location as another chain
+    pub unsafe fn load_descriptor_chain(
+        &mut self,
+        mem: &GuestMemoryMmap,
+        head: DescriptorChain,
+    ) -> Result<(), IoVecError> {
+        self.clear();
+        let _ = self.append_descriptor_chain(mem, head)?;
         Ok(())
     }
 
+    /// Drop descriptor chain from the `IoVecBufferMut` front
+    ///
+    /// This will drop memory described by the `IoVecBufferMut` from the beginning.
+    pub fn drop_chain_front(&mut self, parse_descriptor: &ParsedDescriptorChain) {
+        self.vecs.pop_front(parse_descriptor.nr_iovecs);
+        self.len -= parse_descriptor.length;
+    }
+
+    /// Drop descriptor chain from the `IoVecBufferMut` back
+    ///
+    /// This will drop memory described by the `IoVecBufferMut` from the beginning.
+    pub fn drop_chain_back(&mut self, parse_descriptor: &ParsedDescriptorChain) {
+        self.vecs.pop_back(parse_descriptor.nr_iovecs);
+        self.len -= parse_descriptor.length;
+    }
+
     /// Create an `IoVecBuffer` from a `DescriptorChain`
     ///
     /// # Safety
@@ -276,20 +354,32 @@ impl IoVecBufferMut {
         mem: &GuestMemoryMmap,
         head: DescriptorChain,
     ) -> Result<Self, IoVecError> {
-        let mut new_buffer = Self::default();
+        let mut new_buffer = Self::new()?;
         new_buffer.load_descriptor_chain(mem, head)?;
         Ok(new_buffer)
     }
 
     /// Get the total length of the memory regions covered by this `IoVecBuffer`
-    pub(crate) fn len(&self) -> u32 {
+    #[inline(always)]
+    pub fn len(&self) -> u32 {
         self.len
     }
 
+    /// Returns true if buffer is empty.
+    #[inline(always)]
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    /// Returns a pointer to the memory keeping the `iovec` structs
+    pub fn as_iovec_mut_slice(&mut self) -> &mut [iovec] {
+        self.vecs.as_mut_slice()
+    }
+
     /// Clears the `iovec` array
     pub fn clear(&mut self) {
         self.vecs.clear();
-        self.len = 0u32;
+        self.len = 0;
     }
 
     /// Writes a number of bytes into the `IoVecBufferMut` starting at a given offset.
@@ -337,7 +427,7 @@ impl IoVecBufferMut {
     ) -> Result<usize, VolatileMemoryError> {
         let mut total_bytes_read = 0;
 
-        for iov in &self.vecs {
+        for iov in self.vecs.as_slice() {
             if len == 0 {
                 break;
             }
@@ -381,11 +471,13 @@ impl IoVecBufferMut {
 }
 
 #[cfg(test)]
+#[allow(clippy::cast_possible_truncation)]
 mod tests {
     use libc::{c_void, iovec};
     use vm_memory::VolatileMemoryError;
 
     use super::{IoVecBuffer, IoVecBufferMut};
+    use crate::devices::virtio::iov_deque::IovDeque;
     use crate::devices::virtio::queue::{Queue, VIRTQ_DESC_F_NEXT, VIRTQ_DESC_F_WRITE};
     use crate::devices::virtio::test_utils::VirtQueue;
     use crate::test_utils::multi_region_mem;
@@ -423,16 +515,36 @@ mod tests {
 
     impl From<&mut [u8]> for IoVecBufferMut {
         fn from(buf: &mut [u8]) -> Self {
+            let mut vecs = IovDeque::new().unwrap();
+            vecs.push_back(iovec {
+                iov_base: buf.as_mut_ptr().cast::<c_void>(),
+                iov_len: buf.len(),
+            });
+
             Self {
-                vecs: vec![iovec {
-                    iov_base: buf.as_mut_ptr().cast::<c_void>(),
-                    iov_len: buf.len(),
-                }],
-                len: buf.len().try_into().unwrap(),
+                vecs,
+                len: buf.len() as u32,
             }
         }
     }
 
+    impl From<Vec<&mut [u8]>> for IoVecBufferMut {
+        fn from(buffer: Vec<&mut [u8]>) -> Self {
+            let mut len = 0;
+            let mut vecs = IovDeque::new().unwrap();
+            for slice in buffer {
+                len += slice.len() as u32;
+
+                vecs.push_back(iovec {
+                    iov_base: slice.as_ptr() as *mut c_void,
+                    iov_len: slice.len(),
+                });
+            }
+
+            Self { vecs, len }
+        }
+    }
+
     fn default_mem() -> GuestMemoryMmap {
         multi_region_mem(&[
             (GuestAddress(0), 0x10000),
@@ -521,8 +633,19 @@ mod tests {
         let head = q.pop().unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
-        let iovec = unsafe { IoVecBufferMut::from_descriptor_chain(&mem, head).unwrap() };
+        let mut iovec = unsafe { IoVecBufferMut::from_descriptor_chain(&mem, head).unwrap() };
         assert_eq!(iovec.len(), 4 * 64);
+
+        // We are creating a new queue where we can get descriptors from. Probably, this is not
+        // something that we will ever want to do, as `IoVecBufferMut`s are typically
+        // (concpetually) associated with a single `Queue`. We just do this here to be able to test
+        // the appending logic.
+        let (mut q, _) = write_only_chain(&mem);
+        let head = q.pop().unwrap();
+        // SAFETY: it is actually unsafe, but we just want to check the length of the
+        // `IoVecBufferMut` after appending.
+        let _ = unsafe { iovec.append_descriptor_chain(&mem, head).unwrap() };
+        assert_eq!(iovec.len(), 8 * 64);
     }
 
     #[test]
@@ -680,6 +803,7 @@ mod verification {
     use vm_memory::VolatileSlice;
 
     use super::{IoVecBuffer, IoVecBufferMut};
+    use crate::devices::virtio::iov_deque::IovDeque;
 
     // Maximum memory size to use for our buffers. For the time being 1KB.
     const GUEST_MEMORY_SIZE: usize = 1 << 10;
@@ -695,10 +819,10 @@ mod verification {
         let mut vecs: Vec<iovec> = Vec::with_capacity(nr_descs);
         let mut len = 0u32;
         for _ in 0..nr_descs {
-            // The `IoVecBuffer(Mut)` constructors ensure that the memory region described by every
+            // The `IoVecBuffer` constructors ensure that the memory region described by every
             // `Descriptor` in the chain is a valid, i.e. it is memory with then guest's memory
             // mmap. The assumption, here, that the last address is within the memory object's
-            // bound substitutes these checks that `IoVecBuffer(Mut)::new() performs.`
+            // bound substitutes these checks that `IoVecBuffer::new() performs.`
             let addr: usize = kani::any();
             let iov_len: usize =
                 kani::any_where(|&len| matches!(addr.checked_add(len), Some(x) if x <= size));
@@ -721,6 +845,26 @@ mod verification {
         }
     }
 
+    fn create_iovecs_mut(mem: *mut u8, size: usize, nr_descs: usize) -> (IovDeque, u32) {
+        let mut vecs = IovDeque::new().unwrap();
+        let mut len = 0u32;
+        for _ in 0..nr_descs {
+            // The `IoVecBufferMut` constructors ensure that the memory region described by every
+            // `Descriptor` in the chain is a valid, i.e. it is memory with then guest's memory
+            // mmap. The assumption, here, that the last address is within the memory object's
+            // bound substitutes these checks that `IoVecBufferMut::new() performs.`
+            let addr: usize = kani::any();
+            let iov_len: usize =
+                kani::any_where(|&len| matches!(addr.checked_add(len), Some(x) if x <= size));
+            let iov_base = unsafe { mem.offset(addr.try_into().unwrap()) } as *mut c_void;
+
+            vecs.push_back(iovec { iov_base, iov_len });
+            len += u32::try_from(iov_len).unwrap();
+        }
+
+        (vecs, len)
+    }
+
     impl IoVecBufferMut {
         fn any_of_length(nr_descs: usize) -> Self {
             // We only write into `IoVecBufferMut` objects, so we can simply create a guest memory
@@ -732,8 +876,11 @@ mod verification {
                 ))
             };
 
-            let (vecs, len) = create_iovecs(mem, GUEST_MEMORY_SIZE, nr_descs);
-            Self { vecs, len }
+            let (vecs, len) = create_iovecs_mut(mem, GUEST_MEMORY_SIZE, nr_descs);
+            Self {
+                vecs,
+                len: len.try_into().unwrap(),
+            }
         }
     }
 

src/vmm/src/devices/virtio/rng/device.rs

@@ -13,6 +13,7 @@ use super::metrics::METRICS;
 use super::{RNG_NUM_QUEUES, RNG_QUEUE};
 use crate::devices::virtio::device::{DeviceState, IrqTrigger, IrqType, VirtioDevice};
 use crate::devices::virtio::gen::virtio_rng::VIRTIO_F_VERSION_1;
+use crate::devices::virtio::iov_deque::IovDequeError;
 use crate::devices::virtio::iovec::IoVecBufferMut;
 use crate::devices::virtio::queue::{Queue, FIRECRACKER_MAX_QUEUE_SIZE};
 use crate::devices::virtio::{ActivateError, TYPE_RNG};
@@ -31,6 +32,8 @@ pub enum EntropyError {
     GuestMemory(#[from] GuestMemoryError),
     /// Could not get random bytes: {0}
     Random(#[from] aws_lc_rs::error::Unspecified),
+    /// Underlying IovDeque error: {0}
+    IovDeque(#[from] IovDequeError),
 }
 
 #[derive(Debug)]
@@ -77,7 +80,7 @@ impl Entropy {
             queue_events,
             irq_trigger,
             rate_limiter,
-            buffer: IoVecBufferMut::default(),
+            buffer: IoVecBufferMut::new()?,
         })
     }
 
@@ -111,7 +114,7 @@ impl Entropy {
 
     fn handle_one(&mut self) -> Result<u32, EntropyError> {
         // If guest provided us with an empty buffer just return directly
-        if self.buffer.len() == 0 {
+        if self.buffer.is_empty() {
             return Ok(0);
         }