From 048703551265488206753cc3f1626409556caf76 Mon Sep 17 00:00:00 2001
From: Dylan Reid <dgreid@fb.com>
Date: Tue, 9 Jun 2026 21:02:00 -0700
Subject: [PATCH] vmm: release vIOMMU VFIO container on device eject

Currently ejecting a device leaks its mapping keeping the container fd
open. Remove the mapping so the fd can be closed.

Assisted-by: Claude:Opus-4.8
Signed-off-by: Dylan Reid <dgreid@fb.com>
---
 virtio-devices/src/iommu.rs | 78 +++++++++++++++++++++++++++++++++++++
 vmm/src/device_manager.rs   |  7 ++++
 2 files changed, 85 insertions(+)

diff --git a/virtio-devices/src/iommu.rs b/virtio-devices/src/iommu.rs
index 8eb316b9a..4eb36c2ea 100644
--- a/virtio-devices/src/iommu.rs
+++ b/virtio-devices/src/iommu.rs
@@ -1244,6 +1244,14 @@ impl Iommu {
         self.ext_mapping.lock().unwrap().insert(device_id, mapping);
     }
 
+    /// Removes a mapping added with `add_external_mapping`.
+    pub fn remove_external_mapping(
+        &mut self,
+        device_id: u32,
+    ) -> Option<Arc<dyn ExternalDmaMapping>> {
+        self.ext_mapping.lock().unwrap().remove(&device_id)
+    }
+
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -1362,3 +1370,73 @@ impl Snapshottable for Iommu {
 }
 impl Transportable for Iommu {}
 impl Migratable for Iommu {}
+
+#[cfg(test)]
+mod tests {
+    use std::io;
+    use std::sync::{Arc, Weak};
+
+    use seccompiler::SeccompAction;
+    use vm_device::dma_mapping::ExternalDmaMapping;
+    use vmm_sys_util::eventfd::{EFD_NONBLOCK, EventFd};
+
+    use super::Iommu;
+
+    /// Test stub for VfioDmaMapping.
+    struct MockMapping;
+
+    impl ExternalDmaMapping for MockMapping {
+        fn map(&self, _iova: u64, _gpa: u64, _size: u64) -> Result<(), io::Error> {
+            Ok(())
+        }
+
+        fn unmap(&self, _iova: u64, _size: u64) -> Result<(), io::Error> {
+            Ok(())
+        }
+    }
+
+    fn new_iommu() -> Iommu {
+        let (iommu, _mapping) = Iommu::new(
+            "test-iommu".to_string(),
+            SeccompAction::Allow,
+            EventFd::new(EFD_NONBLOCK).unwrap(),
+            (0, 0),
+            64,
+            false,
+            None,
+        )
+        .unwrap();
+        iommu
+    }
+
+    /// Tests removing a mapping works and releases the ref.
+    #[test]
+    fn remove_external_mapping() {
+        let mut iommu = new_iommu();
+
+        let mapping: Arc<dyn ExternalDmaMapping> = Arc::new(MockMapping);
+        let weak: Weak<dyn ExternalDmaMapping> = Arc::downgrade(&mapping);
+
+        iommu.add_external_mapping(0x100, mapping);
+
+        // Removing the mapping succeeds.
+        let removed = iommu.remove_external_mapping(0x100);
+        assert!(removed.is_some());
+
+        // Dropping the returned Arc drops the last reference.
+        drop(removed);
+        assert!(
+            weak.upgrade().is_none(),
+            "iommu must not retain a reference after removal"
+        );
+
+        // Removing the same id again is a nop.
+        assert!(
+            iommu.remove_external_mapping(0x100).is_none(),
+            "removal is idempotent"
+        );
+
+        // Removing a bogus ID doesn't crash.
+        assert!(iommu.remove_external_mapping(0x999).is_none());
+    }
+}
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index ba279f7b6..93a9c4b68 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -5020,6 +5020,13 @@ impl DeviceManager {
             }
         };
 
+        if let Some(iommu) = &self.iommu_device {
+            iommu
+                .lock()
+                .unwrap()
+                .remove_external_mapping(pci_device_bdf.into());
+        }
+
         if remove_dma_handler {
             for virtio_mem_device in self.virtio_mem_devices.iter() {
                 let source = VirtioMemMappingSource::Device(pci_device_bdf.into());