From 157db33d659597dd3b5395c804c79442eff479bd Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 18 Oct 2022 17:13:36 +0200
Subject: [PATCH] vmm: Refactor hypervisor::Vm creation on restore

This prevents from leaking implementation details to lib.rs, and rather
keep them in vm.rs.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/lib.rs | 28 ++++++++++-----------
 vmm/src/vm.rs  | 66 +++++++++++++++++++++++++++++---------------------
 2 files changed, 51 insertions(+), 43 deletions(-)

diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index e65da51c5..cfe20bb2c 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -27,8 +27,6 @@ use crate::migration::{recv_vm_config, recv_vm_state};
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::vm::{Error as VmError, Vm, VmState};
 use anyhow::anyhow;
-#[cfg(target_arch = "x86_64")]
-use arch::layout::{KVM_IDENTITY_MAP_START, KVM_TSS_START};
 use libc::{EFD_NONBLOCK, SIGINT, SIGTERM};
 use memory_manager::MemoryManagerSnapshotData;
 use pci::PciBdf;
@@ -51,7 +49,6 @@ use std::time::Instant;
 use std::{result, thread};
 use thiserror::Error;
 use tracer::trace_scoped;
-use vm::physical_bits;
 use vm_memory::bitmap::AtomicBitmap;
 use vm_migration::{protocol::*, Migratable};
 use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
@@ -1149,21 +1146,22 @@ impl Vmm {
         let config = vm_migration_config.vm_config.clone();
         self.vm_config = Some(vm_migration_config.vm_config);
 
-        self.hypervisor.check_required_extensions().unwrap();
-        let vm = self.hypervisor.create_vm().unwrap();
+        let vm = Vm::create_hypervisor_vm(
+            &self.hypervisor,
+            #[cfg(feature = "tdx")]
+            false,
+        )
+        .map_err(|e| {
+            MigratableError::MigrateReceive(anyhow!(
+                "Error creating hypervisor VM from snapshot: {:?}",
+                e
+            ))
+        })?;
 
-        #[cfg(target_arch = "x86_64")]
-        {
-            vm.set_identity_map_address(KVM_IDENTITY_MAP_START.0)
-                .unwrap();
-            vm.set_tss_address(KVM_TSS_START.0 as usize).unwrap();
-            vm.enable_split_irq().unwrap();
-        }
-
-        let phys_bits = physical_bits(config.lock().unwrap().cpus.max_phys_bits);
+        let phys_bits = vm::physical_bits(config.lock().unwrap().cpus.max_phys_bits);
 
         let memory_manager = MemoryManager::new(
-            vm.clone(),
+            vm,
             &config.lock().unwrap().memory.clone(),
             None,
             phys_bits,
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 218f55f40..06d21747f 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -730,25 +730,12 @@ impl Vm {
 
         #[cfg(feature = "tdx")]
         let tdx_enabled = config.lock().unwrap().is_tdx_enabled();
-        hypervisor.check_required_extensions().unwrap();
-        #[cfg(feature = "tdx")]
-        let vm = hypervisor
-            .create_vm_with_type(if tdx_enabled {
-                2 // KVM_X86_TDX_VM
-            } else {
-                0 // KVM_X86_LEGACY_VM
-            })
-            .unwrap();
-        #[cfg(not(feature = "tdx"))]
-        let vm = hypervisor.create_vm().unwrap();
 
-        #[cfg(target_arch = "x86_64")]
-        {
-            vm.set_identity_map_address(KVM_IDENTITY_MAP_START.0)
-                .unwrap();
-            vm.set_tss_address(KVM_TSS_START.0 as usize).unwrap();
-            vm.enable_split_irq().unwrap();
-        }
+        let vm = Self::create_hypervisor_vm(
+            &hypervisor,
+            #[cfg(feature = "tdx")]
+            tdx_enabled,
+        )?;
 
         let phys_bits = physical_bits(config.lock().unwrap().cpus.max_phys_bits);
 
@@ -810,16 +797,11 @@ impl Vm {
     ) -> Result<Self> {
         let timestamp = Instant::now();
 
-        hypervisor.check_required_extensions().unwrap();
-        let vm = hypervisor.create_vm().unwrap();
-
-        #[cfg(target_arch = "x86_64")]
-        {
-            vm.set_identity_map_address(KVM_IDENTITY_MAP_START.0)
-                .unwrap();
-            vm.set_tss_address(KVM_TSS_START.0 as usize).unwrap();
-            vm.enable_split_irq().unwrap();
-        }
+        let vm = Self::create_hypervisor_vm(
+            &hypervisor,
+            #[cfg(feature = "tdx")]
+            false,
+        )?;
 
         let memory_manager = if let Some(memory_manager_snapshot) =
             snapshot.snapshots.get(MEMORY_MANAGER_SNAPSHOT_ID)
@@ -856,6 +838,34 @@ impl Vm {
         )
     }
 
+    pub fn create_hypervisor_vm(
+        hypervisor: &Arc<dyn hypervisor::Hypervisor>,
+        #[cfg(feature = "tdx")] tdx_enabled: bool,
+    ) -> Result<Arc<dyn hypervisor::Vm>> {
+        hypervisor.check_required_extensions().unwrap();
+
+        #[cfg(feature = "tdx")]
+        let vm = hypervisor
+            .create_vm_with_type(if tdx_enabled {
+                2 // KVM_X86_TDX_VM
+            } else {
+                0 // KVM_X86_LEGACY_VM
+            })
+            .unwrap();
+        #[cfg(not(feature = "tdx"))]
+        let vm = hypervisor.create_vm().unwrap();
+
+        #[cfg(target_arch = "x86_64")]
+        {
+            vm.set_identity_map_address(KVM_IDENTITY_MAP_START.0)
+                .unwrap();
+            vm.set_tss_address(KVM_TSS_START.0 as usize).unwrap();
+            vm.enable_split_irq().unwrap();
+        }
+
+        Ok(vm)
+    }
+
     fn load_initramfs(&mut self, guest_mem: &GuestMemoryMmap) -> Result<arch::InitramfsConfig> {
         let mut initramfs = self.initramfs.as_ref().unwrap();
         let size: usize = initramfs