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};
 use std::net::Ipv4Addr;
+use std::num::{NonZeroU32, 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 deferred_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()),
-            deferred_descriptor: None,
+            used_descriptors: 0,
+            used_bytes: 0,
         })
     }
 
@@ -141,37 +144,56 @@ 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) -> ParsedDescriptorChain {
-        // 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",
-        );
+    unsafe fn mark_used(&mut self, mut bytes_written: u32, rx_queue: &mut Queue) {
+        self.used_bytes = bytes_written;
+
+        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;
+        };
 
-        self.header_set_num_buffers(1);
-        self.iovec.drop_descriptor_chain(&parsed_dc);
-        parsed_dc.length = size;
-        parsed_dc
+        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
@@ -190,14 +212,24 @@ 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, dc: &ParsedDescriptorChain, rx_queue: &mut Queue) {
-        // 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(dc.head_index, dc.length).unwrap();
+    fn finish_frame(&mut self, rx_queue: &mut Queue) {
+        rx_queue.advance_used_ring(self.used_descriptors);
+        self.used_descriptors = 0;
+        self.used_bytes = 0;
+    }
+
+    /// 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()
     }
 }
 
@@ -260,6 +292,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();
@@ -407,27 +440,35 @@ impl Net {
     // Attempts to copy a single frame into the guest if there is enough
     // rate limiting budget.
     // Returns true on successful frame delivery.
-    fn rate_limited_rx_single_frame(&mut self, dc: &ParsedDescriptorChain) -> bool {
+    fn rate_limited_rx_single_frame(&mut self, frame_size: u32) -> bool {
         let rx_queue = &mut self.queues[RX_INDEX];
-        if !Self::rate_limiter_consume_op(&mut self.rx_rate_limiter, dc.length as u64) {
+        if !Self::rate_limiter_consume_op(&mut self.rx_rate_limiter, frame_size as u64) {
             self.metrics.rx_rate_limiter_throttled.inc();
             return false;
         }
 
-        self.rx_buffer.finish_frame(dc, rx_queue);
+        self.rx_buffer.finish_frame(rx_queue);
         true
     }
 
     /// 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
+            }
         }
     }
 
@@ -442,6 +483,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!(
@@ -528,16 +572,15 @@ impl Net {
     }
 
     // We currently prioritize packets from the MMDS over regular network packets.
-    fn read_from_mmds_or_tap(&mut self) -> Result<Option<ParsedDescriptorChain>, 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() {
+    fn read_from_mmds_or_tap(&mut self) -> Result<Option<NonZeroU32>, NetError> {
+        // 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(None);
             }
         }
@@ -553,29 +596,44 @@ impl Net {
                 self.rx_buffer
                     .iovec
                     .write_all_volatile_at(&self.rx_frame_buf[..vnet_hdr_len() + len], 0)?;
+
+                // SAFETY:
+                // * len is never 0
+                // * Maximum size of the packet is 65562 which fits into u32
+                #[allow(clippy::cast_possible_truncation)]
+                let len = unsafe { NonZeroU32::new_unchecked((vnet_hdr_len() + len) as u32) };
+
                 // SAFETY: This is safe:
                 // * We checked that `rx_buffer` includes at least one `DescriptorChain`
                 // * `rx_frame_buf` has size of `MAX_BUFFER_SIZE` and all `DescriptorChain` objects
                 //   are at least that big.
-                let dc = unsafe {
+                unsafe {
                     self.rx_buffer
-                        .mark_used((vnet_hdr_len() + len).try_into().unwrap())
-                };
-
-                return Ok(Some(dc));
+                        .mark_used(len.get(), &mut self.queues[RX_INDEX]);
+                }
+                return Ok(Some(len));
             }
         }
 
         // SAFETY: this is safe because we ensured that `self.rx_buffer` has at least one
         // DescriptorChain parsed in it.
         let len = unsafe { self.read_tap().map_err(NetError::IO) }?;
 
+        // SAFETY:
+        // * len is never 0
+        // * Maximum size of the packet is 65562 which fits into u32
+        #[allow(clippy::cast_possible_truncation)]
+        let len = unsafe { NonZeroU32::new_unchecked(len as u32) };
+
         // SAFETY: This is safe,
         // * `rx_buffer` has at least one `DescriptorChain`
         // * `read_tap` passes the first `DescriptorChain` to `readv` so we can't have read more
         //   bytes than its capacity.
-        let dc = unsafe { self.rx_buffer.mark_used(len.try_into().unwrap()) };
-        Ok(Some(dc))
+        unsafe {
+            self.rx_buffer
+                .mark_used(len.get(), &mut self.queues[RX_INDEX]);
+        }
+        Ok(Some(len))
     }
 
     /// Read as many frames as possible.
@@ -586,12 +644,11 @@ impl Net {
                     self.metrics.no_rx_avail_buffer.inc();
                     break;
                 }
-                Ok(Some(dc)) => {
+                Ok(Some(bytes)) => {
                     self.metrics.rx_count.inc();
-                    self.metrics.rx_bytes_count.add(dc.length as u64);
+                    self.metrics.rx_bytes_count.add(bytes.get() as u64);
                     self.metrics.rx_packets_count.inc();
-                    if !self.rate_limited_rx_single_frame(&dc) {
-                        self.rx_buffer.deferred_descriptor = Some(dc);
+                    if !self.rate_limited_rx_single_frame(bytes.get()) {
                         break;
                     }
                 }
@@ -619,11 +676,10 @@ impl Net {
 
     fn resume_rx(&mut self) -> Result<(), DeviceError> {
         // First try to handle any deferred frame
-        if let Some(deferred_descriptor) = self.rx_buffer.deferred_descriptor.take() {
+        if self.rx_buffer.used_bytes != 0 {
             // If can't finish sending this frame, re-set it as deferred and return; we can't
             // process any more frames from the TAP.
-            if !self.rate_limited_rx_single_frame(&deferred_descriptor) {
-                self.rx_buffer.deferred_descriptor = Some(deferred_descriptor);
+            if !self.rate_limited_rx_single_frame(self.rx_buffer.used_bytes) {
                 return Ok(());
             }
         }
@@ -689,7 +745,7 @@ impl Net {
                 &self.metrics,
             )
             .unwrap_or(false);
-            if frame_consumed_by_mmds && self.rx_buffer.deferred_descriptor.is_none() {
+            if frame_consumed_by_mmds && self.rx_buffer.used_bytes == 0 {
                 // MMDS consumed this frame/request, let's also try to process the response.
                 process_rx_for_mmds = true;
             }
@@ -774,9 +830,13 @@ 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])
+        let slice = if self.has_feature(VIRTIO_NET_F_MRG_RXBUF as u64) {
+            self.rx_buffer.all_chains_slice_mut()
+        } else {
+            // SAFETY: we only call this if `rx_buffer` is not empty.
+            unsafe { self.rx_buffer.single_chain_slice_mut() }
+        };
+        self.tap.read_iovec(slice)
     }
 
     fn write_tap(tap: &mut Tap, buf: &IoVecBuffer) -> std::io::Result<usize> {
@@ -1815,11 +1875,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.deferred_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,
@@ -1829,7 +1886,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.deferred_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!(