Implement remaining functionality to boot TD

Cargo.toml

@@ -14,7 +14,8 @@ rust-version = "1.71"
 
 [dependencies]
 bitflags = "2.4.2"
-kvm-bindings = "0.7.0"
-kvm-ioctls = "0.16.0"
+kvm-bindings = { git = "https://github.com/jakecorrenti/kvm-bindings.git", branch = "tdx", features = ["fam-wrappers"] }
+kvm-ioctls = { git = "https://github.com/jakecorrenti/kvm-ioctls.git", branch = "tdx" }
+libc = "0.2.155"
 uuid = "1.8.0"
 vmm-sys-util = "0.12.1"

src/launch/linux.rs

@@ -1,20 +1,24 @@
 // SPDX-License-Identifier: Apache-2.0
 
-pub const NR_CPUID_CONFIGS: usize = 12;
+use std::marker::PhantomData;
+
+pub const NR_CPUID_CONFIGS: usize = 24;
 
 /// Trust Domain eXtensions sub-ioctl() commands
 #[repr(u32)]
 pub enum CmdId {
     GetCapabilities,
     InitVm,
     InitVcpu,
+    InitMemRegion,
+    FinalizeVm,
 }
 
 /// Contains information for the sub-ioctl() command to be run. This is
 /// equivalent to `struct kvm_tdx_cmd` in the kernel.
 #[derive(Default)]
 #[repr(C)]
