feat(tdx): initialize TDX VMs with supported CPUID features

Lencerf · Lencerf · commit ca9a64b8636b · 2026-02-22T22:17:13.000-08:00
Signed-off-by: Changyuan Lyu &lt;changyuanl@google.com&gt;
diff --git a/alioth/src/board/board_x86_64/board_x86_64.rs b/alioth/src/board/board_x86_64/board_x86_64.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod sev;
+mod tdx;
 
 use std::arch::x86_64::{__cpuid, CpuidResult};
 use std::collections::HashMap;
@@ -321,7 +322,7 @@ where
         match coco {
             Coco::AmdSev { policy } => self.sev_init(*policy, memory)?,
             Coco::AmdSnp { policy } => self.snp_init(*policy, memory)?,
-            Coco::IntelTdx { attr } => todo!("Intel TDX {attr:?}"),
+            Coco::IntelTdx { attr } => self.tdx_init(*attr, memory)?,
         }
         Ok(())
     }
diff --git a/alioth/src/board/board_x86_64/tdx.rs b/alioth/src/board/board_x86_64/tdx.rs
@@ -0,0 +1,32 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use crate::arch::tdx::TdAttr;
+use crate::board::{Board, Result};
+use crate::hv::{Vm, VmMemory};
+use crate::mem::MarkPrivateMemory;
+
+impl<V> Board<V>
+where
+    V: Vm,
+{
+    pub(crate) fn tdx_init(&self, attr: TdAttr, memory: Arc<dyn VmMemory>) -> Result<()> {
+        self.vm.tdx_init_vm(attr, &self.arch.cpuids)?;
+        let mark_private_memory = Box::new(MarkPrivateMemory { memory });
+        self.memory.register_change_callback(mark_private_memory)?;
+        Ok(())
+    }
+}
diff --git a/alioth/src/hv/hv.rs b/alioth/src/hv/hv.rs
@@ -390,6 +390,9 @@ pub trait Vm {
     #[cfg(target_arch = "x86_64")]
     fn snp_launch_finish(&self) -> Result<()>;
 
+    #[cfg(target_arch = "x86_64")]
+    fn tdx_init_vm(&self, attr: TdAttr, cpuids: &HashMap<CpuidIn, CpuidResult>) -> Result<()>;
+
     #[cfg(target_arch = "aarch64")]
     type GicV2: GicV2;
     #[cfg(target_arch = "aarch64")]
diff --git a/alioth/src/hv/kvm/kvm_x86_64.rs b/alioth/src/hv/kvm/kvm_x86_64.rs
@@ -24,6 +24,26 @@ use crate::sys::kvm::{
     KvmCpuidFeature, KvmX2apicApiFlag, kvm_get_supported_cpuid,
 };
 
+impl From<KvmCpuidEntry2> for (CpuidIn, CpuidResult) {
+    fn from(value: KvmCpuidEntry2) -> Self {
+        let in_ = CpuidIn {
+            func: value.function,
+            index: if value.flags.contains(KvmCpuid2Flag::SIGNIFCANT_INDEX) || value.index > 0 {
+                Some(value.index)
+            } else {
+                None
+            },
+        };
+        let result = CpuidResult {
+            eax: value.eax,
+            ebx: value.ebx,
+            ecx: value.ecx,
+            edx: value.edx,
+        };
+        (in_, result)
+    }
+}
+
 impl Kvm {
     pub fn get_supported_cpuids(&self) -> Result<HashMap<CpuidIn, CpuidResult>> {
         let mut kvm_cpuid2 = KvmCpuid2 {
@@ -32,29 +52,12 @@ impl Kvm {
             entries: [KvmCpuidEntry2::default(); KVM_MAX_CPUID_ENTRIES],
         };
         unsafe { kvm_get_supported_cpuid(&self.fd, &mut kvm_cpuid2) }.context(error::GuestCpuid)?;
-        let map_f = |e: &KvmCpuidEntry2| {
-            let in_ = CpuidIn {
-                func: e.function,
-                index: if e.flags.contains(KvmCpuid2Flag::SIGNIFCANT_INDEX) {
-                    Some(e.index)
-                } else {
-                    None
-                },
-            };
-            let out = CpuidResult {
-                eax: e.eax,
-                ebx: e.ebx,
-                ecx: e.ecx,
-                edx: e.edx,
-            };
-            (in_, out)
-        };
         let mut cpuids: HashMap<_, _> = kvm_cpuid2
             .entries
-            .iter()
+            .into_iter()
             .filter(|e| e.eax != 0 || e.ebx != 0 || e.ecx != 0 || e.edx != 0)
             .take(kvm_cpuid2.nent as usize)
-            .map(map_f)
+            .map(From::from)
             .collect();
 
         let leaf_features = CpuidIn {
diff --git a/alioth/src/hv/kvm/kvm_x86_64_test.rs b/alioth/src/hv/kvm/kvm_x86_64_test.rs
@@ -12,8 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::arch::x86_64::CpuidResult;
+
+use rstest::rstest;
+
+use crate::arch::cpuid::CpuidIn;
 use crate::hv::kvm::{Kvm, KvmConfig};
-use crate::sys::kvm::KVM_CPUID_SIGNATURE;
+use crate::sys::kvm::{KVM_CPUID_SIGNATURE, KvmCpuid2Flag, KvmCpuidEntry2};
 
 #[test]
 #[cfg_attr(not(feature = "test-hv"), ignore)]
@@ -32,3 +37,60 @@ fn test_get_supported_cpuid() {
     }
     assert!(kvm_cpuid_exist);
 }
+
+#[rstest]
+#[case(
+    KvmCpuidEntry2 {
+        function: 0,
+        index: 0,
+        flags: KvmCpuid2Flag::empty(),
+        eax: 0x10,
+        ebx: u32::from_le_bytes(*b"Auth"),
+        ecx: u32::from_le_bytes(*b"cAMD"),
+        edx: u32::from_le_bytes(*b"enti"),
+        padding: [0; 3],
+    },
+    (
+        CpuidIn {
+            func: 0,
+            index: None,
+        },
+        CpuidResult {
+            eax: 0x10,
+            ebx: u32::from_le_bytes(*b"Auth"),
+            ecx: u32::from_le_bytes(*b"cAMD"),
+            edx: u32::from_le_bytes(*b"enti"),
+        }
+    )
+)]
+#[case(
+    KvmCpuidEntry2 {
+        function: 0xb,
+        index: 0,
+        flags: KvmCpuid2Flag::SIGNIFCANT_INDEX,
+        eax: 0x0,
+        ebx: 0x0,
+        ecx: 0x0,
+        edx: 0x6d,
+        padding: [0; 3],
+    },
+    (
+        CpuidIn {
+            func: 0xb,
+            index: Some(0),
+        },
+        CpuidResult {
+            eax: 0x0,
+            ebx: 0x0,
+            ecx: 0x0,
+            edx: 0x6d,
+        }
+    )
+)]
+fn test_convert_cpuid_entry(
+    #[case] value: KvmCpuidEntry2,
+    #[case] expected: (CpuidIn, CpuidResult),
+) {
+    let got: (CpuidIn, CpuidResult) = value.into();
+    assert_eq!(got, expected)
+}
diff --git a/alioth/src/hv/kvm/vm/vm.rs b/alioth/src/hv/kvm/vm/vm.rs
@@ -19,6 +19,8 @@ mod aarch64;
 #[path = "vm_x86_64/vm_x86_64.rs"]
 mod x86_64;
 
