boot shim side changes

Author: babayet2

Cargo.lock

@@ -4975,6 +4975,7 @@ dependencies = [
  "memory_range",
  "minimal_rt",
  "minimal_rt_build",
+ "page_table",
  "safe_intrinsics",
  "sha2",
  "sidecar_defs",

openhcl/openhcl_boot/Cargo.toml

@@ -31,6 +31,7 @@ thiserror.workspace = true
 zerocopy.workspace = true
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
+page_table.workspace = true
 safe_intrinsics.workspace = true
 tdcall.workspace = true
 x86defs.workspace = true

openhcl/openhcl_boot/src/arch/x86_64/memory.rs

@@ -5,21 +5,36 @@
 
 use super::address_space::LocalMap;
 use super::address_space::init_local_map;
+use crate::AddressSpaceManager;
 use crate::ShimParams;
 use crate::arch::TdxHypercallPage;
 use crate::arch::x86_64::address_space::tdx_share_large_page;
 use crate::host_params::PartitionInfo;
 use crate::host_params::shim_params::IsolationType;
 use crate::hypercall::hvcall;
+use crate::memory::AllocationPolicy;
+use crate::memory::AllocationType;
+use crate::off_stack;
+use arrayvec::ArrayVec;
+use loader_defs::shim::MemoryVtlType;
 use memory_range::MemoryRange;
+use page_table::x64::PAGE_TABLE_MAX_BYTES;
+use page_table::x64::PAGE_TABLE_MAX_COUNT;
+use page_table::x64::PageTable;
+use page_table::x64::PageTableBuilder;
 use sha2::Digest;
 use sha2::Sha384;
 use x86defs::X64_LARGE_PAGE_SIZE;
 use x86defs::tdx::TDX_SHARED_GPA_BOUNDARY_ADDRESS_BIT;
+use zerocopy::FromZeros;
 
 /// On isolated systems, transitions all VTL2 RAM to be private and accepted, with the appropriate
 /// VTL permissions applied.
-pub fn setup_vtl2_memory(shim_params: &ShimParams, partition_info: &PartitionInfo) {
+pub fn setup_vtl2_memory(
+    shim_params: &ShimParams,
+    partition_info: &PartitionInfo,
+    address_space: &mut AddressSpaceManager,
+) {
     // Only if the partition is VBS-isolated, accept memory and apply vtl 2 protections here.
     // Non-isolated partitions can undergo servicing, and additional information
     // would be needed to determine whether vtl 2 protections should be applied
@@ -125,6 +140,71 @@ pub fn setup_vtl2_memory(shim_params: &ShimParams, partition_info: &PartitionInf
     // hypercall IO pages. ram_buffer must not be used again beyond this point
     // TODO: find an approach that does not require re-using the ram_buffer
     if shim_params.isolation_type == IsolationType::Tdx {
+        let page_table_region = address_space
+            .allocate(
+                None,
+                PAGE_TABLE_MAX_BYTES as u64,
+                AllocationType::TdxPageTables,
+                AllocationPolicy::LowMemory,
+            )
+            .expect("allocation of space for TDX page tables must succeed");
+
+        const MAX_RANGE_COUNT: usize = 64;
+        let mut ranges = off_stack!(
+            ArrayVec::<(u64, u64), MAX_RANGE_COUNT>,
+            ArrayVec::new_const()
+        );
+
+        let vtl2_ram = address_space
+            .vtl2_ranges()
+            .filter_map(|(range, typ)| match typ {
+                MemoryVtlType::VTL2_RAM => Some((range.start(), range.end())),
+                _ => None,
+            });
+
+        ranges.extend(vtl2_ram);
+
+        let mut page_table_work_buffer =
+            off_stack!(ArrayVec<PageTable, PAGE_TABLE_MAX_COUNT>, ArrayVec::new_const());
+        for _ in 0..PAGE_TABLE_MAX_COUNT {
+            page_table_work_buffer.push(PageTable::new_zeroed());
+        }
+        let mut page_table = off_stack!(ArrayVec<u8, PAGE_TABLE_MAX_BYTES>, ArrayVec::new_const());
+        for _ in 0..PAGE_TABLE_MAX_BYTES {
+            page_table.push(0);
+        }
+
+        let page_table_builder = PageTableBuilder::new(
+            page_table_region.range.start(),
+            page_table_work_buffer.as_mut_slice(),
+            page_table.as_mut_slice(),
+        )
+        .expect("page table builder must return no error")
+        .with_ranges(ranges.as_slice());
+
+        let page_tables = page_table_builder
+            .build()
+            .expect("page table construction must succeed");
+
+        address_space.truncate_range(
+            page_table_region,
+            page_tables.len() as u64,
+            AllocationPolicy::LowMemory,
+        );
+
+        //TODO safety comment
+        unsafe {
+            core::ptr::copy_nonoverlapping(
+                page_tables.as_ptr(),
+                page_table_region.range.start() as *mut u8,
+                page_tables.len(),
+            );
+        }
+
+        crate::arch::tdx::tdx_prepare_ap_trampoline(
+                page_table_region.range.start()
+        );
+
         let free_buffer = ram_buffer.as_mut_ptr() as u64;
         assert!(free_buffer.is_multiple_of(X64_LARGE_PAGE_SIZE));
         // SAFETY: The bottom 2MB region of the ram_buffer is unused by the shim

openhcl/openhcl_boot/src/arch/x86_64/tdx.rs

@@ -182,7 +182,7 @@ pub fn get_tdx_tsc_reftime() -> Option<u64> {
 /// Update the TdxTrampolineContext, setting the necessary control registers for AP startup,
 /// and ensuring that LGDT will be skipped, so the GDT page does not need to be added to the
 /// e820 entries
-pub fn tdx_prepare_ap_trampoline() {
+pub fn tdx_prepare_ap_trampoline(cr3: u64) {
     let context_ptr: *mut TdxTrampolineContext = RESET_VECTOR_PAGE as *mut TdxTrampolineContext;
     // SAFETY: The TdxTrampolineContext is known to be stored at the architectural reset vector address
     let tdxcontext: &mut TdxTrampolineContext = unsafe { context_ptr.as_mut().unwrap() };
@@ -191,13 +191,11 @@ pub fn tdx_prepare_ap_trampoline() {
     tdxcontext.code_selector = 0;
     tdxcontext.task_selector = 0;
     tdxcontext.cr0 |= x86defs::X64_CR0_PG | x86defs::X64_CR0_PE | x86defs::X64_CR0_NE;
+    tdxcontext.cr3 = cr3;
     tdxcontext.cr4 |= x86defs::X64_CR4_PAE | x86defs::X64_CR4_MCE;
 }
 
 pub fn setup_vtl2_vp(partition_info: &PartitionInfo) {
-    // Update the TDX Trampoline Context for AP Startup
-    tdx_prepare_ap_trampoline();
-
     for cpu in 1..partition_info.cpus.len() {
         hvcall()
             .tdx_enable_vp_vtl2(cpu as u32)

openhcl/openhcl_boot/src/host_params/dt/mod.rs

@@ -501,15 +501,6 @@ impl PartitionInfo {
             ));
         }
 
-        // Only specify pagetables as a reserved region on TDX, as they are used
-        // for AP startup via the mailbox protocol. On other platforms, the
-        // memory is free to be reclaimed.
-        if params.isolation_type == IsolationType::Tdx {
-            assert!(params.page_tables.is_some());
-            address_space_builder = address_space_builder
-                .with_page_tables(params.page_tables.expect("always present on tdx"));
-        }
-
         address_space_builder
             .init()
             .expect("failed to initialize address space manager");

openhcl/openhcl_boot/src/host_params/shim_params.rs

@@ -101,8 +101,6 @@ pub struct ShimParams {
     /// Memory used by the shim.
     pub used: MemoryRange,
     pub bounce_buffer: Option<MemoryRange>,
-    /// Page tables region used by the shim.
-    pub page_tables: Option<MemoryRange>,
     /// Log buffer region used by the shim.
     pub log_buffer: MemoryRange,
     /// Memory to be used for the heap.
@@ -133,8 +131,6 @@ impl ShimParams {
             used_end,
             bounce_buffer_start,
             bounce_buffer_size,
-            page_tables_start,
-            page_tables_size,
             log_buffer_start,
             log_buffer_size,
             heap_start_offset,
@@ -150,13 +146,6 @@ impl ShimParams {
             Some(MemoryRange::new(base..base + bounce_buffer_size))
         };
 
-        let page_tables = if page_tables_size == 0 {
-            None
-        } else {
-            let base = shim_base_address.wrapping_add_signed(page_tables_start);
-            Some(MemoryRange::new(base..base + page_tables_size))
-        };
-
         let log_buffer = {
             let base = shim_base_address.wrapping_add_signed(log_buffer_start);
             MemoryRange::new(base..base + log_buffer_size)
@@ -189,7 +178,6 @@ impl ShimParams {
                     ..shim_base_address.wrapping_add_signed(used_end),
             ),
             bounce_buffer,
-            page_tables,
             log_buffer,
             heap,
         }

openhcl/openhcl_boot/src/main.rs

@@ -630,7 +630,7 @@ fn shim_main(shim_params_raw_offset: isize) -> ! {
 
     validate_vp_hw_ids(partition_info);
 
-    setup_vtl2_memory(&p, partition_info);
+    setup_vtl2_memory(&p, partition_info, address_space);
     setup_vtl2_vp(partition_info);
 
     verify_imported_regions_hash(&p);

openhcl/openhcl_boot/src/memory.rs

@@ -170,12 +170,6 @@ impl<'a, I: Iterator<Item = MemoryRange>> AddressSpaceManagerBuilder<'a, I> {
         self
     }
 
-    /// Pagetables that are reported as type [`MemoryVtlType::VTL2_TDX_PAGE_TABLES`].
-    pub fn with_page_tables(mut self, page_tables: MemoryRange) -> Self {
-        self.page_tables = Some(page_tables);
-        self
-    }
-
     /// Log buffer that is reported as type [`MemoryVtlType::VTL2_BOOTSHIM_LOG_BUFFER`].
     pub fn with_log_buffer(mut self, log_buffer: MemoryRange) -> Self {
         self.log_buffer = Some(log_buffer);
@@ -362,6 +356,91 @@ impl AddressSpaceManager {
         allocated
     }
 
+    /// Truncate a memory range, with allocation policy deciding whether
+    /// free the low part or high part.
+    pub fn truncate_range(
+        &mut self,
+        range: AllocatedRange,
+        len: u64,
+        allocation_policy: AllocationPolicy,
+    ) {
+        //TODO babayet2 ASSERT 4k or align
+        //get index of range to modify
+        //let index = self.address_space.iter().position(range).exepct("allocated range must be present");
+
+        let index = self
+            .address_space
+            .iter()
+            .position(|r| r.range.start() == range.range.start())
+            .unwrap();
+
+        let (used, remainder) = match allocation_policy {
+            AllocationPolicy::LowMemory => {
+                // Allocate from the beginning (low addresses)
+                range.range.split_at_offset(len)
+            }
+            AllocationPolicy::HighMemory => {
+                // Allocate from the end (high addresses)
+                let offset = range.range.len() - len;
+                let (remainder, used) = range.range.split_at_offset(offset);
+                (used, remainder)
+            }
+        };
+
+        let remainder = if !remainder.is_empty() {
+            Some(AddressRange {
+                range: remainder,
+                vnode: range.vnode,
+                usage: AddressUsage::Free,
+            })
+        } else {
+            None
+        };
+
+        if let Some(remainder) = remainder {
+            let (free_range, free_index) = match allocation_policy {
+                AllocationPolicy::LowMemory => {
+                    let next = self.address_space.get(index + 1);
+                    if next.is_some()
+                        && next.unwrap().usage == AddressUsage::Free
+                        && next.unwrap().range.start() == remainder.range.end()
+                    {
+                        (MemoryRange::new(remainder.range.start()..next.unwrap().range.end()), index + 1)
+                    } else {
+                        (remainder.range, index + 1)
+                    }
+                }
+                AllocationPolicy::HighMemory => {
+                    // When allocating from high memory, the remainder goes
+                    // before the allocated range
+                    let prev = self.address_space.get(index - 1);
+                    if prev.is_some()
+                        && prev.unwrap().usage == AddressUsage::Free
+                        && prev.unwrap().range.end() == remainder.range.start()
+                    {
+                        (MemoryRange::new(prev.unwrap().range.start()..remainder.range.end()), index - 1)
+                    } else {
+                        (remainder.range, index - 1)
+                    }
+                }
+            };
+
+            let usage = self.address_space[index].usage;
+
+            self.address_space[index] = AddressRange {
+                range: used,
+                vnode: range.vnode,
+                usage,
+            };
+
+            self.address_space[free_index] = AddressRange {
+                range: free_range,
+                vnode: range.vnode,
+                usage: AddressUsage::Free,
+            };
+        }
+    }
+
     /// Allocate a new range of memory with the given type and policy. None is
     /// returned if the allocation was unable to be satisfied.
     ///
@@ -423,6 +502,9 @@ impl AddressSpaceManager {
                     AllocationType::SidecarNode => {
                         AddressUsage::Reserved(ReservedMemoryType::SidecarNode)
                     }
+                    AllocationType::TdxPageTables => {
+                        AddressUsage::Reserved(ReservedMemoryType::TdxPageTables)
+                    }
                 },
                 allocation_policy,
             )
@@ -463,6 +545,7 @@ impl AddressSpaceManager {
 pub enum AllocationType {
     GpaPool,
     SidecarNode,
+    TdxPageTables,
 }
 
 pub enum AllocationPolicy {

vm/loader/loader_defs/src/shim.rs

@@ -53,10 +53,6 @@ pub struct ShimParamsRaw {
     pub bounce_buffer_start: i64,
     /// The size of the bounce buffer range. This is 0 if unavailable.
     pub bounce_buffer_size: u64,
-    /// The offset to the page_tables start address. This is 0 if unavailable.
-    pub page_tables_start: i64,
-    /// The size of the openhcl_boot page tables. This is 0 if unavailable.
-    pub page_tables_size: u64,
     /// The offset to the persisted bootshim log buffer.
     pub log_buffer_start: i64,
     /// The size of the persisted bootshim log buffer.