hyperlight-host 0.17.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.
//! 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;
use hyperlight_common::vmem::PAGE_SIZE;
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<()> {
        Ok(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).
    input_data_size: usize,
    /// Output data buffer size (from SandboxConfiguration).
    output_data_size: usize,
    /// The heap size of this sandbox.
    heap_size: usize,
    /// The size of the guest code section.
    code_size: usize,
    /// The size of the init data section (guest blob).
    init_data_size: usize,
    /// Permission flags for the init data region.
    init_data_permissions: Option<MemoryRegionFlags>,
    /// The size of the scratch region in physical memory.
    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.
    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.
    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 {
    /// 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 = 0x4000;

    // 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)
    }

    pub(crate) fn input_data_size(&self) -> usize {
        self.input_data_size
    }

    pub(crate) fn output_data_size(&self) -> usize {
        self.output_data_size
    }

    pub(crate) fn heap_size(&self) -> usize {
        self.heap_size
    }

    pub(crate) fn code_size(&self) -> usize {
        self.code_size
    }

    pub(crate) fn init_data_size(&self) -> usize {
        self.init_data_size
    }

    pub(crate) fn init_data_permissions(&self) -> Option<MemoryRegionFlags> {
        self.init_data_permissions
    }

    pub(crate) fn get_scratch_size(&self) -> usize {
        self.scratch_size
    }

    /// Guest-visible prefix size of the snapshot blob.
    pub(crate) fn snapshot_size(&self) -> usize {
        self.snapshot_size
    }

    /// Recorded page-table tail size, `None` until page tables are built.
    pub(crate) fn pt_size(&self) -> Option<usize> {
        self.pt_size
    }

    /// Page-table tail size, or 0 if page tables are not yet built.
    pub(crate) fn get_pt_size(&self) -> usize {
        self.pt_size.unwrap_or(0)
    }

    /// 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.
    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(())
    }

    pub(crate) fn set_snapshot_size(&mut self, new_size: usize) {
        self.snapshot_size = new_size;
    }

    /// 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()?)?;

        // This function is primarily used to construct
        // GuestMemoryRegions used to populate initial guest page
        // tables. Therefore, the regions above are aligned based on
        // the guest page size. However, the total final size of the
        // region mapped into the guest needs to be aligned based on
        // the host page size. Therefore, align both of these values
        // to both page sizes before comparing them.
        let final_offset = final_offset.next_multiple_of(page_size::get());
        let expected_final_offset =
            expected_final_offset.next_multiple_of(hyperlight_common::vmem::PAGE_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
    }
}

/// Changes to the below methods is part of Snapshot ABI, and
/// changing any output shifts where the loader
/// reads captured bytes and breaks existing snapshots. Any change here
/// is a snapshot ABI break: see the `layout_offsets_are_pinned` test
/// and docs/snapshot-versioning.md.
impl SandboxMemoryLayout {
    /// Offset of the PEB struct within the snapshot region.
    pub(crate) fn peb_offset(&self) -> usize {
        self.code_size.next_multiple_of(PAGE_SIZE)
    }

