Minor fixes

src/vmm/benches/block_request.rs

@@ -13,6 +13,7 @@ use vmm::devices::virtio::block::virtio::{Request, RequestHeader, VIRTIO_BLK_T_I
 use vmm::devices::virtio::test_utils::VirtQueue;
 use vmm::test_utils::single_region_mem;
 
+#[allow(const_item_mutation)]
 pub fn block_request_benchmark(c: &mut Criterion) {
     let mem = single_region_mem(65562);
     let virt_queue = VirtQueue::new(GuestAddress(0), &mem, 16);
@@ -24,7 +25,7 @@ pub fn block_request_benchmark(c: &mut Criterion) {
     chain.set_header(request_header);
 
     let mut queue = virt_queue.create_queue();
-    let desc = queue.pop().unwrap().unwrap();
+    let desc = queue.pop(&mut u16::MAX).unwrap().unwrap();
 
     c.bench_function("request_parse", |b| {
         b.iter(|| {

src/vmm/benches/queue.rs

@@ -56,14 +56,15 @@ fn set_dtable_many_chains(rxq: &VirtQueue, n: usize) {
     rxq.avail.idx.set(n as u16);
 }
 
+#[allow(const_item_mutation)]
 pub fn queue_benchmark(c: &mut Criterion) {
     let mem = single_region_mem(65562);
     let rxq = VirtQueue::new(GuestAddress(0), &mem, 256);
     let mut queue = rxq.create_queue();
 
     set_dtable_one_chain(&rxq, 16);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop().unwrap().unwrap();
+    let desc = queue.pop(&mut u16::MAX).unwrap().unwrap();
     c.bench_function("next_descriptor_16", |b| {
         b.iter(|| {
             let mut head = Some(desc);
@@ -77,7 +78,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
     c.bench_function("queue_pop_16", |b| {
         b.iter(|| {
             queue.next_avail = Wrapping(0);
-            while let Some(desc) = queue.pop().unwrap() {
+            while let Some(desc) = queue.pop(&mut u16::MAX).unwrap() {
                 std::hint::black_box(desc);
             }
         })

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

@@ -384,6 +384,7 @@ impl Balloon {
         METRICS.inflate_count.inc();
 
         let queue = &mut self.queues[INFLATE_INDEX];
+        let mut queue_bound = queue.size;
         // The pfn buffer index used during descriptor processing.
         let mut pfn_buffer_idx = 0;
         let mut needs_interrupt = false;
@@ -395,7 +396,7 @@ impl Balloon {
             // Internal loop processes descriptors and acummulates the pfns in `pfn_buffer`.
             // Breaks out when there is not enough space in `pfn_buffer` to completely process
             // the next descriptor.
-            while let Some(head) = queue.pop()? {
+            while let Some(head) = queue.pop(&mut queue_bound)? {
                 let len = head.len as usize;
                 let max_len = MAX_PAGES_IN_DESC * SIZE_OF_U32;
                 valid_descs_found = true;
@@ -471,9 +472,10 @@ impl Balloon {
         METRICS.deflate_count.inc();
 
         let queue = &mut self.queues[DEFLATE_INDEX];
+        let mut queue_bound = queue.size;
         let mut needs_interrupt = false;
 
-        while let Some(head) = queue.pop()? {
+        while let Some(head) = queue.pop(&mut queue_bound)? {
             queue.add_used(head.index, 0)?;
             needs_interrupt = true;
         }
@@ -490,8 +492,9 @@ impl Balloon {
         // This is safe since we checked in the event handler that the device is activated.
         let mem = &self.device_state.active_state().unwrap().mem;
         METRICS.stats_updates_count.inc();
+        let mut queue_bound = self.queues[STATS_INDEX].size;
 
-        while let Some(head) = self.queues[STATS_INDEX].pop()? {
+        while let Some(head) = self.queues[STATS_INDEX].pop(&mut queue_bound)? {
             if let Some(prev_stats_desc) = self.stats_desc_index {
                 // We shouldn't ever have an extra buffer if the driver follows
                 // the protocol, but return it if we find one.
@@ -544,12 +547,14 @@ impl Balloon {
             .mem;
 
         let idx = self.free_page_hinting_idx();
-        let queue = &mut self.queues[idx];
         let host_cmd = self.hinting_state.host_cmd;
         let mut needs_interrupt = false;
         let mut complete = false;
 
-        while let Some(head) = queue.pop()? {
+        let queue = &mut self.queues[idx];
+        let mut queue_bound = queue.size;
+
+        while let Some(head) = queue.pop(&mut queue_bound)? {
             let head_index = head.index;
 
             let mut last_desc = Some(head);
@@ -626,10 +631,12 @@ impl Balloon {
             .mem;
 
         let idx = self.free_page_reporting_idx();
-        let queue = &mut self.queues[idx];
         let mut needs_interrupt = false;
 
-        while let Some(head) = queue.pop()? {
+        let queue = &mut self.queues[idx];
+        let mut queue_bound = queue.size;
+
+        while let Some(head) = queue.pop(&mut queue_bound)? {
             let head_index = head.index;
 
             let mut last_desc = Some(head);

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

@@ -391,9 +391,10 @@ impl VirtioBlock {
         let active_state = self.device_state.active_state().unwrap();
 
         let queue = &mut self.queues[queue_index];
+        let mut queue_bound = queue.size;
         let mut used_any = false;
 
-        while let Some(head) = queue.pop_or_enable_notification()? {
+        while let Some(head) = queue.pop_or_enable_notification(&mut queue_bound)? {
             self.metrics.remaining_reqs_count.add(queue.len().into());
             let processing_result =
                 match Request::parse(&head, &active_state.mem, self.disk.nsectors) {

src/vmm/src/devices/virtio/block/virtio/request.rs

@@ -477,21 +477,31 @@ mod tests {
     }
 
     impl RequestDescriptorChain<'_, '_> {
+        #[allow(const_item_mutation)]
         fn check_parse_err(&self, _e: VirtioBlockError) {
             let mut q = self.driver_queue.create_queue();
             let memory = self.driver_queue.memory();
 
             assert!(matches!(
-                Request::parse(&q.pop().unwrap().unwrap(), memory, NUM_DISK_SECTORS),
+                Request::parse(
+                    &q.pop(&mut u16::MAX).unwrap().unwrap(),
+                    memory,
+                    NUM_DISK_SECTORS
+                ),
                 Err(_e)
             ));
         }
 
+        #[allow(const_item_mutation)]
         fn check_parse(&self, check_data: bool) {
             let mut q = self.driver_queue.create_queue();
             let memory = self.driver_queue.memory();
-            let request =
-                Request::parse(&q.pop().unwrap().unwrap(), memory, NUM_DISK_SECTORS).unwrap();
+            let request = Request::parse(
+                &q.pop(&mut u16::MAX).unwrap().unwrap(),
+                memory,
+                NUM_DISK_SECTORS,
+            )
+            .unwrap();
             let expected_header = self.header();
 
             assert_eq!(
@@ -955,10 +965,11 @@ mod tests {
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn parse_random_requests() {
         let cfg = ProptestConfig::with_cases(1000);
         proptest!(cfg, |(mut request in random_request_parse())| {
-            let result = Request::parse(&request.2.pop().unwrap().unwrap(), &request.1, NUM_DISK_SECTORS);
+            let result = Request::parse(&request.2.pop(&mut u16::MAX).unwrap().unwrap(), &request.1, NUM_DISK_SECTORS);
             match result {
                 Ok(r) => prop_assert!(r == request.0.unwrap()),
                 Err(err) => {

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

@@ -612,45 +612,48 @@ mod tests {
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn test_access_mode() {
         let mem = default_mem();
         let (mut q, _) = read_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
         // SAFETY: This descriptor chain is only loaded into one buffer
         unsafe { IoVecBuffer::from_descriptor_chain(&mem, head).unwrap() };
 
         let (mut q, _) = write_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
         // SAFETY: This descriptor chain is only loaded into one buffer
         unsafe { IoVecBuffer::from_descriptor_chain(&mem, head).unwrap_err() };
 
         let (mut q, _) = read_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
         // SAFETY: This descriptor chain is only loaded into one buffer
         unsafe { IoVecBufferMutDefault::from_descriptor_chain(&mem, head).unwrap_err() };
 
         let (mut q, _) = write_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
         // SAFETY: This descriptor chain is only loaded into one buffer
         unsafe { IoVecBufferMutDefault::from_descriptor_chain(&mem, head).unwrap() };
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn test_iovec_length() {
         let mem = default_mem();
         let (mut q, _) = read_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
         let iovec = unsafe { IoVecBuffer::from_descriptor_chain(&mem, head).unwrap() };
         assert_eq!(iovec.len(), 4 * 64);
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn test_iovec_mut_length() {
         let mem = default_mem();
         let (mut q, _) = write_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
         let mut iovec =
@@ -662,18 +665,19 @@ mod tests {
         // (concpetually) associated with a single `Queue`. We just do this here to be able to test
         // the appending logic.
         let (mut q, _) = write_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
         // SAFETY: it is actually unsafe, but we just want to check the length of the
         // `IoVecBufferMut` after appending.
         let _ = unsafe { iovec.append_descriptor_chain(&mem, head).unwrap() };
         assert_eq!(iovec.len(), 8 * 64);
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn test_iovec_read_at() {
         let mem = default_mem();
         let (mut q, _) = read_only_chain(&mem);
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
         let iovec = unsafe { IoVecBuffer::from_descriptor_chain(&mem, head).unwrap() };
@@ -723,12 +727,13 @@ mod tests {
     }
 
     #[test]
+    #[allow(const_item_mutation)]
     fn test_iovec_mut_write_at() {
         let mem = default_mem();
         let (mut q, vq) = write_only_chain(&mem);
 
         // This is a descriptor chain with 4 elements 64 bytes long each.
-        let head = q.pop().unwrap().unwrap();
+        let head = q.pop(&mut u16::MAX).unwrap().unwrap();
 
         // SAFETY: This descriptor chain is only loaded into one buffer
         let mut iovec =

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

@@ -446,7 +446,8 @@ impl VirtioMem {
     }
 
     fn process_mem_queue(&mut self) -> Result<(), VirtioMemError> {
-        while let Some(desc) = self.queues[MEM_QUEUE].pop()? {
+        let mut queue_bound = self.queues[MEM_QUEUE].size;
+        while let Some(desc) = self.queues[MEM_QUEUE].pop(&mut queue_bound)? {
             let index = desc.index;
 
             let (req, resp_addr, used_idx) = self.parse_request(&desc)?;
@@ -539,6 +540,10 @@ impl VirtioMem {
         self.config.plugged_size -= usize_to_u64(self.nb_blocks_to_len(plugged_before));
         self.config.plugged_size += usize_to_u64(self.nb_blocks_to_len(plugged_after));
 
+        // Update kvm slots before doing any discards to prevent guest from re-faulting just
+        // discarded memory.
+        self.update_kvm_slots(range)?;
+
         // If unplugging, discard the range
         if !plug {
             self.guest_memory()
@@ -550,8 +555,7 @@ impl VirtioMem {
                     error!("virtio-mem: Failed to discard memory range: {}", err);
                 });
         }
-
-        self.update_kvm_slots(range)
+        Ok(())
     }
 
     /// Updates the requested size of the virtio-mem device.

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

@@ -486,7 +486,8 @@ impl Net {
         // This is safe since we checked in the event handler that the device is activated.
         let mem = &self.device_state.active_state().unwrap().mem;
         let queue = &mut self.queues[RX_INDEX];
-        while let Some(head) = queue.pop_or_enable_notification()? {
+        let mut queue_bound = queue.size;
+        while let Some(head) = queue.pop_or_enable_notification(&mut queue_bound)? {
             let index = head.index;
             // SAFETY: we are only using this `DescriptorChain` here.
             if let Err(err) = unsafe { self.rx_buffer.add_buffer(mem, head) } {
@@ -540,17 +541,14 @@ impl Net {
         //
         // We do not plan to fix this for a few reasons:
         //
-        // 1. Without MMDS enabled, packets with destination IP 169.254.169.254
-        //    will be forwarded to the TAP without filtering. Operators should
-        //    not rely on MMDS for IMDS access control.
-        // 2. Guest originated traffic is treated as untrusted and Firecracker
-        //    does not filter IPv4 packets. Operators deploying Firecracker
-        //    based services should implement host-level firewall rules to
-        //    restrict guest egress traffic.
-        // 3. Preventing this TOCTOU by copying packets to to a host buffer
-        //    before routing decisions would significantly reduce guest-to-host
-        //    TCP throughput, which is not justifiable given the mitigations
-        //    available at host-level.
+        // 1. Without MMDS enabled, packets with destination IP 169.254.169.254 will be forwarded to
+        //    the TAP without filtering. Operators should not rely on MMDS for IMDS access control.
+        // 2. Guest originated traffic is treated as untrusted and Firecracker does not filter IPv4
+        //    packets. Operators deploying Firecracker based services should implement host-level
+        //    firewall rules to restrict guest egress traffic.
+        // 3. Preventing this TOCTOU by copying packets to to a host buffer before routing decisions
+        //    would significantly reduce guest-to-host TCP throughput, which is not justifiable
+        //    given the mitigations available at host-level.
 
         // Read the frame headers from the IoVecBuffer
         let max_header_len = headers.len();
@@ -737,36 +735,35 @@ impl Net {
         // with the MMDS network stack.
         let mut process_rx_for_mmds = false;
         let mut used_any = false;
-        let tx_queue = &mut self.queues[TX_INDEX];
+        let queue = &mut self.queues[TX_INDEX];
+        let mut queue_bound = queue.size;
 
-        while let Some(head) = tx_queue.pop_or_enable_notification()? {
-            self.metrics
-                .tx_remaining_reqs_count
-                .add(tx_queue.len().into());
+        while let Some(head) = queue.pop_or_enable_notification(&mut queue_bound)? {
+            self.metrics.tx_remaining_reqs_count.add(queue.len().into());
             let head_index = head.index;
             // Parse IoVecBuffer from descriptor head
             // 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(mem, head).is_err() } {
                 self.metrics.tx_fails.inc();
-                tx_queue.add_used(head_index, 0)?;
+                queue.add_used(head_index, 0)?;
                 continue;
             };
 
             // We only handle frames that are up to 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.add_used(head_index, 0)?;
+                queue.add_used(head_index, 0)?;
                 continue;
             }
 
             if !Self::rate_limiter_consume_op(
                 &mut self.tx_rate_limiter,
                 u64::from(self.tx_buffer.len()),
             ) {
-                tx_queue.undo_pop();
+                queue.undo_pop();
                 self.metrics.tx_rate_limiter_throttled.inc();
                 break;
             }
@@ -786,7 +783,7 @@ impl Net {
                 process_rx_for_mmds = true;
             }
 
-            tx_queue.add_used(head_index, 0)?;
+            queue.add_used(head_index, 0)?;
             used_any = true;
         }