From 0f12fe9b3b23424ab0a43374877a3179c1527854 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Thu, 7 Apr 2022 16:33:22 +0200
Subject: [PATCH] vmm: Factorize retrieval of PCI resources

Create a dedicated function for getting the PCI segment, b/d/f and
optional resources. This is meant for handling the potential case of a
restore.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/device_manager.rs | 58 +++++++++++++++++++++++----------------
 1 file changed, 34 insertions(+), 24 deletions(-)

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index b0612e6be..8f458b76e 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -3423,30 +3423,8 @@ impl DeviceManager {
         let mut node = device_node!(id);
         node.children = vec![virtio_device_id.clone()];
 
-        // Look for the id in the device tree. If it can be found, that means
-        // the device is being restored, otherwise it's created from scratch.
-        let (pci_segment_id, pci_device_bdf, resources) = if let Some(node) =
-            self.device_tree.lock().unwrap().get(&id)
-        {
-            info!("Restoring virtio-pci {} resources", id);
-            let pci_device_bdf: PciBdf = node
-                .pci_bdf
-                .ok_or(DeviceManagerError::MissingDeviceNodePciBdf)?;
-            let pci_segment_id = pci_device_bdf.segment();
-
-            self.pci_segments[pci_segment_id as usize]
-                .pci_bus
-                .lock()
-                .unwrap()
-                .get_device_id(pci_device_bdf.device() as usize)
-                .map_err(DeviceManagerError::GetPciDeviceId)?;
-
-            (pci_segment_id, pci_device_bdf, Some(node.resources.clone()))
-        } else {
-            let pci_device_bdf = self.pci_segments[pci_segment_id as usize].next_device_bdf()?;
-
-            (pci_segment_id, pci_device_bdf, None)
-        };
+        let (pci_segment_id, pci_device_bdf, resources) =
+            self.pci_resources(&id, pci_segment_id)?;
 
         // Update the existing virtio node by setting the parent.
         if let Some(node) = self.device_tree.lock().unwrap().get_mut(&virtio_device_id) {
@@ -3572,6 +3550,38 @@ impl DeviceManager {
         Ok(pci_device_bdf)
     }
 
+    fn pci_resources(
+        &self,
+        id: &str,
+        pci_segment_id: u16,
+    ) -> DeviceManagerResult<(u16, PciBdf, Option<Vec<Resource>>)> {
+        // Look for the id in the device tree. If it can be found, that means
+        // the device is being restored, otherwise it's created from scratch.
+        Ok(
+            if let Some(node) = self.device_tree.lock().unwrap().get(id) {
+                info!("Restoring virtio-pci {} resources", id);
+                let pci_device_bdf: PciBdf = node
+                    .pci_bdf
+                    .ok_or(DeviceManagerError::MissingDeviceNodePciBdf)?;
+                let pci_segment_id = pci_device_bdf.segment();
+
+                self.pci_segments[pci_segment_id as usize]
+                    .pci_bus
+                    .lock()
+                    .unwrap()
+                    .get_device_id(pci_device_bdf.device() as usize)
+                    .map_err(DeviceManagerError::GetPciDeviceId)?;
+
+                (pci_segment_id, pci_device_bdf, Some(node.resources.clone()))
+            } else {
+                let pci_device_bdf =
+                    self.pci_segments[pci_segment_id as usize].next_device_bdf()?;
+
+                (pci_segment_id, pci_device_bdf, None)
+            },
+        )
+    }
+
     #[cfg(target_arch = "x86_64")]
     pub fn io_bus(&self) -> &Arc<Bus> {
         &self.address_manager.io_bus