Reuse the IoVecBuffer on TX

On the net virtio device reuse the IoVecBuffer on the TX path

Signed-off-by: Jack Thomson <[email protected]>

Author: JackThomson2

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

@@ -37,19 +37,29 @@ type IoVecVec = SmallVec<[iovec; 4]>;
 /// It describes a 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 IoVecBuffer {
     // container of the memory regions included in this IO vector
     vecs: IoVecVec,
     // Total length of the IoVecBuffer
     len: u32,
 }
 
+// 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 {}
+
 impl IoVecBuffer {
     /// Create an `IoVecBuffer` 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,
+        head: DescriptorChain,
+    ) -> Result<(), IoVecError> {
+        self.clear();
 
         let mut next_descriptor = Some(head);
         while let Some(desc) = next_descriptor {
@@ -66,18 +76,32 @@ impl IoVecBuffer {
                 .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)?;
 
             next_descriptor = desc.next_descriptor();
         }
 
-        Ok(Self { vecs, len })
+        Ok(())
+    }
+
+    /// Create an `IoVecBuffer` 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::default();
+
+        new_buffer.load_descriptor_chain(head)?;
+
+        Ok(new_buffer)
     }
 
     /// Get the total length of the memory regions covered by this `IoVecBuffer`
@@ -95,6 +119,12 @@ impl IoVecBuffer {
         self.vecs.len()
     }
 
+    /// Clears the `iovec` array
+    pub fn clear(&mut self) {
+        self.vecs.clear();
+        self.len = 0u32;
+    }
+
     /// Reads a number of bytes from the `IoVecBuffer` starting at a given offset.
     ///
     /// This will try to fill `buf` reading bytes from the `IoVecBuffer` starting from
@@ -428,11 +458,13 @@ mod tests {
         let mem = default_mem();
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBuffer::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        unsafe { IoVecBuffer::from_descriptor_chain(head).unwrap() };
 
         let (mut q, _) = write_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBuffer::from_descriptor_chain(head).unwrap_err();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        unsafe { IoVecBuffer::from_descriptor_chain(head).unwrap_err() };
 
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
@@ -449,7 +481,8 @@ mod tests {
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
 
-        let iovec = IoVecBuffer::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded once in this test
+        let iovec = unsafe { IoVecBuffer::from_descriptor_chain(head).unwrap() };
         assert_eq!(iovec.len(), 4 * 64);
     }
 
@@ -459,6 +492,7 @@ mod tests {
         let (mut q, _) = write_only_chain(&mem);
         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();
         assert_eq!(iovec.len(), 4 * 64);
     }
@@ -469,7 +503,8 @@ mod tests {
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
 
-        let iovec = IoVecBuffer::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded once in this test
+        let iovec = unsafe { IoVecBuffer::from_descriptor_chain(head).unwrap() };
 
         let mut buf = vec![0u8; 257];
         assert_eq!(

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

@@ -142,6 +142,8 @@ pub struct Net {
     /// Only if MMDS transport has been associated with it.
     pub mmds_ns: Option<MmdsNetworkStack>,
     pub(crate) metrics: Arc<NetDeviceMetrics>,
+
+    tx_buffer: IoVecBuffer,
 }
 
 impl Net {
@@ -199,6 +201,7 @@ impl Net {
             activate_evt: EventFd::new(libc::EFD_NONBLOCK).map_err(NetError::EventFd)?,
             mmds_ns: None,
             metrics: NetMetricsPerDevice::alloc(id),
+            tx_buffer: Default::default(),
         })
     }
 
@@ -597,19 +600,19 @@ impl Net {
                 .add(tx_queue.len(mem).into());
             let head_index = head.index;
             // Parse IoVecBuffer from descriptor head
-            let buffer = match IoVecBuffer::from_descriptor_chain(head) {
-                Ok(buffer) => buffer,
-                Err(_) => {
-                    self.metrics.tx_fails.inc();
-                    tx_queue
-                        .add_used(mem, head_index, 0)
-                        .map_err(DeviceError::QueueError)?;
-                    continue;
-                }
+            // SAFETY: This descriptor chain is only loaded once
+            // virtio requests are handled sequentially so no two IoVecBuffers
+            // are live at the same time, meaning this has exclusive ownership over the memory
+            if unsafe { self.tx_buffer.load_descriptor_chain(head).is_err() } {
+                self.metrics.tx_fails.inc();
+                tx_queue
+                    .add_used(mem, head_index, 0)
+                    .map_err(DeviceError::QueueError)?;
+                continue;
             };
 
             // We only handle frames that are up to MAX_BUFFER_SIZE
-            if buffer.len() as usize > MAX_BUFFER_SIZE {
+            if self.tx_buffer.len() as usize > MAX_BUFFER_SIZE {
                 error!("net: received too big frame from driver");
                 self.metrics.tx_malformed_frames.inc();
                 tx_queue
@@ -618,7 +621,10 @@ impl Net {
                 continue;
             }
 
-            if !Self::rate_limiter_consume_op(&mut self.tx_rate_limiter, u64::from(buffer.len())) {
+            if !Self::rate_limiter_consume_op(
+                &mut self.tx_rate_limiter,
+                u64::from(self.tx_buffer.len()),
+            ) {
                 tx_queue.undo_pop();
                 self.metrics.tx_rate_limiter_throttled.inc();
                 break;
@@ -628,7 +634,7 @@ impl Net {
                 self.mmds_ns.as_mut(),
                 &mut self.tx_rate_limiter,
                 &mut self.tx_frame_headers,
-                &buffer,
+                &self.tx_buffer,
                 &mut self.tap,
                 self.guest_mac,
                 &self.metrics,
@@ -649,6 +655,8 @@ impl Net {
             self.metrics.no_tx_avail_buffer.inc();
         }
 
+        // Cleanup tx_buffer to ensure no two buffers point at the same memory
+        self.tx_buffer.clear();
         self.try_signal_queue(NetQueue::Tx)?;
 
         // An incoming frame for the MMDS may trigger the transmission of a new message.

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

@@ -124,7 +124,10 @@ impl VsockPacket {
     /// - [`VsockError::DescChainTooShortForPacket`] if the contained vsock header describes a vsock
     ///   packet whose length exceeds the descriptor chain's actual total buffer length.
     pub fn from_tx_virtq_head(chain: DescriptorChain) -> Result<Self, VsockError> {
-        let buffer = IoVecBuffer::from_descriptor_chain(chain)?;
+        // SAFETY: This descriptor chain is only loaded once
+        // virtio requests are handled sequentially so no two IoVecBuffers
+        // are live at the same time, meaning this has exclusive ownership over the memory
+        let buffer = unsafe { IoVecBuffer::from_descriptor_chain(chain)? };
 
         let mut hdr = VsockPacketHeader::default();
         match buffer.read_exact_volatile_at(hdr.as_mut_slice(), 0) {