From b8b0dab1ae4a7288d1adaef5fb24ec97a0a05933 Mon Sep 17 00:00:00 2001
From: Rob Bradford <robert.bradford@intel.com>
Date: Thu, 7 Oct 2021 17:47:44 +0100
Subject: [PATCH] vmm: Add segment_id parameter to
 DeviceManager::add_pci_device

Signed-off-by: Rob Bradford <robert.bradford@intel.com>
---
 vmm/src/device_manager.rs | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index a3d60d99e..6b88e9b34 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -2923,6 +2923,7 @@ impl DeviceManager {
         self.add_pci_device(
             vfio_pci_device.clone(),
             vfio_pci_device.clone(),
+            0,
             pci_device_bdf,
         )?;
 
@@ -2958,6 +2959,7 @@ impl DeviceManager {
         &mut self,
         bus_device: Arc<Mutex<dyn BusDevice>>,
         pci_device: Arc<Mutex<dyn PciDevice>>,
+        segment_id: u16,
         bdf: u32,
     ) -> DeviceManagerResult<Vec<(GuestAddress, GuestUsize, PciBarRegionType)>> {
         let bars = pci_device
@@ -2966,7 +2968,10 @@ impl DeviceManager {
             .allocate_bars(&mut self.address_manager.allocator.lock().unwrap())
             .map_err(DeviceManagerError::AllocateBars)?;
 
-        let mut pci_bus = self.pci_segments[0].pci_bus.lock().unwrap();
+        let mut pci_bus = self.pci_segments[segment_id as usize]
+            .pci_bus
+            .lock()
+            .unwrap();
 
         pci_bus
             .add_device(bdf, pci_device)
@@ -3083,6 +3088,7 @@ impl DeviceManager {
         self.add_pci_device(
             vfio_user_pci_device.clone(),
             vfio_user_pci_device.clone(),
+            0,
             pci_device_bdf,
         )?;
 
@@ -3213,6 +3219,7 @@ impl DeviceManager {
         let bars = self.add_pci_device(
             virtio_pci_device.clone(),
             virtio_pci_device.clone(),
+            0,
             pci_device_bdf,
         )?;