refactor: de-duplicate virtio queue validation logic

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

@@ -822,6 +822,7 @@ pub(crate) mod tests {
         // Only initialize the inflate queue to demonstrate invalid request handling.
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);
         balloon.set_queue(INFLATE_INDEX, infq.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, infq.create_queue());
         balloon.activate(mem.clone()).unwrap();
 
         // Fill the second page with non-zero bytes.
@@ -880,6 +881,7 @@ pub(crate) mod tests {
         let mem = default_mem();
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);
         balloon.set_queue(INFLATE_INDEX, infq.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, infq.create_queue());
         balloon.activate(mem.clone()).unwrap();
 
         // Fill the third page with non-zero bytes.
@@ -949,6 +951,7 @@ pub(crate) mod tests {
         let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
         let defq = VirtQueue::new(GuestAddress(0), &mem, 16);
+        balloon.set_queue(INFLATE_INDEX, defq.create_queue());
         balloon.set_queue(DEFLATE_INDEX, defq.create_queue());
         balloon.activate(mem.clone()).unwrap();
 
@@ -997,6 +1000,8 @@ pub(crate) mod tests {
         let mut balloon = Balloon::new(0, true, 1, false).unwrap();
         let mem = default_mem();
         let statsq = VirtQueue::new(GuestAddress(0), &mem, 16);
+        balloon.set_queue(INFLATE_INDEX, statsq.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, statsq.create_queue());
         balloon.set_queue(STATS_INDEX, statsq.create_queue());
         balloon.activate(mem.clone()).unwrap();
 
@@ -1097,6 +1102,9 @@ pub(crate) mod tests {
     fn test_update_stats_interval() {
         let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
+        let q = VirtQueue::new(GuestAddress(0), &mem, 16);
+        balloon.set_queue(INFLATE_INDEX, q.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, q.create_queue());
         balloon.activate(mem).unwrap();
         assert_eq!(
             format!("{:?}", balloon.update_stats_polling_interval(1)),
@@ -1106,6 +1114,10 @@ pub(crate) mod tests {
 
         let mut balloon = Balloon::new(0, true, 1, false).unwrap();
         let mem = default_mem();
+        let q = VirtQueue::new(GuestAddress(0), &mem, 16);
+        balloon.set_queue(INFLATE_INDEX, q.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, q.create_queue());
+        balloon.set_queue(STATS_INDEX, q.create_queue());
         balloon.activate(mem).unwrap();
         assert_eq!(
             format!("{:?}", balloon.update_stats_polling_interval(0)),

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

@@ -146,6 +146,8 @@ pub mod tests {
         let mem = default_mem();
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);
         balloon.set_queue(INFLATE_INDEX, infq.create_queue());
+        balloon.set_queue(DEFLATE_INDEX, infq.create_queue());
+        balloon.set_queue(STATS_INDEX, infq.create_queue());
 
         let balloon = Arc::new(Mutex::new(balloon));
         let _id = event_manager.add_subscriber(balloon.clone());

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

@@ -376,6 +376,7 @@ mod tests {
     use super::*;
     use crate::devices::virtio::block::virtio::device::FileEngineType;
     use crate::devices::virtio::mmio::VIRTIO_MMIO_INT_CONFIG;
+    use crate::devices::virtio::test_utils::VirtQueue;
     use crate::devices::virtio::vhost_user::tests::create_mem;
     use crate::test_utils::create_tmp_socket;
     use crate::vstate::memory::GuestAddress;
@@ -780,6 +781,8 @@ mod tests {
         file.set_len(region_size as u64).unwrap();
         let regions = vec![(GuestAddress(0x0), region_size)];
         let guest_memory = create_mem(file, &regions);
+        let q = VirtQueue::new(GuestAddress(0), &guest_memory, 16);
+        vhost_block.queues[0] = q.create_queue();
 
         // During actiavion of the device features, memory and queues should be set and activated.
         vhost_block.activate(guest_memory).unwrap();

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

@@ -1570,6 +1570,7 @@ mod tests {
 
             let mem = default_mem();
             let vq = VirtQueue::new(GuestAddress(0), &mem, IO_URING_NUM_ENTRIES * 4);
+            block.queues[0] = vq.create_queue();
             block.activate(mem.clone()).unwrap();
 
             // Run scenario that doesn't trigger FullSq BlockError: Add sq_size flush requests.
@@ -1603,6 +1604,7 @@ mod tests {
 
             let mem = default_mem();
             let vq = VirtQueue::new(GuestAddress(0), &mem, IO_URING_NUM_ENTRIES * 4);
+            block.queues[0] = vq.create_queue();
             block.activate(mem.clone()).unwrap();
 
             // Run scenario that triggers FullCqError. Push 2 * IO_URING_NUM_ENTRIES and wait for
@@ -1632,6 +1634,7 @@ mod tests {
 
             let mem = default_mem();
             let vq = VirtQueue::new(GuestAddress(0), &mem, 16);
+            block.queues[0] = vq.create_queue();
             block.activate(mem.clone()).unwrap();
 
             // Add a batch of flush requests.

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

@@ -182,9 +182,9 @@
     }
 
     /// Mark pages used by queues as dirty.
-    fn mark_queue_memory_dirty(&self, mem: &GuestMemoryMmap) -> Result<(), QueueError> {
-        for queue in self.queues() {
-            queue.mark_memory_dirty(mem)?
+    fn mark_queue_memory_dirty(&mut self, mem: &GuestMemoryMmap) -> Result<(), QueueError> {
+        for queue in self.queues_mut() {
+            queue.initialize(mem)?
         }
         Ok(())
     }

src/vmm/src/devices/virtio/mmio.rs

@@ -95,13 +95,6 @@ impl MmioTransport {
         self.device_status & (set | clr) == set
     }
 
-    fn are_queues_valid(&self) -> bool {
-        self.locked_device()
-            .queues()
-            .iter()
-            .all(|q| q.is_valid(&self.mem))
-    }
-
     fn with_queue<U, F>(&self, d: U, f: F) -> U
     where
         F: FnOnce(&Queue) -> U,
@@ -185,7 +178,7 @@ impl MmioTransport {
             DRIVER_OK if self.device_status == (ACKNOWLEDGE | DRIVER | FEATURES_OK) => {
                 self.device_status = status;
                 let device_activated = self.locked_device().is_activated();
-                if !device_activated && self.are_queues_valid() {
+                if !device_activated {
                     // temporary variable needed for borrow checker
                     let activate_result = self.locked_device().activate(self.mem.clone());
                     if let Err(err) = activate_result {
@@ -486,8 +479,6 @@ pub(crate) mod tests {
 
         assert_eq!(d.locked_device().queue_events().len(), 2);
 
-        assert!(!d.are_queues_valid());
-
         d.queue_select = 0;
         assert_eq!(d.with_queue(0, |q| q.max_size), 16);
         assert!(d.with_queue_mut(|q| q.size = 16));
@@ -501,8 +492,6 @@ pub(crate) mod tests {
         d.queue_select = 2;
         assert_eq!(d.with_queue(0, |q| q.max_size), 0);
         assert!(!d.with_queue_mut(|q| q.size = 16));
-
-        assert!(!d.are_queues_valid());
     }
 
     #[test]
@@ -761,7 +750,6 @@ pub(crate) mod tests {
         let m = single_region_mem(0x1000);
         let mut d = MmioTransport::new(m, Arc::new(Mutex::new(DummyDevice::new())), false);
 
-        assert!(!d.are_queues_valid());
         assert!(!d.locked_device().is_activated());
         assert_eq!(d.device_status, device_status::INIT);
 
@@ -800,7 +788,6 @@ pub(crate) mod tests {
             write_le_u32(&mut buf[..], 1);
             d.bus_write(0x44, &buf[..]);
         }
-        assert!(d.are_queues_valid());
         assert!(!d.locked_device().is_activated());
 
         // Device should be ready for activation now.
@@ -860,7 +847,6 @@ pub(crate) mod tests {
             write_le_u32(&mut buf[..], 1);
             d.bus_write(0x44, &buf[..]);
         }
-        assert!(d.are_queues_valid());
         assert_eq!(
             d.locked_device().interrupt_status().load(Ordering::SeqCst),
             0
@@ -910,7 +896,6 @@ pub(crate) mod tests {
             write_le_u32(&mut buf[..], 1);
             d.bus_write(0x44, &buf[..]);
         }
-        assert!(d.are_queues_valid());
         assert!(!d.locked_device().is_activated());
 
         // Device should be ready for activation now.
@@ -937,7 +922,6 @@ pub(crate) mod tests {
         let mut d = MmioTransport::new(m, Arc::new(Mutex::new(DummyDevice::new())), false);
         let mut buf = [0; 4];
 
-        assert!(!d.are_queues_valid());
         assert!(!d.locked_device().is_activated());
         assert_eq!(d.device_status, 0);
         activate_device(&mut d);

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

@@ -184,14 +184,7 @@ impl VirtioDeviceState {
 
         for q in &queues {
             // Sanity check queue size and queue max size.
-            if q.max_size != expected_queue_max_size || q.size > expected_queue_max_size {
-                return Err(PersistError::InvalidInput);
-            }
-            // Snapshot can happen at any time, including during device configuration/activation
-            // when fields are only partially configured.
-            //
-            // Only if the device was activated, check `q.is_valid()`.
-            if self.activated && !q.is_valid(mem) {
+            if q.max_size != expected_queue_max_size {
                 return Err(PersistError::InvalidInput);
             }
         }
@@ -333,6 +326,7 @@ mod tests {
             ..Default::default()
         };
         state.queues = vec![bad_q];
+        state.activated = true;
         state
             .build_queues_checked(&mem, 0, state.queues.len(), max_size)
             .unwrap_err();
@@ -351,6 +345,8 @@ mod tests {
         let mem = default_mem();
 
         let mut queue = Queue::new(128);
+        queue.ready = true;
+        queue.size = queue.max_size;
         queue.initialize(&mem).unwrap();
 
         let mut bytes = vec![0; 4096];