hyperlight_host/mem/
layout.rs

1/*
2Copyright 2024 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*/
16use std::fmt::Debug;
17use std::mem::{offset_of, size_of};
18
19use hyperlight_common::mem::{GuestStackData, HyperlightPEB, RunMode, PAGE_SIZE_USIZE};
20use paste::paste;
21use rand::{rng, RngCore};
22use tracing::{instrument, Span};
23
24use super::memory_region::MemoryRegionType::{
25    BootStack, Code, GuardPage, GuestErrorData, Heap, HostExceptionData, HostFunctionDefinitions,
26    InputData, KernelStack, OutputData, PageTables, PanicContext, Peb, Stack,
27};
28use super::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionVecBuilder};
29use super::mgr::AMOUNT_OF_MEMORY_PER_PT;
30use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory};
31use crate::error::HyperlightError::{GuestOffsetIsInvalid, MemoryRequestTooBig};
32use crate::sandbox::SandboxConfiguration;
33use crate::{log_then_return, new_error, Result};
34
35// +-------------------------------------------+
36// |             Boot Stack (4KiB)             |
37// +-------------------------------------------+
38// |       Kernel Stack Guard Page (4KiB)      |
39// +-------------------------------------------+
40// |             Kernel Stack                  |
41// +-------------------------------------------+
42// |        Guest Stack Guard Page (4KiB)      |
43// +-------------------------------------------+
44// |             Guest (User) Stack            |
45// +-------------------------------------------+
46// |             Guard Page (4KiB)             |
47// +-------------------------------------------+
48// |             Guest Heap                    |
49// +-------------------------------------------+
50// |         Guest Panic Context               |
51// +-------------------------------------------+
52// |             Output Data                   |
53// +-------------------------------------------+
54// |              Input Data                   |
55// +-------------------------------------------+
56// |           Guest Error Log                 |
57// +-------------------------------------------+
58// |        Host Exception Handlers            |
59// +-------------------------------------------+
60// |        Host Function Definitions          |
61// +-------------------------------------------+
62// |                PEB Struct (0x98)          |
63// +-------------------------------------------+
64// |               Guest Code                  |
65// +-------------------------------------------+
66// |                    PT                     |
67// +-------------------------------------------+ 0x203_000
68// |                    PD                     |
69// +-------------------------------------------+ 0x202_000
70// |                   PDPT                    |
71// +-------------------------------------------+ 0x201_000
72// |                   PML4                    |
73// +-------------------------------------------+ 0x200_000
74// |                    ⋮                      |
75// |                 Unmapped                  |
76// |                    ⋮                      |
77// +-------------------------------------------+ 0x0
78
79///
80/// - `HostDefinitions` - the length of this is the `HostFunctionDefinitionSize`
81///   field from `SandboxConfiguration`
82///
83/// - `HostExceptionData` - memory that contains details of any Host Exception that
84///   occurred in outb function. it contains a 32 bit length following by a json
85///   serialisation of any error that occurred. the length of this field is
86///   `HostExceptionSize` from` `SandboxConfiguration`
87///
88/// - `GuestError` - contains a buffer for any guest error that occurred.
89///   the length of this field is `GuestErrorBufferSize` from `SandboxConfiguration`
90///
91/// - `InputData` -  this is a buffer that is used for input data to the host program.
92///   the length of this field is `InputDataSize` from `SandboxConfiguration`
93///
94/// - `OutputData` - this is a buffer that is used for output data from host program.
95///   the length of this field is `OutputDataSize` from `SandboxConfiguration`
96///
97/// - `GuestHeap` - this is a buffer that is used for heap data in the guest. the length
98///   of this field is returned by the `heap_size()` method of this struct
99///
100/// - `GuestStack` - this is a buffer that is used for stack data in the guest. the length
101///   of this field is returned by the `stack_size()` method of this struct. in reality,
102///   the stack might be slightly bigger or smaller than this value since total memory
103///   size is rounded up to the nearest 4K, and there is a 16-byte stack guard written
104///   to the top of the stack. (see below for more details)
105///
106/// - `GuestPanicContext` - contains a buffer for context associated with any guest
107///   panic that occurred.
108///   the length of this field is returned by the `guest_panic_context_size()` fn of this struct.
109///
110/// Boot Stack - this is the stack that is used before the TSS is set up. It is fixed to 4K
111/// Kernel Stack Guard Page is to Guard against boot stack overflow so we dont corrupt the kernel stack
112/// Kernel Stack - this is the stack that is used for kernel mode operations we switch to this early in the initialization function
113/// Guest Stack Guard Page is to Guard against kernel stack overflow so we dont corrupt the user stack
114
115#[derive(Copy, Clone)]
116pub(crate) struct SandboxMemoryLayout {
117    pub(super) sandbox_memory_config: SandboxConfiguration,
118    /// The total stack size of this sandbox.
119    pub(super) stack_size: usize,
120    /// The heap size of this sandbox.
121    pub(super) heap_size: usize,
122
123    /// The following fields are offsets to the actual PEB struct fields.
124    /// They are used when writing the PEB struct itself
125    peb_offset: usize,
126    peb_security_cookie_seed_offset: usize,
127    peb_guest_dispatch_function_ptr_offset: usize, // set by guest in guest entrypoint
128    pub(super) peb_host_function_definitions_offset: usize,
129    pub(crate) peb_host_exception_offset: usize,
130    peb_guest_error_offset: usize,
131    peb_code_and_outb_pointer_offset: usize,
132    peb_runmode_offset: usize,
133    peb_input_data_offset: usize,
134    peb_output_data_offset: usize,
135    peb_guest_panic_context_offset: usize,
136    peb_heap_data_offset: usize,
137    peb_guest_stack_data_offset: usize,
138
139    // The following are the actual values
140    // that are written to the PEB struct
141    pub(crate) host_function_definitions_buffer_offset: usize,
142    pub(crate) host_exception_buffer_offset: usize,
143    pub(super) guest_error_buffer_offset: usize,
144    pub(super) input_data_buffer_offset: usize,
145    pub(super) output_data_buffer_offset: usize,
146    guest_panic_context_buffer_offset: usize,
147    guest_heap_buffer_offset: usize,
148    guard_page_offset: usize,
149    guest_user_stack_buffer_offset: usize, // the lowest address of the user stack
150    user_stack_guard_page_offset: usize,
151    kernel_stack_buffer_offset: usize,
152    kernel_stack_guard_page_offset: usize,
153    #[allow(dead_code)]
154    pub(super) kernel_stack_size_rounded: usize,
155    boot_stack_buffer_offset: usize,
156
157    // other
158    pub(crate) peb_address: usize,
159    code_size: usize,
160    // The total size of the page tables
161    total_page_table_size: usize,
162    // The offset in the sandbox memory where the code starts
163    guest_code_offset: usize,
164}
165
166impl Debug for SandboxMemoryLayout {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        f.debug_struct("SandboxMemoryLayout")
169            .field(
170                "Total Memory Size",
171                &format_args!("{:#x}", self.get_memory_size().unwrap_or(0)),
172            )
173            .field("Stack Size", &format_args!("{:#x}", self.stack_size))
174            .field("Heap Size", &format_args!("{:#x}", self.heap_size))
175            .field("PEB Address", &format_args!("{:#x}", self.peb_address))
176            .field("PEB Offset", &format_args!("{:#x}", self.peb_offset))
177            .field("Code Size", &format_args!("{:#x}", self.code_size))
178            .field(
179                "Security Cookie Seed Offset",
180                &format_args!("{:#x}", self.peb_security_cookie_seed_offset),
181            )
182            .field(
183                "Guest Dispatch Function Pointer Offset",
184                &format_args!("{:#x}", self.peb_guest_dispatch_function_ptr_offset),
185            )
186            .field(
187                "Host Function Definitions Offset",
188                &format_args!("{:#x}", self.peb_host_function_definitions_offset),
189            )
190            .field(
191                "Host Exception Offset",
192                &format_args!("{:#x}", self.peb_host_exception_offset),
193            )
194            .field(
195                "Guest Error Offset",
196                &format_args!("{:#x}", self.peb_guest_error_offset),
197            )
198            .field(
199                "Code and OutB Pointer Offset",
200                &format_args!("{:#x}", self.peb_code_and_outb_pointer_offset),
201            )
202            .field(
203                "Input Data Offset",
204                &format_args!("{:#x}", self.peb_input_data_offset),
205            )
206            .field(
207                "Output Data Offset",
208                &format_args!("{:#x}", self.peb_output_data_offset),
209            )
210            .field(
211                "Guest Panic Context Offset",
212                &format_args!("{:#x}", self.peb_guest_panic_context_offset),
213            )
214            .field(
215                "Guest Heap Offset",
216                &format_args!("{:#x}", self.peb_heap_data_offset),
217            )
218            .field(
219                "Guest Stack Offset",
220                &format_args!("{:#x}", self.peb_guest_stack_data_offset),
221            )
222            .field(
223                "Host Function Definitions Buffer Offset",
224                &format_args!("{:#x}", self.host_function_definitions_buffer_offset),
225            )
226            .field(
227                "Host Exception Buffer Offset",
228                &format_args!("{:#x}", self.host_exception_buffer_offset),
229            )
230            .field(
231                "Guest Error Buffer Offset",
232                &format_args!("{:#x}", self.guest_error_buffer_offset),
233            )
234            .field(
235                "Input Data Buffer Offset",
236                &format_args!("{:#x}", self.input_data_buffer_offset),
237            )
238            .field(
239                "Output Data Buffer Offset",
240                &format_args!("{:#x}", self.output_data_buffer_offset),
241            )
242            .field(
243                "Guest Panic Context Buffer Offset",
244                &format_args!("{:#x}", self.guest_panic_context_buffer_offset),
245            )
246            .field(
247                "Guest Heap Buffer Offset",
248                &format_args!("{:#x}", self.guest_heap_buffer_offset),
249            )
250            .field(
251                "Guard Page Offset",
252                &format_args!("{:#x}", self.guard_page_offset),
253            )
254            .field(
255                "Guest User Stack Buffer Offset",
256                &format_args!("{:#x}", self.guest_user_stack_buffer_offset),
257            )
258            .field(
259                "Page Table Size",
260                &format_args!("{:#x}", self.total_page_table_size),
261            )
262            .field(
263                "Guest Code Offset",
264                &format_args!("{:#x}", self.guest_code_offset),
265            )
266            .field(
267                "User Stack Guard Page Offset",
268                &format_args!("{:#x}", self.user_stack_guard_page_offset),
269            )
270            .field(
271                "Kernel Stack Buffer Offset",
272                &format_args!("{:#x}", self.kernel_stack_buffer_offset),
273            )
274            .field(
275                "Kernel Stack Guard Page Offset",
276                &format_args!("{:#x}", self.kernel_stack_guard_page_offset),
277            )
278            .field(
279                "Boot Stack Buffer Offset",
280                &format_args!("{:#x}", self.boot_stack_buffer_offset),
281            )
282            .finish()
283    }
284}
285
286impl SandboxMemoryLayout {
287    /// The offset into the sandbox's memory where the PML4 Table is located.
288    /// See https://www.pagetable.com/?p=14 for more information.
289    pub(crate) const PML4_OFFSET: usize = 0x0000;
290    /// The offset into the sandbox's memory where the Page Directory Pointer
291    /// Table starts.
292    pub(super) const PDPT_OFFSET: usize = 0x1000;
293    /// The offset into the sandbox's memory where the Page Directory starts.
294    pub(super) const PD_OFFSET: usize = 0x2000;
295    /// The offset into the sandbox's memory where the Page Tables start.
296    pub(super) const PT_OFFSET: usize = 0x3000;
297    /// The address (not the offset) to the start of the page directory
298    pub(super) const PD_GUEST_ADDRESS: usize = Self::BASE_ADDRESS + Self::PD_OFFSET;
299    /// The address (not the offset) into sandbox memory where the Page
300    /// Directory Pointer Table starts
301    pub(super) const PDPT_GUEST_ADDRESS: usize = Self::BASE_ADDRESS + Self::PDPT_OFFSET;
302    /// The address (not the offset) into sandbox memory where the Page
303    /// Tables start
304    pub(super) const PT_GUEST_ADDRESS: usize = Self::BASE_ADDRESS + Self::PT_OFFSET;
305    /// The maximum amount of memory a single sandbox will be allowed.
306    /// The addressable virtual memory with current paging setup is virtual address 0x0 - 0x40000000 (excl.),
307    /// However, the memory up to Self::BASE_ADDRESS is not used.
308    const MAX_MEMORY_SIZE: usize = 0x40000000 - Self::BASE_ADDRESS;
309
310    /// The base address of the sandbox's memory.
311    pub(crate) const BASE_ADDRESS: usize = 0x0200000;
312
313    // the offset into a sandbox's input/output buffer where the stack starts
314    const STACK_POINTER_SIZE_BYTES: u64 = 8;
315
316    /// Create a new `SandboxMemoryLayout` with the given
317    /// `SandboxConfiguration`, code size and stack/heap size.
318    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
319    pub(super) fn new(
320        cfg: SandboxConfiguration,
321        code_size: usize,
322        stack_size: usize,
323        heap_size: usize,
324    ) -> Result<Self> {
325        let total_page_table_size =
326            Self::get_total_page_table_size(cfg, code_size, stack_size, heap_size);
327        let guest_code_offset = total_page_table_size;
328        // The following offsets are to the fields of the PEB struct itself!
329        let peb_offset = total_page_table_size + round_up_to(code_size, PAGE_SIZE_USIZE);
330        let peb_security_cookie_seed_offset =
331            peb_offset + offset_of!(HyperlightPEB, security_cookie_seed);
332        let peb_guest_dispatch_function_ptr_offset =
333            peb_offset + offset_of!(HyperlightPEB, guest_function_dispatch_ptr);
334        let peb_host_function_definitions_offset =
335            peb_offset + offset_of!(HyperlightPEB, hostFunctionDefinitions);
336        let peb_host_exception_offset = peb_offset + offset_of!(HyperlightPEB, hostException);
337        let peb_guest_error_offset = peb_offset + offset_of!(HyperlightPEB, guestErrorData);
338        let peb_code_and_outb_pointer_offset = peb_offset + offset_of!(HyperlightPEB, pCode);
339        let peb_runmode_offset = peb_offset + offset_of!(HyperlightPEB, runMode);
340        let peb_input_data_offset = peb_offset + offset_of!(HyperlightPEB, inputdata);
341        let peb_output_data_offset = peb_offset + offset_of!(HyperlightPEB, outputdata);
342        let peb_guest_panic_context_offset =
343            peb_offset + offset_of!(HyperlightPEB, guestPanicContextData);
344        let peb_heap_data_offset = peb_offset + offset_of!(HyperlightPEB, guestheapData);
345        let peb_guest_stack_data_offset = peb_offset + offset_of!(HyperlightPEB, gueststackData);
346
347        // The following offsets are the actual values that relate to memory layout,
348        // which are written to PEB struct
349        let peb_address = Self::BASE_ADDRESS + peb_offset;
350        // make sure host function definitions buffer starts at 4K boundary
351        let host_function_definitions_buffer_offset = round_up_to(
352            peb_guest_stack_data_offset + size_of::<GuestStackData>(),
353            PAGE_SIZE_USIZE,
354        );
355        // make sure host exception buffer starts at 4K boundary
356        let host_exception_buffer_offset = round_up_to(
357            host_function_definitions_buffer_offset + cfg.get_host_function_definition_size(),
358            PAGE_SIZE_USIZE,
359        );
360        let guest_error_buffer_offset = round_up_to(
361            host_exception_buffer_offset + cfg.get_host_exception_size(),
362            PAGE_SIZE_USIZE,
363        );
364        let input_data_buffer_offset = round_up_to(
365            guest_error_buffer_offset + cfg.get_guest_error_buffer_size(),
366            PAGE_SIZE_USIZE,
367        );
368        let output_data_buffer_offset = round_up_to(
369            input_data_buffer_offset + cfg.get_input_data_size(),
370            PAGE_SIZE_USIZE,
371        );
372        let guest_panic_context_buffer_offset = round_up_to(
373            output_data_buffer_offset + cfg.get_output_data_size(),
374            PAGE_SIZE_USIZE,
375        );
376        // make sure heap buffer starts at 4K boundary
377        let guest_heap_buffer_offset = round_up_to(
378            guest_panic_context_buffer_offset + cfg.get_guest_panic_context_buffer_size(),
379            PAGE_SIZE_USIZE,
380        );
381        // make sure guard page starts at 4K boundary
382        let guard_page_offset = round_up_to(guest_heap_buffer_offset + heap_size, PAGE_SIZE_USIZE);
383        let guest_user_stack_buffer_offset = guard_page_offset + PAGE_SIZE_USIZE;
384        // round up stack size to page size. This is needed for MemoryRegion
385        let stack_size_rounded = round_up_to(stack_size, PAGE_SIZE_USIZE);
386
387        let user_stack_guard_page_offset = guest_user_stack_buffer_offset + stack_size_rounded;
388        let kernel_stack_buffer_offset = user_stack_guard_page_offset + PAGE_SIZE_USIZE;
389        let kernel_stack_size_rounded = round_up_to(cfg.get_kernel_stack_size(), PAGE_SIZE_USIZE);
390        let kernel_stack_guard_page_offset = kernel_stack_buffer_offset + kernel_stack_size_rounded;
391        let boot_stack_buffer_offset = kernel_stack_guard_page_offset + PAGE_SIZE_USIZE;
392
393        Ok(Self {
394            peb_offset,
395            stack_size: stack_size_rounded,
396            heap_size,
397            peb_security_cookie_seed_offset,
398            peb_guest_dispatch_function_ptr_offset,
399            peb_host_function_definitions_offset,
400            peb_host_exception_offset,
401            peb_guest_error_offset,
402            peb_code_and_outb_pointer_offset,
403            peb_runmode_offset,
404            peb_input_data_offset,
405            peb_output_data_offset,
406            peb_guest_panic_context_offset,
407            peb_heap_data_offset,
408            peb_guest_stack_data_offset,
409            guest_error_buffer_offset,
410            sandbox_memory_config: cfg,
411            code_size,
412            host_function_definitions_buffer_offset,
413            host_exception_buffer_offset,
414            input_data_buffer_offset,
415            output_data_buffer_offset,
416            guest_heap_buffer_offset,
417            guest_user_stack_buffer_offset,
418            peb_address,
419            guest_panic_context_buffer_offset,
420            guard_page_offset,
421            total_page_table_size,
422            guest_code_offset,
423            user_stack_guard_page_offset,
424            kernel_stack_buffer_offset,
425            kernel_stack_guard_page_offset,
426            kernel_stack_size_rounded,
427            boot_stack_buffer_offset,
428        })
429    }
430
431    /// Gets the offset in guest memory to the RunMode field in the PEB struct.
432    pub fn get_run_mode_offset(&self) -> usize {
433        self.peb_runmode_offset
434    }
435
436    /// Get the offset in guest memory to the size field in the
437    /// `HostExceptionData` structure.
438    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
439    pub(super) fn get_host_exception_size_offset(&self) -> usize {
440        // The size field is the first field in the `HostExceptionData` struct
441        self.peb_host_exception_offset
442    }
443
444    /// Get the offset in guest memory to the max size of the guest error buffer
445    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
446    pub(super) fn get_guest_error_buffer_size_offset(&self) -> usize {
447        self.peb_guest_error_offset
448    }
449
450    /// Get the offset in guest memory to the error message buffer pointer
451    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
452    fn get_guest_error_buffer_pointer_offset(&self) -> usize {
453        self.peb_guest_error_offset + size_of::<u64>()
454    }
455
456    /// Get the offset in guest memory to the output data size
457    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
458    pub(super) fn get_output_data_size_offset(&self) -> usize {
459        // The size field is the first field in the `OutputData` struct
460        self.peb_output_data_offset
461    }
462
463    /// Get the offset in guest memory to the host function definitions
464    /// size
465    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
466    pub(super) fn get_host_function_definitions_size_offset(&self) -> usize {
467        // The size field is the first field in the `HostFunctions` struct
468        self.peb_host_function_definitions_offset
469    }
470
471    /// Get the offset in guest memory to the host function definitions
472    /// pointer.
473    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
474    fn get_host_function_definitions_pointer_offset(&self) -> usize {
475        // The size field is the field after the size field in the `HostFunctions` struct which is a u64
476        self.peb_host_function_definitions_offset + size_of::<u64>()
477    }
478
479    /// Get the offset in guest memory to the minimum guest stack address.
480    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
481    fn get_min_guest_stack_address_offset(&self) -> usize {
482        // The minimum guest user stack address is the start of the guest stack
483        self.peb_guest_stack_data_offset
484    }
485
486    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
487    pub(super) fn get_guest_stack_size(&self) -> usize {
488        self.stack_size
489    }
490
491    /// Get the offset in guest memory to the start of host errors
492    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
493    pub(super) fn get_host_exception_offset(&self) -> usize {
494        self.host_exception_buffer_offset
495    }
496
497    /// Get the offset in guest memory to the OutB pointer.
498    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
499    pub(super) fn get_outb_pointer_offset(&self) -> usize {
500        // The outb pointer is immediately after the code pointer
501        // in the `CodeAndOutBPointers` struct which is a u64
502        self.peb_code_and_outb_pointer_offset + size_of::<u64>()
503    }
504
505    /// Get the offset in guest memory to the OutB context.
506    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
507    pub(super) fn get_outb_context_offset(&self) -> usize {
508        // The outb context is immediately after the outb pointer
509        // in the `CodeAndOutBPointers` struct which is a u64
510        self.get_outb_pointer_offset() + size_of::<u64>()
511    }
512
513    /// Get the offset in guest memory to the output data pointer.
514    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
515    fn get_output_data_pointer_offset(&self) -> usize {
516        // This field is immediately after the output data size field,
517        // which is a `u64`.
518        self.get_output_data_size_offset() + size_of::<u64>()
519    }
520
521    /// Get the offset in guest memory to the start of output data.
522    ///
523    /// This function exists to accommodate the macro that generates C API
524    /// compatible functions.
525    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
526    pub(crate) fn get_output_data_offset(&self) -> usize {
527        self.output_data_buffer_offset
528    }
529
530    /// Get the offset in guest memory to the input data size.
531    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
532    pub(super) fn get_input_data_size_offset(&self) -> usize {
533        // The input data size is the first field in the `InputData` struct
534        self.peb_input_data_offset
535    }
536
537    /// Get the offset in guest memory to the input data pointer.
538    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
539    fn get_input_data_pointer_offset(&self) -> usize {
540        // The input data pointer is immediately after the input
541        // data size field in the `InputData` struct which is a `u64`.
542        self.get_input_data_size_offset() + size_of::<u64>()
543    }
544
545    /// Get the offset in guest memory to the code pointer
546    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
547    pub(super) fn get_code_pointer_offset(&self) -> usize {
548        // The code pointer is the first field
549        // in the `CodeAndOutBPointers` struct which is a u64
550        self.peb_code_and_outb_pointer_offset
551    }
552
553    /// Get the offset in guest memory to where the guest dispatch function
554    /// pointer is written
555    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
556    pub(super) fn get_dispatch_function_pointer_offset(&self) -> usize {
557        self.peb_guest_dispatch_function_ptr_offset
558    }
559
560    /// Get the offset in guest memory to the PEB address
561    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
562    pub(super) fn get_in_process_peb_offset(&self) -> usize {
563        self.peb_offset
564    }
565
566    /// Get the offset in guest memory to the heap size
567    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
568    fn get_heap_size_offset(&self) -> usize {
569        self.peb_heap_data_offset
570    }
571
572    /// Get the offset of the heap pointer in guest memory,
573    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
574    fn get_heap_pointer_offset(&self) -> usize {
575        // The heap pointer is immediately after the
576        // heap size field in the `GuestHeap` struct which is a `u64`.
577        self.get_heap_size_offset() + size_of::<u64>()
578    }
579
580    /// Get the offset to the top of the stack in guest memory
581    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
582    pub(super) fn get_top_of_user_stack_offset(&self) -> usize {
583        self.guest_user_stack_buffer_offset
584    }
585
586    /// Get the offset of the user stack pointer in guest memory,
587    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
588    fn get_user_stack_pointer_offset(&self) -> usize {
589        // The userStackAddress is immediately after the
590        // minUserStackAddress (top of user stack) field in the `GuestStackData` struct which is a `u64`.
591        self.get_min_guest_stack_address_offset() + size_of::<u64>()
592    }
593
594    /// Get the offset of the kernel stack pointer in guest memory,
595    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
596    fn get_kernel_stack_pointer_offset(&self) -> usize {
597        // The kernelStackAddress is immediately after the
598        // userStackAddress in the `GuestStackData` struct which is a `u64`.
599        self.get_user_stack_pointer_offset() + size_of::<u64>()
600    }
601
602    /// Get the offset of the boot stack pointer in guest memory,
603    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
604    fn get_boot_stack_pointer_offset(&self) -> usize {
605        // The bootStackAddress is immediately after the
606        // kernelStackAddress in the `GuestStackData` struct which is a `u64`.
607        self.get_kernel_stack_pointer_offset() + size_of::<u64>()
608    }
609
610    // Get the offset in guest memory to the start of the guest panic context data
611    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
612    pub(crate) fn get_guest_panic_context_offset(&self) -> usize {
613        self.peb_guest_panic_context_offset
614    }
615
616    // Get the offset to the guest panic context buffer size
617    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
618    pub(crate) fn get_guest_panic_context_size_offset(&self) -> usize {
619        // The size field is the first field in the `GuestPanicContext` data
620        self.peb_guest_panic_context_offset
621    }
622
623    /// Get the offset to the guest panic context buffer pointer
624    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
625    pub(crate) fn get_guest_panic_context_buffer_pointer_offset(&self) -> usize {
626        // The guest panic data pointer is immediately after the guest
627        // panic data size field in the `GuestPanicContext` data which is a `u64`
628        self.get_guest_panic_context_size_offset() + size_of::<u64>()
629    }
630
631    /// Get the offset to the guest panic context buffer pointer
632    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
633    pub(crate) fn get_guest_panic_context_buffer_offset(&self) -> usize {
634        self.guest_panic_context_buffer_offset
635    }
636
637    /// Get the offset to the guest guard page
638    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
639    pub fn get_guard_page_offset(&self) -> usize {
640        self.guard_page_offset
641    }
642
643    /// Get the total size of guest memory in `self`'s memory
644    /// layout.
645    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
646    fn get_unaligned_memory_size(&self) -> usize {
647        self.get_boot_stack_buffer_offset() + PAGE_SIZE_USIZE
648    }
649
650    /// get the code offset
651    /// This is the offset in the sandbox memory where the code starts
652    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
653    pub(super) fn get_guest_code_offset(&self) -> usize {
654        self.guest_code_offset
655    }
656
657    /// Get the guest address of the code section in the sandbox
658    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
659    pub(crate) fn get_guest_code_address(&self) -> usize {
660        Self::BASE_ADDRESS + self.guest_code_offset
661    }
662
663    /// Get the offset in guest memory to the user stack guard page
664    /// This is the offset in the sandbox memory where the user stack guard page starts
665    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
666    pub(super) fn get_user_stack_guard_page_offset(&self) -> usize {
667        self.user_stack_guard_page_offset
668    }
669
670    /// Get the offset in guest memory to the kernel stack buffer
671    /// This is the offset in the sandbox memory where the kernel stack starts
672    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
673    pub(super) fn get_kernel_stack_buffer_offset(&self) -> usize {
674        self.kernel_stack_buffer_offset
675    }
676
677    /// Get the offset in guest memory to the kernel stack guard page
678    /// This is the offset in the sandbox memory where the kernel stack guard page starts
679    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
680    pub(super) fn get_kernel_stack_guard_page_offset(&self) -> usize {
681        self.kernel_stack_guard_page_offset
682    }
683
684    /// Get the offset in guest memory to the boot stack buffer
685    /// This is the offset in the sandbox memory where the boot stack starts
686    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
687    pub(super) fn get_boot_stack_buffer_offset(&self) -> usize {
688        self.boot_stack_buffer_offset
689    }
690
691    #[cfg(test)]
692    /// Get the page table size
693    fn get_page_table_size(&self) -> usize {
694        self.total_page_table_size
695    }
696
697    // This function calculates the page table size for the sandbox
698    // We need enough memory to store the PML4, PDPT, PD and PTs
699    // The size of a single table is 4K, we can map up to 1GB total memory which requires 1 PML4, 1 PDPT, 1 PD and 512 PTs
700    // but we only need enough PTs to map the memory we are using. (In other words we only need 512 PTs to map the memory if the memory size is 1GB)
701    //
702    // Because we always start the physical address space at 0x200_000
703    // we can calculate the amount of memory needed for the PTs by calculating how much memory is needed for the sandbox configuration in total,
704    // then add 0x200_000 to that (as we start at 0x200_000),
705    // and then add 3 * 4K (for the PML4, PDPT and PD)  to that,
706    // then add 2MB to that (the maximum size of memory required for the PTs themselves is 2MB when we map 1GB of memory in 4K pages),
707    // then divide that by 0x200_000 (as we can map 2MB in each PT) and then round the result up by 1 .
708    // This will give us the total size of the PTs required for the sandbox to which we can add the size of the PML4, PDPT and PD.
709    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
710    fn get_total_page_table_size(
711        cfg: SandboxConfiguration,
712        code_size: usize,
713        stack_size: usize,
714        heap_size: usize,
715    ) -> usize {
716        // Get the configured memory size (assume each section is 4K aligned)
717
718        let mut total_mapped_memory_size: usize = round_up_to(code_size, PAGE_SIZE_USIZE);
719        total_mapped_memory_size += round_up_to(stack_size, PAGE_SIZE_USIZE);
720        total_mapped_memory_size += round_up_to(heap_size, PAGE_SIZE_USIZE);
721        total_mapped_memory_size += round_up_to(cfg.get_host_exception_size(), PAGE_SIZE_USIZE);
722        total_mapped_memory_size +=
723            round_up_to(cfg.get_host_function_definition_size(), PAGE_SIZE_USIZE);
724        total_mapped_memory_size += round_up_to(cfg.get_guest_error_buffer_size(), PAGE_SIZE_USIZE);
725        total_mapped_memory_size += round_up_to(cfg.get_input_data_size(), PAGE_SIZE_USIZE);
726        total_mapped_memory_size += round_up_to(cfg.get_output_data_size(), PAGE_SIZE_USIZE);
727        total_mapped_memory_size +=
728            round_up_to(cfg.get_guest_panic_context_buffer_size(), PAGE_SIZE_USIZE);
729        total_mapped_memory_size += round_up_to(size_of::<HyperlightPEB>(), PAGE_SIZE_USIZE);
730
731        // Add the base address of the sandbox
732        total_mapped_memory_size += Self::BASE_ADDRESS;
733
734        // Add the size of  the PML4, PDPT and PD
735        total_mapped_memory_size += 3 * PAGE_SIZE_USIZE;
736
737        // Add the maximum possible size of the PTs
738        total_mapped_memory_size += 512 * PAGE_SIZE_USIZE;
739
740        // Get the number of pages needed for the PTs
741
742        let num_pages: usize = ((total_mapped_memory_size + AMOUNT_OF_MEMORY_PER_PT - 1)
743            / AMOUNT_OF_MEMORY_PER_PT)
744            + 1 // Round up
745            + 3; // PML4, PDPT, PD
746
747        num_pages * PAGE_SIZE_USIZE
748    }
749
750    /// Get the total size of guest memory in `self`'s memory
751    /// layout aligned to page size boundaries.
752    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
753    pub(super) fn get_memory_size(&self) -> Result<usize> {
754        let total_memory = self.get_unaligned_memory_size();
755
756        // Size should be a multiple of page size.
757        let remainder = total_memory % PAGE_SIZE_USIZE;
758        let multiples = total_memory / PAGE_SIZE_USIZE;
759        let size = match remainder {
760            0 => total_memory,
761            _ => (multiples + 1) * PAGE_SIZE_USIZE,
762        };
763
764        if size > Self::MAX_MEMORY_SIZE {
765            Err(MemoryRequestTooBig(size, Self::MAX_MEMORY_SIZE))
766        } else {
767            Ok(size)
768        }
769    }
770
771    /// Returns the memory regions associated with this memory layout,
772    /// suitable for passing to a hypervisor for mapping into memory
773    pub fn get_memory_regions(&self, shared_mem: &GuestSharedMemory) -> Result<Vec<MemoryRegion>> {
774        let mut builder = MemoryRegionVecBuilder::new(Self::BASE_ADDRESS, shared_mem.base_addr());
775
776        // PML4, PDPT, PD
777        let code_offset = builder.push_page_aligned(
778            self.total_page_table_size,
779            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
780            PageTables,
781        );
782
783        if code_offset != self.guest_code_offset {
784            return Err(new_error!(
785                "Code offset does not match expected code offset expected:  {}, actual:  {}",
786                self.guest_code_offset,
787                code_offset
788            ));
789        }
790
791        // code
792        let peb_offset = builder.push_page_aligned(
793            self.code_size,
794            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
795            Code,
796        );
797
798        let expected_peb_offset = TryInto::<usize>::try_into(self.peb_offset)?;
799
800        if peb_offset != expected_peb_offset {
801            return Err(new_error!(
802                "PEB offset does not match expected PEB offset expected:  {}, actual:  {}",
803                expected_peb_offset,
804                peb_offset
805            ));
806        }
807
808        // PEB
809        let host_functions_definitions_offset = builder.push_page_aligned(
810            size_of::<HyperlightPEB>(),
811            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
812            Peb,
813        );
814
815        let expected_host_functions_definitions_offset =
816            TryInto::<usize>::try_into(self.host_function_definitions_buffer_offset)?;
817
818        if host_functions_definitions_offset != expected_host_functions_definitions_offset {
819            return Err(new_error!(
820                "Host Function Definitions offset does not match expected Host Function Definitions offset expected:  {}, actual:  {}",
821                expected_host_functions_definitions_offset,
822                host_functions_definitions_offset
823            ));
824        }
825
826        // host function definitions
827        let host_exception_offset = builder.push_page_aligned(
828            self.sandbox_memory_config
829                .get_host_function_definition_size(),
830            MemoryRegionFlags::READ,
831            HostFunctionDefinitions,
832        );
833
834        let expected_host_exception_offset =
835            TryInto::<usize>::try_into(self.host_exception_buffer_offset)?;
836
837        if host_exception_offset != expected_host_exception_offset {
838            return Err(new_error!(
839                "Host Exception offset does not match expected Host Exception offset expected:  {}, actual:  {}",
840                expected_host_exception_offset,
841                host_exception_offset
842            ));
843        }
844
845        // host exception
846        let guest_error_offset = builder.push_page_aligned(
847            self.sandbox_memory_config.get_host_exception_size(),
848            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
849            HostExceptionData,
850        );
851
852        let expected_guest_error_offset =
853            TryInto::<usize>::try_into(self.guest_error_buffer_offset)?;
854
855        if guest_error_offset != expected_guest_error_offset {
856            return Err(new_error!(
857                "Guest Error offset does not match expected Guest Error offset expected:  {}, actual:  {}",
858                expected_guest_error_offset,
859                guest_error_offset
860            ));
861        }
862
863        // guest error
864        let input_data_offset = builder.push_page_aligned(
865            self.sandbox_memory_config.get_guest_error_buffer_size(),
866            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
867            GuestErrorData,
868        );
869
870        let expected_input_data_offset = TryInto::<usize>::try_into(self.input_data_buffer_offset)?;
871
872        if input_data_offset != expected_input_data_offset {
873            return Err(new_error!(
874                "Input Data offset does not match expected Input Data offset expected:  {}, actual:  {}",
875                expected_input_data_offset,
876                input_data_offset
877            ));
878        }
879
880        // guest input data
881        let output_data_offset = builder.push_page_aligned(
882            self.sandbox_memory_config.get_input_data_size(),
883            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
884            InputData,
885        );
886
887        let expected_output_data_offset =
888            TryInto::<usize>::try_into(self.output_data_buffer_offset)?;
889
890        if output_data_offset != expected_output_data_offset {
891            return Err(new_error!(
892                "Output Data offset does not match expected Output Data offset expected:  {}, actual:  {}",
893                expected_output_data_offset,
894                output_data_offset
895            ));
896        }
897
898        // guest output data
899        let guest_panic_context_offset = builder.push_page_aligned(
900            self.sandbox_memory_config.get_output_data_size(),
901            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
902            OutputData,
903        );
904
905        let expected_guest_panic_context_offset =
906            TryInto::<usize>::try_into(self.guest_panic_context_buffer_offset)?;
907
908        if guest_panic_context_offset != expected_guest_panic_context_offset {
909            return Err(new_error!(
910                "Guest Panic Context offset does not match expected Guest Panic Context offset expected:  {}, actual:  {}",
911                expected_guest_panic_context_offset,
912                guest_panic_context_offset
913            ));
914        }
915
916        // guest panic context
917        let heap_offset = builder.push_page_aligned(
918            self.sandbox_memory_config
919                .get_guest_panic_context_buffer_size(),
920            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
921            PanicContext,
922        );
923
924        let expected_heap_offset = TryInto::<usize>::try_into(self.guest_heap_buffer_offset)?;
925
926        if heap_offset != expected_heap_offset {
927            return Err(new_error!(
928                "Guest Heap offset does not match expected Guest Heap offset expected:  {}, actual:  {}",
929                expected_heap_offset,
930                heap_offset
931            ));
932        }
933
934        // heap
935        #[cfg(feature = "executable_heap")]
936        let guard_page_offset = builder.push_page_aligned(
937            self.heap_size,
938            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
939            Heap,
940        );
941        #[cfg(not(feature = "executable_heap"))]
942        let guard_page_offset = builder.push_page_aligned(
943            self.heap_size,
944            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
945            Heap,
946        );
947
948        let expected_guard_page_offset = TryInto::<usize>::try_into(self.guard_page_offset)?;
949
950        if guard_page_offset != expected_guard_page_offset {
951            return Err(new_error!(
952                "Guard Page offset does not match expected Guard Page offset expected:  {}, actual:  {}",
953                expected_guard_page_offset,
954                guard_page_offset
955            ));
956        }
957
958        // guard page
959        let stack_offset = builder.push_page_aligned(
960            PAGE_SIZE_USIZE,
961            MemoryRegionFlags::READ | MemoryRegionFlags::STACK_GUARD,
962            GuardPage,
963        );
964
965        let expected_stack_offset =
966            TryInto::<usize>::try_into(self.guest_user_stack_buffer_offset)?;
967
968        if stack_offset != expected_stack_offset {
969            return Err(new_error!(
970                "Stack offset does not match expected Stack offset expected:  {}, actual:  {}",
971                expected_stack_offset,
972                stack_offset
973            ));
974        }
975
976        // stack
977        let user_stack_guard_page_offset = builder.push_page_aligned(
978            self.get_guest_stack_size(),
979            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
980            Stack,
981        );
982
983        let expected_user_stack_guard_page_offset =
984            TryInto::<usize>::try_into(self.get_top_of_user_stack_offset())?
985                + self.get_guest_stack_size();
986
987        if user_stack_guard_page_offset != expected_user_stack_guard_page_offset {
988            return Err(new_error!(
989                "User Guard Page offset does not match expected User Guard Page offset expected:  {}, actual:  {}",
990                expected_user_stack_guard_page_offset,
991                user_stack_guard_page_offset
992            ));
993        }
994
995        let kernel_stack_offset = builder.push_page_aligned(
996            PAGE_SIZE_USIZE,
997            MemoryRegionFlags::READ | MemoryRegionFlags::STACK_GUARD,
998            GuardPage,
999        );
1000
1001        let expected_kernel_stack_offset =
1002            TryInto::<usize>::try_into(self.kernel_stack_buffer_offset)?;
1003
1004        if kernel_stack_offset != expected_kernel_stack_offset {
1005            return Err(new_error!(
1006                "Kernel Stack offset does not match expected Kernel Stack offset expected:  {}, actual:  {}",
1007                expected_kernel_stack_offset,
1008                kernel_stack_offset
1009            ));
1010        }
1011
1012        let kernel_stack_guard_page_offset = builder.push_page_aligned(
1013            self.kernel_stack_size_rounded,
1014            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
1015            KernelStack,
1016        );
1017
1018        let expected_kernel_stack_guard_page_offset =
1019            TryInto::<usize>::try_into(self.kernel_stack_guard_page_offset)?;
1020
1021        if kernel_stack_guard_page_offset != expected_kernel_stack_guard_page_offset {
1022            return Err(new_error!(
1023                "Kernel Guard Page offset does not match expected Kernel Guard Page offset expected:  {}, actual:  {}",
1024                expected_kernel_stack_guard_page_offset,
1025                kernel_stack_guard_page_offset
1026            ));
1027        }
1028
1029        let boot_stack_offset = builder.push_page_aligned(
1030            PAGE_SIZE_USIZE,
1031            MemoryRegionFlags::READ | MemoryRegionFlags::STACK_GUARD,
1032            GuardPage,
1033        );
1034
1035        let expected_boot_stack_offset = TryInto::<usize>::try_into(self.boot_stack_buffer_offset)?;
1036
1037        if boot_stack_offset != expected_boot_stack_offset {
1038            return Err(new_error!(
1039                "Boot Stack offset does not match expected Boot Stack offset expected:  {}, actual:  {}",
1040                expected_boot_stack_offset,
1041                boot_stack_offset
1042            ));
1043        }
1044
1045        let final_offset = builder.push_page_aligned(
1046            PAGE_SIZE_USIZE,
1047            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
1048            BootStack,
1049        );
1050
1051        let expected_final_offset = TryInto::<usize>::try_into(self.get_memory_size()?)?;
1052
1053        if final_offset != expected_final_offset {
1054            return Err(new_error!(
1055                "Final offset does not match expected Final offset expected:  {}, actual:  {}",
1056                expected_final_offset,
1057                final_offset
1058            ));
1059        }
1060
1061        Ok(builder.build())
1062    }
1063
1064    /// Write the finished memory layout to `shared_mem` and return
1065    /// `Ok` if successful.
1066    ///
1067    /// Note: `shared_mem` may have been modified, even if `Err` was returned
1068    /// from this function.
1069    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
1070    pub(crate) fn write(
1071        &self,
1072        shared_mem: &mut ExclusiveSharedMemory,
1073        guest_offset: usize,
1074        size: usize,
1075        run_inprocess: bool,
1076    ) -> Result<()> {
1077        macro_rules! get_address {
1078            ($something:ident) => {
1079                paste! {
1080                    if guest_offset == 0 {
1081                        let offset = self.[<$something _offset>];
1082                        let calculated_addr = shared_mem.calculate_address(offset)?;
1083                        u64::try_from(calculated_addr)?
1084                    } else {
1085                        u64::try_from(guest_offset +  self.[<$something _offset>])?
1086                    }
1087                }
1088            };
1089        }
1090
1091        if guest_offset != SandboxMemoryLayout::BASE_ADDRESS
1092            && guest_offset != shared_mem.base_addr()
1093        {
1094            return Err(GuestOffsetIsInvalid(guest_offset));
1095        }
1096
1097        // Start of setting up the PEB. The following are in the order of the PEB fields
1098
1099        // Set up the security cookie seed
1100        let mut security_cookie_seed = [0u8; 8];
1101        rng().fill_bytes(&mut security_cookie_seed);
1102        shared_mem.copy_from_slice(&security_cookie_seed, self.peb_security_cookie_seed_offset)?;
1103
1104        // Skip guest_dispatch_function_ptr_offset because it is set by the guest
1105
1106        // Set up Host Function Definition
1107        shared_mem.write_u64(
1108            self.get_host_function_definitions_size_offset(),
1109            self.sandbox_memory_config
1110                .get_host_function_definition_size()
1111                .try_into()?,
1112        )?;
1113        let addr = get_address!(host_function_definitions_buffer);
1114        shared_mem.write_u64(self.get_host_function_definitions_pointer_offset(), addr)?;
1115
1116        // Set up Host Exception Header
1117        // The peb only needs to include the size, not the actual buffer
1118        // since the the guest wouldn't want to read the buffer anyway
1119        shared_mem.write_u64(
1120            self.get_host_exception_size_offset(),
1121            self.sandbox_memory_config
1122                .get_host_exception_size()
1123                .try_into()?,
1124        )?;
1125
1126        // Set up Guest Error Fields
1127        let addr = get_address!(guest_error_buffer);
1128        shared_mem.write_u64(self.get_guest_error_buffer_pointer_offset(), addr)?;
1129        shared_mem.write_u64(
1130            self.get_guest_error_buffer_size_offset(),
1131            u64::try_from(self.sandbox_memory_config.get_guest_error_buffer_size())?,
1132        )?;
1133
1134        // Skip code, is set when loading binary
1135        // skip outb and outb context, is set when running in_proc
1136
1137        // Set RunMode in PEB
1138        shared_mem.write_u64(
1139            self.get_run_mode_offset(),
1140            match (
1141                run_inprocess,
1142                cfg!(target_os = "windows"),
1143                cfg!(target_os = "linux"),
1144            ) {
1145                (false, _, _) => RunMode::Hypervisor as u64,
1146                (true, true, _) => RunMode::InProcessWindows as u64,
1147                (true, _, true) => RunMode::InProcessLinux as u64,
1148                (true, _, _) => log_then_return!("Unsupported OS for in-process mode"),
1149            },
1150        )?;
1151
1152        // Set up input buffer pointer
1153        shared_mem.write_u64(
1154            self.get_input_data_size_offset(),
1155            self.sandbox_memory_config
1156                .get_input_data_size()
1157                .try_into()?,
1158        )?;
1159        let addr = get_address!(input_data_buffer);
1160        shared_mem.write_u64(self.get_input_data_pointer_offset(), addr)?;
1161
1162        // Set up output buffer pointer
1163        shared_mem.write_u64(
1164            self.get_output_data_size_offset(),
1165            self.sandbox_memory_config
1166                .get_output_data_size()
1167                .try_into()?,
1168        )?;
1169        let addr = get_address!(output_data_buffer);
1170        shared_mem.write_u64(self.get_output_data_pointer_offset(), addr)?;
1171
1172        // Set up the guest panic context buffer
1173        let addr = get_address!(guest_panic_context_buffer);
1174        shared_mem.write_u64(
1175            self.get_guest_panic_context_size_offset(),
1176            self.sandbox_memory_config
1177                .get_guest_panic_context_buffer_size()
1178                .try_into()?,
1179        )?;
1180        shared_mem.write_u64(self.get_guest_panic_context_buffer_pointer_offset(), addr)?;
1181
1182        // Set up heap buffer pointer
1183        let addr = get_address!(guest_heap_buffer);
1184        shared_mem.write_u64(self.get_heap_size_offset(), self.heap_size.try_into()?)?;
1185        shared_mem.write_u64(self.get_heap_pointer_offset(), addr)?;
1186
1187        // Set up user stack pointers
1188
1189        // Set up Min Guest User Stack Address
1190
1191        // The top of the user stack is calculated as the size of the guest memory + the guest offset which gives us the
1192        // address at the bottom of the guest memory.
1193        // we then subtract the size of the stack, the size of the kernel stack,
1194        // the size of the boot stack, the size of the user stack guard page and the size of the kernel stack guard page
1195        // which are all 4K
1196
1197        let bottom = guest_offset + size;
1198        let min_user_stack_address = bottom
1199            - self.stack_size
1200            - self.kernel_stack_size_rounded
1201            - PAGE_SIZE_USIZE
1202            - PAGE_SIZE_USIZE
1203            - PAGE_SIZE_USIZE;
1204
1205        // Top of user stack
1206
1207        shared_mem.write_u64(
1208            self.get_min_guest_stack_address_offset(),
1209            min_user_stack_address.try_into()?,
1210        )?;
1211
1212        // Start of user stack
1213
1214        let start_of_user_stack: u64 = (min_user_stack_address + self.stack_size).try_into()?;
1215
1216        shared_mem.write_u64(self.get_user_stack_pointer_offset(), start_of_user_stack)?;
1217
1218        // Start of kernel stack
1219
1220        // There is a guard page between the user stack and the kernel stack and then we need to add the size of the kernel stack
1221
1222        let start_of_kernel_stack: u64 =
1223            start_of_user_stack + (PAGE_SIZE_USIZE + self.kernel_stack_size_rounded) as u64;
1224
1225        shared_mem.write_u64(
1226            self.get_kernel_stack_pointer_offset(),
1227            start_of_kernel_stack,
1228        )?;
1229
1230        // Start of boot stack
1231
1232        // There is a guard page between the kernel stack and the boot stack and then we need to add the size of the boot stack
1233
1234        let start_of_boot_stack: u64 = start_of_kernel_stack + (PAGE_SIZE_USIZE * 2) as u64;
1235
1236        shared_mem.write_u64(self.get_boot_stack_pointer_offset(), start_of_boot_stack)?;
1237
1238        // End of setting up the PEB
1239
1240        // Initialize the stack pointers of input data and output data
1241        // to point to the ninth (index 8) byte, which is the first free address
1242        // of the each respective stack. The first 8 bytes are the stack pointer itself.
1243        shared_mem.write_u64(
1244            self.input_data_buffer_offset,
1245            Self::STACK_POINTER_SIZE_BYTES,
1246        )?;
1247        shared_mem.write_u64(
1248            self.output_data_buffer_offset,
1249            Self::STACK_POINTER_SIZE_BYTES,
1250        )?;
1251
1252        Ok(())
1253    }
1254}
1255
1256fn round_up_to(value: usize, multiple: usize) -> usize {
1257    (value + multiple - 1) & !(multiple - 1)
1258}
1259
1260#[cfg(test)]
1261mod tests {
1262    use hyperlight_common::mem::PAGE_SIZE_USIZE;
1263
1264    use super::*;
1265
1266    #[test]
1267    fn test_round_up() {
1268        assert_eq!(0, round_up_to(0, 4));
1269        assert_eq!(4, round_up_to(1, 4));
1270        assert_eq!(4, round_up_to(2, 4));
1271        assert_eq!(4, round_up_to(3, 4));
1272        assert_eq!(4, round_up_to(4, 4));
1273        assert_eq!(8, round_up_to(5, 4));
1274        assert_eq!(8, round_up_to(6, 4));
1275        assert_eq!(8, round_up_to(7, 4));
1276        assert_eq!(8, round_up_to(8, 4));
1277        assert_eq!(PAGE_SIZE_USIZE, round_up_to(44, PAGE_SIZE_USIZE));
1278        assert_eq!(PAGE_SIZE_USIZE, round_up_to(4095, PAGE_SIZE_USIZE));
1279        assert_eq!(PAGE_SIZE_USIZE, round_up_to(4096, PAGE_SIZE_USIZE));
1280        assert_eq!(PAGE_SIZE_USIZE * 2, round_up_to(4097, PAGE_SIZE_USIZE));
1281        assert_eq!(PAGE_SIZE_USIZE * 2, round_up_to(8191, PAGE_SIZE_USIZE));
1282    }
1283
1284    // helper func for testing
1285    fn get_expected_memory_size(layout: &SandboxMemoryLayout) -> usize {
1286        let cfg = layout.sandbox_memory_config;
1287        let mut expected_size = 0;
1288        // in order of layout
1289        expected_size += layout.get_page_table_size();
1290        expected_size += layout.code_size;
1291
1292        expected_size += round_up_to(size_of::<HyperlightPEB>(), PAGE_SIZE_USIZE);
1293
1294        expected_size += round_up_to(cfg.get_host_function_definition_size(), PAGE_SIZE_USIZE);
1295
1296        expected_size += round_up_to(cfg.get_host_exception_size(), PAGE_SIZE_USIZE);
1297
1298        expected_size += round_up_to(cfg.get_guest_error_buffer_size(), PAGE_SIZE_USIZE);
1299
1300        expected_size += round_up_to(cfg.get_input_data_size(), PAGE_SIZE_USIZE);
1301
1302        expected_size += round_up_to(cfg.get_output_data_size(), PAGE_SIZE_USIZE);
1303
1304        expected_size += round_up_to(cfg.get_guest_panic_context_buffer_size(), PAGE_SIZE_USIZE);
1305
1306        expected_size += round_up_to(layout.heap_size, PAGE_SIZE_USIZE);
1307
1308        expected_size += PAGE_SIZE_USIZE; // guard page
1309
1310        expected_size += round_up_to(layout.stack_size, PAGE_SIZE_USIZE);
1311
1312        expected_size += PAGE_SIZE_USIZE; // user stack guard page
1313
1314        expected_size += round_up_to(layout.kernel_stack_size_rounded, PAGE_SIZE_USIZE);
1315
1316        expected_size += PAGE_SIZE_USIZE; // kernel stack guard page
1317
1318        expected_size += PAGE_SIZE_USIZE; // boot stack
1319
1320        expected_size
1321    }
1322
1323    #[test]
1324    fn test_get_memory_size() {
1325        let sbox_cfg = SandboxConfiguration::default();
1326        let sbox_mem_layout = SandboxMemoryLayout::new(sbox_cfg, 4096, 2048, 4096).unwrap();
1327        assert_eq!(
1328            sbox_mem_layout.get_memory_size().unwrap(),
1329            get_expected_memory_size(&sbox_mem_layout)
1330        );
1331    }
1332}