td-shim: support loading payload for each vCPU

When feature `multi-payload` is enabled, td-shim will load a payload
binary for each vCPU, and it will also evenly divide the largest
contiguous available memory among each vCPU, generating the e820 table
and HOB for them.

Signed-off-by: Jiaqi Gao <[email protected]>

Author: gaojiaqi7

td-shim/Cargo.toml

@@ -51,6 +51,7 @@ tdx = ["tdx-tdcall", "td-exception/tdx", "td-logger/tdx", "x86"]
 lazy-accept = ["tdx"]
 ring-hash = ["cc-measurement/ring"]
 sha2-hash = ["cc-measurement/sha2"]
+multi-payload = []
 main = [
     "log",
     "td-loader",

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

@@ -8,7 +8,7 @@ use td_shim::e820::{E820Entry, E820Type};
 // Linux BootParam supports 128 e820 entries, so...
 const MAX_E820_ENTRY: usize = 128;
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub struct E820Table {
     entries: [E820Entry; MAX_E820_ENTRY],
     num_entries: usize,

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

@@ -2,6 +2,8 @@
 //
 // SPDX-License-Identifier: BSD-2-Clause-Patent
 
+use core::convert::TryFrom;
+
 use td_layout::memslice;
 use td_loader::elf;
 use td_loader::elf64::ProgramHeader;
@@ -11,9 +13,23 @@ use crate::memory::Memory;
 
 const SIZE_4KB: u64 = 0x00001000u64;
 
+#[repr(u8)]
+#[derive(Debug, Copy, Clone)]
 pub enum ExecutablePayloadType {
-    Elf,
-    PeCoff,
+    Elf = 0,
+    PeCoff = 1,
+}
+
+impl TryFrom<u8> for ExecutablePayloadType {
+    type Error = ();
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(ExecutablePayloadType::Elf),
+            1 => Ok(ExecutablePayloadType::PeCoff),
+            _ => Err(()),
+        }
+    }
 }
 
 pub struct PayloadRelocationInfo {
@@ -79,6 +95,20 @@ pub fn find_and_report_entry_point(
     }
 }
 
