fix: snapshot: disallow snapshots with vhost-user devices

Previously we were incorrectly allow snapshots to be created
even with vhost-user devices by simply ignoring them. Now we
correctly return an error if we detect vhost-user devices.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/vmm/src/device_manager/persist.rs

@@ -37,7 +37,6 @@ use crate::devices::virtio::vsock::persist::{
     VsockConstructorArgs, VsockState, VsockUdsConstructorArgs,
 };
 use crate::devices::virtio::vsock::{Vsock, VsockUnixBackend};
-use crate::logger::warn;
 use crate::mmds::data_store::MmdsVersion;
 use crate::resources::VmResources;
 use crate::snapshot::Persist;
@@ -256,20 +255,17 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                 // Both virtio-block and vhost-user-block share same device type.
                 VirtioDeviceType::Block => {
                     let block = locked_device.as_mut_any().downcast_mut::<Block>().unwrap();
-                    if block.is_vhost_user() {
-                        warn!(
-                            "Skipping vhost-user-block device. VhostUserBlock does not support \
-                             snapshotting yet"
-                        );
-                    } else {
-                        let device_state = block.save();
-                        states.block_devices.push(VirtioDeviceState {
-                            device_id,
-                            device_state,
-                            transport_state,
-                            device_info,
-                        });
-                    }
+                    assert!(
+                        !block.is_vhost_user(),
+                        "vhost-user-block does not support snapshotting yet"
+                    );
+                    let device_state = block.save();
+                    states.block_devices.push(VirtioDeviceState {
+                        device_id,
+                        device_state,
+                        transport_state,
+                        device_info,
+                    });
                 }
                 VirtioDeviceType::Net => {
                     let net = locked_device.as_mut_any().downcast_mut::<Net>().unwrap();

src/vmm/src/lib.rs

@@ -449,6 +449,20 @@ impl Vmm {
         }
     }
 
+    /// Check if the VM can be snapshotted
+    pub fn can_snapshot(&self) -> bool {
+        let mut has_vhost_user_devices = false;
+        self.device_manager
+            .for_each_virtio_device(|device_type, device| {
+                if let VirtioDeviceType::Block = device_type
+                    && let Some(b) = device.as_any().downcast_ref::<Block>()
+                {
+                    has_vhost_user_devices |= b.is_vhost_user();
+                }
+            });
+        !has_vhost_user_devices
+    }
+
     /// Starts the microVM vcpus.
     ///
     /// # Errors
@@ -555,7 +569,12 @@ impl Vmm {
 
     /// Saves the state of a paused Microvm.
     pub fn save_state(&mut self, vm_info: &VmInfo) -> Result<MicrovmState, MicrovmStateError> {
-        use self::MicrovmStateError::SaveVmState;
+        if !self.can_snapshot() {
+            return Err(MicrovmStateError::NotAllowed(
+                "Devices without snapshot support are present".into(),
+            ));
+        }
+
         // We need to save device state before saving KVM state.
         // Some devices, (at the time of writing this comment block device with async engine)
         // might modify the VirtIO transport and send an interrupt to the guest. If we save KVM
@@ -567,13 +586,17 @@ impl Vmm {
         let vm_state = {
             #[cfg(target_arch = "x86_64")]
             {
-                self.vm.save_state().map_err(SaveVmState)?
+                self.vm
+                    .save_state()
+                    .map_err(MicrovmStateError::SaveVmState)?
             }
             #[cfg(target_arch = "aarch64")]
             {
                 let mpidrs = construct_kvm_mpidrs(&vcpu_states);
 
-                self.vm.save_state(&mpidrs).map_err(SaveVmState)?
+                self.vm
+                    .save_state(&mpidrs)
+                    .map_err(MicrovmStateError::SaveVmState)?
             }
         };
 

tests/integration_tests/functional/test_drive_vhost_user.py

@@ -88,7 +88,8 @@ def _check_drives(test_microvm, assert_dict, keys_array):
 def test_vhost_user_block(uvm_vhost_user_booted_ro):
     """
     This test simply tries to boot a VM with
-    vhost-user-block as a root device.
+    vhost-user-block as a root device and then
+    tries to snapshot it.
     """
 
     vm = uvm_vhost_user_booted_ro
@@ -105,6 +106,13 @@ def test_vhost_user_block(uvm_vhost_user_booted_ro):
     _check_drives(vm, assert_dict, assert_dict.keys())
     vhost_user_block_metrics.validate(vm)
 
+    with pytest.raises(
+        RuntimeError, match="Devices without snapshot support are present"
+    ):
+        vm.api.snapshot_create.put(
+            mem_file_path="memfile", snapshot_path="statefile", snapshot_type="Full"
+        )
+
 
 def test_vhost_user_block_read_write(uvm_vhost_user_booted_rw):
     """