hyperlight-host 0.16.0

/*
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.
 */
//! This module describes the virtual and physical addresses of a
//! number of special regions in the hyperlight VM, although we hope
//! to reduce the number of these over time.
//!
//! A snapshot freshly created from an empty VM will result in roughly
//! the following physical layout:
//!
//! +-------------------------------------------+
//! |             Guest Page Tables             |
//! +-------------------------------------------+
//! |              Init Data                    | (GuestBlob size)
//! +-------------------------------------------+
//! |             Guest Heap                    |
//! +-------------------------------------------+
//! |                PEB Struct                 | (HyperlightPEB size)
//! +-------------------------------------------+
//! |               Guest Code                  |
//! +-------------------------------------------+ 0x1_000
//! |              NULL guard page              |
//! +-------------------------------------------+ 0x0_000
//!
//! Everything except for the guest page tables is currently
//! identity-mapped; the guest page tables themselves are mapped at
//! [`hyperlight_common::layout::SNAPSHOT_PT_GVA`] =
//! 0xffff_8000_0000_0000.
//!
//! - `InitData` - some extra data that can be loaded onto the sandbox during
//!   initialization.
//!
//! - `GuestHeap` - this is a buffer that is used for heap data in the guest. the length
//!   of this field is returned by the `heap_size()` method of this struct
//!
//! There is also a scratch region at the top of physical memory,
//! which is mostly laid out as a large undifferentiated blob of
//! memory, although at present the snapshot process specially
//! privileges the statically allocated input and output data regions:
//!
//! +-------------------------------------------+ (top of physical memory)
//! |         Exception Stack, Metadata         |
//! +-------------------------------------------+ (1 page below)
//! |              Scratch Memory               |
//! +-------------------------------------------+
//! |                Output Data                |
//! +-------------------------------------------+
//! |                Input Data                 |
//! +-------------------------------------------+ (scratch size)

use std::fmt::Debug;
use std::mem::size_of;

use hyperlight_common::mem::{HyperlightPEB, PAGE_SIZE_USIZE};
use tracing::{Span, instrument};

use super::memory_region::MemoryRegionType::{Code, Heap, InitData, Peb};
use super::memory_region::{
    DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegion, MemoryRegion_, MemoryRegionFlags, MemoryRegionKind,
    MemoryRegionVecBuilder,
};
#[cfg(readable_shared_mem)]
use super::shared_mem::HostSharedMemory;
use super::shared_mem::{ExclusiveSharedMemory, ReadonlySharedMemory};
use crate::error::HyperlightError::{MemoryRequestTooBig, MemoryRequestTooSmall};
use crate::sandbox::SandboxConfiguration;
use crate::{Result, new_error};

pub(crate) enum BaseGpaRegion<Sn, Sc> {
    Snapshot(Sn),
    Scratch(Sc),
    Mmap(MemoryRegion),
}

// It's an invariant of this type, checked on creation, that the
// offset is in bounds for the base region.
pub(crate) struct ResolvedGpa<Sn, Sc> {
    pub(crate) offset: usize,
    pub(crate) base: BaseGpaRegion<Sn, Sc>,
}

