Skip to main content

hyperlight_host/mem/
layout.rs

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