From 3624573b12405bab4ec7e77672441c19d7441c64 Mon Sep 17 00:00:00 2001
From: Lei Huang <leih@nvidia.com>
Date: Tue, 21 May 2024 15:48:14 -0700
Subject: [PATCH] agent: collect PCI address mapping for both vfio-pci-gk and
 vfio-pci device

The `update_env_pci()` function need the PCI address mapping to
translate the host PCI address to guest PCI address in below
environment variables:
- PCIDEVICE_<prefix>_<resource-name>_INFO
- PCIDEVICE_<prefix>_<resource-name>

So collect PCI address mapping for both vfio-pci-gk and
vfio-pci devices.

Fixes #9614

Signed-off-by: Lei Huang <leih@nvidia.com>
---
 src/agent/src/device.rs | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/agent/src/device.rs b/src/agent/src/device.rs
index f7b7aa085a..a06b54fc36 100644
--- a/src/agent/src/device.rs
+++ b/src/agent/src/device.rs
@@ -896,9 +896,10 @@ async fn vfio_pci_device_handler(
             }
 
             group = Some(devgroup);
-
-            pci_fixups.push((host, guestdev));
         }
+
+        // collect PCI address mapping for both vfio-pci-gk and vfio-pci device
+        pci_fixups.push((host, guestdev));
     }
 
     let dev_update = if vfio_in_guest {
@@ -961,22 +962,22 @@ pub async fn add_devices(
                 ));
             }
 
-            let mut sb = sandbox.lock().await;
-            for (host, guest) in update.pci {
-                if let Some(other_guest) = sb.pcimap.insert(host, guest) {
-                    return Err(anyhow!(
-                        "Conflicting guest address for host device {} ({} versus {})",
-                        host,
-                        guest,
-                        other_guest
-                    ));
-                }
-            }
-
             // Update cgroup to allow all devices added to guest.
             insert_devices_cgroup_rule(spec, &dev_update.info, true, "rwm")
                 .context("Update device cgroup")?;
         }
+
+        let mut sb = sandbox.lock().await;
+        for (host, guest) in update.pci {
+            if let Some(other_guest) = sb.pcimap.insert(host, guest) {
+                return Err(anyhow!(
+                    "Conflicting guest address for host device {} ({} versus {})",
+                    host,
+                    guest,
+                    other_guest
+                ));
+            }
+        }
     }
 
     if let Some(process) = spec.process.as_mut() {