impl AsRef<[u8]> for ExclusiveSharedMemory {
    fn as_ref(&self) -> &[u8] {
        self.as_slice()
    }
}
impl AsRef<[u8]> for ReadonlySharedMemory {
    fn as_ref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl<Sn, Sc> ResolvedGpa<Sn, Sc> {
    pub(crate) fn with_memories<Sn2, Sc2>(self, sn: Sn2, sc: Sc2) -> ResolvedGpa<Sn2, Sc2> {
        ResolvedGpa {
            offset: self.offset,
            base: match self.base {
                BaseGpaRegion::Snapshot(_) => BaseGpaRegion::Snapshot(sn),
                BaseGpaRegion::Scratch(_) => BaseGpaRegion::Scratch(sc),
                BaseGpaRegion::Mmap(r) => BaseGpaRegion::Mmap(r),
            },
        }
    }
}
impl<'a> BaseGpaRegion<&'a [u8], &'a [u8]> {
    pub(crate) fn as_ref<'b>(&'b self) -> &'a [u8] {
        match self {
            BaseGpaRegion::Snapshot(sn) => sn,
            BaseGpaRegion::Scratch(sc) => sc,
            BaseGpaRegion::Mmap(r) => unsafe {
                #[allow(clippy::useless_conversion)]
                let host_region_base: usize = r.host_region.start.into();
                #[allow(clippy::useless_conversion)]
                let host_region_end: usize = r.host_region.end.into();
                let len = host_region_end - host_region_base;
                std::slice::from_raw_parts(host_region_base as *const u8, len)
            },
        }
    }
}
impl<'a> ResolvedGpa<&'a [u8], &'a [u8]> {
    pub(crate) fn as_ref<'b>(&'b self) -> &'a [u8] {
        let base = self.base.as_ref();
        if self.offset > base.len() {
            return &[];
        }
        &self.base.as_ref()[self.offset..]
    }
}
/// A read-only abstraction over the different kinds of backing memory
/// a [`ResolvedGpa`] can point at (the host snapshot mapping, the
/// scratch mapping, or a raw `&[u8]` view of either), letting callers
/// copy guest bytes out without caring which concrete memory type they
/// hold.
///
/// This trait only exists in builds that actually read guest memory
/// through it — see the `readable_shared_mem` cfg alias in `build.rs`
/// for the exact conditions (the `gdb` debug path and the
/// shared-snapshot `mem_profile` path). In every other configuration it
/// is compiled out entirely, so there is no dead code to `#[allow]`.
#[cfg(readable_shared_mem)]
pub(crate) trait ReadableSharedMemory {
    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()>;
}
#[cfg(readable_shared_mem)]
impl ReadableSharedMemory for &HostSharedMemory {
    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
        HostSharedMemory::copy_to_slice(self, slice, offset)
    }
}
/// Coherence workaround for the blanket impl below.
///
/// We want `ReadableSharedMemory` for both `&HostSharedMemory` (above)
/// and for any `T: AsRef<[u8]>` (so that `ExclusiveSharedMemory` /
/// `ReadonlySharedMemory` and their references are covered by a single
/// impl). A naive `impl<T: AsRef<[u8]>> ReadableSharedMemory for T`
/// would *overlap* the `&HostSharedMemory` impl — the compiler can't
/// prove `&HostSharedMemory` never implements `AsRef<[u8]>` — and is
/// rejected with E0119.
///
/// To break the overlap we introduce a private marker trait and
/// implement it *only* for the specific types we want the blanket impl
/// to cover (deliberately excluding `&HostSharedMemory`). The blanket
/// impl is then bounded on this marker rather than on `AsRef<[u8]>`
/// directly, so the two impls provably never overlap.
#[cfg(readable_shared_mem)]
mod coherence_hack {
    use super::{ExclusiveSharedMemory, ReadonlySharedMemory};
    // Used only as a bound on the blanket impl below, so the name reads
    // as unused even though removing it breaks compilation.
    #[allow(unused)]
    pub(super) trait SharedMemoryAsRefMarker: AsRef<[u8]> {}
    impl SharedMemoryAsRefMarker for ExclusiveSharedMemory {}
    impl SharedMemoryAsRefMarker for &ExclusiveSharedMemory {}
    impl SharedMemoryAsRefMarker for ReadonlySharedMemory {}
    impl SharedMemoryAsRefMarker for &ReadonlySharedMemory {}
}
#[cfg(readable_shared_mem)]
impl<T: coherence_hack::SharedMemoryAsRefMarker> ReadableSharedMemory for T {
    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
        let ss: &[u8] = self.as_ref();
        let end = offset + slice.len();
        if end > ss.len() {
            return Err(new_error!(
                "Attempt to read up to {} in memory of size {}",
                offset + slice.len(),
                self.as_ref().len()
            ));
        }
        slice.copy_from_slice(&ss[offset..end]);
        Ok(())
    }
}
/// Copy `slice.len()` bytes out of the resolved guest region.
///
/// Only the `gdb` debug path uses this one-argument convenience (it
/// already carries the offset inside `self`); `mem_profile` reads via
/// the two-argument inherent methods instead. Hence it is gated on the
/// `gdb` cfg alone, even though the [`ReadableSharedMemory`] trait it
/// relies on is available slightly more widely.
#[cfg(gdb)]
impl<Sn: ReadableSharedMemory, Sc: ReadableSharedMemory> ResolvedGpa<Sn, Sc> {
    pub(crate) fn copy_to_slice(&self, slice: &mut [u8]) -> Result<()> {
        match &self.base {
            BaseGpaRegion::Snapshot(sn) => sn.copy_to_slice(slice, self.offset),
            BaseGpaRegion::Scratch(sc) => sc.copy_to_slice(slice, self.offset),
            BaseGpaRegion::Mmap(r) => unsafe {
                #[allow(clippy::useless_conversion)]
                let host_region_base: usize = r.host_region.start.into();
                #[allow(clippy::useless_conversion)]
                let host_region_end: usize = r.host_region.end.into();
                let len = host_region_end - host_region_base;
                // Safety: it's a documented invariant of MemoryRegion
                // that the memory must remain alive as long as the
                // sandbox is alive, and the way this code is used,
                // the lifetimes of the snapshot and scratch memories
                // ensure that the sandbox is still alive. This could
                // perhaps be cleaned up/improved/made harder to
                // misuse significantly, but it would require a much
                // larger rework.
                let ss = std::slice::from_raw_parts(host_region_base as *const u8, len);
                let end = self.offset + slice.len();
                if end > ss.len() {
                    return Err(new_error!(
                        "Attempt to read up to {} in memory of size {}",
                        self.offset + slice.len(),
                        ss.len()
                    ));
                }
                slice.copy_from_slice(&ss[self.offset..end]);
                Ok(())
            },
        }
    }
}

