Add new TDCALL functions to support TD partitioning

td-shim/src/bin/td-shim/td/tdx_mailbox.rs

@@ -10,7 +10,7 @@ use core::cmp::min;
 use core::ops::RangeInclusive;
 use td_exception::idt::DescriptorTablePointer;
 use td_layout::memslice::{get_mem_slice_mut, SliceType};
-use tdx_tdcall::{self, tdx};
+use tdx_tdcall::{self, tdx::*};
 
 use crate::asm::{ap_relocated_func_addr, ap_relocated_func_size};
 
@@ -23,8 +23,6 @@ static AP_TEMP_STACK: [u8; AP_TEMP_STACK_TOTAL_SIZE] = [0; AP_TEMP_STACK_TOTAL_S
 
 const ACCEPT_CHUNK_SIZE: u64 = 0x2000000;
 const ACCEPT_PAGE_SIZE: u64 = 0x200000;
-const PAGE_SIZE_2M: u64 = 0x200000;
-const PAGE_SIZE_4K: u64 = 0x1000;
 const MAILBOX_SIZE: usize = 0x1000;
 
 #[derive(Debug)]
@@ -169,32 +167,6 @@ pub fn ap_assign_work(cpu_index: u32, stack_top: u64, entry: u32) {
     wait_for_ap_response(&mut mail_box);
 }
 
-fn td_accept_pages(address: u64, pages: u64, page_size: u64) {
-    for i in 0..pages {
-        let mut accept_addr = address + i * page_size;
-        let accept_level = if page_size == PAGE_SIZE_2M { 1 } else { 0 };
-        match tdx::tdcall_accept_page(accept_addr | accept_level) {
-            Ok(()) => {}
-            Err(e) => {
-                if let tdx_tdcall::TdCallError::LeafSpecific(error_code) = e {
-                    if error_code == tdx_tdcall::TDCALL_STATUS_PAGE_SIZE_MISMATCH {
-                        if page_size == PAGE_SIZE_2M {
-                            td_accept_pages(accept_addr, 512, PAGE_SIZE_4K);
-                            continue;
-                        }
-                    } else if error_code == tdx_tdcall::TDCALL_STATUS_PAGE_ALREADY_ACCEPTED {
-                        continue;
-                    }
-                }
-                panic!(
-                    "Accept Page Error: 0x{:x}, page_size: {}\n",
-                    accept_addr, page_size
-                );
-            }
-        }
-    }
-}
-
 extern "win64" fn parallel_accept_memory(cpu_index: u64) {
     // Safety:
     // During this state, all the BSPs/APs are accessing the mailbox in shared immutable mode.

tdx-tdcall/src/asm/tdcall.asm

@@ -30,6 +30,7 @@ asm_td_call:
         push rbx
         push rsi
         push rdi
+        push rcx
 
         # Use RDI to save RCX value
         mov rdi, rcx
@@ -38,7 +39,7 @@ asm_td_call:
         test rdi, rdi
         jz td_call_exit
 
-        # Copy the input operands from memory to registers 
+        # Copy the input operands from memory to registers
         mov rax, [rdi + TDCALL_ARG_RAX]
         mov rcx, [rdi + TDCALL_ARG_RCX]
         mov rdx, [rdi + TDCALL_ARG_RDX]
@@ -52,9 +53,11 @@ asm_td_call:
         # tdcall
         .byte 0x66,0x0f,0x01,0xcc
 
-        # Exit if tdcall reports failure.
-        test rax, rax
-        jnz td_call_exit
+        # On TDVM exit, RDI loses the RCX value (TdVmcallArgs struct) that was saved earlier.
+        # so pop saved RCX back to RDI.
+        # Note: TDG.VP.ENTER uses RDI to store CS base address but since it is not used, RDI is
+        # repurposed here.
+        pop rdi
 
         # Copy the output operands from registers to the struct
         mov [rdi + TDCALL_ARG_RAX], rax
@@ -68,6 +71,8 @@ asm_td_call:
         mov [rdi + TDCALL_ARG_R13], r13
 
 td_call_exit:
+        # Restore saved RCX value
+        mov rcx, rdi
         # Pop out saved registers from stack
         pop rdi
         pop rsi

tdx-tdcall/src/lib.rs

@@ -39,9 +39,17 @@ const TDCALL_TDEXTENDRTMR: u64 = 2;
 const TDCALL_TDGETVEINFO: u64 = 3;
 const TDCALL_TDREPORT: u64 = 4;
 const TDCALL_TDACCEPTPAGE: u64 = 6;
+const TDCALL_VM_RD: u64 = 7;
+const TDCALL_VM_WR: u64 = 8;
+const TDCALL_VP_RD: u64 = 9;
+const TDCALL_VP_WR: u64 = 10;
 const TDCALL_SYS_RD: u64 = 11;
 const TDCALL_SERVTD_RD: u64 = 18;
 const TDCALL_SERVTD_WR: u64 = 20;
+const TDCALL_MEM_PAGE_ATTR_WR: u64 = 24;
+const TDCALL_VP_ENTER: u64 = 25;
+const TDCALL_VP_INVEPT: u64 = 26;
+const TDCALL_VP_INVVPID: u64 = 27;
 
 // GTDG.VP.VMCALL leaf sub-function numbers
 const TDVMCALL_CPUID: u64 = 0x0000a;
@@ -57,6 +65,7 @@ const TDVMCALL_SERVICE: u64 = 0x10005;
 
 // TDCALL completion status code
 const TDCALL_STATUS_SUCCESS: u64 = 0;
+
 // leaf-specific completion status code
 pub const TDCALL_STATUS_PAGE_ALREADY_ACCEPTED: u64 = 0x00000B0A00000000;
 pub const TDCALL_STATUS_PAGE_SIZE_MISMATCH: u64 = 0xC0000B0B00000001;

tdx-tdcall/src/tdx.rs

@@ -21,6 +21,8 @@ use crate::*;
 const IO_READ: u64 = 0;
 const IO_WRITE: u64 = 1;
 const TARGET_TD_UUID_NUM: usize = 4;
+pub const PAGE_SIZE_4K: u64 = 0x1000;
+pub const PAGE_SIZE_2M: u64 = 0x200000;
 
 /// SHA384 digest value extended to RTMR
 /// Both alignment and size are 64 bytes.
@@ -575,13 +577,83 @@ pub fn tdcall_accept_page(address: u64) -> Result<(), TdCallError> {
         ..Default::default()
     };
 
-    let ret = td_call(&mut args);
+    const MAX_RETRIES_ACCEPT_PAGE: usize = 5;
+    let mut retry_counter = 0;
+    let mut ret = 0;
 
-    if ret != TDCALL_STATUS_SUCCESS {
-        return Err(ret.into());
+    while retry_counter < MAX_RETRIES_ACCEPT_PAGE {
+        ret = td_call(&mut args);
+
+        if ret == TDCALL_STATUS_SUCCESS {
+            return Ok(());
+        } else if TdCallError::TdxExitReasonOperandBusy != ret.into() {
+            return Err(ret.into());
+        }
+
+        retry_counter += 1;
     }
 
-    Ok(())
+    return Err(ret.into());
+}
+
+/// Accept a range of private pages and initialize the pages to zeros using the TD ephemeral
+/// private key.
+///
+/// This function is a wrapper to `tdcall_accept_page()`.
+pub fn td_accept_pages(address: u64, pages: u64, page_size: u64) {
+    for i in 0..pages {
+        let accept_addr = address + i * page_size;
+        let accept_level = if page_size == PAGE_SIZE_2M { 1 } else { 0 };
+        match tdcall_accept_page(accept_addr | accept_level) {
+            Ok(()) => {}
+            Err(e) => {
+                if let TdCallError::LeafSpecific(error_code) = e {
+                    if error_code == TDCALL_STATUS_PAGE_SIZE_MISMATCH {
+                        if page_size == PAGE_SIZE_2M {
+                            td_accept_pages(accept_addr, 512, PAGE_SIZE_4K);
+                            continue;
+                        }
+                    } else if error_code == TDCALL_STATUS_PAGE_ALREADY_ACCEPTED {
+                        continue;
+                    }
+                }
+                panic!(
+                    "Accept Page Error: 0x{:x}, page_size: {}, err {:?}\n",
+                    accept_addr, page_size, e
+                );
+            }
+        }
+    }
+}
+
+/// Accept a range of either 4K normal pages or 2M huge pages. This is basically a wrapper over
+/// td_accept_pages and initializes the pages to zero using the TD ephemeral private key.
+pub fn td_accept_memory(address: u64, len: u64) {
+    let mut start = address;
+    let end = address + len;
+
+    while start < end {
+        let remaining = end - start;
+
+        // Try larger accepts first to keep 1G/2M Secure EPT entries
+        // where possible and speeds up process by cutting number of
+        // tdcalls (if successful).
+        if remaining >= PAGE_SIZE_2M && (start & (PAGE_SIZE_2M - 1)) == 0 {
+            let npages = remaining >> 21;
+            td_accept_pages(start, npages, PAGE_SIZE_2M);
+            start += npages << 21;
+        } else if remaining >= PAGE_SIZE_4K && (start & (PAGE_SIZE_4K - 1)) == 0 {
+            let mut npages = remaining >> 12;
+            // Try to consume in 4K chunks until 2M aligned.
+            if remaining >= PAGE_SIZE_2M {
+                npages = (PAGE_SIZE_2M - (start & (PAGE_SIZE_2M - 1))) >> 12;
+            }
+            td_accept_pages(start, npages, PAGE_SIZE_4K);
+            start += npages << 12;
+        } else {
+            panic!("Accept Memory Error: 0x{:x}, length: {}\n", address, len);
+        }
+    }
 }
 
 /// Get the guest physical address (GPA) width via TDG.VP.INFO
@@ -695,6 +767,177 @@ pub fn tdcall_sys_rd(field_identifier: u64) -> core::result::Result<(u64, u64),
     Ok((args.rdx, args.r8))
 }
 
+/// Read a VCPU-scope metadata field (control structure field) of a TD.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VP.RD Leaf'.
+pub fn tdcall_vp_read(field: u64) -> Result<(u64, u64), TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VP_RD,
+        rdx: field,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok((args.rdx, args.r8))
+}
+
+/// Write a VCPU-scope metadata field (control structure field) of a TD.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VP.WR Leaf'.
+pub fn tdcall_vp_write(field: u64, value: u64, mask: u64) -> Result<u64, TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VP_WR,
+        rdx: field,
+        r8: value,
+        r9: mask,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok(args.r8)
+}
+
+/// Invalidate mappings in the translation lookaside buffers (TLBs) and paging-structure caches
+/// for a specified L2 VM and a specified list of 4KB page linear addresses.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VP.INVVPID Leaf'.
+pub fn tdcall_vp_invvpid(flags: u64, gla: u64) -> Result<u64, TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VP_INVVPID,
+        rcx: flags,
+        rdx: gla,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok(args.rdx)
+}
+
+/// Invalidate cached EPT translations for selected L2 VM.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VP.INVEPT Leaf'.
+pub fn tdcall_vp_invept(vm_flags: u64) -> Result<(), TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VP_INVEPT,
+        rcx: vm_flags,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok(())
+}
+
+/// Enter L2 VCPU operation.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VP.ENTER Leaf'.
+pub fn tdcall_vp_enter(vm_flags: u64, gpa: u64) -> TdcallArgs {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VP_ENTER,
+        rcx: vm_flags,
+        rdx: gpa,
+        ..Default::default()
+    };
+
+    td_call(&mut args);
+
+    args
+}
+
+/// Read a TD-scope metadata field (control structure field) of a TD.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VM.RD Leaf'.
+pub fn tdcall_vm_read(field: u64, version: u8) -> Result<(u64, u64), TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VM_RD | (version as u64) << 16,
+        rdx: field,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok((args.rdx, args.r8))
+}
+
+/// Write a TD-scope metadata field (control structure field) of a TD.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.VM.WR Leaf'.
+pub fn tdcall_vm_write(field: u64, value: u64, mask: u64) -> Result<u64, TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_VM_WR,
+        rdx: field,
+        r8: value,
+        r9: mask,
+        ..Default::default()
+    };
+
+    let ret = td_call(&mut args);
+
+    if ret != TDCALL_STATUS_SUCCESS {
+        return Err(ret.into());
+    }
+
+    Ok(args.r8)
+}
+
+/// Write the attributes of a private page.  Create or remove L2 page aliases as required.
+///
+/// Details can be found in TDX Module v1.5 ABI spec section 'TDG.MEM.PAGE.ATTR.WR Leaf'.
+pub fn tdcall_mem_page_attr_wr(
+    gpa_mapping: u64,
+    gpa_attr: u64,
+    attr_flags: u64,
+) -> Result<(u64, u64), TdCallError> {
+    let mut args = TdcallArgs {
+        rax: TDCALL_MEM_PAGE_ATTR_WR,
+        rcx: gpa_mapping,
+        rdx: gpa_attr,
+        r8: attr_flags,
+        ..Default::default()
+    };
+
+    const MAX_RETRIES_ATTR_WR: usize = 5;
+    let mut retry_counter = 0;
+    let mut ret = 0;
+
+    while retry_counter < MAX_RETRIES_ATTR_WR {
+        ret = td_call(&mut args);
+
+        if ret == TDCALL_STATUS_SUCCESS {
+            return Ok((args.rcx, args.rdx));
+        } else if TdCallError::TdxExitReasonOperandBusy != ret.into() {
+            return Err(ret.into());
+        }
+
+        retry_counter += 1;
+    }
+
+    return Err(ret.into());
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;