From 1ea63f50a1ed9283e6a744610cb8a61d4fa40d78 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Mon, 27 Sep 2021 13:43:04 +0200
Subject: [PATCH] vmm: Move MemoryRangeTable creation to the MemoryManager

The function memory_range_table() will be reused by the MemoryManager in
a following patch to describe all the ranges that we should snapshot.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/memory_manager.rs | 22 ++++++++++++++++++++--
 vmm/src/vm.rs             | 32 ++++++++------------------------
 2 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 6d34b202b..cafdac26c 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -41,8 +41,8 @@ use vm_memory::{
     GuestMemory, GuestMemoryAtomic, GuestMemoryError, GuestMemoryRegion, GuestUsize, MmapRegion,
 };
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
-    SnapshotDataSection, Snapshottable, Transportable, VersionMapped,
+    protocol::MemoryRange, protocol::MemoryRangeTable, Migratable, MigratableError, Pausable,
+    Snapshot, SnapshotDataSection, Snapshottable, Transportable, VersionMapped,
 };
 
 #[cfg(feature = "acpi")]
@@ -1508,6 +1508,24 @@ impl MemoryManager {
     pub fn memory_zones(&self) -> &MemoryZones {
         &self.memory_zones
     }
+
+    pub fn memory_range_table(&self) -> std::result::Result<MemoryRangeTable, MigratableError> {
+        let mut table = MemoryRangeTable::default();
+
+        for memory_zone in self.memory_zones.values() {
+            for region in memory_zone.regions() {
+                table.push(MemoryRange {
+                    gpa: region.start_addr().raw_value(),
+                    length: region.len() as u64,
+                });
+            }
+            if let Some(virtio_mem_zone) = memory_zone.virtio_mem_zone() {
+                table.extend(virtio_mem_zone.plugged_ranges());
+            }
+        }
+
+        Ok(table)
+    }
 }
 
 #[cfg(feature = "acpi")]
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 35ccd1a94..360cde873 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -66,15 +66,14 @@ use std::{result, str, thread};
 use vm_device::Bus;
 #[cfg(target_arch = "x86_64")]
 use vm_device::BusDevice;
+#[cfg(any(target_arch = "aarch64", feature = "tdx"))]
+use vm_memory::Address;
+use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
 #[cfg(feature = "tdx")]
-use vm_memory::GuestMemory;
-use vm_memory::{
-    Address, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryRegion,
-};
+use vm_memory::{GuestMemory, GuestMemoryRegion};
 use vm_migration::{
-    protocol::{MemoryRange, MemoryRangeTable},
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
+    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    SnapshotDataSection, Snapshottable, Transportable,
 };
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::signal::unblock_signal;
@@ -2189,22 +2188,7 @@ impl Vm {
     }
 
     pub fn memory_range_table(&self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        let mut table = MemoryRangeTable::default();
-        let mm = self.memory_manager.lock().unwrap();
-
-        for memory_zone in mm.memory_zones().values() {
-            for region in memory_zone.regions() {
-                table.push(MemoryRange {
-                    gpa: region.start_addr().raw_value(),
-                    length: region.len() as u64,
-                });
-            }
-            if let Some(virtio_mem_zone) = memory_zone.virtio_mem_zone() {
-                table.extend(virtio_mem_zone.plugged_ranges());
-            }
-        }
-
-        Ok(table)
+        self.memory_manager.lock().unwrap().memory_range_table()
     }
 
     pub fn device_tree(&self) -> Arc<Mutex<DeviceTree>> {
@@ -2670,7 +2654,7 @@ mod tests {
 #[test]
 pub fn test_vm() {
     use hypervisor::VmExit;
-    use vm_memory::{GuestMemory, GuestMemoryRegion};
+    use vm_memory::{Address, GuestMemory, GuestMemoryRegion};
     // This example based on https://lwn.net/Articles/658511/
     let code = [
         0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */