feat: impl basic iovec ability for IoVecBufferMut

Add load_descriptor_chain and other essential methods for
IoVecBufferMut.

Signed-off-by: ihciah <[email protected]>

Author: ihciah

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

@@ -45,6 +45,16 @@ pub struct IoVecBuffer {
     len: u32,
 }
 
+impl IoVecBuffer {
+    /// Create an empty `IoVecBuffer` with a given capacity.
+    pub fn with_capacity(cap: usize) -> IoVecBuffer {
+        IoVecBuffer {
+            vecs: IoVecVec::with_capacity(cap),
+            len: 0,
+        }
+    }
+}
+
 // SAFETY: `IoVecBuffer` doesn't allow for interior mutability and no shared ownership is possible
 // as it doesn't implement clone
 unsafe impl Send for IoVecBuffer {}
@@ -218,21 +228,39 @@ impl IoVecBuffer {
 /// 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)]
+#[derive(Debug, Default)]
 pub struct IoVecBufferMut {
     // container of the memory regions included in this IO vector
     vecs: IoVecVec,
     // Total length of the IoVecBufferMut
     len: u32,
 }
 
+impl IoVecBufferMut {
+    /// Create an empty `IoVecBufferMut` with a given capacity.
+    pub fn with_capacity(cap: usize) -> IoVecBufferMut {
+        IoVecBufferMut {
+            vecs: IoVecVec::with_capacity(cap),
+            len: 0,
+        }
+    }
+}
+
+// SAFETY: iovec pointers are safe to send across threads.
+unsafe impl Send for IoVecBufferMut {}
+
 impl IoVecBufferMut {
     /// Create an `IoVecBufferMut` from a `DescriptorChain`
-    pub fn from_descriptor_chain(head: DescriptorChain) -> Result<Self, IoVecError> {
-        let mut vecs = IoVecVec::new();
-        let mut len = 0u32;
+    /// # Safety
+    /// The descriptor chain cannot be referencing the same memory location as another chain.
+    pub unsafe fn load_descriptor_chain(
+        &mut self,
+        mut desc: DescriptorChain,
+        max_size: Option<u32>,
+    ) -> Result<(), IoVecError> {
+        self.clear();
 
-        for desc in head {
+        loop {
             if !desc.is_write_only() {
                 return Err(IoVecError::ReadOnlyDescriptor);
             }
@@ -248,23 +276,57 @@ impl IoVecBufferMut {
             slice.bitmap().mark_dirty(0, desc.len as usize);
 
             let iov_base = slice.ptr_guard_mut().as_ptr().cast::<c_void>();
-            vecs.push(iovec {
+            self.vecs.push(iovec {
                 iov_base,
                 iov_len: desc.len as size_t,
             });
-            len = len
+            self.len = self
+                .len
                 .checked_add(desc.len)
                 .ok_or(IoVecError::OverflowedDescriptor)?;
+            if matches!(max_size, Some(max) if self.len >= max) {
+                break;
+            }
+            if desc.load_next_descriptor().is_none() {
+                break;
+            }
         }
 
-        Ok(Self { vecs, len })
+        Ok(())
+    }
+
+    /// Create an `IoVecBufferMut` from a `DescriptorChain`
+    /// # Safety
+    /// The descriptor chain cannot be referencing the same memory location as another chain.
+    pub unsafe fn from_descriptor_chain(head: DescriptorChain) -> Result<Self, IoVecError> {
+        let mut new_buffer = Self::default();
+
+        Self::load_descriptor_chain(&mut new_buffer, head, None)?;
+
+        Ok(new_buffer)
     }
 
     /// Get the total length of the memory regions covered by this `IoVecBuffer`
     pub(crate) fn len(&self) -> u32 {
         self.len
     }
 
+    /// Returns a pointer to the memory keeping the `iovec` structs
+    pub fn as_iovec_ptr(&self) -> *const iovec {
+        self.vecs.as_ptr()
+    }
+
+    /// Returns the length of the `iovec` array.
+    pub fn iovec_count(&self) -> usize {
+        self.vecs.len()
+    }
+
+    /// Clears the `iovec` array
+    pub fn clear(&mut self) {
+        self.vecs.clear();
+        self.len = 0u32;
+    }
+
     /// Writes a number of bytes into the `IoVecBufferMut` starting at a given offset.
     ///
     /// This will try to fill `IoVecBufferMut` writing bytes from the `buf` starting from
@@ -469,11 +531,13 @@ mod tests {
 
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBufferMut::from_descriptor_chain(head).unwrap_err();
+        // SAFETY: This descriptor chain is only loaded into one buffer.
+        unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap_err() };
 
         let (mut q, _) = write_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer.
+        unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
     }
 
     #[test]
@@ -494,7 +558,7 @@ mod tests {
         let head = q.pop(&mem).unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
-        let iovec = IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        let iovec = unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
         assert_eq!(iovec.len(), 4 * 64);
     }
 
@@ -559,7 +623,8 @@ mod tests {
         // This is a descriptor chain with 4 elements 64 bytes long each.
         let head = q.pop(&mem).unwrap();
 
-        let mut iovec = IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer.
+        let mut iovec = unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
         let buf = vec![0u8, 1, 2, 3, 4];
 
         // One test vector for each part of the chain

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

@@ -132,7 +132,8 @@ impl Entropy {
             let index = desc.index;
             METRICS.entropy_event_count.inc();
 
-            let bytes = match IoVecBufferMut::from_descriptor_chain(desc) {
+            // SAFETY: This descriptor chain is only loaded into one buffer.
+            let bytes = match unsafe { IoVecBufferMut::from_descriptor_chain(desc) } {
                 Ok(mut iovec) => {
                     debug!(
                         "entropy: guest request for {} bytes of entropy",
@@ -428,13 +429,15 @@ mod tests {
         // This should succeed, we just added two descriptors
         let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
         assert!(matches!(
-            IoVecBufferMut::from_descriptor_chain(desc),
+            // SAFETY: This descriptor chain is only loaded into one buffer.
+            unsafe { IoVecBufferMut::from_descriptor_chain(desc) },
             Err(crate::devices::virtio::iovec::IoVecError::ReadOnlyDescriptor)
         ));
 
         // This should succeed, we should have one more descriptor
         let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
-        let mut iovec = IoVecBufferMut::from_descriptor_chain(desc).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer.
+        let mut iovec = unsafe { IoVecBufferMut::from_descriptor_chain(desc).unwrap() };
         entropy_dev.handle_one(&mut iovec).unwrap();
     }
 

src/vmm/src/devices/virtio/vsock/packet.rs

@@ -161,7 +161,8 @@ impl VsockPacket {
     /// Returns [`VsockError::DescChainTooShortForHeader`] if the descriptor chain's total buffer
     /// length is insufficient to hold the 44 byte vsock header
     pub fn from_rx_virtq_head(chain: DescriptorChain) -> Result<Self, VsockError> {
-        let buffer = IoVecBufferMut::from_descriptor_chain(chain)?;
+        // SAFETY: This descriptor chain is only loaded into one buffer.
+        let buffer = unsafe { IoVecBufferMut::from_descriptor_chain(chain)? };
 
         if buffer.len() < VSOCK_PKT_HDR_SIZE {
             return Err(VsockError::DescChainTooShortForHeader(buffer.len() as usize));