Modify guest physical page allocator to allocate from the scratch region

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_guest/src/arch/amd64/layout.rs

@@ -0,0 +1,34 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+pub const MAIN_STACK_TOP_GVA: usize = 0xffff_feff_ffff_f000;
+
+pub fn scratch_size() -> u64 {
+    let addr = crate::layout::scratch_size_gva();
+    let x: u64;
+    unsafe {
+        core::arch::asm!("mov {x}, [{addr}]", x = out(reg) x, addr = in(reg) addr);
+    }
+    x
+}
+
+pub fn scratch_base_gpa() -> u64 {
+    hyperlight_common::layout::scratch_base_gpa(scratch_size() as usize)
+}
+
+pub fn scratch_base_gva() -> u64 {
+    hyperlight_common::layout::scratch_base_gva(scratch_size() as usize)
+}

src/hyperlight_guest/src/arch/amd64/prim_alloc.rs

@@ -0,0 +1,36 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// There are no notable architecture-specific safety considerations
+// here, and the general conditions are documented in the
+// architecture-independent re-export in prim_alloc.rs
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
+    let addr = crate::layout::allocator_gva();
+    let nbytes = n * hyperlight_common::vmem::PAGE_SIZE as u64;
+    let mut x = nbytes;
+    unsafe {
+        core::arch::asm!(
+            "lock xadd qword ptr [{addr}], {x}",
+            addr = in(reg) addr,
+            x = inout(reg) x
+        );
+    }
+    if x.checked_add(nbytes).is_none() {
+        panic!("Out of physical memory!")
+    }
+    x
+}

src/hyperlight_guest/src/layout.rs

@@ -0,0 +1,29 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+#[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")]
+mod arch;
+
+pub use arch::MAIN_STACK_TOP_GVA;
+pub fn scratch_size_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SIZE_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_SIZE_OFFSET + 1) as *mut u64
+}
+pub fn allocator_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_ALLOCATOR_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_ALLOCATOR_OFFSET + 1) as *mut u64
+}
+pub use arch::{scratch_base_gpa, scratch_base_gva};

src/hyperlight_guest/src/lib.rs

@@ -23,6 +23,8 @@ extern crate alloc;
 // Modules
 pub mod error;
 pub mod exit;
+pub mod layout;
+pub mod prim_alloc;
 
 pub mod guest_handle {
     pub mod handle;

src/hyperlight_guest/src/prim_alloc.rs

@@ -0,0 +1,36 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+#[cfg_attr(target_arch = "x86_64", path = "arch/amd64/prim_alloc.rs")]
+mod arch;
+
+/// Allocate n contiguous physical pages and return the physical
+/// addresses of the pages in question.
+/// # Safety
+/// Since this reads and writes specific allocator state addresses, it
+/// is only safe when the allocator has been set up properly. It may
+/// become less safe in the future.
+///
+/// # Panics
+/// This function will panic if memory allocation fails
+///
+/// This is defined in an arch-specific module because it reads and
+/// writes the actual allocator state with inline assembly in order to
+/// access it atomically according to the architecture memory model
+/// rather than the Rust memory model: the stronger constraints of the
+/// latter cannot be perfectly satisfied due to the lack of per-byte
+/// atomic memcpy in the host.
+pub use arch::alloc_phys_pages;

src/hyperlight_guest_bin/src/paging.rs

@@ -14,12 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use alloc::alloc::Layout;
 use core::arch::asm;
 
+use hyperlight_guest::prim_alloc::alloc_phys_pages;
 use tracing::{Span, instrument};
 
-use crate::OS_PAGE_SIZE;
 
 /// Convert a physical address in main memory to a virtual address
 /// through the pysmap
@@ -137,28 +136,6 @@ pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
     }
 }
 
