Skip to main content

nub_host_kvm/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//! `nub_host_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::{offset_of, size_of};
65
66use nub_host_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_, MemoryRegionFlags, MemoryRegionKind,
72    MemoryRegionVecBuilder,
73};
74#[cfg(any(gdb, feature = "mem_profile"))]
75use super::shared_mem::HostSharedMemory;
76use crate::error::HyperlightError::{MemoryRequestTooBig, MemoryRequestTooSmall};
77use crate::sandbox::SandboxConfiguration;
78use crate::{Result, new_error};
79
80#[cfg(any(gdb, feature = "mem_profile"))]
81#[allow(unused)] // may be unused when i686-guest is also enabled
82pub(crate) trait ReadableSharedMemory {
83    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()>;
84}
85#[cfg(any(gdb, feature = "mem_profile"))]
86impl ReadableSharedMemory for &HostSharedMemory {
87    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
88        HostSharedMemory::copy_to_slice(self, slice, offset)
89    }
90}
91#[cfg(any(gdb, feature = "mem_profile"))]
92mod coherence_hack {
93    use super::{ExclusiveSharedMemory, ReadonlySharedMemory};
94    #[allow(unused)] // it actually is; see the impl below
95    pub(super) trait SharedMemoryAsRefMarker: AsRef<[u8]> {}
96    impl SharedMemoryAsRefMarker for ExclusiveSharedMemory {}
97    impl SharedMemoryAsRefMarker for &ExclusiveSharedMemory {}
98    impl SharedMemoryAsRefMarker for ReadonlySharedMemory {}
99    impl SharedMemoryAsRefMarker for &ReadonlySharedMemory {}
100}
101#[cfg(any(gdb, feature = "mem_profile"))]
102impl<T: coherence_hack::SharedMemoryAsRefMarker> ReadableSharedMemory for T {
103    fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
104        let ss: &[u8] = self.as_ref();
105        let end = offset + slice.len();
106        if end > ss.len() {
107            return Err(new_error!(
108                "Attempt to read up to {} in memory of size {}",
109                offset + slice.len(),
110                self.as_ref().len()
111            ));
112        }
113        slice.copy_from_slice(&ss[offset..end]);
114        Ok(())
115    }
116}
117#[cfg(any(gdb, feature = "mem_profile"))]
118impl<Sn: ReadableSharedMemory, Sc: ReadableSharedMemory> ResolvedGpa<Sn, Sc> {
119    #[allow(unused)] // may be unused when i686-guest is also enabled
120    pub(crate) fn copy_to_slice(&self, slice: &mut [u8]) -> Result<()> {
121        match &self.base {
122            BaseGpaRegion::Snapshot(sn) => sn.copy_to_slice(slice, self.offset),
123            BaseGpaRegion::Scratch(sc) => sc.copy_to_slice(slice, self.offset),
124            BaseGpaRegion::Mmap(r) => unsafe {
125                #[allow(clippy::useless_conversion)]
126                let host_region_base: usize = r.host_region.start.into();
127                #[allow(clippy::useless_conversion)]
128                let host_region_end: usize = r.host_region.end.into();
129                let len = host_region_end - host_region_base;
130                // Safety: it's a documented invariant of MemoryRegion
131                // that the memory must remain alive as long as the
132                // sandbox is alive, and the way this code is used,
133                // the lifetimes of the snapshot and scratch memories
134                // ensure that the sandbox is still alive. This could
135                // perhaps be cleaned up/improved/made harder to
136                // misuse significantly, but it would require a much
137                // larger rework.
138                let ss = std::slice::from_raw_parts(host_region_base as *const u8, len);
139                let end = self.offset + slice.len();
140                if end > ss.len() {
141                    return Err(new_error!(
142                        "Attempt to read up to {} in memory of size {}",
143                        self.offset + slice.len(),
144                        ss.len()
145                    ));
146                }
147                slice.copy_from_slice(&ss[self.offset..end]);
148                Ok(())
149            },
150        }
151    }
152}
153
154#[derive(Copy, Clone)]
155pub(crate) struct SandboxMemoryLayout {
156    pub(super) sandbox_memory_config: SandboxConfiguration,
157    /// The heap size of this sandbox.
158    pub(super) heap_size: usize,
159    init_data_size: usize,
160
161    /// The following fields are offsets to the actual PEB struct fields.
162    /// They are used when writing the PEB struct itself
163    peb_offset: usize,
164    peb_input_data_offset: usize,
165    peb_output_data_offset: usize,
166    peb_init_data_offset: usize,
167    peb_heap_data_offset: usize,
168    #[cfg(feature = "nanvix-unstable")]
169    peb_file_mappings_offset: usize,
170
171    guest_heap_buffer_offset: usize,
172    init_data_offset: usize,
173    pt_size: Option<usize>,
174
175    // other
176    pub(crate) peb_address: usize,
177    code_size: usize,
178    // The offset in the sandbox memory where the code starts
179    guest_code_offset: usize,
180    #[cfg_attr(feature = "i686-guest", allow(unused))]
181    pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
182
183    // The size of the scratch region in physical memory; note that
184    // this will appear under the top of physical memory.
185    scratch_size: usize,
186    // The guest-visible size of the snapshot region in physical
187    // memory. After compaction this may be smaller than the full
188    // snapshot blob (which also contains a PT tail that is only
189    // host-accessible).
190    snapshot_size: usize,
191}
192
193impl Debug for SandboxMemoryLayout {
194    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        let mut ff = f.debug_struct("SandboxMemoryLayout");
196        ff.field(
197            "Total Memory Size",
198            &format_args!("{:#x}", self.get_memory_size().unwrap_or(0)),
199        )
200        .field("Heap Size", &format_args!("{:#x}", self.heap_size))
201        .field(
202            "Init Data Size",
203            &format_args!("{:#x}", self.init_data_size),
204        )
205        .field("PEB Address", &format_args!("{:#x}", self.peb_address))
206        .field("PEB Offset", &format_args!("{:#x}", self.peb_offset))
207        .field("Code Size", &format_args!("{:#x}", self.code_size))
208        .field(
209            "Input Data Offset",
210            &format_args!("{:#x}", self.peb_input_data_offset),
211        )
212        .field(
213            "Output Data Offset",
214            &format_args!("{:#x}", self.peb_output_data_offset),
215        )
216        .field(
217            "Init Data Offset",
218            &format_args!("{:#x}", self.peb_init_data_offset),
219        )
220        .field(
221            "Guest Heap Offset",
222            &format_args!("{:#x}", self.peb_heap_data_offset),
223        );
224        #[cfg(feature = "nanvix-unstable")]
225        ff.field(
226            "File Mappings Offset",
227            &format_args!("{:#x}", self.peb_file_mappings_offset),
228        );
229        ff.field(
230            "Guest Heap Buffer Offset",
231            &format_args!("{:#x}", self.guest_heap_buffer_offset),
232        )
233        .field(
234            "Init Data Offset",
235            &format_args!("{:#x}", self.init_data_offset),
236        )
237        .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0)))
238        .field(
239            "Guest Code Offset",
240            &format_args!("{:#x}", self.guest_code_offset),
241        )
242        .field(
243            "Scratch region size",
244            &format_args!("{:#x}", self.scratch_size),
245        )
246        .finish()
247    }
248}
249
250impl SandboxMemoryLayout {
251    /// The maximum amount of memory a single sandbox will be allowed.
252    ///
253    /// Both the scratch region and the snapshot region are bounded by
254    /// this size. The value is arbitrary but chosen to be large enough
255    /// for most workloads while preventing accidental resource exhaustion.
256    const MAX_MEMORY_SIZE: usize = (16 * 1024 * 1024 * 1024) - Self::BASE_ADDRESS; // 16 GiB - BASE_ADDRESS
257
258    /// The base address of the sandbox's memory.
259    pub(crate) const BASE_ADDRESS: usize = 0x1000;
260
261    /// Virtual-address base where the kernel is mapped.
262    ///
263    /// Matches the linker base in
264    /// [`rust/nub-arch-x86/link.x`](../../../nub-arch-x86/link.x); the
265    /// initial PT maps `[BASE_ADDRESS, ...) GPAs` to
266    /// `[KERNEL_HIGH_BASE, ...) GVAs` via a constant offset, so
267    /// `kernel_gva = KERNEL_HIGH_BASE + (gpa - BASE_ADDRESS)`.
268    ///
269    /// Now equal to `GUEST_VA_BASE_DEFAULT + KERNEL_OFFSET`
270    /// (see `nub_host_common::layout`). Lives in canonical low-half
271    /// so the host process can mmap-shadow this region at the same VA.
272    ///
273    /// TODO: rename to `KERNEL_BASE` — the name `KERNEL_HIGH_BASE` is
274    /// no longer accurate now that the kernel is in low-half. Kept
275    /// for this commit to minimise churn.
276    pub(crate) const KERNEL_HIGH_BASE: u64 = 0x5001_4000_0000;
277
278    // the offset into a sandbox's input/output buffer where the stack starts
279    pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8;
280
281    /// Create a new `SandboxMemoryLayout` with the given
282    /// `SandboxConfiguration`, code size and stack/heap size.
283    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
284    pub(crate) fn new(
285        cfg: SandboxConfiguration,
286        code_size: usize,
287        init_data_size: usize,
288        init_data_permissions: Option<MemoryRegionFlags>,
289    ) -> Result<Self> {
290        let heap_size = usize::try_from(cfg.get_heap_size())?;
291        let scratch_size = cfg.get_scratch_size();
292        if scratch_size > Self::MAX_MEMORY_SIZE {
293            return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE));
294        }
295        let min_scratch_size = nub_host_common::layout::min_scratch_size(
296            cfg.get_input_data_size(),
297            cfg.get_output_data_size(),
298        );
299        if scratch_size < min_scratch_size {
300            return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size));
301        }
302
303        let guest_code_offset = 0;
304        // The following offsets are to the fields of the PEB struct itself!
305        let peb_offset = code_size.next_multiple_of(PAGE_SIZE_USIZE);
306        let peb_input_data_offset = peb_offset + offset_of!(HyperlightPEB, input_stack);
307        let peb_output_data_offset = peb_offset + offset_of!(HyperlightPEB, output_stack);
308        let peb_init_data_offset = peb_offset + offset_of!(HyperlightPEB, init_data);
309        let peb_heap_data_offset = peb_offset + offset_of!(HyperlightPEB, guest_heap);
310        #[cfg(feature = "nanvix-unstable")]
311        let peb_file_mappings_offset = peb_offset + offset_of!(HyperlightPEB, file_mappings);
312
313        // The following offsets are the actual values that relate to memory layout,
314        // which are written to PEB struct
315        let peb_address = Self::BASE_ADDRESS + peb_offset;
316        // make sure heap buffer starts at 4K boundary.
317        // The FileMappingInfo array is stored immediately after the PEB struct.
318        // We statically reserve space for MAX_FILE_MAPPINGS entries so that
319        // the heap never overlaps the array, even when all slots are used.
320        // The host writes file mapping metadata here via write_file_mapping_entry;
321        // the guest only reads the entries. We don't know at layout time how
322        // many file mappings the host will register, so we reserve space for
323        // the maximum number.
324        // The heap starts at the next page boundary after this reserved area.
325        #[cfg(feature = "nanvix-unstable")]
326        let file_mappings_array_end = peb_offset
327            + size_of::<HyperlightPEB>()
328            + nub_host_common::mem::MAX_FILE_MAPPINGS
329                * size_of::<nub_host_common::mem::FileMappingInfo>();
330        #[cfg(feature = "nanvix-unstable")]
331        let guest_heap_buffer_offset = file_mappings_array_end.next_multiple_of(PAGE_SIZE_USIZE);
332        #[cfg(not(feature = "nanvix-unstable"))]
333        let guest_heap_buffer_offset =
334            (peb_offset + size_of::<HyperlightPEB>()).next_multiple_of(PAGE_SIZE_USIZE);
335
336        // make sure init data starts at 4K boundary
337        let init_data_offset =
338            (guest_heap_buffer_offset + heap_size).next_multiple_of(PAGE_SIZE_USIZE);
339        let mut ret = Self {
340            peb_offset,
341            heap_size,
342            peb_input_data_offset,
343            peb_output_data_offset,
344            peb_init_data_offset,
345            peb_heap_data_offset,
346            #[cfg(feature = "nanvix-unstable")]
347            peb_file_mappings_offset,
348            sandbox_memory_config: cfg,
349            code_size,
350            guest_heap_buffer_offset,
351            peb_address,
352            guest_code_offset,
353            init_data_offset,
354            init_data_size,
355            init_data_permissions,
356            pt_size: None,
357            scratch_size,
358            snapshot_size: 0,
359        };
360        ret.set_snapshot_size(ret.get_memory_size()?);
361        Ok(ret)
362    }
363
364    /// Get the offset in guest memory to the output data size
365    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
366    pub(super) fn get_output_data_size_offset(&self) -> usize {
367        // The size field is the first field in the `OutputData` struct
368        self.peb_output_data_offset
369    }
370
371    /// Get the offset in guest memory to the init data size
372    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
373    pub(super) fn get_init_data_size_offset(&self) -> usize {
374        // The init data size is the first field in the `GuestMemoryRegion` struct
375        self.peb_init_data_offset
376    }
377
378    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
379    pub(crate) fn get_scratch_size(&self) -> usize {
380        self.scratch_size
381    }
382
383    /// Get the offset in guest memory to the output data pointer.
384    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
385    fn get_output_data_pointer_offset(&self) -> usize {
386        // This field is immediately after the output data size field,
387        // which is a `u64`.
388        self.get_output_data_size_offset() + size_of::<u64>()
389    }
390
391    /// Get the offset in guest memory to the init data pointer.
392    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
393    pub(super) fn get_init_data_pointer_offset(&self) -> usize {
394        // The init data pointer is immediately after the init data size field,
395        // which is a `u64`.
396        self.get_init_data_size_offset() + size_of::<u64>()
397    }
398
399    /// Get the guest virtual address of the start of output data.
400    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
401    pub(crate) fn get_output_data_buffer_gva(&self) -> u64 {
402        nub_host_common::layout::scratch_base_gva(self.scratch_size)
403            + self.sandbox_memory_config.get_input_data_size() as u64
404    }
405
406    /// Get the offset into the host scratch buffer of the start of
407    /// the output data.
408    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
409    pub(crate) fn get_output_data_buffer_scratch_host_offset(&self) -> usize {
410        self.sandbox_memory_config.get_input_data_size()
411    }
412
413    /// Get the offset in guest memory to the input data size.
414    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
415    pub(super) fn get_input_data_size_offset(&self) -> usize {
416        // The input data size is the first field in the input stack's `GuestMemoryRegion` struct
417        self.peb_input_data_offset
418    }
419
420    /// Get the offset in guest memory to the input data pointer.
421    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
422    fn get_input_data_pointer_offset(&self) -> usize {
423        // The input data pointer is immediately after the input
424        // data size field in the input data `GuestMemoryRegion` struct which is a `u64`.
425        self.get_input_data_size_offset() + size_of::<u64>()
426    }
427
428    /// Get the guest virtual address of the start of input data
429    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
430    fn get_input_data_buffer_gva(&self) -> u64 {
431        nub_host_common::layout::scratch_base_gva(self.scratch_size)
432    }
433
434    /// Get the offset into the host scratch buffer of the start of
435    /// the input data
436    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
437    pub(crate) fn get_input_data_buffer_scratch_host_offset(&self) -> usize {
438        0
439    }
440
441    /// Get the offset from the beginning of the scratch region to the
442    /// location where page tables will be eagerly copied on restore
443    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
444    pub(crate) fn get_pt_base_scratch_offset(&self) -> usize {
445        (self.sandbox_memory_config.get_input_data_size()
446            + self.sandbox_memory_config.get_output_data_size())
447        .next_multiple_of(nub_host_common::vmem::PAGE_SIZE)
448    }
449
450    /// Get the base GPA to which the page tables will be eagerly
451    /// copied on restore
452    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
453    pub(crate) fn get_pt_base_gpa(&self) -> u64 {
454        nub_host_common::layout::scratch_base_gpa(self.scratch_size)
455            + self.get_pt_base_scratch_offset() as u64
456    }
457
458    /// Get the first GPA of the scratch region that the host hasn't
459    /// used for something else
460    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
461    pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 {
462        self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64
463    }
464
465    /// Get the offset in guest memory to the heap size
466    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
467    fn get_heap_size_offset(&self) -> usize {
468        self.peb_heap_data_offset
469    }
470
471    /// Get the offset in guest memory to the file_mappings count field
472    /// (the `size` field of the `GuestMemoryRegion` in the PEB).
473    #[cfg(feature = "nanvix-unstable")]
474    pub(crate) fn get_file_mappings_size_offset(&self) -> usize {
475        self.peb_file_mappings_offset
476    }
477
478    /// Get the offset in guest memory to the file_mappings pointer field.
479    #[cfg(feature = "nanvix-unstable")]
480    fn get_file_mappings_pointer_offset(&self) -> usize {
481        self.get_file_mappings_size_offset() + size_of::<u64>()
482    }
483
484    /// Get the offset in snapshot memory where the FileMappingInfo array starts
485    /// (immediately after the PEB struct, within the same page).
486    #[cfg(feature = "nanvix-unstable")]
487    pub(crate) fn get_file_mappings_array_offset(&self) -> usize {
488        self.peb_offset + size_of::<HyperlightPEB>()
489    }
490
491    /// Get the guest address of the FileMappingInfo array.
492    #[cfg(feature = "nanvix-unstable")]
493    fn get_file_mappings_array_gva(&self) -> u64 {
494        (Self::BASE_ADDRESS + self.get_file_mappings_array_offset()) as u64
495    }
496
497    /// Get the offset of the heap pointer in guest memory,
498    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
499    fn get_heap_pointer_offset(&self) -> usize {
500        // The heap pointer is immediately after the
501        // heap size field in the guest heap's `GuestMemoryRegion` struct which is a `u64`.
502        self.get_heap_size_offset() + size_of::<u64>()
503    }
504
505    /// Get the total size of guest memory in `self`'s memory
506    /// layout.
507    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
508    fn get_unaligned_memory_size(&self) -> usize {
509        self.init_data_offset + self.init_data_size
510    }
511
512    /// get the code offset
513    /// This is the offset in the sandbox memory where the code starts
514    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
515    pub(crate) fn get_guest_code_offset(&self) -> usize {
516        self.guest_code_offset
517    }
518
519    /// Get the guest address of the code section in the sandbox
520    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
521    pub(crate) fn get_guest_code_address(&self) -> usize {
522        Self::BASE_ADDRESS + self.guest_code_offset
523    }
524
525    /// Get the total size of guest memory in `self`'s memory
526    /// layout aligned to page size boundaries.
527    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
528    pub(crate) fn get_memory_size(&self) -> Result<usize> {
529        let total_memory = self.get_unaligned_memory_size();
530
531        // Size should be a multiple of page size.
532        let remainder = total_memory % PAGE_SIZE_USIZE;
533        let multiples = total_memory / PAGE_SIZE_USIZE;
534        let size = match remainder {
535            0 => total_memory,
536            _ => (multiples + 1) * PAGE_SIZE_USIZE,
537        };
538
539        if size > Self::MAX_MEMORY_SIZE {
540            Err(MemoryRequestTooBig(size, Self::MAX_MEMORY_SIZE))
541        } else {
542            Ok(size)
543        }
544    }
545
546    /// Sets the size of the memory region used for page tables
547    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
548    pub(crate) fn set_pt_size(&mut self, size: usize) -> Result<()> {
549        let min_fixed_scratch = nub_host_common::layout::min_scratch_size(
550            self.sandbox_memory_config.get_input_data_size(),
551            self.sandbox_memory_config.get_output_data_size(),
552        );
553        let min_scratch = min_fixed_scratch + size;
554        if self.scratch_size < min_scratch {
555            return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch));
556        }
557        let old_pt_size = self.pt_size.unwrap_or(0);
558        self.snapshot_size = self.snapshot_size - old_pt_size + size;
559        self.pt_size = Some(size);
560        Ok(())
561    }
562
563    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
564    pub(crate) fn set_snapshot_size(&mut self, new_size: usize) {
565        self.snapshot_size = new_size;
566    }
567
568    /// Get the size of the memory region used for page tables
569    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
570    pub(crate) fn get_pt_size(&self) -> usize {
571        self.pt_size.unwrap_or(0)
572    }
573
574    /// Returns the memory regions associated with this memory layout,
575    /// suitable for passing to a hypervisor for mapping into memory
576    #[cfg_attr(feature = "i686-guest", allow(unused))]
577    pub(crate) fn get_memory_regions_<K: MemoryRegionKind>(
578        &self,
579        host_base: K::HostBaseType,
580    ) -> Result<Vec<MemoryRegion_<K>>> {
581        let mut builder = MemoryRegionVecBuilder::new(Self::BASE_ADDRESS, host_base);
582
583        // code
584        let peb_offset = builder.push_page_aligned(
585            self.code_size,
586            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
587            Code,
588        );
589
590        let expected_peb_offset = TryInto::<usize>::try_into(self.peb_offset)?;
591
592        if peb_offset != expected_peb_offset {
593            return Err(new_error!(
594                "PEB offset does not match expected PEB offset expected:  {}, actual:  {}",
595                expected_peb_offset,
596                peb_offset
597            ));
598        }
599
600        // PEB + preallocated FileMappingInfo array
601        #[cfg(feature = "nanvix-unstable")]
602        let heap_offset = {
603            let peb_and_array_size = size_of::<HyperlightPEB>()
604                + nub_host_common::mem::MAX_FILE_MAPPINGS
605                    * size_of::<nub_host_common::mem::FileMappingInfo>();
606            builder.push_page_aligned(
607                peb_and_array_size,
608                MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
609                Peb,
610            )
611        };
612        #[cfg(not(feature = "nanvix-unstable"))]
613        let heap_offset =
614            builder.push_page_aligned(size_of::<HyperlightPEB>(), MemoryRegionFlags::READ, Peb);
615
616        let expected_heap_offset = TryInto::<usize>::try_into(self.guest_heap_buffer_offset)?;
617
618        if heap_offset != expected_heap_offset {
619            return Err(new_error!(
620                "Guest Heap offset does not match expected Guest Heap offset expected:  {}, actual:  {}",
621                expected_heap_offset,
622                heap_offset
623            ));
624        }
625
626        // heap
627        #[cfg(feature = "executable_heap")]
628        let init_data_offset = builder.push_page_aligned(
629            self.heap_size,
630            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
631            Heap,
632        );
633        #[cfg(not(feature = "executable_heap"))]
634        let init_data_offset = builder.push_page_aligned(
635            self.heap_size,
636            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
637            Heap,
638        );
639
640        let expected_init_data_offset = TryInto::<usize>::try_into(self.init_data_offset)?;
641
642        if init_data_offset != expected_init_data_offset {
643            return Err(new_error!(
644                "Init Data offset does not match expected Init Data offset expected:  {}, actual:  {}",
645                expected_init_data_offset,
646                init_data_offset
647            ));
648        }
649
650        // init data
651        let after_init_offset = if self.init_data_size > 0 {
652            let mem_flags = self
653                .init_data_permissions
654                .unwrap_or(DEFAULT_GUEST_BLOB_MEM_FLAGS);
655            builder.push_page_aligned(self.init_data_size, mem_flags, InitData)
656        } else {
657            init_data_offset
658        };
659
660        let final_offset = after_init_offset;
661
662        let expected_final_offset = TryInto::<usize>::try_into(self.get_memory_size()?)?;
663
664        if final_offset != expected_final_offset {
665            return Err(new_error!(
666                "Final offset does not match expected Final offset expected:  {}, actual:  {}",
667                expected_final_offset,
668                final_offset
669            ));
670        }
671
672        Ok(builder.build())
673    }
674
675    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
676    pub(crate) fn write_init_data(&self, out: &mut [u8], bytes: &[u8]) -> Result<()> {
677        out[self.init_data_offset..self.init_data_offset + self.init_data_size]
678            .copy_from_slice(bytes);
679        Ok(())
680    }
681
682    /// Write the finished memory layout to `mem` and return `Ok` if
683    /// successful.
684    ///
685    /// Note: `mem` may have been modified, even if `Err` was returned
686    /// from this function.
687    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
688    pub(crate) fn write_peb(&self, mem: &mut [u8]) -> Result<()> {
689        // (No `guest_offset` here — Stage F kernel relocation replaced
690        // the legacy `get_address!` macro with `get_gva!` for the
691        // kernel-half PEB pointers, and the input/output buffer
692        // pointers compute their GVAs via `get_*_gva()`.)
693
694        fn write_u64(mem: &mut [u8], offset: usize, value: u64) -> Result<()> {
695            if offset + 8 > mem.len() {
696                return Err(new_error!(
697                    "Cannot write to offset {} in slice of len {}",
698                    offset,
699                    mem.len()
700                ));
701            }
702            mem[offset..offset + 8].copy_from_slice(&u64::to_ne_bytes(value));
703            Ok(())
704        }
705
706        // Returns the kernel-half GVA for a PEB-pointer field. Used for
707        // `init_data.ptr` / `guest_heap.ptr` — the guest dereferences
708        // these and expects high-VA pointers after Stage F kernel
709        // relocation. The legacy `get_address!` macro (GPA-returning)
710        // is gone with the file-mapping API.
711        macro_rules! get_gva {
712            ($something:ident) => {
713                Self::KERNEL_HIGH_BASE + (self.$something as u64)
714            };
715        }
716
717        // Start of setting up the PEB. The following are in the order of the PEB fields
718
719        // Set up input buffer pointer
720        write_u64(
721            mem,
722            self.get_input_data_size_offset(),
723            self.sandbox_memory_config
724                .get_input_data_size()
725                .try_into()?,
726        )?;
727        write_u64(
728            mem,
729            self.get_input_data_pointer_offset(),
730            self.get_input_data_buffer_gva(),
731        )?;
732
733        // Set up output buffer pointer
734        write_u64(
735            mem,
736            self.get_output_data_size_offset(),
737            self.sandbox_memory_config
738                .get_output_data_size()
739                .try_into()?,
740        )?;
741        write_u64(
742            mem,
743            self.get_output_data_pointer_offset(),
744            self.get_output_data_buffer_gva(),
745        )?;
746
747        // Set up init data pointer (high GVA — guest dereferences via
748        // the kernel-half mapping).
749        write_u64(
750            mem,
751            self.get_init_data_size_offset(),
752            (self.get_unaligned_memory_size() - self.init_data_offset).try_into()?,
753        )?;
754        let addr = get_gva!(init_data_offset);
755        write_u64(mem, self.get_init_data_pointer_offset(), addr)?;
756
757        // Set up heap buffer pointer (high GVA — talc's `claim` reads
758        // this directly from `(*peb).guest_heap.ptr`).
759        let addr = get_gva!(guest_heap_buffer_offset);
760        write_u64(mem, self.get_heap_size_offset(), self.heap_size.try_into()?)?;
761        write_u64(mem, self.get_heap_pointer_offset(), addr)?;
762
763        // Set up the file_mappings descriptor in the PEB.
764        // - The `size` field holds the number of valid FileMappingInfo
765        //   entries currently written (initially 0 — entries are added
766        //   later by map_file_cow / evolve).
767        // - The `ptr` field holds the guest address of the preallocated
768        //   FileMappingInfo array
769        #[cfg(feature = "nanvix-unstable")]
770        write_u64(mem, self.get_file_mappings_size_offset(), 0)?;
771        #[cfg(feature = "nanvix-unstable")]
772        write_u64(
773            mem,
774            self.get_file_mappings_pointer_offset(),
775            self.get_file_mappings_array_gva(),
776        )?;
777
778        // End of setting up the PEB
779
780        // The input and output data regions do not have their layout
781        // initialised here, because they are in the scratch
782        // region---they are instead set in
783        // [`SandboxMemoryManager::update_scratch_bookkeeping`].
784
785        Ok(())
786    }
787}
788
789#[cfg(test)]
790mod tests {
791    use nub_host_common::mem::PAGE_SIZE_USIZE;
792
793    use super::*;
794
795    // helper func for testing
796    fn get_expected_memory_size(layout: &SandboxMemoryLayout) -> usize {
797        let mut expected_size = 0;
798        // in order of layout
799        expected_size += layout.code_size;
800
801        // PEB + preallocated FileMappingInfo array
802        #[cfg(feature = "nanvix-unstable")]
803        let peb_and_array = size_of::<HyperlightPEB>()
804            + nub_host_common::mem::MAX_FILE_MAPPINGS
805                * size_of::<nub_host_common::mem::FileMappingInfo>();
806        #[cfg(not(feature = "nanvix-unstable"))]
807        let peb_and_array = size_of::<HyperlightPEB>();
808        expected_size += peb_and_array.next_multiple_of(PAGE_SIZE_USIZE);
809
810        expected_size += layout.heap_size.next_multiple_of(PAGE_SIZE_USIZE);
811
812        expected_size
813    }
814
815    #[test]
816    fn test_get_memory_size() {
817        let sbox_cfg = SandboxConfiguration::default();
818        let sbox_mem_layout = SandboxMemoryLayout::new(sbox_cfg, 4096, 0, None).unwrap();
819        assert_eq!(
820            sbox_mem_layout.get_memory_size().unwrap(),
821            get_expected_memory_size(&sbox_mem_layout)
822        );
823    }
824
825    #[test]
826    fn test_max_memory_sandbox() {
827        let mut cfg = SandboxConfiguration::default();
828        // scratch_size exceeds 16 GiB limit
829        cfg.set_scratch_size(17 * 1024 * 1024 * 1024);
830        cfg.set_input_data_size(16 * 1024 * 1024 * 1024);
831        let layout = SandboxMemoryLayout::new(cfg, 4096, 4096, None);
832        assert!(matches!(layout.unwrap_err(), MemoryRequestTooBig(..)));
833    }
834}