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

@@ -24,6 +24,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};
 
 pub(crate) struct ElfInfo {
@@ -73,15 +74,26 @@ impl ElfInfo {
             .unwrap();
         (max_phdr.p_vaddr + max_phdr.p_memsz - self.get_base_va()) as usize
     }
-    pub(crate) fn load_at(&self, load_addr: usize, target: &mut [u8]) -> Result<()> {
+    pub(crate) fn load_at(
+        &self,
+        load_addr: usize,
+        guest_code_offset: usize,
+        excl: &mut ExclusiveSharedMemory,
+    ) -> Result<()> {
         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
@@ -104,8 +116,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

@@ -20,6 +20,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
@@ -71,10 +72,15 @@ 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(&mut self, load_addr: usize, target: &mut [u8]) -> Result<()> {
+    pub fn load(
+        &mut self,
+        load_addr: usize,
+        guest_code_offset: usize,
+        target: &mut ExclusiveSharedMemory,
+    ) -> Result<()> {
         match self {
             ExeInfo::Elf(elf) => {
-                elf.load_at(load_addr, target)?;
+                elf.load_at(load_addr, guest_code_offset, target)?;
             }
         }
         Ok(())

src/hyperlight_host/src/mem/mgr.rs

@@ -341,7 +341,8 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
 
         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(Self::new(layout, shared_mem, load_addr, entrypoint_offset))

src/hyperlight_host/src/mem/shared_mem.rs

@@ -576,7 +576,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()) }
     }