fix: Prevent balloon to be used with huge pages via snapshot mod

Insert a check into the snapshot restore routine that checks that a
snapshot file cannot indicate both huge pages being enabled, while also
listing a balloon device. If this condition is detected, refuse to load
the snapshot and print an error suggesting that the snapshot file is
corrupt.

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

src/vmm/src/device_manager/persist.rs

@@ -39,7 +39,7 @@ use crate::devices::virtio::vsock::{
 };
 use crate::devices::virtio::{TYPE_BALLOON, TYPE_BLOCK, TYPE_NET, TYPE_RNG};
 use crate::mmds::data_store::MmdsVersion;
-use crate::resources::VmResources;
+use crate::resources::{ResourcesError, VmResources};
 use crate::snapshot::Persist;
 use crate::vmm_config::mmds::MmdsConfigError;
 use crate::vstate::memory::GuestMemoryMmap;
@@ -69,6 +69,8 @@ pub enum DevicePersistError {
     MmdsConfig(#[from] MmdsConfigError),
     /// Entropy: {0}
     Entropy(#[from] EntropyError),
+    /// Resource misconfiguration: {0}. Is the snapshot file corrupted?
+    ResourcesError(#[from] ResourcesError),
 }
 
 /// Holds the state of a balloon device connected to the MMIO space.
@@ -461,7 +463,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Balloon(device.clone()));
+                .update_from_restored_device(SharedDeviceType::Balloon(device.clone()))?;
 
             restore_helper(
                 device.clone(),
@@ -482,7 +484,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::VirtioBlock(device.clone()));
+                .update_from_restored_device(SharedDeviceType::VirtioBlock(device.clone()))?;
 
             restore_helper(
                 device.clone(),
@@ -527,7 +529,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Network(device.clone()));
+                .update_from_restored_device(SharedDeviceType::Network(device.clone()))?;
 
             restore_helper(
                 device.clone(),
@@ -555,7 +557,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Vsock(device.clone()));
+                .update_from_restored_device(SharedDeviceType::Vsock(device.clone()))?;
 
             restore_helper(
                 device.clone(),
@@ -578,7 +580,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Entropy(device.clone()));
+                .update_from_restored_device(SharedDeviceType::Entropy(device.clone()))?;
 
             restore_helper(
                 device.clone(),

src/vmm/src/resources.rs

@@ -202,7 +202,10 @@ impl VmResources {
 
     /// Updates the resources from a restored device (used for configuring resources when
     /// restoring from a snapshot).
-    pub fn update_from_restored_device(&mut self, device: SharedDeviceType) {
+    pub fn update_from_restored_device(
+        &mut self,
+        device: SharedDeviceType,
+    ) -> Result<(), ResourcesError> {
         match device {
             SharedDeviceType::VirtioBlock(block) => {
                 self.block.add_virtio_device(block);
@@ -214,6 +217,10 @@ impl VmResources {
 
             SharedDeviceType::Balloon(balloon) => {
                 self.balloon.set_device(balloon);
+
+                if self.vm_config.huge_pages != HugePageConfig::None {
+                    return Err(ResourcesError::BalloonDevice(BalloonConfigError::HugePages));
+                }
             }
 
             SharedDeviceType::Vsock(vsock) => {
@@ -223,6 +230,8 @@ impl VmResources {
                 self.entropy.set_device(entropy);
             }
         }
+
+        Ok(())
     }
 
     /// Returns whether dirty page tracking is enabled or not.
@@ -498,6 +507,7 @@ mod tests {
 
     use super::*;
     use crate::cpu_config::templates::{CpuTemplateType, StaticCpuTemplate};
+    use crate::devices::virtio::balloon::Balloon;
     use crate::devices::virtio::block::virtio::VirtioBlockError;
     use crate::devices::virtio::block::{BlockError, CacheType};
     use crate::devices::virtio::vsock::VSOCK_DEV_ID;
@@ -1442,6 +1452,33 @@ mod tests {
             .unwrap_err();
     }
 
+    #[test]
+    fn test_negative_restore_balloon_device_with_huge_pages() {
+        if KernelVersion::get().unwrap() >= KernelVersion::new(4, 16, 0) {
+            let mut vm_resources = default_vm_resources();
+            vm_resources.balloon = BalloonBuilder::new();
+            vm_resources
+                .update_vm_config(&MachineConfigUpdate {
+                    huge_pages: Some(HugePageConfig::Hugetlbfs2M),
+                    ..Default::default()
+                })
+                .unwrap();
+            let err = vm_resources
+                .update_from_restored_device(SharedDeviceType::Balloon(Arc::new(Mutex::new(
+                    Balloon::new(128, false, 0, true).unwrap(),
+                ))))
+                .unwrap_err();
+            assert!(
+                matches!(
+                    err,
+                    ResourcesError::BalloonDevice(BalloonConfigError::HugePages)
+                ),
+                "{:?}",
+                err
+            );
+        }
+    }
+
     #[test]
     fn test_set_entropy_device() {
         let mut vm_resources = default_vm_resources();