    /// 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)
    }

    /// 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)
    }

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

    /// Guest address of the code section in the sandbox.
    pub(crate) fn get_guest_code_address(&self) -> usize {
        Self::BASE_ADDRESS + self.guest_code_offset()
    }

    /// Guest virtual address of the start of output data.
    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
    }

    /// Offset into the host scratch buffer of the start of the output data.
    pub(crate) fn get_output_data_buffer_scratch_host_offset(&self) -> usize {
        self.input_data_size
    }

    /// Guest virtual address of the start of input data.
    fn get_input_data_buffer_gva(&self) -> u64 {
        hyperlight_common::layout::scratch_base_gva(self.scratch_size)
    }

    /// Offset into the host scratch buffer of the start of the input data.
    pub(crate) fn get_input_data_buffer_scratch_host_offset(&self) -> usize {
        0
    }

    /// Offset from the beginning of the scratch region to the location
    /// where page tables are eagerly copied on restore.
    pub(crate) fn get_pt_base_scratch_offset(&self) -> usize {
        (self.input_data_size + self.output_data_size).next_multiple_of(PAGE_SIZE)
    }

    /// Base GPA to which the page tables are eagerly copied on restore.
    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
    }

    /// First GPA of the scratch region the host has not used for
    /// something else.
    pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 {
        self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64
    }

    /// Total size of guest memory in `self`'s memory layout.
    fn get_unaligned_memory_size(&self) -> usize {
        self.init_data_offset() + self.init_data_size
    }

    /// Total size of guest memory in `self`'s memory layout, aligned
    /// to page size boundaries.
    pub(crate) fn get_memory_size(&self) -> Result<usize> {
        let total_memory = self.get_unaligned_memory_size();

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

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

#[cfg(test)]
mod tests {
    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);

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

        expected_size.next_multiple_of(page_size::get())
    }

    #[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(..)));
    }

    /// Pinned region offsets. These methods place every region that a
    /// restored snapshot is interpreted against, so a change shifts
    /// where the loader reads captured bytes and breaks existing
    /// snapshots. Treat a failure as an ABI change: follow
    /// docs/snapshot-versioning.md rather than editing the constants.
    #[test]
    fn layout_offsets_are_pinned() {
        /// `assert_eq!` carrying the shared snapshot-ABI failure
        /// message, mirroring `abi_assert!` in the snapshot tripwires.
        macro_rules! pin_eq {
            ($left:expr, $right:expr) => {
                assert_eq!(
                    $left, $right,
                    "snapshot ABI changed: this breaks loading of existing snapshots. \
                     Do not just update the expected value to make this compile. \
                     See docs/snapshot-versioning.md."
                );
            };
        }

        // The scratch region's top GVA and GPA are baked into snapshot
        // page tables by `map_specials`, so they are part of the
        // snapshot ABI. The pins below fix offsets from this base.
        #[cfg(target_arch = "x86_64")]
        {
            pin_eq!(
                hyperlight_common::layout::SCRATCH_TOP_GVA,
                0xffff_ffff_ffff_efff
            );
            pin_eq!(
                hyperlight_common::layout::SCRATCH_TOP_GPA,
                0x0000_000f_ffff_ffff
            );
        }
        #[cfg(target_arch = "aarch64")]
        {
            pin_eq!(
                hyperlight_common::layout::SCRATCH_TOP_GVA,
                0x0000_ffff_ffff_dfff
            );
            pin_eq!(
                hyperlight_common::layout::SCRATCH_TOP_GPA,
                0x0000_000f_ffff_bfff
            );
        }

        // `map_specials` bakes the IO page mapping into snapshot page
        // tables, so its address is part of the snapshot ABI. amd64 has
        // no IO page. aarch64 maps one fixed GPA to one fixed GVA.
        #[cfg(target_arch = "x86_64")]
        pin_eq!(hyperlight_common::layout::io_page(), None);
        #[cfg(target_arch = "aarch64")]
        pin_eq!(
            hyperlight_common::layout::io_page(),
            Some((0x0000_000f_ffff_f000, 0x0000_ffff_ffff_e000))
        );

        let mut cfg = SandboxConfiguration::default();
        cfg.set_input_data_size(0x2000);
        cfg.set_output_data_size(0x2000);
        cfg.set_heap_size(0x2000);
        cfg.set_scratch_size(0x10000);
        let layout = SandboxMemoryLayout::new(cfg, 0x1000, 0, None).unwrap();

        pin_eq!(layout.guest_code_offset(), 0);
        pin_eq!(layout.peb_offset(), 0x1000);
        pin_eq!(layout.peb_address(), 0x5000);
        pin_eq!(layout.guest_heap_buffer_offset(), 0x2000);
        pin_eq!(layout.init_data_offset(), 0x4000);
        pin_eq!(layout.get_memory_size().unwrap(), 0x4000);

        pin_eq!(layout.get_scratch_size(), 0x10000);
        pin_eq!(layout.get_pt_size(), 0);

        pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0);
        pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x2000);
        pin_eq!(layout.get_pt_base_scratch_offset(), 0x4000);

        // The output buffer sits one input buffer past the input
        // buffer in the guest's scratch view.
        pin_eq!(
            layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(),
            0x2000
        );

        // The input buffer sits at the scratch base. The page tables
        // sit `get_pt_base_scratch_offset` above it. With the
        // `SCRATCH_TOP` pins above, these fix the absolute addresses.
        pin_eq!(
            layout.get_input_data_buffer_gva()
                - hyperlight_common::layout::scratch_base_gva(0x10000),
            0
        );
        pin_eq!(
            layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x10000),
            0x4000
        );
        // pt_size is zero here, so the first free scratch GPA equals
        // the page table base.
        pin_eq!(
            layout.get_first_free_scratch_gpa(),
            layout.get_pt_base_gpa()
        );

        // A second config with different sizes shifts the offsets off
        // the first config's page boundaries.
        let mut cfg = SandboxConfiguration::default();
        cfg.set_input_data_size(0x4000);
        cfg.set_output_data_size(0x2000);
        cfg.set_heap_size(0x5000);
        cfg.set_scratch_size(0x20000);
        let layout = SandboxMemoryLayout::new(cfg, 0x3000, 0, None).unwrap();

        pin_eq!(layout.guest_code_offset(), 0);
        pin_eq!(layout.peb_offset(), 0x3000);
        pin_eq!(layout.peb_address(), 0x7000);
        pin_eq!(layout.guest_heap_buffer_offset(), 0x4000);
        pin_eq!(layout.init_data_offset(), 0x9000);
        pin_eq!(
            layout.get_memory_size().unwrap(),
            0x9000_usize.next_multiple_of(page_size::get())
        );

        pin_eq!(layout.get_scratch_size(), 0x20000);
        pin_eq!(layout.get_pt_size(), 0);

        pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0);
        pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x4000);
        pin_eq!(layout.get_pt_base_scratch_offset(), 0x6000);

        pin_eq!(
            layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(),
            0x4000
        );

        pin_eq!(
            layout.get_input_data_buffer_gva()
                - hyperlight_common::layout::scratch_base_gva(0x20000),
            0
        );
        pin_eq!(
            layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000),
            0x6000
        );
        pin_eq!(
            layout.get_first_free_scratch_gpa(),
            layout.get_pt_base_gpa()
        );
    }
}