Skip to main content

hyperlight_host/mem/
layout.rs

1/*
2Copyright 2025  The Hyperlight Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15 */
16//! This module describes the virtual and physical addresses of a
17//! number of special regions in the hyperlight VM, although we hope
18//! to reduce the number of these over time.
19//!
20//! A snapshot freshly created from an empty VM will result in roughly
21//! the following physical layout:
22//!
23//! +-------------------------------------------+
24//! |             Guest Page Tables             |
25//! +-------------------------------------------+
26//! |              Init Data                    | (GuestBlob size)
27//! +-------------------------------------------+
28//! |             Guest Heap                    |
29//! +-------------------------------------------+
30//! |                PEB Struct                 | (HyperlightPEB size)
31//! +-------------------------------------------+
32//! |               Guest Code                  |
33//! +-------------------------------------------+ 0x1_000
34//! |              NULL guard page              |
35//! +-------------------------------------------+ 0x0_000
36//!
37//! Everything except for the guest page tables is currently
38//! identity-mapped; the guest page tables themselves are mapped at
39//! [`hyperlight_common::layout::SNAPSHOT_PT_GVA`] =
40//! 0xffff_8000_0000_0000.
41//!
42//! - `InitData` - some extra data that can be loaded onto the sandbox during
43//!   initialization.
44//!
45//! - `GuestHeap` - this is a buffer that is used for heap data in the guest. the length
46//!   of this field is returned by the `heap_size()` method of this struct
47//!
48//! There is also a scratch region at the top of physical memory,
49//! which is mostly laid out as a large undifferentiated blob of
50//! memory, although at present the snapshot process specially
51//! privileges the statically allocated input and output data regions:
52//!
53//! +-------------------------------------------+ (top of physical memory)
54//! |         Exception Stack, Metadata         |
55//! +-------------------------------------------+ (1 page below)
56//! |              Scratch Memory               |
57//! +-------------------------------------------+
58//! |                Output Data                |
59//! +-------------------------------------------+
60//! |                Input Data                 |
61//! +-------------------------------------------+ (scratch size)
62
63use std::fmt::Debug;
64use std::mem::size_of;
65
66use hyperlight_common::mem::{HyperlightPEB, PAGE_SIZE_USIZE};
67use tracing::{Span, instrument};
68
69use super::memory_region::MemoryRegionType::{Code, Heap, InitData, Peb};
70use super::memory_region::{
71    DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegion, MemoryRegion_, MemoryRegionFlags, MemoryRegionKind,
72    MemoryRegionVecBuilder,
73};
74#[cfg(readable_shared_mem)]
75use super::shared_mem::HostSharedMemory;
76use super::shared_mem::{ExclusiveSharedMemory, ReadonlySharedMemory};
77use crate::error::HyperlightError::{MemoryRequestTooBig, MemoryRequestTooSmall};
78use crate::sandbox::SandboxConfiguration;
79use crate::{Result, new_error};
80
81pub(crate) enum BaseGpaRegion<Sn, Sc> {
82    Snapshot(Sn),
83    Scratch(Sc),
84    Mmap(MemoryRegion),
85}
86
87// It's an invariant of this type, checked on creation, that the
88// offset is in bounds for the base region.
89pub(crate) struct ResolvedGpa<Sn, Sc> {
90    pub(crate) offset: usize,
91    pub(crate) base: BaseGpaRegion<Sn, Sc>,
92}
93
94impl AsRef<[u8]> for ExclusiveSharedMemory {
95    fn as_ref(&self) -> &[u8] {
96        self.as_slice()
97    }
98}
99impl AsRef<[u8]> for ReadonlySharedMemory {
100    fn as_ref(&self) -> &[u8] {
101        self.as_slice()
102    }
103}
104
105impl<Sn, Sc> ResolvedGpa<Sn, Sc> {
106    pub(crate) fn with_memories<Sn2, Sc2>(self, sn: Sn2, sc: Sc2) -> ResolvedGpa<Sn2, Sc2> {
107        ResolvedGpa {
108            offset: self.offset,
109            base: match self.base {
110                BaseGpaRegion::Snapshot(_) => BaseGpaRegion::Snapshot(sn),
111                BaseGpaRegion::Scratch(_) => BaseGpaRegion::Scratch(sc),
112                BaseGpaRegion::Mmap(r) => BaseGpaRegion::Mmap(r),
113            },
114        }
115    }
116}
117impl<'a> BaseGpaRegion<&'a [u8], &'a [u8]> {
118    pub(crate) fn as_ref<'b>(&'b self) -> &'a [u8] {
119        match self {
120            BaseGpaRegion::Snapshot(sn) => sn,
121            BaseGpaRegion::Scratch(sc) => sc,
122            BaseGpaRegion::Mmap(r) => unsafe {
123                #[allow(clippy::useless_conversion)]
124                let host_region_base: usize = r.host_region.start.into();
125                #[allow(clippy::useless_conversion)]
126                let host_region_end: usize = r.host_region.end.into();
127                let len = host_region_end - host_region_base;
128                std::slice::from_raw_parts(host_region_base as *const u8, len)
129            },
130        }
131    }
132}
133impl<'a> ResolvedGpa<&'a [u8], &'a [u8]> {
134    pub(crate) fn as_ref<'b>(&'b self) -> &'a [u8] {
135        let base = self.base.as_ref();
136        if self.offset > base.len() {
137            return &[];
138        }
139        &self.base.as_ref()[self.offset..]
140    }
141}
142/// A read-only abstraction over the different kinds of backing memory
143/// a [`ResolvedGpa`] can point at (the host snapshot mapping, the
144/// scratch mapping, or a raw `&[u8]` view of either), letting callers
145/// copy guest bytes out without caring which concrete memory type they
146/// hold.
147///
148/// This trait only exists in builds that actually read guest memory
149/// through it — see the `readable_shared_mem` cfg alias in `build.rs`
150/// for the exact conditions (the `gdb` debug path and the
151/// shared-snapshot `mem_profile` path). In every other configuration it
152/// is compiled out entirely, so there is no dead code to `#[allow]`.
153#[cfg(readable_shared_mem)]
154pub(crate) trait ReadableSharedMemory {
155    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()>;
156}
157#[cfg(readable_shared_mem)]
158impl ReadableSharedMemory for &HostSharedMemory {
159    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
160        HostSharedMemory::copy_to_slice(self, slice, offset)
161    }
162}
163/// Coherence workaround for the blanket impl below.
164///
165/// We want `ReadableSharedMemory` for both `&HostSharedMemory` (above)
166/// and for any `T: AsRef<[u8]>` (so that `ExclusiveSharedMemory` /
167/// `ReadonlySharedMemory` and their references are covered by a single
168/// impl). A naive `impl<T: AsRef<[u8]>> ReadableSharedMemory for T`
169/// would *overlap* the `&HostSharedMemory` impl — the compiler can't
170/// prove `&HostSharedMemory` never implements `AsRef<[u8]>` — and is
171/// rejected with E0119.
172///
173/// To break the overlap we introduce a private marker trait and
174/// implement it *only* for the specific types we want the blanket impl
175/// to cover (deliberately excluding `&HostSharedMemory`). The blanket
176/// impl is then bounded on this marker rather than on `AsRef<[u8]>`
177/// directly, so the two impls provably never overlap.
178#[cfg(readable_shared_mem)]
179mod coherence_hack {
180    use super::{ExclusiveSharedMemory, ReadonlySharedMemory};
181    // Used only as a bound on the blanket impl below, so the name reads
182    // as unused even though removing it breaks compilation.
183    #[allow(unused)]
184    pub(super) trait SharedMemoryAsRefMarker: AsRef<[u8]> {}
185    impl SharedMemoryAsRefMarker for ExclusiveSharedMemory {}
186    impl SharedMemoryAsRefMarker for &ExclusiveSharedMemory {}
187    impl SharedMemoryAsRefMarker for ReadonlySharedMemory {}
188    impl SharedMemoryAsRefMarker for &ReadonlySharedMemory {}
189}
190#[cfg(readable_shared_mem)]
191impl<T: coherence_hack::SharedMemoryAsRefMarker> ReadableSharedMemory for T {
192    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
193        let ss: &[u8] = self.as_ref();
194        let end = offset + slice.len();
195        if end > ss.len() {
196            return Err(new_error!(
197                "Attempt to read up to {} in memory of size {}",
198                offset + slice.len(),
199                self.as_ref().len()
200            ));
201        }
202        slice.copy_from_slice(&ss[offset..end]);
203        Ok(())
204    }
205}
206/// Copy `slice.len()` bytes out of the resolved guest region.
207///
208/// Only the `gdb` debug path uses this one-argument convenience (it
209/// already carries the offset inside `self`); `mem_profile` reads via
210/// the two-argument inherent methods instead. Hence it is gated on the
211/// `gdb` cfg alone, even though the [`ReadableSharedMemory`] trait it
212/// relies on is available slightly more widely.
213#[cfg(gdb)]
214impl<Sn: ReadableSharedMemory, Sc: ReadableSharedMemory> ResolvedGpa<Sn, Sc> {
215    pub(crate) fn copy_to_slice(&self, slice: &mut [u8]) -> Result<()> {
216        match &self.base {
217            BaseGpaRegion::Snapshot(sn) => sn.copy_to_slice(slice, self.offset),
218            BaseGpaRegion::Scratch(sc) => sc.copy_to_slice(slice, self.offset),
219            BaseGpaRegion::Mmap(r) => unsafe {
220                #[allow(clippy::useless_conversion)]
221                let host_region_base: usize = r.host_region.start.into();
222                #[allow(clippy::useless_conversion)]
223                let host_region_end: usize = r.host_region.end.into();
224                let len = host_region_end - host_region_base;
225                // Safety: it's a documented invariant of MemoryRegion
226                // that the memory must remain alive as long as the
227                // sandbox is alive, and the way this code is used,
228                // the lifetimes of the snapshot and scratch memories
229                // ensure that the sandbox is still alive. This could
230                // perhaps be cleaned up/improved/made harder to
231                // misuse significantly, but it would require a much
232                // larger rework.
233                let ss = std::slice::from_raw_parts(host_region_base as *const u8, len);
234                let end = self.offset + slice.len();
235                if end > ss.len() {
236                    return Err(new_error!(
237                        "Attempt to read up to {} in memory of size {}",
238                        self.offset + slice.len(),
239                        ss.len()
240                    ));
241                }
242                slice.copy_from_slice(&ss[self.offset..end]);
243                Ok(())
244            },
245        }
246    }
247}
248
249#[derive(Copy, Clone)]
250pub(crate) struct SandboxMemoryLayout {
251    /// Input data buffer size (from SandboxConfiguration).
252    pub(crate) input_data_size: usize,
253    /// Output data buffer size (from SandboxConfiguration).
254    pub(crate) output_data_size: usize,
255    /// The heap size of this sandbox.
256    pub(crate) heap_size: usize,
257    /// The size of the guest code section.
258    pub(crate) code_size: usize,
259    /// The size of the init data section (guest blob).
260    pub(crate) init_data_size: usize,
261    /// Permission flags for the init data region.
262    pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
263    /// The size of the scratch region in physical memory.
264    pub(crate) scratch_size: usize,
265    /// Size of the primary guest memory region at `BASE_ADDRESS`
266    /// (code, PEB, heap, init data). For a snapshot-backed layout
267    /// this is also the guest-visible prefix of the host snapshot
268    /// mapping.
269    pub(crate) snapshot_size: usize,
270    /// Size of the page-table region. Sits at the tail of the host
271    /// snapshot mapping but is never mapped to the guest from there.
272    /// On restore the host copies it into scratch, where the guest
273    /// sees it at `SNAPSHOT_PT_GVA`. `None` until page tables are built.
274    pub(crate) pt_size: Option<usize>,
275}
276
277impl Debug for SandboxMemoryLayout {
278    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
279        let mut ff = f.debug_struct("SandboxMemoryLayout");
280        ff.field(
281            "Total Memory Size",
282            &format_args!("{:#x}", self.get_memory_size().unwrap_or(0)),
283        )
284        .field("Code Size", &format_args!("{:#x}", self.code_size))
285        .field("Heap Size", &format_args!("{:#x}", self.heap_size))
286        .field(
287            "Init Data Size",
288            &format_args!("{:#x}", self.init_data_size),
289        )
290        .field(
291            "Input Data Size",
292            &format_args!("{:#x}", self.input_data_size),
293        )
294        .field(
295            "Output Data Size",
296            &format_args!("{:#x}", self.output_data_size),
297        )
298        .field("Scratch Size", &format_args!("{:#x}", self.scratch_size))
299        .field("Snapshot Size", &format_args!("{:#x}", self.snapshot_size))
300        .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0)))
301        .field(
302            "Guest Code Offset",
303            &format_args!("{:#x}", self.guest_code_offset()),
304        )
305        .field("PEB Offset", &format_args!("{:#x}", self.peb_offset()))
306        .field("PEB Address", &format_args!("{:#x}", self.peb_address()));
307        ff.field(
308            "Guest Heap Buffer Offset",
309            &format_args!("{:#x}", self.guest_heap_buffer_offset()),
310        )
311        .field(
312            "Init Data Offset",
313            &format_args!("{:#x}", self.init_data_offset()),
314        )
315        .finish()
316    }
317}
318
319impl SandboxMemoryLayout {
320    /// Whether `other` has the same layout configuration as `self`,
321    /// i.e. the fields that come from the guest binary and the
322    /// `SandboxConfiguration`. `snapshot_size` and `pt_size` are
323    /// excluded because they are outputs of building a snapshot blob
324    /// (the compacted data size and the size of the rebuilt
325    /// page-table tail), not configuration inputs, so they differ
326    /// between the sandbox's live layout and any snapshot taken
327    /// from it.
328    ///
329    /// TODO: separate/remove snapshot_size and pt_size from this struct.
330    pub(crate) fn is_compatible_with(&self, other: &Self) -> bool {
331        // Exhaustive destructure so adding a field to
332        // `SandboxMemoryLayout` fails to compile here, forcing the
333        // author to decide whether it participates in compatibility.
334        let Self {
335            input_data_size,
336            output_data_size,
337            heap_size,
338            code_size,
339            init_data_size,
340            init_data_permissions,
341            scratch_size,
342            snapshot_size: _,
343            pt_size: _,
344        } = self;
345        *input_data_size == other.input_data_size
346            && *output_data_size == other.output_data_size
347            && *heap_size == other.heap_size
348            && *code_size == other.code_size
349            && *init_data_size == other.init_data_size
350            && *init_data_permissions == other.init_data_permissions
351            && *scratch_size == other.scratch_size
352    }
353
354    /// The maximum amount of memory a single sandbox will be allowed.
355    ///
356    /// Both the scratch region and the snapshot region are bounded by
357    /// this size. The value is arbitrary but chosen to be large enough
358    /// for most workloads while preventing accidental resource exhaustion.
359    pub(crate) const MAX_MEMORY_SIZE: usize = (16 * 1024 * 1024 * 1024) - Self::BASE_ADDRESS; // 16 GiB - BASE_ADDRESS
360
361    /// The base address of the sandbox's memory.
362    pub(crate) const BASE_ADDRESS: usize = 0x1000;
363
364    // the offset into a sandbox's input/output buffer where the stack starts
365    pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8;
366
367    /// Create a new `SandboxMemoryLayout` with the given
368    /// `SandboxConfiguration`, code size and stack/heap size.
369    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
370    pub(crate) fn new(
371        cfg: SandboxConfiguration,
372        code_size: usize,
373        init_data_size: usize,
374        init_data_permissions: Option<MemoryRegionFlags>,
375    ) -> Result<Self> {
376        let heap_size = usize::try_from(cfg.get_heap_size())?;
377        let scratch_size = cfg.get_scratch_size();
378        if scratch_size > Self::MAX_MEMORY_SIZE {
379            return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE));
380        }
381        let input_data_size = cfg.get_input_data_size();
382        let output_data_size = cfg.get_output_data_size();
383        let min_scratch_size =
384            hyperlight_common::layout::min_scratch_size(input_data_size, output_data_size);
385        if scratch_size < min_scratch_size {
386            return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size));
387        }
388
389        let mut ret = Self {
390            input_data_size,
391            output_data_size,
392            heap_size,
393            code_size,
394            init_data_size,
395            init_data_permissions,
396            pt_size: None,
397            scratch_size,
398            snapshot_size: 0,
399        };
400        ret.set_snapshot_size(ret.get_memory_size()?);
401        Ok(ret)
402    }
403
404    /// Offset of the PEB struct within the snapshot region.
405    pub(crate) fn peb_offset(&self) -> usize {
406        self.code_size.next_multiple_of(PAGE_SIZE_USIZE)
407    }
408
409    /// Guest physical address of the PEB.
410    pub(crate) fn peb_address(&self) -> usize {
411        Self::BASE_ADDRESS + self.peb_offset()
412    }
413
414    /// Offset of the guest heap buffer within the snapshot region.
415    pub(crate) fn guest_heap_buffer_offset(&self) -> usize {
416        (self.peb_offset() + size_of::<HyperlightPEB>()).next_multiple_of(PAGE_SIZE_USIZE)
417    }
418
419    /// Offset of the init data section within the snapshot region.
420    pub(crate) fn init_data_offset(&self) -> usize {
421        (self.guest_heap_buffer_offset() + self.heap_size).next_multiple_of(PAGE_SIZE_USIZE)
422    }
423
424    /// The code offset is always 0.
425    pub(crate) fn guest_code_offset(&self) -> usize {
426        0
427    }
428
429    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
430    pub(crate) fn get_scratch_size(&self) -> usize {
431        self.scratch_size
432    }
433
434    /// Get the guest virtual address of the start of output data.
435    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
436    pub(crate) fn get_output_data_buffer_gva(&self) -> u64 {
437        hyperlight_common::layout::scratch_base_gva(self.scratch_size) + self.input_data_size as u64
438    }
439
440    /// Get the offset into the host scratch buffer of the start of
441    /// the output data.
442    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
443    pub(crate) fn get_output_data_buffer_scratch_host_offset(&self) -> usize {
444        self.input_data_size
445    }
446
447    /// Get the guest virtual address of the start of input data
448    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
449    fn get_input_data_buffer_gva(&self) -> u64 {
450        hyperlight_common::layout::scratch_base_gva(self.scratch_size)
451    }
452
453    /// Get the offset into the host scratch buffer of the start of
454    /// the input data
455    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
456    pub(crate) fn get_input_data_buffer_scratch_host_offset(&self) -> usize {
457        0
458    }
459
460    /// Get the offset from the beginning of the scratch region to the
461    /// location where page tables will be eagerly copied on restore
462    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
463    pub(crate) fn get_pt_base_scratch_offset(&self) -> usize {
464        (self.input_data_size + self.output_data_size)
465            .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE)
466    }
467
468    /// Get the base GPA to which the page tables will be eagerly
469    /// copied on restore
470    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
471    pub(crate) fn get_pt_base_gpa(&self) -> u64 {
472        hyperlight_common::layout::scratch_base_gpa(self.scratch_size)
473            + self.get_pt_base_scratch_offset() as u64
474    }
475
476    /// Get the first GPA of the scratch region that the host hasn't
477    /// used for something else
478    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
479    pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 {
480        self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64
481    }
482
483    /// Get the total size of guest memory in `self`'s memory
484    /// layout.
485    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
486    fn get_unaligned_memory_size(&self) -> usize {
487        self.init_data_offset() + self.init_data_size
488    }
489
490    /// get the code offset
491    /// This is the offset in the sandbox memory where the code starts
492    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
493    pub(crate) fn get_guest_code_offset(&self) -> usize {
494        self.guest_code_offset()
495    }
496
497    /// Get the guest address of the code section in the sandbox
498    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
499    pub(crate) fn get_guest_code_address(&self) -> usize {
500        Self::BASE_ADDRESS + self.guest_code_offset()
501    }
502
503    /// Get the total size of guest memory in `self`'s memory
504    /// layout aligned to page size boundaries.
505    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
506    pub(crate) fn get_memory_size(&self) -> Result<usize> {
507        let total_memory = self.get_unaligned_memory_size();
508
509        // Size should be a multiple of page size.
510        let remainder = total_memory % PAGE_SIZE_USIZE;
511        let multiples = total_memory / PAGE_SIZE_USIZE;
512        let size = match remainder {
513            0 => total_memory,
514            _ => (multiples + 1) * PAGE_SIZE_USIZE,
515        };
516
517        if size > Self::MAX_MEMORY_SIZE {
518            Err(MemoryRequestTooBig(size, Self::MAX_MEMORY_SIZE))
519        } else {
520            Ok(size)
521        }
522    }
523
524    /// Record the size of the page-table tail appended to the
525    /// snapshot blob. The PT bytes live at the end of the blob and
526    /// the host mapping, outside the guest mapping of the snapshot
527    /// region, and are copied into the scratch region on restore.
528    /// `snapshot_size` (the guest-visible prefix of the blob) is an
529    /// independent field and must be set separately.
530    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
531    pub(crate) fn set_pt_size(&mut self, size: usize) -> Result<()> {
532        let min_fixed_scratch = hyperlight_common::layout::min_scratch_size(
533            self.input_data_size,
534            self.output_data_size,
535        );
536        let min_scratch = min_fixed_scratch + size;
537        if self.scratch_size < min_scratch {
538            return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch));
539        }
540        self.pt_size = Some(size);
541        Ok(())
542    }
543
544    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
545    pub(crate) fn set_snapshot_size(&mut self, new_size: usize) {
546        self.snapshot_size = new_size;
547    }
548
549    /// Get the size of the memory region used for page tables
550    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
551    pub(crate) fn get_pt_size(&self) -> usize {
552        self.pt_size.unwrap_or(0)
553    }
554
555    /// Returns the memory regions associated with this memory layout,
556    /// suitable for passing to a hypervisor for mapping into memory
557    pub(crate) fn get_memory_regions_<K: MemoryRegionKind>(
558        &self,
559        host_base: K::HostBaseType,
560    ) -> Result<Vec<MemoryRegion_<K>>> {
561        let mut builder = MemoryRegionVecBuilder::new(Self::BASE_ADDRESS, host_base);
562
563        // code
564        let peb_offset = builder.push_page_aligned(
565            self.code_size,
566            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
567            Code,
568        );
569
570        let expected_peb_offset = TryInto::<usize>::try_into(self.peb_offset())?;
571
572        if peb_offset != expected_peb_offset {
573            return Err(new_error!(
574                "PEB offset does not match expected PEB offset expected:  {}, actual:  {}",
575                expected_peb_offset,
576                peb_offset
577            ));
578        }
579
580        // PEB
581        let heap_offset =
582            builder.push_page_aligned(size_of::<HyperlightPEB>(), MemoryRegionFlags::READ, Peb);
583
584        let expected_heap_offset = TryInto::<usize>::try_into(self.guest_heap_buffer_offset())?;
585
586        if heap_offset != expected_heap_offset {
587            return Err(new_error!(
588                "Guest Heap offset does not match expected Guest Heap offset expected:  {}, actual:  {}",
589                expected_heap_offset,
590                heap_offset
591            ));
592        }
593
594        // heap
595        #[cfg(feature = "executable_heap")]
596        let init_data_offset = builder.push_page_aligned(
597            self.heap_size,
598            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
599            Heap,
600        );
601        #[cfg(not(feature = "executable_heap"))]
602        let init_data_offset = builder.push_page_aligned(
603            self.heap_size,
604            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
605            Heap,
606        );
607
608        let expected_init_data_offset = TryInto::<usize>::try_into(self.init_data_offset())?;
609
610        if init_data_offset != expected_init_data_offset {
611            return Err(new_error!(
612                "Init Data offset does not match expected Init Data offset expected:  {}, actual:  {}",
613                expected_init_data_offset,
614                init_data_offset
615            ));
616        }
617
618        // init data
619        let after_init_offset = if self.init_data_size > 0 {
620            let mem_flags = self
621                .init_data_permissions
622                .unwrap_or(DEFAULT_GUEST_BLOB_MEM_FLAGS);
623            builder.push_page_aligned(self.init_data_size, mem_flags, InitData)
624        } else {
625            init_data_offset
626        };
627
628        let final_offset = after_init_offset;
629
630        let expected_final_offset = TryInto::<usize>::try_into(self.get_memory_size()?)?;
631
632        if final_offset != expected_final_offset {
633            return Err(new_error!(
634                "Final offset does not match expected Final offset expected:  {}, actual:  {}",
635                expected_final_offset,
636                final_offset
637            ));
638        }
639
640        Ok(builder.build())
641    }
642
643    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
644    pub(crate) fn write_init_data(&self, out: &mut [u8], bytes: &[u8]) -> Result<()> {
645        out[self.init_data_offset()..self.init_data_offset() + self.init_data_size]
646            .copy_from_slice(bytes);
647        Ok(())
648    }
649
650    /// Write the finished memory layout to `mem` and return `Ok` if
651    /// successful.
652    ///
653    /// Note: `mem` may have been modified, even if `Err` was returned
654    /// from this function.
655    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
656    pub(crate) fn write_peb(&self, mem: &mut [u8]) -> Result<()> {
657        use hyperlight_common::mem::GuestMemoryRegion;
658
659        let guest_base = Self::BASE_ADDRESS as u64;
660
661        let peb = HyperlightPEB {
662            input_stack: GuestMemoryRegion {
663                size: self.input_data_size as u64,
664                ptr: self.get_input_data_buffer_gva(),
665            },
666            output_stack: GuestMemoryRegion {
667                size: self.output_data_size as u64,
668                ptr: self.get_output_data_buffer_gva(),
669            },
670            init_data: GuestMemoryRegion {
671                size: (self.get_unaligned_memory_size() - self.init_data_offset()) as u64,
672                ptr: guest_base + self.init_data_offset() as u64,
673            },
674            guest_heap: GuestMemoryRegion {
675                size: self.heap_size as u64,
676                ptr: guest_base + self.guest_heap_buffer_offset() as u64,
677            },
678        };
679
680        let offset = self.peb_offset();
681        let bytes = bytemuck::bytes_of(&peb);
682        let end = offset + bytes.len();
683        let mem_len = mem.len();
684        let dst = mem.get_mut(offset..end).ok_or_else(|| {
685            new_error!(
686                "memory too small to write PEB: need {} bytes at offset {:#x}, have {} bytes",
687                bytes.len(),
688                offset,
689                mem_len
690            )
691        })?;
692        dst.copy_from_slice(bytes);
693
694        // The input and output data regions do not have their layout
695        // initialised here, because they are in the scratch
696        // region---they are instead set in
697        // [`SandboxMemoryManager::update_scratch_bookkeeping`].
698
699        Ok(())
700    }
701
702    /// Determine what region this gpa is in, and its offset into that region
703    pub(crate) fn resolve_gpa(
704        &self,
705        gpa: u64,
706        mmap_regions: &[MemoryRegion],
707    ) -> Option<ResolvedGpa<(), ()>> {
708        let scratch_base = hyperlight_common::layout::scratch_base_gpa(self.scratch_size);
709        if gpa >= scratch_base && gpa < scratch_base + self.scratch_size as u64 {
710            return Some(ResolvedGpa {
711                offset: (gpa - scratch_base) as usize,
712                base: BaseGpaRegion::Scratch(()),
713            });
714        } else if gpa >= SandboxMemoryLayout::BASE_ADDRESS as u64
715            && gpa < SandboxMemoryLayout::BASE_ADDRESS as u64 + self.snapshot_size as u64
716        {
717            return Some(ResolvedGpa {
718                offset: gpa as usize - SandboxMemoryLayout::BASE_ADDRESS,
719                base: BaseGpaRegion::Snapshot(()),
720            });
721        }
722        for rgn in mmap_regions {
723            if gpa >= rgn.guest_region.start as u64 && gpa < rgn.guest_region.end as u64 {
724                return Some(ResolvedGpa {
725                    offset: gpa as usize - rgn.guest_region.start,
726                    base: BaseGpaRegion::Mmap(rgn.clone()),
727                });
728            }
729        }
730        None
731    }
732}
733
734#[cfg(test)]
735mod tests {
736    use hyperlight_common::mem::PAGE_SIZE_USIZE;
737
738    use super::*;
739
740    // helper func for testing
741    fn get_expected_memory_size(layout: &SandboxMemoryLayout) -> usize {
742        let mut expected_size = 0;
743        // in order of layout
744        expected_size += layout.code_size;
745
746        // PEB
747        let peb_and_array = size_of::<HyperlightPEB>();
748        expected_size += peb_and_array.next_multiple_of(PAGE_SIZE_USIZE);
749
750        expected_size += layout.heap_size.next_multiple_of(PAGE_SIZE_USIZE);
751
752        expected_size
753    }
754
755    #[test]
756    fn test_get_memory_size() {
757        let sbox_cfg = SandboxConfiguration::default();
758        let sbox_mem_layout = SandboxMemoryLayout::new(sbox_cfg, 4096, 0, None).unwrap();
759        assert_eq!(
760            sbox_mem_layout.get_memory_size().unwrap(),
761            get_expected_memory_size(&sbox_mem_layout)
762        );
763    }
764
765    #[test]
766    fn test_max_memory_sandbox() {
767        let mut cfg = SandboxConfiguration::default();
768        // scratch_size exceeds 16 GiB limit
769        cfg.set_scratch_size(17 * 1024 * 1024 * 1024);
770        cfg.set_input_data_size(16 * 1024 * 1024 * 1024);
771        let layout = SandboxMemoryLayout::new(cfg, 4096, 4096, None);
772        assert!(matches!(layout.unwrap_err(), MemoryRequestTooBig(..)));
773    }
774
775    #[test]
776    fn is_compatible_with_identical_layouts() {
777        let cfg = SandboxConfiguration::default();
778        let a = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
779        let b = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
780        assert!(a.is_compatible_with(&b));
781        assert!(b.is_compatible_with(&a));
782    }
783
784    #[test]
785    fn is_compatible_with_ignores_snapshot_size_and_pt_size() {
786        // `snapshot_size` and `pt_size` are outputs of building a
787        // snapshot blob, not configuration inputs, so flipping
788        // them must not break compatibility.
789        let cfg = SandboxConfiguration::default();
790        let a = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
791        let mut b = a;
792        b.snapshot_size = a.snapshot_size + PAGE_SIZE_USIZE;
793        b.set_pt_size(PAGE_SIZE_USIZE).unwrap();
794        assert!(a.is_compatible_with(&b));
795        assert!(b.is_compatible_with(&a));
796    }
797
798    #[test]
799    fn is_compatible_with_rejects_each_configured_field() {
800        let cfg = SandboxConfiguration::default();
801        let base = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap();
802
803        // Each mutation must independently break compatibility.
804        let mutators: &[fn(&mut SandboxMemoryLayout)] = &[
805            |l| l.input_data_size += PAGE_SIZE_USIZE,
806            |l| l.output_data_size += PAGE_SIZE_USIZE,
807            |l| l.heap_size += PAGE_SIZE_USIZE,
808            |l| l.code_size += PAGE_SIZE_USIZE,
809            |l| l.init_data_size += PAGE_SIZE_USIZE,
810            |l| l.scratch_size += PAGE_SIZE_USIZE,
811            |l| {
812                l.init_data_permissions = Some(MemoryRegionFlags::READ);
813            },
814        ];
815        for mutate in mutators {
816            let mut other = base;
817            mutate(&mut other);
818            assert!(
819                !base.is_compatible_with(&other),
820                "mutation should have broken compatibility: {:?} vs {:?}",
821                base,
822                other,
823            );
824        }
825    }
826}