From 49f19e061b643f3285cb5ea29c9d4fd515522356 Mon Sep 17 00:00:00 2001
From: Rob Bradford <robert.bradford@intel.com>
Date: Fri, 8 Oct 2021 09:27:13 +0100
Subject: [PATCH] vmm: Use device's segment when removing a device

The segment ID has been stored in the DeviceTree.

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

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 4c6f3bc6d..73d8a85c8 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -3435,7 +3435,11 @@ impl DeviceManager {
 
         let pci_device_bdf = pci_device_node
             .pci_bdf
-            .ok_or(DeviceManagerError::MissingPciDevice)?;
+            .ok_or(DeviceManagerError::MissingDeviceNodePciBdf)?;
+        let pci_segment_id = pci_device_node
+            .pci_segment_id
+            .ok_or(DeviceManagerError::MissingDeviceNodePciSegmentId)?;
+
         let pci_device_handle = pci_device_node
             .pci_device_handle
             .as_ref()
@@ -3462,7 +3466,7 @@ impl DeviceManager {
         }
 
         // Update the PCID bitmap
-        self.pci_segments[0].pci_devices_down |= 1 << (pci_device_bdf >> 3);
+        self.pci_segments[pci_segment_id as usize].pci_devices_down |= 1 << (pci_device_bdf >> 3);
 
         Ok(())
     }