feat(net): add MRG_RXBUF support to virtio-net device

Now virtio-net device VIRTIO_NET_F_MRG_RXBUF feature which allows it
to write single packet into multiple descriptor chains.
The amount of descriptor chains (also known as heads) is written into
the `virtio_net_hdr_v1` structure which is located at the very begging
of the packet.

Signed-off-by: Egor Lazarchuk <[email protected]>

Author: ShadowCurse

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

@@ -367,10 +367,15 @@ impl IoVecBufferMut {
     /// In contrast to the equivalent [`IoVecBuffer::len()`] which returns `u32`, this one returns
     /// `usize` since the buffer can contain multiple `DescriptorChain` objects, so we don't have
     /// the limit that the length of a buffer is limited by `u32`.
-    pub(crate) fn len(&self) -> usize {
+    pub fn len(&self) -> usize {
         self.len
     }
 
+    /// Returns true if there is buffer is empty.
+    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()

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

@@ -6,11 +6,12 @@
 // found in the THIRD-PARTY file.
 
 use std::collections::VecDeque;
-use std::mem;
+use std::mem::{self, offset_of};
 use std::net::Ipv4Addr;
+use std::num::Wrapping;
 use std::sync::{Arc, Mutex};
 
-use libc::EAGAIN;
+use libc::{iovec, EAGAIN};
 use log::error;
 use vmm_sys_util::eventfd::EventFd;
 
@@ -19,7 +20,7 @@ use crate::devices::virtio::gen::virtio_blk::VIRTIO_F_VERSION_1;
 use crate::devices::virtio::gen::virtio_net::{
     virtio_net_hdr_v1, VIRTIO_NET_F_CSUM, VIRTIO_NET_F_GUEST_CSUM, VIRTIO_NET_F_GUEST_TSO4,
     VIRTIO_NET_F_GUEST_TSO6, VIRTIO_NET_F_GUEST_UFO, VIRTIO_NET_F_HOST_TSO4,
-    VIRTIO_NET_F_HOST_TSO6, VIRTIO_NET_F_HOST_UFO, VIRTIO_NET_F_MAC,
+    VIRTIO_NET_F_HOST_TSO6, VIRTIO_NET_F_HOST_UFO, VIRTIO_NET_F_MAC, VIRTIO_NET_F_MRG_RXBUF,
 };
 use crate::devices::virtio::gen::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
 use crate::devices::virtio::iovec::{
@@ -108,7 +109,8 @@ pub struct RxBuffers {
     // A map of which part of the memory belongs to which `DescriptorChain` object
     pub parsed_descriptors: VecDeque<ParsedDescriptorChain>,
     // Buffers that we have used and they are ready to be given back to the guest.
-    pub used_descriptor: Option<ParsedDescriptorChain>,
+    pub used_descriptors: u16,
+    pub used_bytes: u32,
 }
 
 impl RxBuffers {
@@ -118,7 +120,8 @@ impl RxBuffers {
             min_buffer_size: 0,
             iovec: IoVecBufferMut::new()?,
             parsed_descriptors: VecDeque::with_capacity(FIRECRACKER_MAX_QUEUE_SIZE.into()),
-            used_descriptor: None,
+            used_descriptors: 0,
+            used_bytes: 0,
         })
     }
 
@@ -141,37 +144,61 @@ impl RxBuffers {
         Ok(())
     }
 
-    /// Returns the number of available `iovec` objects.
+    /// Returns the total size of available space in the buffer.
     #[inline(always)]
-    fn len(&self) -> usize {
+    fn capacity(&self) -> usize {
         self.iovec.len()
     }
 
-    /// Returns `true` if there aren't any available `iovec` objects.
-    #[inline(always)]
-    fn is_empty(&self) -> bool {
-        self.len() == 0
-    }
-
     /// Mark the first `size` bytes of available memory as used.
     ///
     /// # Safety:
     ///
     /// * The `RxBuffers` should include at least one parsed `DescriptorChain`.
     /// * `size` needs to be smaller or equal to total length of the first `DescriptorChain` stored
     ///   in the `RxBuffers`.
-    unsafe fn mark_used(&mut self, size: u32) {
+    unsafe fn mark_used(&mut self, mut bytes_written: u32, rx_queue: &mut Queue) {
         // Since we were able to write a frame in guest memory, we should have at least one
         // descriptor chain here. If not, we have a bug, so fail fast, since the device is
         // fundamentally broken.
-        let mut parsed_dc = self.parsed_descriptors.pop_front().expect(
-            "net: internal bug. Mismatch between written frame size and available descriptors",
-        );
+        debug_assert!(!self.iovec.is_empty());
+
+        self.used_bytes = bytes_written;
 
-        self.header_set_num_buffers(1);
-        self.iovec.drop_descriptor_chain(&parsed_dc);
-        parsed_dc.length = size;
-        self.used_descriptor = Some(parsed_dc);
+        let mut used_heads: u16 = 0;
+        let mut write_used = |head_index: u16, bytes_written: u32, rx_queue: &mut Queue| {
+            if let Err(err) = rx_queue.write_used_element(
+                (rx_queue.next_used + Wrapping(self.used_descriptors)).0,
+                head_index,
+                bytes_written,
+            ) {
+                error!(
+                    "net: Failed to add used descriptor {} of length {} to RX queue: {err}",
+                    head_index, bytes_written
+                );
+            }
+            used_heads += 1;
+            self.used_descriptors += 1;
+        };
+
+        for parsed_dc in self.parsed_descriptors.iter() {
+            if bytes_written <= parsed_dc.length {
+                write_used(parsed_dc.head_index, bytes_written, rx_queue);
+                break;
+            } else {
+                write_used(parsed_dc.head_index, parsed_dc.length, rx_queue);
+                bytes_written -= parsed_dc.length;
+            }
+        }
+
+        self.header_set_num_buffers(used_heads);
+        for _ in 0..used_heads {
+            let parsed_dc = self
+                .parsed_descriptors
+                .pop_front()
+                .expect("This should never happen if write to the buffer succeded.");
+            self.iovec.drop_descriptor_chain(&parsed_dc);
+        }
     }
 
     /// Write the number of descriptors used in VirtIO header
@@ -191,25 +218,28 @@ impl RxBuffers {
     /// This will let the guest know that about all the `DescriptorChain` object that has been
     /// used to receive a frame from the TAP.
     fn finish_frame(&mut self, rx_queue: &mut Queue) {
-        if let Some(used_dc) = self.used_descriptor.take() {
-            // It is fine to `.unrap()` here. The only reason why `add_used` can fail is if the
-            // `head_index` is not a valid descriptor id. `head_index` here is a valid
-            // `DescriptorChain` index. We got it from `queue.pop_or_enable_notification()` which
-            // checks for its validity. In other words, if this unwrap() fails there's a bug in our
-            // emulation logic which, most likely, we can't recover from. So, let's crash here
-            // instead of logging an error and continuing.
-            rx_queue
-                .add_used(used_dc.head_index, used_dc.length)
-                .unwrap();
-        }
+        rx_queue.advance_used_ring(self.used_descriptors);
+        self.used_descriptors = 0;
+        self.used_bytes = 0;
     }
 
     /// Returns the number of bytes that have been used from the buffer
     fn used_bytes(&self) -> u32 {
-        match self.used_descriptor {
-            Some(ref dc) => dc.length,
-            None => 0,
-        }
+        self.used_bytes
+    }
+
+    /// Return a slice of iovecs for the first slice in the buffer.
+    ///
+    /// # Safety
+    /// Buffer needs to have at least one descriptor chain parsed.
+    unsafe fn single_chain_slice_mut(&mut self) -> &mut [iovec] {
+        let nr_iovecs = self.parsed_descriptors[0].nr_iovecs as usize;
+        &mut self.iovec.as_iovec_mut_slice()[..nr_iovecs]
+    }
+
+    /// Return a slice of iovecs for all descriptor chains in the buffer.
+    fn all_chains_slice_mut(&mut self) -> &mut [iovec] {
+        self.iovec.as_iovec_mut_slice()
     }
 }
 
@@ -272,6 +302,7 @@ impl Net {
             | 1 << VIRTIO_NET_F_HOST_TSO6
             | 1 << VIRTIO_NET_F_HOST_UFO
             | 1 << VIRTIO_F_VERSION_1
+            | 1 << VIRTIO_NET_F_MRG_RXBUF
             | 1 << VIRTIO_RING_F_EVENT_IDX;
 
         let mut config_space = ConfigSpace::default();
@@ -433,13 +464,21 @@ impl Net {
     /// Returns the minimum size of buffer we expect the guest to provide us depending on the
     /// features we have negotiated with it
     fn minimum_rx_buffer_size(&self) -> u32 {
-        if self.has_feature(VIRTIO_NET_F_GUEST_TSO4 as u64)
-            || self.has_feature(VIRTIO_NET_F_GUEST_TSO6 as u64)
-            || self.has_feature(VIRTIO_NET_F_GUEST_UFO as u64)
-        {
-            65562
+        if !self.has_feature(VIRTIO_NET_F_MRG_RXBUF as u64) {
+            if self.has_feature(VIRTIO_NET_F_GUEST_TSO4 as u64)
+                || self.has_feature(VIRTIO_NET_F_GUEST_TSO6 as u64)
+                || self.has_feature(VIRTIO_NET_F_GUEST_UFO as u64)
+            {
+                65562
+            } else {
+                1526
+            }
         } else {
-            1526
+            // header is 12 bytes long
+            #[allow(clippy::cast_possible_truncation)]
+            {
+                vnet_hdr_len() as u32
+            }
         }
     }
 
@@ -454,6 +493,9 @@ impl Net {
             if let Err(err) = unsafe { self.rx_buffer.add_buffer(mem, head) } {
                 self.metrics.rx_fails.inc();
                 error!("net: Could not parse an RX descriptor: {err}");
+                // Notify queue about ready frames. We need this
+                // to bring queue into up to date state.
+                self.rx_buffer.finish_frame(queue);
                 // Try to add the bad descriptor to the used ring.
                 if let Err(err) = queue.add_used(index, 0) {
                     error!(
@@ -541,15 +583,14 @@ impl Net {
 
     // We currently prioritize packets from the MMDS over regular network packets.
     fn read_from_mmds_or_tap(&mut self) -> Result<usize, NetError> {
-        // If we don't have any buffers available try to parse more from the RX queue. There might
-        // be some buffers we didn't get the chance to process, because we got to handle the TAP
-        // event before the RX queue event.
-        if self.rx_buffer.is_empty() {
+        // We only want to read from TAP (or mmds) if we have at least 64K of available capacity as
+        // this is the max size of 1 packet.
+        if self.rx_buffer.capacity() < u16::MAX as usize {
             self.parse_rx_descriptors();
 
-            // If after parsing the RX queue we still don't have any buffers stop processing RX
+            // If after parsing the RX queue we still don't have enough capacity, stop processing RX
             // frames.
-            if self.rx_buffer.is_empty() {
+            if self.rx_buffer.capacity() < u16::MAX as usize {
                 return Ok(0);
             }
         }
@@ -570,8 +611,10 @@ impl Net {
                 // * `rx_frame_buf` has size of `MAX_BUFFER_SIZE` and all `DescriptorChain` objects
                 //   are at least that big.
                 unsafe {
-                    self.rx_buffer
-                        .mark_used((vnet_hdr_len() + len).try_into().unwrap());
+                    self.rx_buffer.mark_used(
+                        (vnet_hdr_len() + len).try_into().unwrap(),
+                        &mut self.queues[RX_INDEX],
+                    );
                 }
                 return Ok(vnet_hdr_len() + len);
             }
@@ -586,7 +629,8 @@ impl Net {
         // * `read_tap` passes the first `DescriptorChain` to `readv` so we can't have read more
         //   bytes than its capacity.
         unsafe {
-            self.rx_buffer.mark_used(len.try_into().unwrap());
+            self.rx_buffer
+                .mark_used(len.try_into().unwrap(), &mut self.queues[RX_INDEX]);
         }
         Ok(len)
     }
@@ -630,13 +674,15 @@ impl Net {
     }
 
     fn has_deferred_frame(&self) -> bool {
-        self.rx_buffer.used_descriptor.is_some()
+        self.rx_buffer.used_descriptors != 0
     }
 
     // Process the deferred frame first, then continue reading from tap.
     fn handle_deferred_frame(&mut self) -> Result<(), DeviceError> {
         let used_bytes = self.rx_buffer.used_bytes();
         if self.rate_limited_rx_single_frame(used_bytes as usize) {
+            // Finish with rate limitted packet.
+            self.rx_buffer.finish_frame(&mut self.queues[RX_INDEX]);
             // process_rx() was interrupted possibly before consuming all
             // packets in the tap; try continuing now.
             return self.process_rx();
@@ -797,9 +843,12 @@ impl Net {
     ///
     /// `self.rx_buffer` needs to have at least one descriptor chain parsed
     pub unsafe fn read_tap(&mut self) -> std::io::Result<usize> {
-        let nr_iovecs = self.rx_buffer.parsed_descriptors[0].nr_iovecs as usize;
-        self.tap
-            .read_iovec(&mut self.rx_buffer.iovec.as_iovec_mut_slice()[..nr_iovecs])
+        if self.has_feature(VIRTIO_NET_F_MRG_RXBUF as u64) {
+            self.tap.read_iovec(self.rx_buffer.all_chains_slice_mut())
+        } else {
+            // SAFETY: we only call this if `rx_buffer` is not empty.
+            unsafe { self.tap.read_iovec(self.rx_buffer.single_chain_slice_mut()) }
+        }
     }
 
     fn write_tap(tap: &mut Tap, buf: &IoVecBuffer) -> std::io::Result<usize> {
@@ -1846,11 +1895,8 @@ pub mod tests {
         unsafe { libc::close(th.net.lock().unwrap().tap.as_raw_fd()) };
 
         // The RX queue is empty and there is a deferred frame.
-        th.net().rx_buffer.used_descriptor = Some(ParsedDescriptorChain {
-            head_index: 1,
-            length: 100,
-            nr_iovecs: 1,
-        });
+        th.net().rx_buffer.used_descriptors = 1;
+        th.net().rx_buffer.used_bytes = 100;
         check_metric_after_block!(
             th.net().metrics.no_rx_avail_buffer,
             1,
@@ -1860,7 +1906,8 @@ pub mod tests {
         // We need to set this here to false, otherwise the device will try to
         // handle a deferred frame, it will fail and will never try to read from
         // the tap.
-        th.net().rx_buffer.used_descriptor = None;
+        th.net().rx_buffer.used_descriptors = 0;
+        th.net().rx_buffer.used_bytes = 0;
 
         th.add_desc_chain(NetQueue::Rx, 0, &[(0, 4096, VIRTQ_DESC_F_WRITE)]);
         check_metric_after_block!(

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

@@ -12,7 +12,6 @@ use serde::{Deserialize, Serialize};
 use super::device::{Net, RxBuffers};
 use super::{TapError, NET_NUM_QUEUES, RX_INDEX};
 use crate::devices::virtio::device::DeviceState;
-use crate::devices::virtio::iovec::ParsedDescriptorChain;
 use crate::devices::virtio::persist::{PersistError as VirtioStateError, VirtioDeviceState};
 use crate::devices::virtio::queue::FIRECRACKER_MAX_QUEUE_SIZE;
 use crate::devices::virtio::TYPE_NET;
@@ -37,14 +36,15 @@ pub struct NetConfigSpaceState {
 pub struct RxBufferState {
     // Number of iovecs we have parsed from the guest
     parsed_descriptor_chains_nr: u16,
-    used_descriptor: Option<ParsedDescriptorChain>,
+    // Number of deferred_descriptors
+    deferred_descriptors: u16,
 }
 
 impl RxBufferState {
     fn from_rx_buffers(rx_buffer: &RxBuffers) -> Self {
         RxBufferState {
             parsed_descriptor_chains_nr: rx_buffer.parsed_descriptors.len().try_into().unwrap(),
-            used_descriptor: rx_buffer.used_descriptor.clone(),
+            deferred_descriptors: rx_buffer.used_descriptors,
         }
     }
 }
@@ -162,9 +162,7 @@ impl Persist<'_> for Net {
             // rolling back `next_avail` in the RX queue and call `parse_rx_descriptors`.
             net.queues[RX_INDEX].next_avail -= state.rx_buffers_state.parsed_descriptor_chains_nr;
             net.parse_rx_descriptors();
-            net.rx_buffer
-                .used_descriptor
-                .clone_from(&state.rx_buffers_state.used_descriptor);
+            net.rx_buffer.used_descriptors = state.rx_buffers_state.deferred_descriptors;
         }
 
         Ok(net)