#[derive(Copy, Clone)]
pub(crate) struct SandboxMemoryLayout {
    /// Input data buffer size (from SandboxConfiguration).
    pub(crate) input_data_size: usize,
    /// Output data buffer size (from SandboxConfiguration).
    pub(crate) output_data_size: usize,
    /// The heap size of this sandbox.
    pub(crate) heap_size: usize,
    /// The size of the guest code section.
    pub(crate) code_size: usize,
    /// The size of the init data section (guest blob).
    pub(crate) init_data_size: usize,
    /// Permission flags for the init data region.
    pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
    /// The size of the scratch region in physical memory.
    pub(crate) scratch_size: usize,
    /// Size of the primary guest memory region at `BASE_ADDRESS`
    /// (code, PEB, heap, init data). For a snapshot-backed layout
    /// this is also the guest-visible prefix of the host snapshot
    /// mapping.
    pub(crate) snapshot_size: usize,
    /// Size of the page-table region. Sits at the tail of the host
    /// snapshot mapping but is never mapped to the guest from there.
    /// On restore the host copies it into scratch, where the guest
    /// sees it at `SNAPSHOT_PT_GVA`. `None` until page tables are built.
    pub(crate) pt_size: Option<usize>,
}

impl Debug for SandboxMemoryLayout {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut ff = f.debug_struct("SandboxMemoryLayout");
        ff.field(
            "Total Memory Size",
            &format_args!("{:#x}", self.get_memory_size().unwrap_or(0)),
        )
        .field("Code Size", &format_args!("{:#x}", self.code_size))
        .field("Heap Size", &format_args!("{:#x}", self.heap_size))
        .field(
            "Init Data Size",
            &format_args!("{:#x}", self.init_data_size),
        )
        .field(
            "Input Data Size",
            &format_args!("{:#x}", self.input_data_size),
        )
        .field(
            "Output Data Size",
            &format_args!("{:#x}", self.output_data_size),
        )
        .field("Scratch Size", &format_args!("{:#x}", self.scratch_size))
        .field("Snapshot Size", &format_args!("{:#x}", self.snapshot_size))
        .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0)))
        .field(
            "Guest Code Offset",
            &format_args!("{:#x}", self.guest_code_offset()),
        )
        .field("PEB Offset", &format_args!("{:#x}", self.peb_offset()))
        .field("PEB Address", &format_args!("{:#x}", self.peb_address()));
        ff.field(
            "Guest Heap Buffer Offset",
            &format_args!("{:#x}", self.guest_heap_buffer_offset()),
        )
        .field(
            "Init Data Offset",
            &format_args!("{:#x}", self.init_data_offset()),
        )
        .finish()
    }
}