+pub fn jump_to_payload_entry(entry: u64, image_type: u8, hob_addr: u64, payload_addr: u64) {
+    let image_type = ExecutablePayloadType::try_from(image_type).expect("Unknown image type");
+    match image_type {
+        ExecutablePayloadType::Elf => {
+            let entry = unsafe { core::mem::transmute::<u64, extern "sysv64" fn(u64, u64)>(entry) };
+            entry(hob_addr, payload_addr);
+        }
+        ExecutablePayloadType::PeCoff => {
+            let entry = unsafe { core::mem::transmute::<u64, extern "win64" fn(u64, u64)>(entry) };
+            entry(hob_addr, payload_addr);
+        }
+    };
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -7,15 +7,20 @@
 #![cfg_attr(not(test), no_main)]
 #![allow(unused_imports)]
 
+#[macro_use]
 extern crate alloc;
 use alloc::boxed::Box;
 use alloc::vec::Vec;
 use asm::{empty_exception_handler, empty_exception_handler_size};
 use core::ffi::c_void;
 use core::mem::size_of;
 use core::panic::PanicInfo;
+use e820::E820Table;
 use memory::Memory;
+use td::ap_set_payload;
 use td_exception::idt::{load_idtr, DescriptorTablePointer, Idt, IdtEntry};
+use td_layout::runtime::exec::PAYLOAD_SIZE;
+use td_shim::e820::E820Type;
 use td_shim::event_log::CCEL_CC_TYPE_TDX;
 use x86_64::instructions::hlt;
 use x86_64::registers::segmentation::{Segment, CS};
@@ -28,7 +33,7 @@ use zerocopy::{AsBytes, ByteSlice, FromBytes};
 use cc_measurement::{log::CcEventLogWriter, EV_EFI_HANDOFF_TABLES2, EV_PLATFORM_CONFIG_FLAGS};
 use td_layout::build_time::{self, *};
 use td_layout::memslice::{self, SliceType};
-use td_layout::RuntimeMemoryLayout;
+use td_layout::{mailbox, RuntimeMemoryLayout};
 use td_shim::event_log::{log_hob_list, log_payload_binary, log_payload_parameter};
 use td_shim::{
     speculation_barrier, PayloadInfo, TdPayloadInfoHobType, TD_ACPI_TABLE_HOB_GUID,
@@ -151,6 +156,7 @@ pub extern "win64" fn _start(
         &mut mem,
         &mut td_event_log,
         &dynamic_info.acpi_tables,
+        num_vcpus,
     );
 
     panic!("payload entry() should not return here, deadloop!!!");
@@ -207,6 +213,7 @@ fn boot_builtin_payload(
     mem: &mut memory::Memory,
     event_log: &mut CcEventLogWriter,
     acpi_tables: &Vec<&[u8]>,
+    num_vcpus: u32,
 ) {
     // Get and parse image file from the payload firmware volume.
     let fv_buffer = memslice::get_mem_slice(memslice::SliceType::ShimPayload);
@@ -230,6 +237,10 @@ fn boot_builtin_payload(
     // Create an EV_SEPARATOR event to mark the end of the td-shim events
     event_log.create_seperator();
 
+    if cfg!(feature = "multi-payload") {
+        mp::launch_payload_per_cpu(mem, payload_bin, acpi_tables, num_vcpus);
+    }
+
     let payload = mem.get_dynamic_mem_slice_mut(memslice::SliceType::Payload);
     let relocation_info = ipl::find_and_report_entry_point(mem, payload_bin, payload)
         .expect("Entry point not found!");
@@ -241,8 +252,11 @@ fn boot_builtin_payload(
         payload_hob_region.size as u64,
     );
 
+    let payload_hob = mem.get_dynamic_mem_slice_mut(SliceType::Acpi);
+    let e820 = mem.create_e820();
     // Prepare the HOB list to run the image
-    payload_hob::build_payload_hob(acpi_tables, &mem).expect("Fail to create payload HOB");
+    payload_hob::build_payload_hob(payload_hob, acpi_tables, &e820)
+        .expect("Fail to create payload HOB");
 
     // Finally let's switch stack and jump to the image entry point...
     log::info!(
@@ -258,30 +272,12 @@ fn boot_builtin_payload(
     // Relocate Mailbox along side with the AP function
     td::relocate_mailbox(mailbox);
 
-    match relocation_info.image_type {
-        ExecutablePayloadType::Elf => {
-            let entry = unsafe {
-                core::mem::transmute::<u64, extern "sysv64" fn(u64, u64)>(
-                    relocation_info.entry_point,
-                )
-            };
-            entry(
-                payload_hob_region.base_address as u64,
-                payload.as_ptr() as u64,
-            );
-        }
-        ExecutablePayloadType::PeCoff => {
-            let entry = unsafe {
-                core::mem::transmute::<u64, extern "win64" fn(u64, u64)>(
-                    relocation_info.entry_point,
-                )
-            };
-            entry(
-                payload_hob_region.base_address as u64,
-                payload.as_ptr() as u64,
-            );
-        }
-    };
+    ipl::jump_to_payload_entry(
+        relocation_info.entry_point,
+        relocation_info.image_type as u8,
+        payload_hob_region.base_address as u64,
+        payload.as_ptr() as u64,
+    );
 }
 
 // Reuse mailbox for the IDT used before payload setup its own one.

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

@@ -2,12 +2,22 @@
 //
 // SPDX-License-Identifier: BSD-2-Clause-Patent
 
+use alloc::vec::Vec;
 use core::convert::TryInto;
 use core::mem::size_of;
+use td_layout::runtime::exec::{ACPI_SIZE, PAYLOAD_SIZE};
+use td_shim::e820::E820Type;
 use zerocopy::{AsBytes, FromBytes, FromZeroes};
 
 use td_shim_interface::acpi::{self, GenericSdtHeader};
 
+use crate::{
+    e820::E820Table,
+    ipl::{self, jump_to_payload_entry, ExecutablePayloadType},
+    memory::Memory,
+    payload_hob, td,
+};
+
 // 255 vCPUs needs 2278 bytes, refer to create_madt().
 const MADT_MAX_SIZE: usize = 0xc00;
 const NUM_8259_IRQS: usize = 16;
@@ -144,6 +154,80 @@ pub fn create_madt_default(cpu_num: u32, mailbox_base: u64) -> Option<Madt> {
     Some(madt)
 }
 
+pub fn launch_payload_per_cpu(
+    mem: &mut Memory,
+    payload_bin: &[u8],
+    acpi_tables: &Vec<&[u8]>,
+    num_vcpus: u32,
+) {
+    const TEMP_STACK_SIZE: u64 = 0x4000;
+
+    let e820 = mem.create_e820();
+
+    // Get the largest usable memory index in e820 table
+    let (max_addr, max_usable) = e820
+        .as_slice()
+        .iter()
+        .enumerate()
+        .filter(|&(_, e)| e.r#type == E820Type::Memory as u32)
+        .max_by_key(|&(_, e)| e.size)
+        .map(|(_, e)| (e.addr, e.size))
+        .expect("No usable memory found!");
+
+    let size_per_cpu = max_usable / num_vcpus as u64;
+    let mut bsp_relocation_info = None;
+    let mut bsp_hob_addr = 0;
+    let mut bsp_payload_addr = 0;
+    for cpu_idx in 0..num_vcpus {
+        let start = cpu_idx as u64 * size_per_cpu + max_addr;
+        let temp_stack = start;
+        let payload_addr = temp_stack + TEMP_STACK_SIZE;
+        let hob_addr = payload_addr + PAYLOAD_SIZE as u64;
+
+        let mut table = E820Table::new();
+        table.add_range(E820Type::Memory, start, size_per_cpu);
+        table.convert_range(E820Type::Reserved, temp_stack, TEMP_STACK_SIZE);
+        table.convert_range(E820Type::Reserved, payload_addr, PAYLOAD_SIZE as u64);
+        table.convert_range(E820Type::Acpi, hob_addr, ACPI_SIZE as u64);
+        let payload =
+            unsafe { core::slice::from_raw_parts_mut(payload_addr as *mut u8, PAYLOAD_SIZE) };
+        let hob = unsafe { core::slice::from_raw_parts_mut(hob_addr as *mut u8, ACPI_SIZE) };
+
+        let relocation_info = ipl::find_and_report_entry_point(mem, payload_bin, payload)
+            .expect("Entry point not found!");
+        // Prepare the HOB list to run the image
+        payload_hob::build_payload_hob(hob, acpi_tables, &table)
+            .expect("Fail to create payload HOB");
+
+        // Skip the BSP
+        if cpu_idx != 0 {
+            td::ap_set_payload(
+                cpu_idx,
+                temp_stack + TEMP_STACK_SIZE,
+                relocation_info.entry_point,
+                relocation_info.image_type as u8,
+                hob_addr,
+                payload_addr,
+            );
+        }
+        // Save the relocation info for BSP paylaod
+        if cpu_idx == 0 {
+            bsp_relocation_info = Some(relocation_info);
+            bsp_hob_addr = hob_addr;
+            bsp_payload_addr = payload_addr;
+        }
+    }
+
+    // For BSP
+    let bsp_relocation_info = bsp_relocation_info.expect("Payload entry point not found for BSP!");
+    jump_to_payload_entry(
+        bsp_relocation_info.entry_point,
+        bsp_relocation_info.image_type as u8,
+        bsp_hob_addr,
+        bsp_payload_addr,
+    );
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -23,13 +23,13 @@ pub enum PayloadHobError {
     OutOfResource,
 }
 
-pub struct PayloadHob {
-    memory: &'static mut [u8],
+pub struct PayloadHob<'a> {
+    memory: &'a mut [u8],
     end: usize,
 }
 
-impl PayloadHob {
-    pub fn new(memory: &'static mut [u8]) -> Option<Self> {
+impl<'a> PayloadHob<'a> {
+    pub fn new(memory: &'a mut [u8]) -> Option<Self> {
         if memory.len() < size_of::<pi::hob::HandoffInfoTable>() {
             return None;
         }
@@ -168,10 +168,13 @@ impl PayloadHob {
     }
 }
 
-pub fn build_payload_hob(acpi_tables: &Vec<&[u8]>, memory: &Memory) -> Option<PayloadHob> {
+pub fn build_payload_hob<'a>(
+    hob: &'a mut [u8],
+    acpi_tables: &Vec<&[u8]>,
+    e820: &E820Table,
+) -> Option<PayloadHob<'a>> {
     // Reuse the ACPI memory to build the payload HOB.
-    let mut payload_hob =
-        PayloadHob::new(memory.get_dynamic_mem_slice_mut(memslice::SliceType::Acpi))?;
+    let mut payload_hob = PayloadHob::new(hob)?;
 
     payload_hob.add_cpu(memory::cpu_get_memory_space_size(), 16);
     payload_hob.add_fv(TD_SHIM_PAYLOAD_BASE as u64, TD_SHIM_PAYLOAD_SIZE as u64);
@@ -182,9 +185,6 @@ pub fn build_payload_hob(acpi_tables: &Vec<&[u8]>, memory: &Memory) -> Option<Pa
             .ok()?;
     }
 
-    let mut e820 = memory.create_e820();
-
-    log::info!("e820 table: {:x?}\n", e820.as_slice());
     payload_hob
         .add_guided_data(&TD_E820_TABLE_HOB_GUID, e820.as_bytes())
         .ok()?;

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

@@ -6,6 +6,8 @@ use cc_measurement::log::CcEventLogError;
 use td_exception::idt::DescriptorTablePointer;
 use tdx_tdcall::tdx;
 
+pub use super::tdx_mailbox::ap_set_payload;
+
 extern "win64" {
     fn asm_read_msr64(index: u32) -> u64;
     fn asm_write_msr64(index: u32, value: u64) -> u64;