+#[cfg(target_arch = "x86_64")]
+use std::arch::x86_64::CpuidResult;
 use std::collections::HashMap;
 use std::fmt::{self, Display, Formatter};
 use std::io::ErrorKind;
@@ -35,8 +37,12 @@ use libc::{EFD_CLOEXEC, EFD_NONBLOCK, SIGRTMIN, eventfd};
 use parking_lot::{Mutex, RwLock};
 use snafu::ResultExt;
 
+#[cfg(target_arch = "x86_64")]
+use crate::arch::cpuid::CpuidIn;
 #[cfg(target_arch = "x86_64")]
 use crate::arch::sev::{SevPolicy, SnpPageType, SnpPolicy};
+#[cfg(target_arch = "x86_64")]
+use crate::arch::tdx::TdAttr;
 use crate::ffi;
 #[cfg(not(target_arch = "x86_64"))]
 use crate::hv::IrqSender;
@@ -755,6 +761,11 @@ impl Vm for KvmVm {
         KvmVm::snp_launch_finish(self)
     }
 
+    #[cfg(target_arch = "x86_64")]
+    fn tdx_init_vm(&self, attr: TdAttr, cpuids: &HashMap<CpuidIn, CpuidResult>) -> Result<()> {
+        KvmVm::tdx_init_vm(self, attr, cpuids)
+    }
+
     #[cfg(target_arch = "aarch64")]
     fn create_gic_v2(&self, distributor_base: u64, cpu_interface_base: u64) -> Result<Self::GicV2> {
         aarch64::KvmGicV2::new(self, distributor_base, cpu_interface_base)
diff --git a/alioth/src/hv/kvm/vm/vm_x86_64/tdx.rs b/alioth/src/hv/kvm/vm/vm_x86_64/tdx.rs
@@ -12,9 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::arch::x86_64::CpuidResult;
+use std::collections::HashMap;
+use std::mem::MaybeUninit;
+
+use crate::arch::cpuid::CpuidIn;
+use crate::arch::tdx::TdAttr;
 use crate::hv::Result;
+use crate::hv::kvm::tdx::tdx_op;
 use crate::hv::kvm::vm::KvmVm;
-use crate::sys::kvm::{KvmCap, KvmHypercall};
+use crate::sys::kvm::{KvmCap, KvmCpuid2Flag, KvmCpuidEntry2, KvmHypercall};
+use crate::sys::tdx::{KvmTdxCapabilities, KvmTdxCmdId, KvmTdxInitVm};
 
 impl KvmVm {
     pub fn tdx_init(&self) -> Result<()> {
@@ -23,4 +31,55 @@ impl KvmVm {
         self.vm.enable_cap(KvmCap::X86_APIC_BUS_CYCLES_NS, 40)?;
         Ok(())
     }
+
+    fn tdx_get_capabilities(&self) -> Result<Box<KvmTdxCapabilities>> {
+        let mut caps: Box<KvmTdxCapabilities> =
+            Box::new(unsafe { MaybeUninit::zeroed().assume_init() });
+        caps.cpuid.nent = caps.cpuid.entries.len() as u32;
+        tdx_op(&self.vm.fd, KvmTdxCmdId::CAPABILITIES, 0, Some(&mut *caps))?;
+        Ok(caps)
+    }
+
+    pub fn tdx_init_vm(&self, attr: TdAttr, cpuids: &HashMap<CpuidIn, CpuidResult>) -> Result<()> {
+        let mut init: Box<KvmTdxInitVm> = Box::new(unsafe { MaybeUninit::zeroed().assume_init() });
+        init.attributes = attr;
+
+        let caps = self.tdx_get_capabilities()?;
+        let convert = |e: &KvmCpuidEntry2| {
+            let (mut in_, out) = From::from(*e);
+            if in_.index.is_none() {
+                in_.index = Some(0)
+            }
+            (in_, out)
+        };
+        let caps_cpuid = caps.cpuid.entries.iter().take(caps.cpuid.nent as usize);
+        let caps_cpuid: HashMap<_, _> = caps_cpuid.map(convert).collect();
+        for (in_, out) in cpuids {
+            let cap_cpuid_in = CpuidIn {
+                func: in_.func,
+                index: in_.index.or(Some(0)),
+            };
+            let Some(cap_out) = caps_cpuid.get(&cap_cpuid_in) else {
+                continue;
+            };
+
+            let entry = &mut init.cpuid.entries[init.cpuid.nent as usize];
+            entry.function = in_.func;
+            entry.index = in_.index.unwrap_or(0);
+            entry.flags = if in_.index.is_some() {
+                KvmCpuid2Flag::SIGNIFCANT_INDEX
+            } else {
+                KvmCpuid2Flag::empty()
+            };
+            entry.eax = out.eax & cap_out.eax;
+            entry.ebx = out.ebx & cap_out.ebx;
+            entry.ecx = out.ecx & cap_out.ecx;
+            entry.edx = out.edx & cap_out.edx;
+
+            init.cpuid.nent += 1;
+        }
+
+        tdx_op(&self.vm.fd, KvmTdxCmdId::INIT_VM, 0, Some(&mut *init))?;
+        Ok(())
+    }
 }
diff --git a/alioth/src/sys/linux/tdx.rs b/alioth/src/sys/linux/tdx.rs
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use crate::arch::tdx::TdAttr;
-use crate::sys::kvm::KvmCpuid2;
+use crate::sys::kvm::{KVM_MAX_CPUID_ENTRIES, KvmCpuid2};
 use crate::{bitflags, consts};
 
 consts! {
@@ -39,7 +39,7 @@ pub struct KvmTdxCmd {
 
 #[repr(C)]
 #[derive(Debug, Clone)]
-pub struct KvmTdxCapabilities<const N: usize> {
+pub struct KvmTdxCapabilities<const N: usize = KVM_MAX_CPUID_ENTRIES> {
     pub supported_attrs: TdAttr,
     pub supported_xfam: u64,
     pub kernel_tdvmcallinfo_1_r11: u64,
@@ -52,8 +52,8 @@ pub struct KvmTdxCapabilities<const N: usize> {
 
 #[repr(C)]
 #[derive(Debug, Clone)]
-pub struct KvmTdxInitVm<const N: usize> {
-    pub attributes: u64,
+pub struct KvmTdxInitVm<const N: usize = KVM_MAX_CPUID_ENTRIES> {
+    pub attributes: TdAttr,
     pub xfam: u64,
     pub mrconfigid: [u8; 48],
     pub mrowner: [u8; 48],

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@`
`13`	`13`	`// limitations under the License.`
`14`	`14`
`15`	`15`	`mod sev;`
	`16`	`+mod tdx;`
`16`	`17`
`17`	`18`	`use std::arch::x86_64::{__cpuid, CpuidResult};`
`18`	`19`	`use std::collections::HashMap;`
`@@ -321,7 +322,7 @@ where`
`321`	`322`	`match coco {`
`322`	`323`	`Coco::AmdSev { policy } => self.sev_init(*policy, memory)?,`
`323`	`324`	`Coco::AmdSnp { policy } => self.snp_init(*policy, memory)?,`
`324`		`- Coco::IntelTdx { attr } => todo!("Intel TDX {attr:?}"),`
	`325`	`+ Coco::IntelTdx { attr } => self.tdx_init(*attr, memory)?,`
`325`	`326`	`}`
`326`	`327`	`Ok(())`
`327`	`328`	`}`