From eb8b10add056b03f9f85d7282d6c9b436b4cce05 Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@meta.com>
Date: Sun, 26 Apr 2026 22:15:54 +0100
Subject: [PATCH] virtio-devices: iommu: Cap number of domains per device

The per-domain mapping cap bounds memory inside one domain, but a
guest can still grow the domains map indefinitely with ATTACH
requests for distinct domain IDs. Reject ATTACH with
VIRTIO_IOMMU_S_NOMEM at 64K domains.

Signed-off-by: Rob Bradford <rbradford@meta.com>
Assisted-by: Claude:claude-opus-4-7
---
 virtio-devices/src/iommu.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/virtio-devices/src/iommu.rs b/virtio-devices/src/iommu.rs
index 165dfdaf7..b56d44b20 100644
--- a/virtio-devices/src/iommu.rs
+++ b/virtio-devices/src/iommu.rs
@@ -77,6 +77,10 @@ const VIRTIO_IOMMU_PAGE_GRANULE: u64 = 1u64 << VIRTIO_IOMMU_PAGE_SIZE_MASK.trail
 // ~64 MiB at ~64 bytes/entry, well above any legitimate workload.
 const MAX_MAPPINGS_PER_DOMAIN: usize = 1 << 20;
 
+// Bound the per-device domain count so a guest cannot grow the
+// domain map indefinitely with ATTACH-only requests.
+const MAX_DOMAINS: usize = 1 << 16;
+
 #[derive(Copy, Clone, Debug, Default)]
 #[repr(C, packed)]
 #[allow(dead_code)]
@@ -321,6 +325,8 @@ enum Error {
     InvalidUnmapRequestPartialOverlap,
     #[error("Per-domain mapping count cap exceeded")]
     MappingCountExceeded,
+    #[error("Per-device domain count cap exceeded: current:{0}, max:{MAX_DOMAINS}")]
+    DomainCountExceeded(usize),
     #[error("Guest sent us invalid PROBE request")]
     InvalidProbeRequest,
     #[error("Failed to performing external mapping")]
@@ -419,6 +425,17 @@ impl Request {
                         return Err(Error::InvalidAttachRequest);
                     }
 
+                    // Refuse before mutating any state so a failed ATTACH
+                    // does not leave a phantom endpoint pointing at a
+                    // domain that was never created.
+                    {
+                        let domains = mapping.domains.read().unwrap();
+                        if !domains.contains_key(&domain_id) && domains.len() >= MAX_DOMAINS {
+                            status = VIRTIO_IOMMU_S_NOMEM;
+                            return Err(Error::DomainCountExceeded(domains.len()));
+                        }
+                    }
+
                     let mut old_domain_id = domain_id;
                     if let Some(&id) = mapping.endpoints.read().unwrap().get(&endpoint) {
                         old_domain_id = id;
@@ -449,6 +466,18 @@ impl Request {
 
                     // Add new domain with no mapping if the entry didn't exist yet
                     let mut domains = mapping.domains.write().unwrap();
+                    if !domains.contains_key(&domain_id) && domains.len() >= MAX_DOMAINS {
+                        // Defensive re-check under write lock. Single-threaded
+                        // today, but keeps the invariant if processing ever
+                        // becomes concurrent. Drop the domains write lock
+                        // before taking the endpoints write lock so we never
+                        // hold both simultaneously.
+                        let count = domains.len();
+                        drop(domains);
+                        mapping.endpoints.write().unwrap().remove(&endpoint);
+                        status = VIRTIO_IOMMU_S_NOMEM;
+                        return Err(Error::DomainCountExceeded(count));
+                    }
                     let domain = Domain {
                         mappings: BTreeMap::new(),
                         bypass,