refactor(virtio): create VirtioDeviceType enum

Instead of using `u32` values directly from generated `virtio_ids`
create an enum backed by u8 with all types we actually use and start using
them.

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

Author: ShadowCurse

src/vmm/src/builder.rs

@@ -753,7 +753,7 @@ pub(crate) mod tests {
     use super::*;
     use crate::device_manager::tests::default_device_manager;
     use crate::devices::virtio::block::CacheType;
-    use crate::devices::virtio::generated::virtio_ids;
+    use crate::devices::virtio::device::VirtioDeviceType;
     use crate::devices::virtio::rng::device::ENTROPY_DEV_ID;
     use crate::devices::virtio::vsock::VSOCK_DEV_ID;
     use crate::mmds::data_store::{Mmds, MmdsVersion};
@@ -955,7 +955,7 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .get_virtio_device(virtio_ids::VIRTIO_ID_VSOCK, &vsock_dev_id)
+                .get_virtio_device(VirtioDeviceType::Vsock, &vsock_dev_id)
                 .is_some()
         );
     }
@@ -980,7 +980,7 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .get_virtio_device(virtio_ids::VIRTIO_ID_RNG, ENTROPY_DEV_ID)
+                .get_virtio_device(VirtioDeviceType::Rng, ENTROPY_DEV_ID)
                 .is_some()
         );
     }
@@ -1044,7 +1044,7 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
+                .get_virtio_device(VirtioDeviceType::Balloon, BALLOON_DEV_ID)
                 .is_some()
         );
     }
@@ -1095,7 +1095,7 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda ro"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1116,7 +1116,7 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1138,7 +1138,7 @@ pub(crate) mod tests {
             assert!(!cmdline_contains(&cmdline, "root=/dev/vda"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1175,17 +1175,17 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "root")
+                    .get_virtio_device(VirtioDeviceType::Block, "root")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "secondary")
+                    .get_virtio_device(VirtioDeviceType::Block, "secondary")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, "third")
+                    .get_virtio_device(VirtioDeviceType::Block, "third")
                     .is_some()
             );
 
@@ -1214,7 +1214,7 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1235,7 +1235,7 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 ro"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1256,7 +1256,7 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .get_virtio_device(virtio_ids::VIRTIO_ID_BLOCK, drive_id.as_str())
+                    .get_virtio_device(VirtioDeviceType::Block, drive_id.as_str())
                     .is_some()
             );
         }
@@ -1279,7 +1279,7 @@ pub(crate) mod tests {
         assert!(cmdline_contains(&cmdline, "root=/dev/pmem0 ro"));
         assert!(
             vmm.device_manager
-                .get_virtio_device(virtio_ids::VIRTIO_ID_PMEM, id.as_str())
+                .get_virtio_device(VirtioDeviceType::Pmem, id.as_str())
                 .is_some()
         );
     }

src/vmm/src/device_manager/mmio.rs

@@ -25,6 +25,7 @@ use crate::arch::{RTC_MEM_START, SERIAL_MEM_START};
 #[cfg(target_arch = "aarch64")]
 use crate::devices::legacy::{RTCDevice, SerialDevice};
 use crate::devices::pseudo::BootTimer;
+use crate::devices::virtio::device::VirtioDeviceType;
 use crate::devices::virtio::transport::mmio::MmioTransport;
 use crate::vstate::bus::{Bus, BusError};
 #[cfg(target_arch = "x86_64")]
