From 2cd0bc0a2c8887d57c1019f8e88c7cc5f6e3ae7c Mon Sep 17 00:00:00 2001
From: Samuel Ortiz <sameo@linux.intel.com>
Date: Tue, 17 Mar 2020 00:48:12 +0100
Subject: [PATCH] vmm: Create initial VM from its snapshot

The MemoryManager is somehow a special case, as its restore() function
was not implemented as part of the Snapshottable trait. Instead, and
because restoring memory regions rely both on vm.json and every memory
region snapshot file, the memory manager is restored at creation time.
This makes the restore path slightly different from CpuManager, Vcpu,
DeviceManager and Vm, but achieve the correct restoration of the
MemoryManager along with its memory regions filled with the correct
content.

Signed-off-by: Samuel Ortiz <sameo@linux.intel.com>
---
 vmm/src/memory_manager.rs | 63 ++++++++++++++++++++++++++++++++++++++-
 vmm/src/vm.rs             | 39 +++++++++++++++++++++++-
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index b8c123b78..6ab80cbec 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -24,7 +24,7 @@ use vm_memory::guest_memory::FileOffset;
 use vm_memory::{
     mmap::MmapRegionError, Address, Bytes, Error as MmapError, GuestAddress, GuestAddressSpace,
     GuestMemory, GuestMemoryAtomic, GuestMemoryLoadGuard, GuestMemoryMmap, GuestMemoryRegion,
-    GuestRegionMmap, GuestUsize, MmapRegion,
+    GuestRegionMmap, GuestUsize, MemoryRegionAddress, MmapRegion,
 };
 use vm_migration::{
     Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
@@ -105,6 +105,9 @@ pub enum Error {
     /// Failed to virtio-mem resize
     VirtioMemResizeFail(vm_virtio::mem::Error),
 
+    /// Cannot restore VM
+    Restore(MigratableError),
+
     /// Cannot create the system allocator
     CreateSystemAllocator,
 }
@@ -350,6 +353,64 @@ impl MemoryManager {
         Ok(memory_manager)
     }
 
+    pub fn new_from_snapshot(
+        snapshot: &Snapshot,
+        fd: Arc<VmFd>,
+        config: &MemoryConfig,
+        source_url: &str,
+    ) -> Result<Arc<Mutex<MemoryManager>>, Error> {
+        let memory_manager = MemoryManager::new(fd, config)?;
+
+        let url = Url::parse(source_url).unwrap();
+        /* url must be valid dir which is verified in recv_vm_snapshot() */
+        let vm_snapshot_path = url.to_file_path().unwrap();
+
+        if let Some(mem_section) = snapshot
+            .snapshot_data
+            .get(&format!("{}-section", MEMORY_MANAGER_SNAPSHOT_ID))
+        {
+            let mem_snapshot: MemoryManagerSnapshotData =
+                match serde_json::from_slice(&mem_section.snapshot) {
+                    Ok(snapshot) => snapshot,
+                    Err(error) => {
+                        return Err(Error::Restore(MigratableError::Restore(anyhow!(
+                            "Could not deserialize MemoryManager {}",
+                            error
+                        ))))
+                    }
+                };
+
+            memory_manager
+                .lock()
+                .unwrap()
+                .guest_memory
+                .memory()
+                .with_regions(|index, region| {
+                    let mut memory_region_path = vm_snapshot_path.clone();
+                    memory_region_path
+                        .push(mem_snapshot.memory_regions[index].backing_file.clone());
+
+                    // Create the snapshot file for the region
+                    let mut memory_region_file = OpenOptions::new()
+                        .read(true)
+                        .open(memory_region_path)
+                        .map_err(|e| Error::Restore(MigratableError::MigrateReceive(e.into())))?;
+
+                    region
+                        .read_from(
+                            MemoryRegionAddress(0),
+                            &mut memory_region_file,
+                            region.len().try_into().unwrap(),
+                        )
+                        .map_err(|e| Error::Restore(MigratableError::MigrateReceive(e.into())))?;
+
+                    Ok(())
+                })?;
+        }
+
+        Ok(memory_manager)
+    }
+
     fn create_ram_region(
         backing_file: &Option<PathBuf>,
         start_addr: GuestAddress,
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 3a2547ceb..4a363c2c6 100755
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -29,7 +29,7 @@ use crate::config::{DeviceConfig, DiskConfig, HotplugMethod, NetConfig, PmemConf
 use crate::cpu;
 use crate::device_manager::{get_win_size, Console, DeviceManager, DeviceManagerError};
 use crate::memory_manager::{Error as MemoryManagerError, MemoryManager};
-use crate::migration::{url_to_path, VM_SNAPSHOT_FILE};
+use crate::migration::{url_to_path, vm_config_from_snapshot, VM_SNAPSHOT_FILE};
 use crate::{CPU_MANAGER_SNAPSHOT_ID, DEVICE_MANAGER_SNAPSHOT_ID, MEMORY_MANAGER_SNAPSHOT_ID};
 use anyhow::anyhow;
 use arch::{BootProtocol, EntryPoint};
@@ -374,6 +374,43 @@ impl Vm {
         )
     }
 
+    pub fn new_from_snapshot(
+        snapshot: &Snapshot,
+        exit_evt: EventFd,
+        reset_evt: EventFd,
+        vmm_path: PathBuf,
+        source_url: &str,
+    ) -> Result<Self> {
+        let (kvm, fd) = Vm::kvm_new()?;
+        let config = vm_config_from_snapshot(snapshot).map_err(Error::Restore)?;
+
+        let memory_manager = if let Some(memory_manager_snapshot) =
+            snapshot.snapshots.get(MEMORY_MANAGER_SNAPSHOT_ID)
+        {
+            MemoryManager::new_from_snapshot(
+                memory_manager_snapshot,
+                fd.clone(),
+                &config.lock().unwrap().memory.clone(),
+                source_url,
+            )
+            .map_err(Error::MemoryManager)?
+        } else {
+            return Err(Error::Restore(MigratableError::Restore(anyhow!(
+                "Missing memory manager snapshot"
+            ))));
+        };
+
+        Vm::new_from_memory_manager(
+            config,
+            memory_manager,
+            fd,
+            kvm,
+            exit_evt,
+            reset_evt,
+            vmm_path,
+        )
+    }
+
     fn load_initramfs(&mut self, guest_mem: &GuestMemoryMmap) -> Result<arch::InitramfsConfig> {
         let mut initramfs = self.initramfs.as_ref().unwrap();
         let size: usize = initramfs