-pub struct Cmd {
+pub struct Cmd<'a, T: 'a> {
     /// TDX command identifier
     pub id: u32,
 
@@ -31,6 +35,21 @@ pub struct Cmd {
 
     /// Reserved.
     pub _unused: u64,
+
+    _phantom: PhantomData<&'a T>,
+}
+
+impl<'a, T: 'a> Cmd<'a, T> {
+    pub fn from(id: CmdId, data: &'a T) -> Self {
+        Self {
+            id: id as u32,
+            flags: 0,
+            data: data as *const T as _,
+            error: 0,
+            _unused: 0,
+            _phantom: PhantomData,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -159,18 +178,6 @@ impl Default for Capabilities {
     }
 }
 
-impl From<&Capabilities> for Cmd {
-    fn from(caps: &Capabilities) -> Self {
-        Self {
-            id: CmdId::GetCapabilities as u32,
-            flags: 0,
-            data: caps as *const Capabilities as _,
-            error: 0,
-            _unused: 0,
-        }
-    }
-}
-
 /// TDX specific VM initialization information
 #[derive(Debug)]
 #[repr(C)]
@@ -227,14 +234,15 @@ impl Default for InitVm {
     }
 }
 
-impl From<&InitVm> for Cmd {
-    fn from(init_vm: &InitVm) -> Self {
-        Self {
-            id: CmdId::InitVm as u32,
-            flags: 0,
-            data: init_vm as *const InitVm as _,
-            error: 0,
-            _unused: 0,
-        }
-    }
+#[repr(C)]
+#[derive(Debug)]
+pub struct TdxInitMemRegion {
+    /// Host physical address of the target page to be added to the TD
+    pub source_addr: u64,
+
+    /// Guest physical address to be mapped
+    pub gpa: u64,
+
+    /// Number of pages to be mapped
+    pub nr_pages: u64,
 }

src/launch/mod.rs

@@ -2,26 +2,21 @@
 
 mod linux;
 
-use kvm_bindings::{kvm_enable_cap, KVM_CAP_MAX_VCPUS, KVM_CAP_SPLIT_IRQCHIP};
-use linux::{Capabilities, Cmd, CpuidConfig, InitVm, TdxError};
+use kvm_bindings::{kvm_enable_cap, KVM_CAP_MAX_VCPUS};
+use linux::{Capabilities, Cmd, CmdId, CpuidConfig, InitVm, TdxError};
 
 use bitflags::bitflags;
-use kvm_ioctls::{Kvm, VmFd};
-use std::arch::x86_64;
+use kvm_ioctls::VmFd;
 
 // Defined in linux/arch/x86/include/uapi/asm/kvm.h
-const KVM_X86_TDX_VM: u64 = 2;
+pub const KVM_X86_TDX_VM: u64 = 5;
 
 /// Handle to the TDX VM file descriptor
-pub struct TdxVm {
-    pub fd: VmFd,
-}
+pub struct TdxVm {}
 
 impl TdxVm {
     /// Create a new TDX VM with KVM
-    pub fn new(kvm_fd: &Kvm, max_vcpus: u64) -> Result<Self, TdxError> {
-        let vm_fd = kvm_fd.create_vm_with_type(KVM_X86_TDX_VM)?;
-
+    pub fn new(vm_fd: &VmFd, max_vcpus: u64) -> Result<Self, TdxError> {
         // TDX requires that MAX_VCPUS and SPLIT_IRQCHIP be set
         let mut cap: kvm_enable_cap = kvm_enable_cap {
             cap: KVM_CAP_MAX_VCPUS,
@@ -30,20 +25,20 @@ impl TdxVm {
         cap.args[0] = max_vcpus;
         vm_fd.enable_cap(&cap).unwrap();
 
-        cap.cap = KVM_CAP_SPLIT_IRQCHIP;
-        cap.args[0] = 24;
+        cap.cap = kvm_bindings::KVM_CAP_X2APIC_API;
+        cap.args[0] = (1 << 0) | (1 << 1);
         vm_fd.enable_cap(&cap).unwrap();
 
-        Ok(Self { fd: vm_fd })
+        Ok(Self {})
     }
 
     /// Retrieve information about the Intel TDX module
-    pub fn get_capabilities(&self) -> Result<TdxCapabilities, TdxError> {
+    pub fn get_capabilities(&self, fd: &VmFd) -> Result<TdxCapabilities, TdxError> {
         let caps = Capabilities::default();
-        let mut cmd: Cmd = Cmd::from(&caps);
+        let mut cmd: Cmd<Capabilities> = Cmd::from(CmdId::GetCapabilities, &caps);
 
         unsafe {
-            self.fd.encrypt_op(&mut cmd)?;
+            fd.encrypt_op(&mut cmd)?;
         }
 
         Ok(TdxCapabilities {
@@ -61,57 +56,58 @@ impl TdxVm {
     }
 
     /// Do additional VM initialization that is specific to Intel TDX
-    pub fn init_vm(&self, kvm_fd: &Kvm, caps: &TdxCapabilities) -> Result<(), TdxError> {
-        let cpuid = kvm_fd
-            .get_supported_cpuid(kvm_bindings::KVM_MAX_CPUID_ENTRIES)
-            .unwrap();
+    pub fn init_vm(&self, fd: &VmFd, cpuid: kvm_bindings::CpuId) -> Result<(), TdxError> {
         let mut cpuid_entries: Vec<kvm_bindings::kvm_cpuid_entry2> = cpuid.as_slice().to_vec();
 
         // resize to 256 entries to make sure that InitVm is 8KB
         cpuid_entries.resize(256, kvm_bindings::kvm_cpuid_entry2::default());
 
-        // hex for Ob1100000001011111111 based on the XSAVE state-components architecture
-        let xcr0_mask = 0x602ff;
-        // hex for 0b11111110100000000 based on the XSAVE state-components architecture
-        let xss_mask = 0x1FD00;
-
-        let xfam_fixed0 = caps.xfam.fixed0.bits();
-        let xfam_fixed1 = caps.xfam.fixed1.bits();
-
-        // patch cpuid
-        for entry in cpuid_entries.as_mut_slice() {
-            // mandatory patches for TDX based on XFAM values reported by TdxCapabilities
-            match entry.index {
-                // XSAVE features and state-components
-                0xD => {
-                    if entry.index == 0 {
-                        // XSAVE XCR0 LO
-                        entry.eax &= (xfam_fixed0 as u32) & (xcr0_mask as u32);
-                        entry.eax |= (xfam_fixed1 as u32) & (xcr0_mask as u32);
-                        // XSAVE XCR0 HI
-                        entry.edx &= ((xfam_fixed0 & xcr0_mask) >> 32) as u32;
-                        entry.edx |= ((xfam_fixed1 & xcr0_mask) >> 32) as u32;
-                    } else if entry.index == 1 {
-                        // XSAVE XCR0 LO
-                        entry.ecx &= (xfam_fixed0 as u32) & (xss_mask as u32);
-                        entry.ecx |= (xfam_fixed1 as u32) & (xss_mask as u32);
-                        // XSAVE XCR0 HI
-                        entry.edx &= ((xfam_fixed0 & xss_mask) >> 32) as u32;
-                        entry.edx |= ((xfam_fixed1 & xss_mask) >> 32) as u32;
-                    }
-                }
-                0x8000_0008 => {
-                    // host physical address bits supported
-                    let phys_bits = unsafe { x86_64::__cpuid(0x8000_0008).eax } & 0xff;
-                    entry.eax = (entry.eax & 0xffff_ff00) | (phys_bits & 0xff);
-                }
-                _ => (),
-            }
+        let init_vm = InitVm::new(&cpuid_entries);
+        let mut cmd: Cmd<InitVm> = Cmd::from(CmdId::InitVm, &init_vm);
+        unsafe {
+            fd.encrypt_op(&mut cmd)?;
         }
 
-        let mut cmd = Cmd::from(&InitVm::new(&cpuid_entries));
+        Ok(())
+    }
+
+    /// Encrypt a memory continuous region
+    pub fn init_mem_region(
+        &self,
+        fd: &VmFd,
+        gpa: u64,
+        nr_pages: u64,
+        attributes: u32,
+        source_addr: u64,
+    ) -> Result<(), TdxError> {
+        const TDVF_SECTION_ATTRIBUTES_MR_EXTEND: u32 = 1u32 << 0;
+        let mem_region = linux::TdxInitMemRegion {
+            source_addr,
+            gpa,
+            nr_pages,
+        };
+
+        let mut cmd: Cmd<linux::TdxInitMemRegion> = Cmd::from(CmdId::InitMemRegion, &mem_region);
+
+        // determines if we also extend the measurement
+        cmd.flags = if attributes & TDVF_SECTION_ATTRIBUTES_MR_EXTEND > 0 {
+            1
+        } else {
+            0
+        };
+
         unsafe {
-            self.fd.encrypt_op(&mut cmd)?;
+            fd.encrypt_op(&mut cmd)?;
+        }
+
+        Ok(())
+    }
+
+    /// Complete measurement of the initial TD contents and mark it ready to run
+    pub fn finalize(&self, fd: &VmFd) -> Result<(), TdxError> {
+        let mut cmd: Cmd<u64> = Cmd::from(CmdId::FinalizeVm, &0);
+        unsafe {
+            fd.encrypt_op(&mut cmd)?;
         }
 
         Ok(())
@@ -274,20 +270,12 @@ ioctl_iowr_nr!(
     std::os::raw::c_ulong
 );
 
-pub struct TdxVcpu<'a> {
-    pub fd: &'a mut kvm_ioctls::VcpuFd,
-}
+pub struct TdxVcpu {}
 
-impl<'a> TdxVcpu<'a> {
-    pub fn init(&self, hob_address: u64) -> Result<(), TdxError> {
-        let mut cmd = Cmd {
-            id: linux::CmdId::InitVcpu as u32,
-            flags: 0,
-            data: hob_address as *const u64 as _,
-            error: 0,
-            _unused: 0,
-        };
-        let ret = unsafe { ioctl::ioctl_with_mut_ptr(self.fd, KVM_MEMORY_ENCRYPT_OP(), &mut cmd) };
+impl TdxVcpu {
+    pub fn init(fd: &kvm_ioctls::VcpuFd, hob_address: u64) -> Result<(), TdxError> {
+        let mut cmd: Cmd<u64> = Cmd::from(CmdId::InitVcpu, &hob_address);
+        let ret = unsafe { ioctl::ioctl_with_mut_ptr(fd, KVM_MEMORY_ENCRYPT_OP(), &mut cmd) };
         if ret < 0 {
             // can't return `ret` because it will just return -1 and not give the error
             // code. `cmd.error` will also just be 0.
@@ -296,24 +284,3 @@ impl<'a> TdxVcpu<'a> {
         Ok(())
     }
 }
-
-impl<'a> TryFrom<(&'a mut kvm_ioctls::VcpuFd, &'a mut kvm_ioctls::Kvm)> for TdxVcpu<'a> {
-    type Error = TdxError;
-
-    fn try_from(
-        value: (&'a mut kvm_ioctls::VcpuFd, &'a mut kvm_ioctls::Kvm),
-    ) -> Result<Self, Self::Error> {
-        // need to enable the X2APIC bit for CPUID[0x1] so that the kernel can call
-        // KVM_SET_MSRS(MSR_IA32_APIC_BASE) without failing
-        let mut cpuid = value
-            .1
-            .get_supported_cpuid(kvm_bindings::KVM_MAX_CPUID_ENTRIES)?;
-        for entry in cpuid.as_mut_slice().iter_mut() {
-            if entry.index == 0x1 {
-                entry.ecx &= 1 << 21;
-            }
-        }
-        value.0.set_cpuid2(&cpuid)?;
-        Ok(Self { fd: value.0 })
-    }
-}