Don't use shared_mem.as_mut_slice() since it doesn't dirty touched pages in the host

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/mem/elf.rs

@@ -27,6 +27,7 @@ use goblin::elf32::program_header::PT_LOAD;
 #[cfg(feature = "init-paging")]
 use goblin::elf64::program_header::PT_LOAD;
 
+use super::shared_mem::ExclusiveSharedMemory;
 use crate::{Result, log_then_return, new_error};
 
 #[cfg(feature = "unwind_guest")]
@@ -164,16 +165,23 @@ impl ElfInfo {
     pub(crate) fn load_at(
         self,
         load_addr: usize,
-        target: &mut [u8],
+        guest_code_offset: usize,
+        excl: &mut ExclusiveSharedMemory,
     ) -> Result<super::exe::LoadInfo> {
         let base_va = self.get_base_va();
         for phdr in self.phdrs.iter().filter(|phdr| phdr.p_type == PT_LOAD) {
             let start_va = (phdr.p_vaddr - base_va) as usize;
             let payload_offset = phdr.p_offset as usize;
             let payload_len = phdr.p_filesz as usize;
-            target[start_va..start_va + payload_len]
-                .copy_from_slice(&self.payload[payload_offset..payload_offset + payload_len]);
-            target[start_va + payload_len..start_va + phdr.p_memsz as usize].fill(0);
+            excl.copy_from_slice(
+                &self.payload[payload_offset..payload_offset + payload_len],
+                guest_code_offset + start_va,
+            )?;
+
+            excl.zero_fill(
+                guest_code_offset + start_va + payload_len,
+                phdr.p_memsz as usize - payload_len,
+            )?;
         }
         let get_addend = |name, r: &Reloc| {
             r.r_addend
@@ -196,8 +204,10 @@ impl ElfInfo {
             match r.r_type {
                 R_X86_64_RELATIVE => {
                     let addend = get_addend("R_X86_64_RELATIVE", r)?;
-                    target[r.r_offset as usize..r.r_offset as usize + 8]
-                        .copy_from_slice(&(load_addr as i64 + addend).to_le_bytes());
+                    excl.copy_from_slice(
+                        &(load_addr as i64 + addend).to_le_bytes(),
+                        guest_code_offset + r.r_offset as usize,
+                    )?;
                 }
                 R_X86_64_NONE => {}
                 _ => {

src/hyperlight_host/src/mem/exe.rs

@@ -22,6 +22,7 @@ use std::vec::Vec;
 
 use super::elf::ElfInfo;
 use super::ptr_offset::Offset;
+use super::shared_mem::ExclusiveSharedMemory;
 use crate::Result;
 
 // This is used extremely infrequently, so being unusually large for PE
@@ -108,9 +109,14 @@ impl ExeInfo {
     // copying into target, but the PE loader chooses to apply
     // relocations in its owned representation of the PE contents,
     // which requires it to be &mut.
-    pub fn load(self, load_addr: usize, target: &mut [u8]) -> Result<LoadInfo> {
+    pub fn load(
+        self,
+        load_addr: usize,
+        guest_code_offset: usize,
+        target: &mut ExclusiveSharedMemory,
+    ) -> Result<LoadInfo> {
         match self {
-            ExeInfo::Elf(elf) => elf.load_at(load_addr, target),
+            ExeInfo::Elf(elf) => elf.load_at(load_addr, guest_code_offset, target),
         }
     }
 }

src/hyperlight_host/src/mem/mgr.rs

@@ -346,7 +346,8 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         #[allow(clippy::let_unit_value)]
         let load_info = exe_info.load(
             load_addr.clone().try_into()?,
-            &mut shared_mem.as_mut_slice()[layout.get_guest_code_offset()..],
+            layout.get_guest_code_offset(),
+            &mut shared_mem,
         )?;
 
         Ok((

src/hyperlight_host/src/mem/shared_mem.rs

@@ -577,7 +577,10 @@ impl ExclusiveSharedMemory {
     ///   the safety documentation of pointer::offset.
     ///
     ///   This is ensured by a check in ::new()
-    pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
+    ///
+    /// Additionally, writes to the returned slice will not mark pages as dirty.
+    /// User must call `mark_pages_dirty` manually to mark pages as dirty.
+    fn as_mut_slice(&mut self) -> &mut [u8] {
         unsafe { std::slice::from_raw_parts_mut(self.base_ptr(), self.mem_size()) }
     }