impl SandboxMemoryLayout {
    /// Whether `other` has the same layout configuration as `self`,
    /// i.e. the fields that come from the guest binary and the
    /// `SandboxConfiguration`. `snapshot_size` and `pt_size` are
    /// excluded because they are outputs of building a snapshot blob
    /// (the compacted data size and the size of the rebuilt
    /// page-table tail), not configuration inputs, so they differ
    /// between the sandbox's live layout and any snapshot taken
    /// from it.
    ///
    /// TODO: separate/remove snapshot_size and pt_size from this struct.
    pub(crate) fn is_compatible_with(&self, other: &Self) -> bool {
        // Exhaustive destructure so adding a field to
        // `SandboxMemoryLayout` fails to compile here, forcing the
        // author to decide whether it participates in compatibility.
        let Self {
            input_data_size,
            output_data_size,
            heap_size,
            code_size,
            init_data_size,
            init_data_permissions,
            scratch_size,
            snapshot_size: _,
            pt_size: _,
        } = self;
        *input_data_size == other.input_data_size
            && *output_data_size == other.output_data_size
            && *heap_size == other.heap_size
            && *code_size == other.code_size
            && *init_data_size == other.init_data_size
            && *init_data_permissions == other.init_data_permissions
            && *scratch_size == other.scratch_size
    }

    /// The maximum amount of memory a single sandbox will be allowed.
    ///
    /// Both the scratch region and the snapshot region are bounded by
    /// this size. The value is arbitrary but chosen to be large enough
    /// for most workloads while preventing accidental resource exhaustion.
    pub(crate) const MAX_MEMORY_SIZE: usize = (16 * 1024 * 1024 * 1024) - Self::BASE_ADDRESS; // 16 GiB - BASE_ADDRESS

    /// The base address of the sandbox's memory.
    pub(crate) const BASE_ADDRESS: usize = 0x1000;

    // the offset into a sandbox's input/output buffer where the stack starts
    pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8;

