From 7c3c2ee0d6967dd4b4bc51c47df1456830c635b1 Mon Sep 17 00:00:00 2001
From: Bo Chen <chen.bo@intel.com>
Date: Mon, 11 Dec 2023 20:04:55 +0000
Subject: [PATCH] vmm: tdx: Error out early for TD migration

Signed-off-by: Bo Chen <chen.bo@intel.com>
(cherry picked from commit e64b66054ee0af45a3e1e97a039ea429c8b3487b)
---
 vmm/src/lib.rs | 18 +++++++++++++++---
 vmm/src/vm.rs  |  7 ++-----
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index e230c7902..b105d7a8d 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -1640,7 +1640,12 @@ impl Vmm {
         #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
         let common_cpuid = {
             #[cfg(feature = "tdx")]
-            let tdx_enabled = vm_config.lock().unwrap().is_tdx_enabled();
+            if vm_config.lock().unwrap().is_tdx_enabled() {
+                return Err(MigratableError::MigrateSend(anyhow!(
+                    "Live Migration is not supported when TDX is enabled"
+                )));
+            };
+
             let phys_bits =
                 vm::physical_bits(&hypervisor, vm_config.lock().unwrap().cpus.max_phys_bits);
             arch::generate_common_cpuid(
@@ -1650,7 +1655,7 @@ impl Vmm {
                 phys_bits,
                 vm_config.lock().unwrap().cpus.kvm_hyperv,
                 #[cfg(feature = "tdx")]
-                tdx_enabled,
+                false,
             )
             .map_err(|e| {
                 MigratableError::MigrateReceive(anyhow!("Error generating common cpuid': {:?}", e))
@@ -1827,6 +1832,13 @@ impl Vmm {
         src_vm_config: &Arc<Mutex<VmConfig>>,
         src_vm_cpuid: &[hypervisor::arch::x86::CpuIdEntry],
     ) -> result::Result<(), MigratableError> {
+        #[cfg(feature = "tdx")]
+        if src_vm_config.lock().unwrap().is_tdx_enabled() {
+            return Err(MigratableError::MigrateReceive(anyhow!(
+                "Live Migration is not supported when TDX is enabled"
+            )));
+        };
+
         // We check the `CPUID` compatibility of between the source vm and destination, which is
         // mostly about feature compatibility and "topology/sgx" leaves are not relevant.
         let dest_cpuid = &{
@@ -1840,7 +1852,7 @@ impl Vmm {
                 phys_bits,
                 vm_config.cpus.kvm_hyperv,
                 #[cfg(feature = "tdx")]
-                vm_config.is_tdx_enabled(),
+                false,
             )
             .map_err(|e| {
                 MigratableError::MigrateReceive(anyhow!("Error generating common cpuid: {:?}", e))
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 2b829df52..cef7bbbb3 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2374,12 +2374,9 @@ impl Snapshottable for Vm {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         event!("vm", "snapshotting");
 
-        #[cfg(feature = "tdx")]
-        let tdx_enabled = self.config.lock().unwrap().is_tdx_enabled();
-
         #[cfg(feature = "tdx")]
         {
-            if tdx_enabled {
+            if self.config.lock().unwrap().is_tdx_enabled() {
                 return Err(MigratableError::Snapshot(anyhow!(
                     "Snapshot not possible with TDX VM"
                 )));
@@ -2406,7 +2403,7 @@ impl Snapshottable for Vm {
                 phys_bits,
                 self.config.lock().unwrap().cpus.kvm_hyperv,
                 #[cfg(feature = "tdx")]
-                tdx_enabled,
+                false,
             )
             .map_err(|e| {
                 MigratableError::MigrateReceive(anyhow!("Error generating common cpuid: {:?}", e))