Let windows decide at which address to map shared memory in surrogate process

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -14,14 +14,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use core::ffi::c_void;
 use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::string::String;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 
-use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use log::LevelFilter;
 use tracing::{Span, instrument};
 use windows::Win32::System::Hypervisor::{
@@ -55,10 +53,8 @@ use crate::{Result, debug, new_error};
 
 /// A Hypervisor driver for HyperV-on-Windows.
 pub(crate) struct HypervWindowsDriver {
-    size: usize, // this is the size of the memory region, excluding the 2 surrounding guard pages
     processor: VMProcessor,
     _surrogate_process: SurrogateProcess, // we need to keep a reference to the SurrogateProcess for the duration of the driver since otherwise it will dropped and the memory mapping will be unmapped and the surrogate process will be returned to the pool
-    source_address: *mut c_void,          // this points into the first guard page
     entrypoint: u64,
     orig_rsp: GuestPtr,
     mem_regions: Vec<MemoryRegion>,
@@ -80,7 +76,6 @@ impl HypervWindowsDriver {
     pub(crate) fn new(
         mem_regions: Vec<MemoryRegion>,
         raw_size: usize,
-        raw_source_address: *mut c_void,
         pml4_address: u64,
         entrypoint: u64,
         rsp: u64,
@@ -94,23 +89,18 @@ impl HypervWindowsDriver {
         // with guard pages setup
         let surrogate_process = {
             let mgr = get_surrogate_process_manager()?;
-            mgr.get_surrogate_process(raw_size, raw_source_address, mmap_file_handle)
+            mgr.get_surrogate_process(raw_size, mmap_file_handle)
         }?;
 
-        partition.map_gpa_range(&mem_regions, surrogate_process.process_handle)?;
+        partition.map_gpa_range(&mem_regions, &surrogate_process)?;
 
         let mut proc = VMProcessor::new(partition)?;
         Self::setup_initial_sregs(&mut proc, pml4_address)?;
         let partition_handle = proc.get_partition_hdl();
 
-        // subtract 2 pages for the guard pages, since when we copy memory to and from surrogate process,
-        // we don't want to copy the guard pages themselves (that would cause access violation)
-        let mem_size = raw_size - 2 * PAGE_SIZE_USIZE;
         Ok(Self {
-            size: mem_size,
             processor: proc,
             _surrogate_process: surrogate_process,
-            source_address: raw_source_address,
             entrypoint,
             orig_rsp: GuestPtr::try_from(RawPtr::from(rsp))?,
             mem_regions,
@@ -177,9 +167,7 @@ impl Debug for HypervWindowsDriver {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         let mut fs = f.debug_struct("HyperV Driver");
 
-        fs.field("Size", &self.size)
-            .field("Source Address", &self.source_address)
-            .field("Entrypoint", &self.entrypoint)
+        fs.field("Entrypoint", &self.entrypoint)
             .field("Original RSP", &self.orig_rsp);
 
         for region in &self.mem_regions {

src/hyperlight_host/src/hypervisor/surrogate_process.rs

@@ -30,6 +30,7 @@ use super::wrappers::HandleWrapper;
 #[derive(Debug)]
 pub(super) struct SurrogateProcess {
     /// The address of memory allocated in the surrogate process to be mapped to the VM.
+    /// This includes the first guard page
     pub(crate) allocated_address: *mut c_void,
     /// The handle to the surrogate process.
     pub(crate) process_handle: HandleWrapper,

src/hyperlight_host/src/hypervisor/surrogate_process_manager.rs

@@ -144,7 +144,6 @@ impl SurrogateProcessManager {
     pub(super) fn get_surrogate_process(
         &self,
         raw_size: usize,
-        raw_source_address: *const c_void,
         mmap_file_handle: HandleWrapper,
     ) -> Result<SurrogateProcess> {
         let surrogate_process_handle: HANDLE = self.process_receiver.recv()?.into();
@@ -160,7 +159,7 @@ impl SurrogateProcessManager {
                 mapping_file_handle,
                 surrogate_process_handle,
                 0,
-                Some(raw_source_address),
+                None,
                 raw_size,
                 0,
                 PAGE_READWRITE.0,
@@ -171,19 +170,7 @@ impl SurrogateProcessManager {
         if allocated_address.Value.is_null() {
             // Safety: `MapViewOfFileNuma2` will set the last error code if it fails.
             let error = unsafe { windows::Win32::Foundation::GetLastError() };
-            log_then_return!(
-                "MapViewOfFileNuma2 failed with error code: {:?} for mem address {:?} ",
-                error,
-                raw_source_address
-            );
-        }
-
-        if allocated_address.Value as *const c_void != raw_source_address {
-            log_then_return!(
-                "Address Mismatch Allocated: {:?} Requested: {:?}",
-                allocated_address.Value,
-                raw_source_address
-            );
+            log_then_return!("MapViewOfFileNuma2 failed with error code: {:?}", error);
         }
 
         // set up guard pages
@@ -193,7 +180,7 @@ impl SurrogateProcessManager {
         let mut unused_out_old_prot_flags = PAGE_PROTECTION_FLAGS(0);
 
         // the first page of the raw_size is the guard page
-        let first_guard_page_start = raw_source_address;
+        let first_guard_page_start = allocated_address.Value;
         if let Err(e) = unsafe {
             VirtualProtectEx(
                 surrogate_process_handle,
@@ -207,7 +194,8 @@ impl SurrogateProcessManager {
         }
 
         // the last page of the raw_size is the guard page
-        let last_guard_page_start = unsafe { raw_source_address.add(raw_size - PAGE_SIZE_USIZE) };
+        let last_guard_page_start =
+            unsafe { first_guard_page_start.add(raw_size - PAGE_SIZE_USIZE) };
         if let Err(e) = unsafe {
             VirtualProtectEx(
                 surrogate_process_handle,
@@ -485,11 +473,8 @@ mod tests {
 
                     let timer = Instant::now();
                     let surrogate_process = {
-                        let res = surrogate_process_manager.get_surrogate_process(
-                            size,
-                            addr.Value,
-                            HandleWrapper::from(handle),
-                        )?;
+                        let res = surrogate_process_manager
+                            .get_surrogate_process(size, HandleWrapper::from(handle))?;
                         let elapsed = timer.elapsed();
                         // Print out the time it took to get the process if its greater than 150ms (this is just to allow us to see that threads are blocking on the process queue)
                         if (elapsed.as_millis() as u64) > 150 {
@@ -579,7 +564,6 @@ mod tests {
         let process = mgr
             .get_surrogate_process(
                 mem.raw_mem_size(),
-                mem.raw_ptr() as *mut c_void,
                 HandleWrapper::from(mem.get_mmap_file_handle()),
             )
             .unwrap();

src/hyperlight_host/src/hypervisor/windows_hypervisor_platform.rs

@@ -16,14 +16,15 @@ limitations under the License.
 
 use core::ffi::c_void;
 
+use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use tracing::{Span, instrument};
 use windows::Win32::Foundation::{FreeLibrary, HANDLE};
 use windows::Win32::System::Hypervisor::*;
 use windows::Win32::System::LibraryLoader::*;
 use windows::core::s;
 use windows_result::HRESULT;
 
-use super::wrappers::HandleWrapper;
+use super::surrogate_process::SurrogateProcess;
 #[cfg(crashdump)]
 use crate::HyperlightError;
 use crate::hypervisor::wrappers::{WHvFPURegisters, WHvGeneralRegisters, WHvSpecialRegisters};
@@ -89,9 +90,28 @@ impl VMPartition {
     pub(super) fn map_gpa_range(
         &mut self,
         regions: &[MemoryRegion],
-        process_handle: HandleWrapper,
+        surrogate_process: &SurrogateProcess,
     ) -> Result<()> {
-        let process_handle: HANDLE = process_handle.into();
+        let process_handle: HANDLE = surrogate_process.process_handle.into();
+        // this is the address in the surrogate process where shared memory starts.
+        // We add page-size because we don't care about the first guard page
+        let surrogate_address = surrogate_process.allocated_address as usize + PAGE_SIZE_USIZE;
+        if regions.is_empty() {
+            return Err(new_error!("No memory regions to map"));
+        }
+        // this is the address in the main process where the shared memory starts
+        let host_address = regions[0].host_region.start;
+
+        let surrogate_address_bigger;
+        // offset between the surrogate process and the host process address of start of shared memory
+        let offset = if surrogate_address > host_address {
+            surrogate_address_bigger = true;
+            surrogate_address - host_address
+        } else {
+            surrogate_address_bigger = false;
+            host_address - surrogate_address
+        };
+
         // The function pointer to WHvMapGpaRange2 is resolved dynamically to allow us to detect
         // when we are running on older versions of windows that do not support this API and
         // return a more informative error message, rather than failing with an error about a missing entrypoint
@@ -118,12 +138,18 @@ impl VMPartition {
                 .iter()
                 .fold(WHvMapGpaRangeFlagNone, |acc, flag| acc | *flag); // collect using bitwise OR
 
+            let surrogate_region_start_address = if surrogate_address_bigger {
+                region.host_region.start + offset
+            } else {
+                region.host_region.start - offset
+            };
+
             let res = whvmapgparange2_func(
                 self.0,
                 process_handle,
-                region.host_region.start as *const c_void,
+                surrogate_region_start_address as *const c_void,
                 region.guest_region.start as u64,
-                (region.guest_region.end - region.guest_region.start) as u64,
+                region.guest_region.len() as u64,
                 flags,
             );
             if res.is_err() {

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -238,8 +238,6 @@ pub(crate) fn set_up_hypervisor_partition(
 
         #[cfg(target_os = "windows")]
         Some(HypervisorType::Whp) => {
-            use std::ffi::c_void;
-
             use crate::hypervisor::wrappers::HandleWrapper;
 
             let mmap_file_handle = mgr
@@ -248,7 +246,6 @@ pub(crate) fn set_up_hypervisor_partition(
             let hv = crate::hypervisor::hyperv_windows::HypervWindowsDriver::new(
                 regions,
                 mgr.shared_mem.raw_mem_size(), // we use raw_* here because windows driver requires 64K aligned addresses,
-                mgr.shared_mem.raw_ptr() as *mut c_void, // and instead convert it to base_addr where needed in the driver itself
                 pml4_ptr.absolute()?,
                 entrypoint_ptr.absolute()?,
                 rsp_ptr.absolute()?,