virtio-devices, vmm: Register a DMA handler to VirtioPciDevice

Given that some virtio device might need some DMA handling, we provide a
way to store this through the VirtioPciDevice layer, so that it can be
accessed when the PCI device is removed.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

virtio-devices/src/transport/pci_device.rs

@@ -31,6 +31,7 @@ use versionize::{VersionMap, Versionize, VersionizeResult};
 use versionize_derive::Versionize;
 use virtio_queue::{Error as QueueError, Queue};
 use vm_allocator::{AddressAllocator, SystemAllocator};
+use vm_device::dma_mapping::ExternalDmaMapping;
 use vm_device::interrupt::{
     InterruptIndex, InterruptManager, InterruptSourceGroup, MsiIrqGroupConfig,
 };
@@ -338,6 +339,9 @@ pub struct VirtioPciDevice {
 
     // Barrier that is used to wait on for activation
     activate_barrier: Arc<Barrier>,
+
+    // Optional DMA handler
+    dma_handler: Option<Arc<dyn ExternalDmaMapping>>,
 }
 
 impl VirtioPciDevice {
@@ -353,6 +357,7 @@ impl VirtioPciDevice {
         pci_device_bdf: u32,
         activate_evt: EventFd,
         use_64bit_bar: bool,
+        dma_handler: Option<Arc<dyn ExternalDmaMapping>>,
     ) -> Result<Self> {
         let device_clone = device.clone();
         let mut locked_device = device_clone.lock().unwrap();
@@ -454,6 +459,7 @@ impl VirtioPciDevice {
             bar_regions: vec![],
             activate_evt,
             activate_barrier: Arc::new(Barrier::new(2)),
+            dma_handler,
         };
 
         if let Some(msix_config) = &virtio_pci_device.msix_config {
@@ -701,6 +707,10 @@ impl VirtioPciDevice {
     fn needs_activation(&self) -> bool {
         !self.device_activated.load(Ordering::SeqCst) && self.is_driver_ready()
     }
+
+    pub fn dma_handler(&self) -> Option<&Arc<dyn ExternalDmaMapping>> {
+        self.dma_handler.as_ref()
+    }
 }
 
 impl VirtioTransport for VirtioPciDevice {

vmm/src/device_manager.rs

@@ -3295,7 +3295,7 @@ impl DeviceManager {
         iommu_mapping: &Option<Arc<IommuMapping>>,
         virtio_device_id: String,
         pci_segment_id: u16,
-        _dma_handler: Option<Arc<dyn ExternalDmaMapping>>,
+        dma_handler: Option<Arc<dyn ExternalDmaMapping>>,
     ) -> DeviceManagerResult<PciBdf> {
         let id = format!("{}-{}", VIRTIO_PCI_DEVICE_NAME_PREFIX, virtio_device_id);
 
@@ -3385,6 +3385,7 @@ impl DeviceManager {
             // to firmware without requiring excessive identity mapping.
             // The exception being if not on the default PCI segment.
             pci_segment_id > 0 || device_type != VirtioDeviceType::Block as u32,
+            dma_handler,
         )
         .map_err(DeviceManagerError::VirtioDevice)?;