From c44de53f0bc313a2a6753d89f1c82e966c10e640 Mon Sep 17 00:00:00 2001
From: "Liang, Zhou" <liang1.zhou@intel.com>
Date: Fri, 4 Sep 2026 20:15:29 +0530
Subject: [PATCH] tdx: mask CPUID.0xD using TDX supported_xfam

Signed-off-by: Liang, Zhou <liang1.zhou@intel.com>
---
 arch/src/x86_64/mod.rs | 49 +++++++++++++++++++++++++++++++++++++++---
 vmm/src/cpu.rs         |  2 ++
 vmm/src/lib.rs         |  6 ++++++
 vmm/src/vm.rs          |  3 +++
 4 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/arch/src/x86_64/mod.rs b/arch/src/x86_64/mod.rs
index 8c1f1417b..ecc1dfd8e 100644
--- a/arch/src/x86_64/mod.rs
+++ b/arch/src/x86_64/mod.rs
@@ -564,6 +564,7 @@ impl CpuidFeatureEntry {
 pub fn generate_common_cpuid(
     hypervisor: &dyn hypervisor::Hypervisor,
     config: &CpuidConfig,
+    #[cfg(feature = "tdx")] vm: Option<&dyn hypervisor::Vm>,
 ) -> super::Result<Vec<CpuIdEntry>> {
     #[allow(unused_unsafe)]
     // SAFETY: cpuid called with valid leaves
@@ -634,6 +635,16 @@ pub fn generate_common_cpuid(
 
     CpuidPatch::patch_cpuid(&mut cpuid, &cpuid_patches);
 
+    #[cfg(feature = "tdx")]
+    if config.tdx {
+        let tdx_vm = vm.ok_or_else(|| {
+            Error::TdxCapabilities(HypervisorVmError::InitializeTdx(std::io::Error::other(
+                "Missing VM instance for TDX CPUID generation",
+            )))
+        })?;
+        common_cpuid_tdx_configuration(&mut cpuid, tdx_vm)?;
+    }
+
     // Update some existing CPUID
     for entry in cpuid.as_mut_slice().iter_mut() {
         #[allow(unused_unsafe)]
@@ -648,9 +659,6 @@ pub fn generate_common_cpuid(
                     entry.edx &= !(1 << AMX_COMPLEX);
                 }
             }
-            0xd =>
-            {
-            }
             // Tile Information (purely AMX related).
             0x1d if !config.amx => {
                 entry.eax = 0;
@@ -792,6 +800,41 @@ pub fn generate_common_cpuid(
     Ok(cpuid)
 }
 
+#[cfg(feature = "tdx")]
+fn common_cpuid_tdx_configuration(
+    cpuid: &mut [CpuIdEntry],
+    vm: &dyn hypervisor::Vm,
+) -> super::Result<()> {
+    let caps = vm.tdx_capabilities().map_err(Error::TdxCapabilities)?;
+    debug!("TDX capabilities supported_attrs={:#x} supported_xfam={:#x}",
+        caps.supported_attrs, caps.supported_xfam);
+
+    // XCR0-managed (user) XSAVE state components. Must include the AMX
+    // tile state components (XTILECFG bit 17, XTILEDATA bit 18); otherwise
+    // AMX is stripped from CPUID.0xD.0 and never makes it into the TD XFAM.
+    // Bits: x87(0) SSE(1) AVX(2) BNDREG(3) BNDCSR(4) OPMASK(5) ZMM_Hi256(6)
+    //       Hi16_ZMM(7) PKRU(9) XTILECFG(17) XTILEDATA(18).
+    let xcr0_mask: u64 = 0x602ff;
+    // IA32_XSS-managed (supervisor) XSAVE state components currently defined
+    // by the architecture. Listed explicitly (rather than `!xcr0_mask`) so
+    // reserved/undefined bits (19-63) are masked to zero instead of relying
+    // on `supported_xfam` to already be zero there.
+    // Bits: PT(8) ENQCMD/PASID(10) CET_U(11) CET_S(12) HDC(13) UINTR(14)
+    //       LBR(15) HWP(16).
+    let xss_mask: u64 = 0x1fd00;
+    for entry in cpuid.iter_mut().filter(|entry| entry.function == 0xd) {
+        if entry.index == 0 {
+            entry.eax &= (caps.supported_xfam as u32) & (xcr0_mask as u32);
+            entry.edx &= ((caps.supported_xfam & xcr0_mask) >> 32) as u32;
+        } else if entry.index == 1 {
+            entry.ecx &= (caps.supported_xfam as u32) & (xss_mask as u32);
+            entry.edx &= ((caps.supported_xfam & xss_mask) >> 32) as u32;
+        }
+    }
+
+    Ok(())
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn configure_vcpu(
     vcpu: &dyn hypervisor::Vcpu,
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 187979249..1669ca274 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -1109,6 +1109,8 @@ impl CpuManager {
                     tdx,
                     amx: self.config.features.amx,
                 },
+                #[cfg(feature = "tdx")]
+                vm,
             )
             .map_err(Error::CommonCpuId)?
         };
diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index a23187d2e..c0a89dc61 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -1331,6 +1331,9 @@ impl Vmm {
                     tdx: false,
                     amx,
                 },
+                #[cfg(feature = "tdx")]
+                // Live Migration is not supported when TDX is enabled
+                None,
             )
             .map_err(|e| {
                 MigratableError::MigrateSend(anyhow!("Error generating common cpuid': {e:?}"))
@@ -1440,6 +1443,9 @@ impl Vmm {
                     tdx: false,
                     amx: vm_config.cpus.features.amx,
                 },
+                #[cfg(feature = "tdx")]
+                // Live Migration is not supported when TDX is enabled
+                None,
             )
             .map_err(|e| {
                 MigratableError::MigrateReceive(anyhow!("Error generating common cpuid: {e:?}"))
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 74180c406..d7ead06d5 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -3395,6 +3395,9 @@ impl Snapshottable for Vm {
                     tdx: false,
                     amx,
                 },
+                #[cfg(feature = "tdx")]
+                // Snapshot not possible with TDX VM
+                None,
             )
             .map_err(|e| {
                 MigratableError::MigrateReceive(anyhow!("Error generating common cpuid: {e:?}"))