@@ -118,7 +119,7 @@ pub struct MMIODevice<T> {
 #[derive(Debug, Default)]
 pub struct MMIODeviceManager {
     /// VirtIO devices using an MMIO transport layer
-    pub(crate) virtio_devices: HashMap<(u32, String), MMIODevice<MmioTransport>>,
+    pub(crate) virtio_devices: HashMap<(VirtioDeviceType, String), MMIODevice<MmioTransport>>,
     /// Boot timer device
     pub(crate) boot_timer: Option<MMIODevice<BootTimer>>,
     #[cfg(target_arch = "aarch64")]
@@ -382,20 +383,20 @@ impl MMIODeviceManager {
     /// Gets the specified device.
     pub fn get_virtio_device(
         &self,
-        virtio_type: u32,
+        device_type: VirtioDeviceType,
         device_id: &str,
     ) -> Option<&MMIODevice<MmioTransport>> {
         self.virtio_devices
-            .get(&(virtio_type, device_id.to_string()))
+            .get(&(device_type, device_id.to_string()))
     }
 
     /// Run fn for each registered virtio device.
     pub fn for_each_virtio_device<F, E: Debug>(&self, mut f: F) -> Result<(), E>
     where
-        F: FnMut(&u32, &String, &MMIODevice<MmioTransport>) -> Result<(), E>,
+        F: FnMut(&VirtioDeviceType, &String, &MMIODevice<MmioTransport>) -> Result<(), E>,
     {
-        for ((virtio_type, device_id), mmio_device) in &self.virtio_devices {
-            f(virtio_type, device_id, mmio_device)?;
+        for ((device_type, device_id), mmio_device) in &self.virtio_devices {
+            f(device_type, device_id, mmio_device)?;
         }
         Ok(())
     }
@@ -430,7 +431,7 @@ pub(crate) mod tests {
 
     use super::*;
     use crate::devices::virtio::ActivateError;
-    use crate::devices::virtio::device::VirtioDevice;
+    use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
     use crate::devices::virtio::queue::Queue;
     use crate::devices::virtio::transport::VirtioInterrupt;
     use crate::devices::virtio::transport::mmio::IrqTrigger;
@@ -491,7 +492,7 @@ pub(crate) mod tests {
     }
 
     impl VirtioDevice for DummyDevice {
-        impl_device_type!(0);
+        impl_device_type!(VirtioDeviceType::Net);
 
         fn avail_features(&self) -> u64 {
             0
@@ -575,15 +576,21 @@ pub(crate) mod tests {
             )
             .unwrap();
 
-        assert!(device_manager.get_virtio_device(0, "foo").is_none());
-        let dev = device_manager.get_virtio_device(0, "dummy").unwrap();
+        assert!(
+            device_manager
+                .get_virtio_device(VirtioDeviceType::Net, "foo")
+                .is_none()
+        );
+        let dev = device_manager
+            .get_virtio_device(VirtioDeviceType::Net, "dummy")
+            .unwrap();
         assert_eq!(dev.resources.addr, arch::MEM_32BIT_DEVICES_START);
         assert_eq!(dev.resources.len, MMIO_LEN);
         assert_eq!(dev.resources.gsi, Some(arch::GSI_LEGACY_START));
 
         device_manager
-            .for_each_virtio_device(|virtio_type, device_id, mmio_device| {
-                assert_eq!(*virtio_type, 0);
+            .for_each_virtio_device(|device_type, device_id, mmio_device| {
+                assert_eq!(*device_type, VirtioDeviceType::Net);
                 assert_eq!(device_id, "dummy");
                 assert_eq!(mmio_device.resources.addr, arch::MEM_32BIT_DEVICES_START);
                 assert_eq!(mmio_device.resources.len, MMIO_LEN);
@@ -642,7 +649,7 @@ pub(crate) mod tests {
     #[test]
     fn test_dummy_device() {
         let dummy = DummyDevice::new();
-        assert_eq!(dummy.device_type(), 0);
+        assert_eq!(dummy.device_type(), VirtioDeviceType::Net);
         assert_eq!(dummy.queues().len(), QUEUE_SIZES.len());
     }
 

src/vmm/src/device_manager/mod.rs

@@ -31,7 +31,7 @@ use crate::devices::legacy::RTCDevice;
 use crate::devices::legacy::serial::SerialOut;
 use crate::devices::legacy::{IER_RDA_BIT, IER_RDA_OFFSET, SerialDevice};
 use crate::devices::pseudo::BootTimer;
-use crate::devices::virtio::device::VirtioDevice;
+use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
 use crate::devices::virtio::transport::mmio::{IrqTrigger, MmioTransport};
 use crate::resources::VmResources;
 use crate::snapshot::Persist;
@@ -332,11 +332,11 @@ impl DeviceManager {
     /// Get a VirtIO device of type `virtio_type` with ID `device_id`
     pub fn get_virtio_device(
         &self,
-        virtio_type: u32,
+        device_type: VirtioDeviceType,
         device_id: &str,
     ) -> Option<Arc<Mutex<dyn VirtioDevice>>> {
         if self.pci_devices.pci_segment.is_some() {
-            let pci_device = self.pci_devices.get_virtio_device(virtio_type, device_id)?;
+            let pci_device = self.pci_devices.get_virtio_device(device_type, device_id)?;
             Some(
                 pci_device
                     .lock()
@@ -347,7 +347,7 @@ impl DeviceManager {
         } else {
             let mmio_device = self
                 .mmio_devices
-                .get_virtio_device(virtio_type, device_id)?;
+                .get_virtio_device(device_type, device_id)?;
             Some(
                 mmio_device
                     .inner

src/vmm/src/device_manager/pci_mngr.rs

@@ -16,8 +16,7 @@ use crate::devices::virtio::balloon::Balloon;
 use crate::devices::virtio::balloon::persist::{BalloonConstructorArgs, BalloonState};
 use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
-use crate::devices::virtio::device::VirtioDevice;
-use crate::devices::virtio::generated::virtio_ids;
+use crate::devices::virtio::device::{VirtioDevice, VirtioDeviceType};
 use crate::devices::virtio::mem::VirtioMem;
 use crate::devices::virtio::mem::persist::{VirtioMemConstructorArgs, VirtioMemState};
 use crate::devices::virtio::net::Net;
@@ -48,7 +47,7 @@ pub struct PciDevices {
     /// PCIe segment of the VMM, if PCI is enabled. We currently support a single PCIe segment.
     pub pci_segment: Option<PciSegment>,
     /// All VirtIO PCI devices of the system
-    pub virtio_devices: HashMap<(u32, String), Arc<Mutex<VirtioPciDevice>>>,
+    pub virtio_devices: HashMap<(VirtioDeviceType, String), Arc<Mutex<VirtioPciDevice>>>,
 }
 
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -120,7 +119,7 @@ impl PciDevices {
         debug!("Allocating BDF: {pci_device_bdf:?} for device");
         let mem = vm.guest_memory().clone();
 
-        let device_type: u32 = device.lock().expect("Poisoned lock").device_type();
+        let device_type = device.lock().expect("Poisoned lock").device_type();
 
         // Allocate one MSI vector per queue, plus one for configuration
         let msix_num =
@@ -172,7 +171,7 @@ impl PciDevices {
     ) -> Result<(), PciManagerError> {
         // We should only be reaching this point if PCI is enabled
         let pci_segment = self.pci_segment.as_ref().unwrap();
-        let device_type: u32 = device.lock().expect("Poisoned lock").device_type();
+        let device_type = device.lock().expect("Poisoned lock").device_type();
 
         let virtio_device = Arc::new(Mutex::new(VirtioPciDevice::new_from_state(
             device_id.to_string(),
@@ -207,7 +206,7 @@ impl PciDevices {
     /// Gets the specified device.
     pub fn get_virtio_device(
         &self,
-        device_type: u32,
+        device_type: VirtioDeviceType,
         device_id: &str,
     ) -> Option<&Arc<Mutex<VirtioPciDevice>>> {
         self.virtio_devices
@@ -290,7 +289,7 @@ impl<'a> Persist<'a> for PciDevices {
             let pci_device_bdf = transport_state.pci_device_bdf.into();
 
             match locked_virtio_dev.device_type() {
-                virtio_ids::VIRTIO_ID_BALLOON => {
+                VirtioDeviceType::Balloon => {
                     let balloon_device = locked_virtio_dev
                         .as_any()
                         .downcast_ref::<Balloon>()
@@ -305,7 +304,7 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     });
                 }
-                virtio_ids::VIRTIO_ID_BLOCK => {
+                VirtioDeviceType::Block => {
                     let block_dev = locked_virtio_dev
                         .as_mut_any()
                         .downcast_mut::<Block>()
@@ -326,7 +325,7 @@ impl<'a> Persist<'a> for PciDevices {
                         });
                     }
                 }
-                virtio_ids::VIRTIO_ID_NET => {
+                VirtioDeviceType::Net => {
                     let net_dev = locked_virtio_dev
                         .as_mut_any()
                         .downcast_mut::<Net>()
@@ -348,7 +347,7 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     })
                 }
-                virtio_ids::VIRTIO_ID_VSOCK => {
+                VirtioDeviceType::Vsock => {
                     let vsock_dev = locked_virtio_dev
                         .as_mut_any()
                         // Currently, VsockUnixBackend is the only implementation of VsockBackend.
@@ -379,7 +378,7 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     });
                 }
-                virtio_ids::VIRTIO_ID_RNG => {
+                VirtioDeviceType::Rng => {
                     let rng_dev = locked_virtio_dev
                         .as_mut_any()
                         .downcast_mut::<Entropy>()
@@ -393,7 +392,7 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     })
                 }
-                virtio_ids::VIRTIO_ID_PMEM => {
+                VirtioDeviceType::Pmem => {
                     let pmem_dev = locked_virtio_dev
                         .as_mut_any()
                         .downcast_mut::<Pmem>()
@@ -406,7 +405,7 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     });
                 }
-                virtio_ids::VIRTIO_ID_MEM => {
+                VirtioDeviceType::Mem => {
                     let mem_dev = locked_virtio_dev
                         .as_mut_any()
                         .downcast_mut::<VirtioMem>()
@@ -420,7 +419,6 @@ impl<'a> Persist<'a> for PciDevices {
                         transport_state,
                     })
                 }
-                _ => unreachable!(),
             }
         }