chore: add unit tests for virtio-net with VIRTIO_NET_F_MRG_RXBUF

Add new unit tests for virtio-net with enabled VIRTIO_NET_F_MRG_RXBUF
feature.

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

Author: ShadowCurse

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

@@ -1048,8 +1048,8 @@ pub mod tests {
     };
     use crate::devices::virtio::net::test_utils::test::TestHelper;
     use crate::devices::virtio::net::test_utils::{
-        default_net, if_index, inject_tap_tx_frame, set_mac, NetEvent, NetQueue, ReadTapMock,
-        TapTrafficSimulator, WriteTapMock,
+        default_net, if_index, inject_tap_tx_frame, mock_frame_set_num_buffers, set_mac, NetEvent,
+        NetQueue, ReadTapMock, TapTrafficSimulator, WriteTapMock,
     };
     use crate::devices::virtio::net::NET_QUEUE_SIZES;
     use crate::devices::virtio::queue::VIRTQ_DESC_F_WRITE;
@@ -1290,7 +1290,40 @@ pub mod tests {
     }
 
     #[test]
-    fn test_rx_partial_write() {
+    fn test_rx_mrg_buf_multiple_short_writable_descriptors() {
+        let mut th = TestHelper::get_default();
+        th.activate_net();
+        th.net().tap.mocks.set_read_tap(ReadTapMock::TapFrame);
+
+        // VIRTIO_NET_F_MRG_RXBUF is not enabled by default
+        th.net().acked_features = 1 << VIRTIO_NET_F_MRG_RXBUF;
+
+        th.add_desc_chain(NetQueue::Rx, 0, &[(0, 500, VIRTQ_DESC_F_WRITE)]);
+        th.add_desc_chain(NetQueue::Rx, 500, &[(1, 500, VIRTQ_DESC_F_WRITE)]);
+
+        // There will be 2 heads used.
+        let mut frame = inject_tap_tx_frame(&th.net(), 1000);
+        mock_frame_set_num_buffers(&mut frame, 2);
+
+        check_metric_after_block!(
+            th.net().metrics.rx_packets_count,
+            1,
+            th.event_manager.run_with_timeout(100).unwrap()
+        );
+
+        assert_eq!(th.rxq.used.idx.get(), 2);
+        assert!(th.net().irq_trigger.has_pending_irq(IrqType::Vring));
+        assert!(!th.net().rx_deferred_frame);
+
+        th.rxq.check_used_elem(0, 0, 500);
+        th.rxq.check_used_elem(1, 1, 500);
+
+        th.rxq.dtable[0].check_data(&frame[..500]);
+        th.rxq.dtable[1].check_data(&frame[500..]);
+    }
+
+    #[test]
+    fn test_rx_invalid_desc_chain() {
         let mut th = TestHelper::get_default();
         th.activate_net();
 
@@ -1312,6 +1345,79 @@ pub mod tests {
         th.check_rx_queue_resume(&frame);
     }
 
+    #[test]
+    fn test_rx_mrg_buf_invalid_desc_chain() {
+        let mut th = TestHelper::get_default();
+        th.activate_net();
+
+        // VIRTIO_NET_F_MRG_RXBUF is not enabled by default
+        th.net().acked_features = 1 << VIRTIO_NET_F_MRG_RXBUF;
+
+        // The descriptor chain is created so that the last descriptor doesn't fit in the
+        // guest memory.
+        let offset = th.mem.last_addr().raw_value() - th.data_addr() - 300;
+        th.add_desc_chain(
+            NetQueue::Rx,
+            offset,
+            &[
+                (0, 100, VIRTQ_DESC_F_WRITE),
+                (1, 50, VIRTQ_DESC_F_WRITE),
+                (2, 4096, VIRTQ_DESC_F_WRITE),
+            ],
+        );
+        let frame = th.check_rx_deferred_frame(1000);
+        th.rxq.check_used_elem(0, 0, 0);
+
+        th.check_rx_queue_resume(&frame);
+    }
+
+    #[test]
+    fn test_rx_mrg_buf_partial_write() {
+        let mut th = TestHelper::get_default();
+        th.activate_net();
+        th.net().tap.mocks.set_read_tap(ReadTapMock::TapFrame);
+
+        // VIRTIO_NET_F_MRG_RXBUF is not enabled by default
+        th.net().acked_features = 1 << VIRTIO_NET_F_MRG_RXBUF;
+
+        // Add descriptor that is not big enough to store the
+        // whole packet.
+        th.add_desc_chain(NetQueue::Rx, 0, &[(0, 500, VIRTQ_DESC_F_WRITE)]);
+
+        // There will be 2 heads used.
+        let mut frame = inject_tap_tx_frame(&th.net(), 1000);
+        mock_frame_set_num_buffers(&mut frame, 2);
+
+        // For now only 1 descriptor chain is used,
+        // but the packet is not fully written yet.
+        check_metric_after_block!(
+            th.net().metrics.rx_packets_count,
+            0,
+            th.event_manager.run_with_timeout(100).unwrap()
+        );
+        th.rxq.check_used_elem(0, 0, 500);
+
+        // The write was converted to partial write
+        assert!(th.net().rx_partial_write.is_some());
+
+        // Continuing writing.
+        th.add_desc_chain(NetQueue::Rx, 500, &[(1, 500, VIRTQ_DESC_F_WRITE)]);
+        check_metric_after_block!(
+            th.net().metrics.rx_packets_count,
+            1,
+            th.event_manager.run_with_timeout(100).unwrap()
+        );
+        assert!(th.net().rx_partial_write.is_none());
+        assert_eq!(th.rxq.used.idx.get(), 2);
+        assert!(th.net().irq_trigger.has_pending_irq(IrqType::Vring));
+
+        th.rxq.check_used_elem(0, 0, 500);
+        th.rxq.check_used_elem(1, 1, 500);
+
+        th.rxq.dtable[0].check_data(&frame[..500]);
+        th.rxq.dtable[1].check_data(&frame[500..]);
+    }
+
     #[test]
     fn test_rx_retry() {
         let mut th = TestHelper::get_default();
@@ -1363,6 +1469,62 @@ pub mod tests {
         th.rxq.dtable[5].check_data(&frame);
     }
 
+    #[test]
+    fn test_rx_mrg_buf_retry() {
+        let mut th = TestHelper::get_default();
+        th.activate_net();
+        th.net().tap.mocks.set_read_tap(ReadTapMock::TapFrame);
+
+        // VIRTIO_NET_F_MRG_RXBUF is not enabled by default
+        th.net().acked_features = 1 << VIRTIO_NET_F_MRG_RXBUF;
+
+        // Add invalid descriptor chain - read only descriptor.
+        th.add_desc_chain(
+            NetQueue::Rx,
+            0,
+            &[
+                (0, 100, VIRTQ_DESC_F_WRITE),
+                (1, 100, 0),
+                (2, 1000, VIRTQ_DESC_F_WRITE),
+            ],
+        );
+        // Add valid descriptor chain but too short. This one will be used
+        // and the write will be converted to partial write.
+        th.add_desc_chain(NetQueue::Rx, 1200, &[(3, 100, VIRTQ_DESC_F_WRITE)]);
+        // Add invalid descriptor chain - invalid memory offset.
+        // The partial write stated with previous descriptor should halt here.
+        th.add_desc_chain(
+            NetQueue::Rx,
+            th.mem.last_addr().raw_value(),
+            &[(4, 1000, VIRTQ_DESC_F_WRITE)],
+        );
+
+        // Add valid descriptor chain.
+        th.add_desc_chain(NetQueue::Rx, 1300, &[(5, 1000, VIRTQ_DESC_F_WRITE)]);
+
+        // Inject frame to tap and run epoll.
+        let frame = inject_tap_tx_frame(&th.net(), 1000);
+        check_metric_after_block!(
+            th.net().metrics.rx_packets_count,
+            1,
+            th.event_manager.run_with_timeout(100).unwrap()
+        );
+
+        // Check that the used queue has advanced.
+        assert_eq!(th.rxq.used.idx.get(), 4);
+        assert!(&th.net().irq_trigger.has_pending_irq(IrqType::Vring));
+        // Check that the invalid descriptor chains have been discarded
+        th.rxq.check_used_elem(0, 0, 0);
+        th.rxq.check_used_elem(1, 3, 0);
+        th.rxq.check_used_elem(2, 4, 0);
+        // Check that the frame wasn't deferred.
+        assert!(!th.net().rx_deferred_frame);
+        // Check that the frame has been written successfully to the valid Rx descriptor chain.
+        th.rxq
+            .check_used_elem(3, 5, frame.len().try_into().unwrap());
+        th.rxq.dtable[5].check_data(&frame);
+    }
+
     #[test]
     fn test_rx_complex_desc_chain() {
         let mut th = TestHelper::get_default();

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

@@ -93,6 +93,16 @@ pub fn mock_frame(len: usize) -> Vec<u8> {
     mock_frame
 }
 
+pub fn mock_frame_set_num_buffers(frame: &mut [u8], num_buffers: u16) {
+    assert!(std::mem::size_of::<virtio_net_hdr_v1>() <= frame.len());
+    // SAFETY:
+    // Frame is bigger than the header.
+    unsafe {
+        let hdr = &mut *frame.as_mut_ptr().cast::<virtio_net_hdr_v1>();
+        hdr.num_buffers = num_buffers;
+    }
+}
+
 #[derive(Debug)]
 pub enum ReadTapMock {
     Failure,