    /// Create a new `SandboxMemoryLayout` with the given
    /// `SandboxConfiguration`, code size and stack/heap size.
    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn new(
        cfg: SandboxConfiguration,
        code_size: usize,
        init_data_size: usize,
        init_data_permissions: Option<MemoryRegionFlags>,
    ) -> Result<Self> {
        let heap_size = usize::try_from(cfg.get_heap_size())?;
        let scratch_size = cfg.get_scratch_size();
        if scratch_size > Self::MAX_MEMORY_SIZE {
            return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE));
        }
        let input_data_size = cfg.get_input_data_size();
        let output_data_size = cfg.get_output_data_size();
        let min_scratch_size =
            hyperlight_common::layout::min_scratch_size(input_data_size, output_data_size);
        if scratch_size < min_scratch_size {
            return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size));
        }

        let mut ret = Self {
            input_data_size,
            output_data_size,
            heap_size,
            code_size,
            init_data_size,
            init_data_permissions,
            pt_size: None,
            scratch_size,
            snapshot_size: 0,
        };
        ret.set_snapshot_size(ret.get_memory_size()?);
        Ok(ret)
    }

    /// Offset of the PEB struct within the snapshot region.
    pub(crate) fn peb_offset(&self) -> usize {
        self.code_size.next_multiple_of(PAGE_SIZE_USIZE)
    }

    /// Guest physical address of the PEB.
    pub(crate) fn peb_address(&self) -> usize {
        Self::BASE_ADDRESS + self.peb_offset()
    }

    /// Offset of the guest heap buffer within the snapshot region.
    pub(crate) fn guest_heap_buffer_offset(&self) -> usize {
        (self.peb_offset() + size_of::<HyperlightPEB>()).next_multiple_of(PAGE_SIZE_USIZE)
    }

    /// Offset of the init data section within the snapshot region.
    pub(crate) fn init_data_offset(&self) -> usize {
        (self.guest_heap_buffer_offset() + self.heap_size).next_multiple_of(PAGE_SIZE_USIZE)
    }

    /// The code offset is always 0.
    pub(crate) fn guest_code_offset(&self) -> usize {
        0
    }

    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_scratch_size(&self) -> usize {
        self.scratch_size
    }

    /// Get the guest virtual address of the start of output data.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_output_data_buffer_gva(&self) -> u64 {
        hyperlight_common::layout::scratch_base_gva(self.scratch_size) + self.input_data_size as u64
    }

    /// Get the offset into the host scratch buffer of the start of
    /// the output data.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_output_data_buffer_scratch_host_offset(&self) -> usize {
        self.input_data_size
    }

    /// Get the guest virtual address of the start of input data
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    fn get_input_data_buffer_gva(&self) -> u64 {
        hyperlight_common::layout::scratch_base_gva(self.scratch_size)
    }

    /// Get the offset into the host scratch buffer of the start of
    /// the input data
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_input_data_buffer_scratch_host_offset(&self) -> usize {
        0
    }

    /// Get the offset from the beginning of the scratch region to the
    /// location where page tables will be eagerly copied on restore
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_pt_base_scratch_offset(&self) -> usize {
        (self.input_data_size + self.output_data_size)
            .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE)
    }

    /// Get the base GPA to which the page tables will be eagerly
    /// copied on restore
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_pt_base_gpa(&self) -> u64 {
        hyperlight_common::layout::scratch_base_gpa(self.scratch_size)
            + self.get_pt_base_scratch_offset() as u64
    }

    /// Get the first GPA of the scratch region that the host hasn't
    /// used for something else
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 {
        self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64
    }

    /// Get the total size of guest memory in `self`'s memory
    /// layout.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    fn get_unaligned_memory_size(&self) -> usize {
        self.init_data_offset() + self.init_data_size
    }

    /// get the code offset
    /// This is the offset in the sandbox memory where the code starts
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_guest_code_offset(&self) -> usize {
        self.guest_code_offset()
    }

    /// Get the guest address of the code section in the sandbox
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_guest_code_address(&self) -> usize {
        Self::BASE_ADDRESS + self.guest_code_offset()
    }

    /// Get the total size of guest memory in `self`'s memory
    /// layout aligned to page size boundaries.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_memory_size(&self) -> Result<usize> {
        let total_memory = self.get_unaligned_memory_size();

        // Size should be a multiple of page size.
        let remainder = total_memory % PAGE_SIZE_USIZE;
        let multiples = total_memory / PAGE_SIZE_USIZE;
        let size = match remainder {
            0 => total_memory,
            _ => (multiples + 1) * PAGE_SIZE_USIZE,
        };

        if size > Self::MAX_MEMORY_SIZE {
            Err(MemoryRequestTooBig(size, Self::MAX_MEMORY_SIZE))
        } else {
            Ok(size)
        }
    }

    /// Record the size of the page-table tail appended to the
    /// snapshot blob. The PT bytes live at the end of the blob and
    /// the host mapping, outside the guest mapping of the snapshot
    /// region, and are copied into the scratch region on restore.
    /// `snapshot_size` (the guest-visible prefix of the blob) is an
    /// independent field and must be set separately.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn set_pt_size(&mut self, size: usize) -> Result<()> {
        let min_fixed_scratch = hyperlight_common::layout::min_scratch_size(
            self.input_data_size,
            self.output_data_size,
        );
        let min_scratch = min_fixed_scratch + size;
        if self.scratch_size < min_scratch {
            return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch));
        }
        self.pt_size = Some(size);
        Ok(())
    }

    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn set_snapshot_size(&mut self, new_size: usize) {
        self.snapshot_size = new_size;
    }

    /// Get the size of the memory region used for page tables
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn get_pt_size(&self) -> usize {
        self.pt_size.unwrap_or(0)
    }

    /// Returns the memory regions associated with this memory layout,
    /// suitable for passing to a hypervisor for mapping into memory
    pub(crate) fn get_memory_regions_<K: MemoryRegionKind>(
        &self,
        host_base: K::HostBaseType,
    ) -> Result<Vec<MemoryRegion_<K>>> {
        let mut builder = MemoryRegionVecBuilder::new(Self::BASE_ADDRESS, host_base);

        // code
        let peb_offset = builder.push_page_aligned(
            self.code_size,
            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
            Code,
        );

        let expected_peb_offset = TryInto::<usize>::try_into(self.peb_offset())?;

        if peb_offset != expected_peb_offset {
            return Err(new_error!(
                "PEB offset does not match expected PEB offset expected:  {}, actual:  {}",
                expected_peb_offset,
                peb_offset
            ));
        }

        // PEB
        let heap_offset =
            builder.push_page_aligned(size_of::<HyperlightPEB>(), MemoryRegionFlags::READ, Peb);

        let expected_heap_offset = TryInto::<usize>::try_into(self.guest_heap_buffer_offset())?;

        if heap_offset != expected_heap_offset {
            return Err(new_error!(
                "Guest Heap offset does not match expected Guest Heap offset expected:  {}, actual:  {}",
                expected_heap_offset,
                heap_offset
            ));
        }

        // heap
        #[cfg(feature = "executable_heap")]
        let init_data_offset = builder.push_page_aligned(
            self.heap_size,
            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
            Heap,
        );
        #[cfg(not(feature = "executable_heap"))]
        let init_data_offset = builder.push_page_aligned(
            self.heap_size,
            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
            Heap,
        );

        let expected_init_data_offset = TryInto::<usize>::try_into(self.init_data_offset())?;

        if init_data_offset != expected_init_data_offset {
            return Err(new_error!(
                "Init Data offset does not match expected Init Data offset expected:  {}, actual:  {}",
                expected_init_data_offset,
                init_data_offset
            ));
        }

        // init data
        let after_init_offset = if self.init_data_size > 0 {
            let mem_flags = self
                .init_data_permissions
                .unwrap_or(DEFAULT_GUEST_BLOB_MEM_FLAGS);
            builder.push_page_aligned(self.init_data_size, mem_flags, InitData)
        } else {
            init_data_offset
        };

        let final_offset = after_init_offset;

        let expected_final_offset = TryInto::<usize>::try_into(self.get_memory_size()?)?;

        if final_offset != expected_final_offset {
            return Err(new_error!(
                "Final offset does not match expected Final offset expected:  {}, actual:  {}",
                expected_final_offset,
                final_offset
            ));
        }

        Ok(builder.build())
    }

    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn write_init_data(&self, out: &mut [u8], bytes: &[u8]) -> Result<()> {
        out[self.init_data_offset()..self.init_data_offset() + self.init_data_size]
            .copy_from_slice(bytes);
        Ok(())
    }

    /// Write the finished memory layout to `mem` and return `Ok` if
    /// successful.
    ///
    /// Note: `mem` may have been modified, even if `Err` was returned
    /// from this function.
    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub(crate) fn write_peb(&self, mem: &mut [u8]) -> Result<()> {
        use hyperlight_common::mem::GuestMemoryRegion;

        let guest_base = Self::BASE_ADDRESS as u64;

        let peb = HyperlightPEB {
            input_stack: GuestMemoryRegion {
                size: self.input_data_size as u64,
                ptr: self.get_input_data_buffer_gva(),
            },
            output_stack: GuestMemoryRegion {
                size: self.output_data_size as u64,
                ptr: self.get_output_data_buffer_gva(),
            },
            init_data: GuestMemoryRegion {
                size: (self.get_unaligned_memory_size() - self.init_data_offset()) as u64,
                ptr: guest_base + self.init_data_offset() as u64,
            },
            guest_heap: GuestMemoryRegion {
                size: self.heap_size as u64,
                ptr: guest_base + self.guest_heap_buffer_offset() as u64,
            },
        };

        let offset = self.peb_offset();
        let bytes = bytemuck::bytes_of(&peb);
        let end = offset + bytes.len();
        let mem_len = mem.len();
        let dst = mem.get_mut(offset..end).ok_or_else(|| {
            new_error!(
                "memory too small to write PEB: need {} bytes at offset {:#x}, have {} bytes",
                bytes.len(),
                offset,
                mem_len
            )
        })?;
        dst.copy_from_slice(bytes);

        // The input and output data regions do not have their layout
        // initialised here, because they are in the scratch
        // region---they are instead set in
        // [`SandboxMemoryManager::update_scratch_bookkeeping`].

        Ok(())
    }

    /// Determine what region this gpa is in, and its offset into that region
    pub(crate) fn resolve_gpa(
        &self,
        gpa: u64,
        mmap_regions: &[MemoryRegion],
    ) -> Option<ResolvedGpa<(), ()>> {
        let scratch_base = hyperlight_common::layout::scratch_base_gpa(self.scratch_size);
        if gpa >= scratch_base && gpa < scratch_base + self.scratch_size as u64 {
            return Some(ResolvedGpa {
                offset: (gpa - scratch_base) as usize,
                base: BaseGpaRegion::Scratch(()),
            });
        } else if gpa >= SandboxMemoryLayout::BASE_ADDRESS as u64
            && gpa < SandboxMemoryLayout::BASE_ADDRESS as u64 + self.snapshot_size as u64
        {
            return Some(ResolvedGpa {
                offset: gpa as usize - SandboxMemoryLayout::BASE_ADDRESS,
                base: BaseGpaRegion::Snapshot(()),
            });
        }
        for rgn in mmap_regions {
            if gpa >= rgn.guest_region.start as u64 && gpa < rgn.guest_region.end as u64 {
                return Some(ResolvedGpa {
                    offset: gpa as usize - rgn.guest_region.start,
                    base: BaseGpaRegion::Mmap(rgn.clone()),
                });
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use hyperlight_common::mem::PAGE_SIZE_USIZE;

    use super::*;

    // helper func for testing
    fn get_expected_memory_size(layout: &SandboxMemoryLayout) -> usize {
        let mut expected_size = 0;
        // in order of layout
        expected_size += layout.code_size;

        // PEB
        let peb_and_array = size_of::<HyperlightPEB>();
        expected_size += peb_and_array.next_multiple_of(PAGE_SIZE_USIZE);

        expected_size += layout.heap_size.next_multiple_of(PAGE_SIZE_USIZE);

        expected_size
    }

    #[test]
    fn test_get_memory_size() {
        let sbox_cfg = SandboxConfiguration::default();
        let sbox_mem_layout = SandboxMemoryLayout::new(sbox_cfg, 4096, 0, None).unwrap();
        assert_eq!(
            sbox_mem_layout.get_memory_size().unwrap(),
            get_expected_memory_size(&sbox_mem_layout)
        );
    }

    #[test]
    fn test_max_memory_sandbox() {
        let mut cfg = SandboxConfiguration::default();
        // scratch_size exceeds 16 GiB limit
        cfg.set_scratch_size(17 * 1024 * 1024 * 1024);
        cfg.set_input_data_size(16 * 1024 * 1024 * 1024);
        let layout = SandboxMemoryLayout::new(cfg, 4096, 4096, None);
        assert!(matches!(layout.unwrap_err(), MemoryRequestTooBig(..)));
    }

    #[test]
    fn is_compatible_with_identical_layouts() {
        let cfg = SandboxConfiguration::default();
        let a = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
        let b = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
        assert!(a.is_compatible_with(&b));
        assert!(b.is_compatible_with(&a));
    }

    #[test]
    fn is_compatible_with_ignores_snapshot_size_and_pt_size() {
        // `snapshot_size` and `pt_size` are outputs of building a
        // snapshot blob, not configuration inputs, so flipping
        // them must not break compatibility.
        let cfg = SandboxConfiguration::default();
        let a = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
        let mut b = a;
        b.snapshot_size = a.snapshot_size + PAGE_SIZE_USIZE;
        b.set_pt_size(PAGE_SIZE_USIZE).unwrap();
        assert!(a.is_compatible_with(&b));
        assert!(b.is_compatible_with(&a));
    }

    #[test]
    fn is_compatible_with_rejects_each_configured_field() {
        let cfg = SandboxConfiguration::default();
        let base = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();

        // Each mutation must independently break compatibility.
        let mutators: &[fn(&mut SandboxMemoryLayout)] = &[
            |l| l.input_data_size += PAGE_SIZE_USIZE,
            |l| l.output_data_size += PAGE_SIZE_USIZE,
            |l| l.heap_size += PAGE_SIZE_USIZE,
            |l| l.code_size += PAGE_SIZE_USIZE,
            |l| l.init_data_size += PAGE_SIZE_USIZE,
            |l| l.scratch_size += PAGE_SIZE_USIZE,
            |l| {
                l.init_data_permissions = Some(MemoryRegionFlags::READ);
            },
        ];
        for mutate in mutators {
            let mut other = base;
            mutate(&mut other);
            assert!(
                !base.is_compatible_with(&other),
                "mutation should have broken compatibility: {:?} vs {:?}",
                base,
                other,
            );
        }
    }
}