td-payload: add new shared memory init function with private shadow

td-payload/src/arch/x86_64/init.rs

@@ -7,8 +7,12 @@ use crate::{
     arch::{gdt, idt},
     hob::{self, get_hob},
     mm::{
-        get_usable, heap::init_heap, init_ram, layout::RuntimeLayout,
-        page_table::init_pt_frame_allocator, shared::init_shared_memory,
+        get_usable,
+        heap::init_heap,
+        init_ram,
+        layout::RuntimeLayout,
+        page_table::init_pt_frame_allocator,
+        shared::{init_shared_memory, init_shared_memory_with_shadow},
     },
 };
 
@@ -22,7 +26,7 @@ use super::{
     idt::{PAGE_FAULT_EXCEPTION, PAGE_FAULT_IST},
 };
 
-pub fn pre_init(hob: u64, layout: &RuntimeLayout) {
+pub fn pre_init(hob: u64, layout: &RuntimeLayout, use_shared_shadow: bool) {
     let hob = hob::init(hob).expect("Invalid payload HOB");
     let memory_map = init_ram(hob).expect("Failed to parse E820 table from payload HOB");
 
@@ -35,7 +39,13 @@ pub fn pre_init(hob: u64, layout: &RuntimeLayout) {
     init_heap(heap, layout.heap_size);
 
     let shared = get_usable(layout.shared_memory_size).expect("Failed to allocate shared memory");
-    init_shared_memory(shared, layout.shared_memory_size);
+    if use_shared_shadow {
+        let shadow =
+            get_usable(layout.shared_memory_size).expect("Failed to allocate shared shadow");
+        init_shared_memory_with_shadow(shared, layout.shared_memory_size, shadow);
+    } else {
+        init_shared_memory(shared, layout.shared_memory_size);
+    }
 
     // Init Global Descriptor Table and Task State Segment
     gdt::init_gdt();

td-payload/src/lib.rs

@@ -33,7 +33,7 @@ pub extern "C" fn _start(hob: u64, _payload: u64) -> ! {
 
     let layout = RuntimeLayout::default();
 
-    arch::init::pre_init(hob, &layout);
+    arch::init::pre_init(hob, &layout, false);
     arch::init::init(&layout, main);
 }
 

td-payload/src/mm/shared.rs

@@ -10,35 +10,38 @@ use super::SIZE_4K;
 use crate::arch::shared::decrypt;
 
 static SHARED_MEMORY_ALLOCATOR: LockedHeap = LockedHeap::empty();
-static SHADOW_OFFSET: Once<usize> = Once::new();
+static SHARED_START: Once<usize> = Once::new();
+static SHADOW_START: Once<usize> = Once::new();
 
 pub fn init_shared_memory(start: u64, size: usize) {
     if size % SIZE_4K != 0 {
         panic!("Failed to initialize shared memory: size needs to be aligned with 0x1000");
     }
-    let shared_size = size / 2;
 
     // Set the shared memory region to be shared
-    decrypt(start, shared_size);
+    decrypt(start, size);
     // Initialize the shared memory allocator
     unsafe {
-        SHARED_MEMORY_ALLOCATOR
-            .lock()
-            .init(start as *mut u8, shared_size);
+        SHARED_MEMORY_ALLOCATOR.lock().init(start as *mut u8, size);
     }
-    SHADOW_OFFSET.call_once(|| shared_size);
+}
+
+pub fn init_shared_memory_with_shadow(start: u64, size: usize, shadow_start: u64) {
+    init_shared_memory(start, size);
+    SHARED_START.call_once(|| start as usize);
+    SHADOW_START.call_once(|| shadow_start as usize);
 }
 
 pub struct SharedMemory {
     addr: usize,
-    shadow_addr: usize,
+    shadow_addr: Option<usize>,
     size: usize,
 }
 
 impl SharedMemory {
     pub fn new(num_page: usize) -> Option<Self> {
         let addr = unsafe { alloc_shared_pages(num_page)? };
-        let shadow_addr = alloc_private_shadow_pages(addr)?;
+        let shadow_addr = alloc_private_shadow_pages(addr);
 
         Some(Self {
             addr,
@@ -47,12 +50,13 @@ impl SharedMemory {
         })
     }
 
-    pub fn copy_to_private_shadow(&mut self) -> &[u8] {
-        let shadow =
-            unsafe { core::slice::from_raw_parts_mut(self.shadow_addr as *mut u8, self.size) };
-        shadow.copy_from_slice(self.as_bytes());
+    pub fn copy_to_private_shadow(&mut self) -> Option<&[u8]> {
+        self.shadow_addr.map(|addr| {
+            let shadow = unsafe { core::slice::from_raw_parts_mut(addr as *mut u8, self.size) };
+            shadow.copy_from_slice(self.as_bytes());
 
-        shadow
+            &shadow[..]
+        })
     }
 
     pub fn as_bytes(&self) -> &[u8] {
@@ -110,5 +114,6 @@ pub unsafe fn free_shared_page(addr: usize) {
 }
 
 fn alloc_private_shadow_pages(shared_addr: usize) -> Option<usize> {
-    Some(shared_addr + SHADOW_OFFSET.get()?)
+    let offset = shared_addr.checked_sub(*SHARED_START.get()?)?;
+    Some(SHADOW_START.get()? + offset)
 }

tests/test-td-payload/src/main.rs

@@ -132,7 +132,7 @@ pub extern "C" fn _start(hob: u64, _payload: u64) -> ! {
         shadow_stack_size: layout::DEFAULT_SHADOW_STACK_SIZE,
     };
 
-    arch::init::pre_init(hob, &layout);
+    arch::init::pre_init(hob, &layout, false);
     arch::init::init(&layout, main);
 }