From 2b150ac2ea2dff26e4f6c2cc8726c640ecfd5f47 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 11 Oct 2022 15:49:42 +0200
Subject: [PATCH] pci, virtio-devices: Restore proper BAR type

When restoring a VM, the BAR type can be found directly from the
snapshot resources. It is more reliable than the previous method which
was using self.use_64bit_bar from VirtioPciDevice because at the time
the BARs are allocated, the VirtioDevice hasn't been restored yet,
meaning the way to determine the value of use_64bit_bar is wrong for a
device like vDPA. At this time, the device type is not known and relying
on the stored resources is the only reliable way.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 pci/src/device.rs                          |  3 +++
 virtio-devices/src/transport/pci_device.rs | 23 ++++++++++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/pci/src/device.rs b/pci/src/device.rs
index b4743d741..1a49beb18 100644
--- a/pci/src/device.rs
+++ b/pci/src/device.rs
@@ -21,6 +21,8 @@ pub enum Error {
     IoRegistrationFailed(u64, configuration::Error),
     /// Expected resource not found.
     MissingResource,
+    /// Invalid resource.
+    InvalidResource(Resource),
 }
 pub type Result<T> = std::result::Result<T, Error>;
 
@@ -37,6 +39,7 @@ impl Display for Error {
                 write!(f, "failed to register an IO BAR, addr={} err={}", addr, e)
             }
             MissingResource => write!(f, "failed to find expected resource"),
+            InvalidResource(r) => write!(f, "invalid resource {:?}", r),
         }
     }
 }
diff --git a/virtio-devices/src/transport/pci_device.rs b/virtio-devices/src/transport/pci_device.rs
index 748d4328d..7654930ca 100644
--- a/virtio-devices/src/transport/pci_device.rs
+++ b/virtio-devices/src/transport/pci_device.rs
@@ -36,7 +36,7 @@ use vm_device::dma_mapping::ExternalDmaMapping;
 use vm_device::interrupt::{
     InterruptIndex, InterruptManager, InterruptSourceGroup, MsiIrqGroupConfig,
 };
-use vm_device::{BusDevice, Resource};
+use vm_device::{BusDevice, PciBarType, Resource};
 use vm_memory::{Address, ByteValued, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, Le32};
 use vm_migration::{
     Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, VersionMapped,
@@ -882,11 +882,22 @@ impl PciDevice for VirtioPciDevice {
         let device = device_clone.lock().unwrap();
 
         let mut settings_bar_addr = None;
-        if let Some(resources) = &resources {
+        let mut use_64bit_bar = self.use_64bit_bar;
+        if let Some(resources) = resources {
             for resource in resources {
-                if let Resource::PciBar { index, base, .. } = resource {
-                    if *index == VIRTIO_COMMON_BAR_INDEX {
-                        settings_bar_addr = Some(GuestAddress(*base));
+                if let Resource::PciBar {
+                    index, base, type_, ..
+                } = resource
+                {
+                    if index == VIRTIO_COMMON_BAR_INDEX {
+                        settings_bar_addr = Some(GuestAddress(base));
+                        use_64bit_bar = match type_ {
+                            PciBarType::Io => {
+                                return Err(PciDeviceError::InvalidResource(resource))
+                            }
+                            PciBarType::Mmio32 => false,
+                            PciBarType::Mmio64 => true,
+                        };
                         break;
                     }
                 }
@@ -899,7 +910,7 @@ impl PciDevice for VirtioPciDevice {
 
         // Allocate the virtio-pci capability BAR.
         // See http://docs.oasis-open.org/virtio/virtio/v1.0/cs04/virtio-v1.0-cs04.html#x1-740004
-        let (virtio_pci_bar_addr, region_type) = if self.use_64bit_bar {
+        let (virtio_pci_bar_addr, region_type) = if use_64bit_bar {
             let region_type = PciBarRegionType::Memory64BitRegion;
             let addr = mmio_allocator
                 .allocate(