From 036b7787dd96272832f740ffd7e448b2c112efb1 Mon Sep 17 00:00:00 2001
From: Archana Shinde <archana.m.shinde@intel.com>
Date: Sat, 21 Oct 2023 06:28:54 -0700
Subject: [PATCH] runtime-rs: Use PCI path from hypervisor for vfio devices

Remove earlier functionality that tries to assign PCI path to vfio
devices from the host assuming pci slots to start from 1.
Get this from the hypervisor instead.

Signed-off-by: Archana Shinde <archana.m.shinde@intel.com>
---
 .../crates/hypervisor/src/ch/inner_device.rs  |  8 ++-
 .../hypervisor/src/device/driver/vfio.rs      | 53 ++++++++++++-------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/src/runtime-rs/crates/hypervisor/src/ch/inner_device.rs b/src/runtime-rs/crates/hypervisor/src/ch/inner_device.rs
index 482a23a9b..f7509bd8b 100644
--- a/src/runtime-rs/crates/hypervisor/src/ch/inner_device.rs
+++ b/src/runtime-rs/crates/hypervisor/src/ch/inner_device.rs
@@ -169,7 +169,7 @@ impl CloudHypervisorInner {
     }
 
     async fn handle_vfio_device(&mut self, device: VfioDevice) -> Result<DeviceType> {
-        let vfio_device: VfioDevice = device.clone();
+        let mut vfio_device: VfioDevice = device.clone();
 
         // A device with multi-funtions, or a IOMMU group with one more
         // devices, the Primary device is selected to be passed to VM.
@@ -210,6 +210,12 @@ impl CloudHypervisorInner {
                 serde_json::from_str(detail.as_str()).map_err(|e| anyhow!(e))?;
             self.device_ids
                 .insert(device.device_id.clone(), dev_info.id);
+
+            // Update PCI path for the vfio host device. It is safe to directly access the slice element
+            // here as we have already checked if it exists.
+            // Todo: Handle vfio-ap mediated devices - return error for them.
+            vfio_device.devices[0].guest_pci_path =
+                Some(Self::clh_pci_info_to_path(&dev_info.bdf)?);
         }
 
         Ok(DeviceType::Vfio(vfio_device))
diff --git a/src/runtime-rs/crates/hypervisor/src/device/driver/vfio.rs b/src/runtime-rs/crates/hypervisor/src/device/driver/vfio.rs
index 132535dbd..a689bcb35 100644
--- a/src/runtime-rs/crates/hypervisor/src/device/driver/vfio.rs
+++ b/src/runtime-rs/crates/hypervisor/src/device/driver/vfio.rs
@@ -389,7 +389,7 @@ impl VfioDevice {
             .get_vfio_device_vendor(&dev_bdf)
             .context("get property device and vendor failed")?;
 
-        let mut vfio_dev = HostDevice {
+        let vfio_dev = HostDevice {
             bus_slot_func: dev_bdf.clone(),
             device_vendor: Some(dev_vendor),
             sysfs_path: vfio_dev_details.1,
@@ -397,18 +397,6 @@ impl VfioDevice {
             ..Default::default()
         };
 
-        // when vfio pci, kata-agent handles with device_options, and its
-        // format: "DDDD:BB:DD.F=<pcipath>"
-        // DDDD:BB:DD.F is the device's PCI address on host
-        // <pcipath> is the device's PCI path in the guest
-        if self.bus_mode == VfioBusMode::PCI {
-            let pci_path =
-                generate_guest_pci_path(dev_bdf.clone()).context("generate pci path failed")?;
-            vfio_dev.guest_pci_path = Some(pci_path.clone());
-            self.device_options
-                .push(format!("0000:{}={}", dev_bdf, pci_path.convert_to_string()));
-        }
-
         Ok(vfio_dev)
     }
 
@@ -493,13 +481,42 @@ impl Device for VfioDevice {
         }
 
         // do add device for vfio deivce
-        if let Err(e) = h.add_device(DeviceType::Vfio(self.clone())).await {
-            self.decrease_attach_count().await?;
+        match h.add_device(DeviceType::Vfio(self.clone())).await {
+            Ok(dev) => {
+                // Update device info with the one received from device attach
+                if let DeviceType::Vfio(vfio) = dev {
+                    self.config = vfio.config;
+                    self.devices = vfio.devices;
+                }
 
-            return Err(e);
+                if self.bus_mode == VfioBusMode::PCI {
+                    for hostdev in self.devices.iter_mut() {
+                        if hostdev.guest_pci_path.is_none() {
+                            // guest_pci_path may be empty for certain hypervisors such as
+                            // dragonball
+                            hostdev.guest_pci_path = Some(
+                                generate_guest_pci_path(hostdev.bus_slot_func.clone())
+                                    .map_err(|e| anyhow!("generate pci path failed: {:?}", e))?,
+                            );
+                        }
+
+                        // Safe to call unwrap here because of previous assignment.
+                        let pci_path = hostdev.guest_pci_path.clone().unwrap();
+                        self.device_options.push(format!(
+                            "0000:{}={}",
+                            hostdev.bus_slot_func.clone(),
+                            pci_path.convert_to_string()
+                        ));
+                    }
+                }
+
+                Ok(())
+            }
+            Err(e) => {
+                self.decrease_attach_count().await?;
+                return Err(e);
+            }
         }
-
-        Ok(())
     }
 
     async fn detach(&mut self, h: &dyn hypervisor) -> Result<Option<u64>> {