From a5de49558e9cee7a947c1b34b7af57bca0914309 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Mon, 27 Apr 2020 16:55:52 +0200
Subject: [PATCH] vmm: Only allow removal of specific types of virtio device

Now that all virtio devices are assigned with identifiers, they could
all be removed from the VM. This is not something that we want to allow
because it does not make sense for some devices. That's why based on the
device type, we remove the device or we return an error to the user.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
Signed-off-by: Rob Bradford <robert.bradford@intel.com>
---
 vm-virtio/src/lib.rs      |  4 ++--
 vmm/src/device_manager.rs | 42 +++++++++++++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/vm-virtio/src/lib.rs b/vm-virtio/src/lib.rs
index 68958c3c3..46f17a7ae 100755
--- a/vm-virtio/src/lib.rs
+++ b/vm-virtio/src/lib.rs
@@ -68,11 +68,11 @@ const VIRTIO_F_IOMMU_PLATFORM: u32 = 33;
 const VIRTIO_F_IN_ORDER: u32 = 35;
 
 // Types taken from linux/virtio_ids.h
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 #[allow(dead_code)]
 #[allow(non_camel_case_types)]
 #[repr(C)]
-enum VirtioDeviceType {
+pub enum VirtioDeviceType {
     TYPE_NET = 1,
     TYPE_BLOCK = 2,
     TYPE_CONSOLE = 3,
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index b83b8dd3f..96efc12eb 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -63,7 +63,7 @@ use vm_virtio::transport::VirtioPciDevice;
 use vm_virtio::transport::VirtioTransport;
 use vm_virtio::vhost_user::VhostUserConfig;
 #[cfg(feature = "pci_support")]
-use vm_virtio::{DmaRemapping, IommuMapping, VirtioIommuRemapping};
+use vm_virtio::{DmaRemapping, IommuMapping, VirtioDeviceType, VirtioIommuRemapping};
 use vm_virtio::{VirtioSharedMemory, VirtioSharedMemoryList};
 use vmm_sys_util::eventfd::EventFd;
 
@@ -254,9 +254,17 @@ pub enum DeviceManagerError {
     /// Failed removing a bus device from the MMIO bus.
     RemoveDeviceFromMmioBus(devices::BusError),
 
-    /// Failed to find VFIO device corresponding to the given identifier.
+    /// Failed to find the device corresponding to a specific PCI b/d/f.
     #[cfg(feature = "pci_support")]
-    UnknownVfioDeviceId(String),
+    UnknownPciBdf(u32),
+
+    /// Not allowed to remove this type of device from the VM.
+    #[cfg(feature = "pci_support")]
+    RemovalNotAllowed(vm_virtio::VirtioDeviceType),
+
+    /// Failed to find device corresponding to the given identifier.
+    #[cfg(feature = "pci_support")]
+    UnknownDeviceId(String),
 
     /// Failed to find an available PCI device ID.
     #[cfg(feature = "pci_support")]
@@ -2150,12 +2158,38 @@ impl DeviceManager {
     #[cfg(feature = "pci_support")]
     pub fn remove_device(&mut self, id: String) -> DeviceManagerResult<()> {
         if let Some(pci_device_bdf) = self.pci_id_list.get(&id) {
+            if let Some(any_device) = self.pci_devices.get(&pci_device_bdf) {
+                if let Ok(virtio_pci_device) =
+                    Arc::clone(any_device).downcast::<Mutex<VirtioPciDevice>>()
+                {
+                    let device_type = VirtioDeviceType::from(
+                        virtio_pci_device
+                            .lock()
+                            .unwrap()
+                            .virtio_device()
+                            .lock()
+                            .unwrap()
+                            .device_type(),
+                    );
+                    match device_type {
+                        VirtioDeviceType::TYPE_NET
+                        | VirtioDeviceType::TYPE_BLOCK
+                        | VirtioDeviceType::TYPE_PMEM
+                        | VirtioDeviceType::TYPE_FS
+                        | VirtioDeviceType::TYPE_VSOCK => {}
+                        _ => return Err(DeviceManagerError::RemovalNotAllowed(device_type)),
+                    }
+                }
+            } else {
+                return Err(DeviceManagerError::UnknownPciBdf(*pci_device_bdf));
+            }
+
             // Update the PCID bitmap
             self.pci_devices_down |= 1 << (*pci_device_bdf >> 3);
 
             Ok(())
         } else {
-            Err(DeviceManagerError::UnknownVfioDeviceId(id))
+            Err(DeviceManagerError::UnknownDeviceId(id))
         }
     }