-/// Allocate n contiguous physical pages and return the physical
-/// addresses of the pages in question.
-/// # Safety
-/// This function is not inherently unsafe but will likely become so in the future
-/// when a real physical page allocator is implemented.
-/// # Panics
-/// This function will panic if:
-/// - The Layout creation fails
-/// - Memory allocation fails
-pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
-    // Currently, since all of main memory is idmap'd, we can just
-    // allocate any appropriately aligned section of memory.
-    unsafe {
-        let v = alloc::alloc::alloc_zeroed(
-            Layout::from_size_align(n as usize * OS_PAGE_SIZE as usize, OS_PAGE_SIZE as usize)
-                .expect("could not create physical page allocation layout"),
-        );
-        if v.is_null() {
-            panic!("could not allocate a physical page");
-        }
-        v as u64
-    }
 }
 
 pub fn flush_tlb() {

src/hyperlight_host/src/hypervisor/gdb/mod.rs

@@ -497,7 +497,7 @@ mod tests {
                 .inspect_err(|_| unsafe {
                     libc::munmap(mapped_mem, size);
                 })?;
-            let (mem_mgr, _) = sandbox.mgr.build();
+            let (mem_mgr, _) = sandbox.mgr.build()?;
 
             // Create the memory access struct
             let mem_access = DebugMemoryAccess {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -525,7 +525,7 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mut mem_mgr, gshm) = sandbox.mgr.build();
+        let (mut mem_mgr, gshm) = sandbox.mgr.build().unwrap();
         let mut vm = set_up_hypervisor_partition(
             gshm,
             &config,

src/hyperlight_host/src/mem/mgr.rs

@@ -95,13 +95,18 @@ impl vmem::TableOps for GuestPageTableBuffer {
             .unwrap_or(0)
     }
 
-    unsafe fn write_entry(&self, addr: (usize, usize), x: PageTableEntry) {
+    unsafe fn write_entry(
+        &self,
+        addr: (usize, usize),
+        entry: PageTableEntry,
+    ) -> Option<(usize, usize)> {
         let mut b = self.buffer.borrow_mut();
         let byte_offset =
             (addr.0 - self.phys_base / PAGE_TABLE_SIZE) * PAGE_TABLE_SIZE + addr.1 * 8;
         if let Some(slice) = b.get_mut(byte_offset..byte_offset + 8) {
-            slice.copy_from_slice(&x.to_ne_bytes());
+            slice.copy_from_slice(&entry.to_ne_bytes());
         }
+        None
     }
 
     fn to_phys(addr: (usize, usize)) -> PhysAddr {
@@ -229,36 +234,45 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
     }
 
     /// Wraps ExclusiveSharedMemory::build
+    // Morally, this should not have to be a Result: this operation is
+    // infallible. The source of the Result is
+    // update_scratch_bookkeeping(), which calls functions that can
+    // fail due to bounds checks (which are statically known to be ok
+    // in this situation) or due to failing to take the scratch shared
+    // memory lock, but the scratch shared memory is built in this
+    // function, its lock does not escape before the end of the
+    // function, and the lock is taken by no other code path, so we
+    // know it is not contended.
     pub fn build(
         self,
-    ) -> (
+    ) -> Result<(
         SandboxMemoryManager<HostSharedMemory>,
         SandboxMemoryManager<GuestSharedMemory>,
-    ) {
+    )> {
         let (hshm, gshm) = self.shared_mem.build();
         let (hscratch, gscratch) = self.scratch_mem.build();
-        (
-            SandboxMemoryManager {
-                shared_mem: hshm,
-                scratch_mem: hscratch,
-                layout: self.layout,
-                load_addr: self.load_addr.clone(),
-                entrypoint_offset: self.entrypoint_offset,
-                mapped_rgns: self.mapped_rgns,
-                stack_cookie: self.stack_cookie,
-                abort_buffer: self.abort_buffer,
-            },
-            SandboxMemoryManager {
-                shared_mem: gshm,
-                scratch_mem: gscratch,
-                layout: self.layout,
-                load_addr: self.load_addr.clone(),
-                entrypoint_offset: self.entrypoint_offset,
-                mapped_rgns: self.mapped_rgns,
-                stack_cookie: self.stack_cookie,
-                abort_buffer: Vec::new(), // Guest doesn't need abort buffer
-            },
-        )
+        let mut host_mgr = SandboxMemoryManager {
+            shared_mem: hshm,
+            scratch_mem: hscratch,
+            layout: self.layout,
+            load_addr: self.load_addr.clone(),
+            entrypoint_offset: self.entrypoint_offset,
+            mapped_rgns: self.mapped_rgns,
+            stack_cookie: self.stack_cookie,
+            abort_buffer: self.abort_buffer,
+        };
+        let guest_mgr = SandboxMemoryManager {
+            shared_mem: gshm,
+            scratch_mem: gscratch,
+            layout: self.layout,
+            load_addr: self.load_addr.clone(),
+            entrypoint_offset: self.entrypoint_offset,
+            mapped_rgns: self.mapped_rgns,
+            stack_cookie: self.stack_cookie,
+            abort_buffer: Vec::new(), // Guest doesn't need abort buffer
+        };
+        host_mgr.update_scratch_bookkeeping()?;
+        Ok((host_mgr, guest_mgr))
     }
 }
 
@@ -400,9 +414,9 @@ impl SandboxMemoryManager<HostSharedMemory> {
         }
         self.shared_mem.restore_from_snapshot(snapshot)?;
         let new_scratch_size = snapshot.layout().get_scratch_size();
-        if new_scratch_size == self.scratch_mem.mem_size() {
+        let gscratch = if new_scratch_size == self.scratch_mem.mem_size() {
             self.scratch_mem.zero()?;
-            Ok(None)
+            None
         } else {
             let new_scratch_mem = ExclusiveSharedMemory::new(new_scratch_size)?;
             let (hscratch, gscratch) = new_scratch_mem.build();
@@ -413,8 +427,28 @@ impl SandboxMemoryManager<HostSharedMemory> {
             // has been unmapped from the VM.
             self.scratch_mem = hscratch;
 
-            Ok(Some(gscratch))
-        }
+            Some(gscratch)
+        };
+        self.update_scratch_bookkeeping()?;
+        Ok(gscratch)
+    }
+
+    fn update_scratch_bookkeeping(&mut self) -> Result<()> {
+        let scratch_size = self.scratch_mem.mem_size();
+
+        let size_offset =
+            scratch_size - hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET as usize;
+        self.scratch_mem
+            .write::<u64>(size_offset, scratch_size as u64)?;
+
+        let alloc_offset =
+            scratch_size - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize;
+        self.scratch_mem.write::<u64>(
+            alloc_offset,
+            hyperlight_common::layout::scratch_base_gpa(scratch_size),
+        )?;
+
+        Ok(())
     }
 }
 

src/hyperlight_host/src/sandbox/outb.rs

@@ -287,7 +287,7 @@ mod tests {
             layout
                 .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                 .unwrap();
-            let (hmgr, _) = mgr.build();
+            let (hmgr, _) = mgr.build().unwrap();
             hmgr
         };
         {
@@ -399,7 +399,7 @@ mod tests {
                 layout
                     .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                     .unwrap();
-                let (hmgr, _) = mgr.build();
+                let (hmgr, _) = mgr.build().unwrap();
                 hmgr
             };