nub_host_kvm/mem/
mgr.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 */
16use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
17use nub_arch_x86_abi::ParallelInvokeSlot;
18use nub_host_common::vmem::{self, PAGE_TABLE_SIZE};
19use std::mem::{align_of, offset_of, size_of};
20use tracing::{Span, instrument};
21
22use super::layout::SandboxMemoryLayout;
23use super::shared_mem::{
24    ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, ReadonlySharedMemory, SharedMemory,
25};
26use crate::Result;
27use crate::sandbox::snapshot::{NextAction, Snapshot};
28
29// `SnapshotSharedMemory<S: SharedMemory>` is unconditionally
30// `ReadonlySharedMemory`, but it is expressed through an associated
31// type with an unused type parameter `S`. rustc rejects an unused type
32// parameter on a plain type alias, so this module wraps it in a trait
33// to placate the compiler.
34mod unused_hack {
35    use crate::mem::shared_mem::ReadonlySharedMemory;
36    use crate::mem::shared_mem::SharedMemory;
37    pub trait SnapshotSharedMemoryT {
38        type T<S: SharedMemory>;
39    }
40    pub struct SnapshotSharedMemory_;
41    impl SnapshotSharedMemoryT for SnapshotSharedMemory_ {
42        type T<S: SharedMemory> = ReadonlySharedMemory;
43    }
44    pub type SnapshotSharedMemory<S> = <SnapshotSharedMemory_ as SnapshotSharedMemoryT>::T<S>;
45}
46impl ReadonlySharedMemory {
47    pub(crate) fn to_mgr_snapshot_mem(
48        &self,
49    ) -> Result<SnapshotSharedMemory<ExclusiveSharedMemory>> {
50        let ret = self.clone();
51        Ok(ret)
52    }
53}
54pub(crate) use unused_hack::SnapshotSharedMemory;
55/// A struct that is responsible for laying out and managing the memory
56/// for a given `Sandbox`.
57#[derive(Clone)]
58pub(crate) struct SandboxMemoryManager<S: SharedMemory> {
59    /// Shared memory for the Sandbox
60    pub(crate) shared_mem: SnapshotSharedMemory<S>,
61    /// Scratch memory for the Sandbox
62    pub(crate) scratch_mem: S,
63    /// The memory layout of the underlying shared memory
64    pub(crate) layout: SandboxMemoryLayout,
65    /// Offset for the execution entrypoint from `load_addr`
66    pub(crate) entrypoint: NextAction,
67    /// How many memory regions were mapped after sandbox creation
68    pub(crate) mapped_rgns: u64,
69    /// Buffer for accumulating guest abort messages
70    pub(crate) abort_buffer: Vec<u8>,
71    /// Generation counter: how many snapshots have been taken from
72    /// this sandbox's execution path from init to here. Incremented
73    /// on each `snapshot` call; on `restore_snapshot` we inherit the
74    /// restored snapshot's own generation number so the guest-visible
75    /// counter tracks which snapshot the sandbox is a clone of.
76    pub(crate) snapshot_count: u64,
77}
78
79/// Buffer for building guest page tables during snapshot creation.
80/// `TableAddr` is an absolute GPA (u64) so the same address space is
81/// used regardless of entry size.
82pub(crate) struct GuestPageTableBuffer {
83    buffer: std::cell::RefCell<Vec<u8>>,
84    phys_base: usize,
85    /// Absolute GPA of the currently-active root table. For
86    /// multi-root guests, `set_root` switches which root subsequent
87    /// `vmem::map` / `vmem::space_aware_map` calls target — typically
88    /// to an address previously returned by `alloc_table`.
89    root: std::cell::Cell<u64>,
90}
91
92impl vmem::TableReadOps for GuestPageTableBuffer {
93    type TableAddr = u64;
94
95    fn entry_addr(addr: u64, offset: u64) -> u64 {
96        addr + offset
97    }
98
99    unsafe fn read_entry(&self, addr: u64) -> vmem::PageTableEntry {
100        let buffer = self.buffer.borrow();
101        let byte_offset = addr as usize - self.phys_base;
102        let pte_size = core::mem::size_of::<vmem::PageTableEntry>();
103        let Some(bytes) = buffer.get(byte_offset..byte_offset + pte_size) else {
104            return 0;
105        };
106        let mut buf = [0u8; 8];
107        buf[..pte_size].copy_from_slice(bytes);
108        vmem::PageTableEntry::from_le_bytes(buf[..pte_size].try_into().unwrap_or_default())
109    }
110
111    fn to_phys(addr: u64) -> vmem::PhysAddr {
112        addr as vmem::PhysAddr
113    }
114
115    fn from_phys(addr: vmem::PhysAddr) -> u64 {
116        #[allow(clippy::unnecessary_cast)]
117        {
118            addr as u64
119        }
120    }
121
122    fn root_table(&self) -> u64 {
123        self.root.get()
124    }
125}
126
127impl vmem::TableOps for GuestPageTableBuffer {
128    type TableMovability = vmem::MayNotMoveTable;
129
130    unsafe fn alloc_table(&self) -> u64 {
131        let mut b = self.buffer.borrow_mut();
132        let offset = b.len();
133        b.resize(offset + PAGE_TABLE_SIZE, 0);
134        (self.phys_base + offset) as u64
135    }
136
137    unsafe fn write_entry(&self, addr: u64, entry: vmem::PageTableEntry) -> Option<vmem::Void> {
138        let mut b = self.buffer.borrow_mut();
139        let byte_offset = addr as usize - self.phys_base;
140        let pte_size = core::mem::size_of::<vmem::PageTableEntry>();
141        if let Some(slice) = b.get_mut(byte_offset..byte_offset + pte_size) {
142            slice.copy_from_slice(&entry.to_le_bytes()[..pte_size]);
143        }
144        None
145    }
146
147    unsafe fn update_root(&self, impossible: vmem::Void) {
148        match impossible {}
149    }
150}
151
152impl core::convert::AsRef<GuestPageTableBuffer> for GuestPageTableBuffer {
153    fn as_ref(&self) -> &Self {
154        self
155    }
156}
157
158impl GuestPageTableBuffer {
159    /// Create a new buffer with an initial zeroed root table at
160    /// `phys_base`. The returned buffer's current root is `phys_base`;
161    /// additional roots can be obtained by calling `alloc_table`.
162    pub(crate) fn new(phys_base: usize) -> Self {
163        GuestPageTableBuffer {
164            buffer: std::cell::RefCell::new(vec![0u8; PAGE_TABLE_SIZE]),
165            phys_base,
166            root: std::cell::Cell::new(phys_base as u64),
167        }
168    }
169
170    pub(crate) fn into_bytes(self) -> Box<[u8]> {
171        self.buffer.into_inner().into_boxed_slice()
172    }
173}
174
175impl<S> SandboxMemoryManager<S>
176where
177    S: SharedMemory,
178{
179    /// Create a new `SandboxMemoryManager` with the given parameters
180    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
181    pub(crate) fn new(
182        layout: SandboxMemoryLayout,
183        shared_mem: SnapshotSharedMemory<S>,
184        scratch_mem: S,
185        entrypoint: NextAction,
186    ) -> Self {
187        Self {
188            layout,
189            shared_mem,
190            scratch_mem,
191            entrypoint,
192            mapped_rgns: 0,
193            abort_buffer: Vec::new(),
194            snapshot_count: 0,
195        }
196    }
197
198    /// Get mutable access to the abort buffer
199    pub(crate) fn get_abort_buffer_mut(&mut self) -> &mut Vec<u8> {
200        &mut self.abort_buffer
201    }
202}
203
204impl SandboxMemoryManager<ExclusiveSharedMemory> {
205    pub(crate) fn from_snapshot(s: &Snapshot) -> Result<Self> {
206        let layout = *s.layout();
207        let shared_mem = s.memory().to_mgr_snapshot_mem()?;
208        let scratch_mem = ExclusiveSharedMemory::new(s.layout().get_scratch_size())?;
209        let entrypoint = s.entrypoint();
210        Ok(Self::new(layout, shared_mem, scratch_mem, entrypoint))
211    }
212
213    /// Wraps ExclusiveSharedMemory::build
214    // Morally, this should not have to be a Result: this operation is
215    // infallible. The source of the Result is
216    // update_scratch_bookkeeping(), which calls functions that can
217    // fail due to bounds checks (which are statically known to be ok
218    // in this situation) or due to failing to take the scratch shared
219    // memory lock, but the scratch shared memory is built in this
220    // function, its lock does not escape before the end of the
221    // function, and the lock is taken by no other code path, so we
222    // know it is not contended.
223    pub fn build(
224        self,
225    ) -> Result<(
226        SandboxMemoryManager<HostSharedMemory>,
227        SandboxMemoryManager<GuestSharedMemory>,
228    )> {
229        let (hshm, gshm) = self.shared_mem.build();
230        let (hscratch, gscratch) = self.scratch_mem.build();
231        let mut host_mgr = SandboxMemoryManager {
232            shared_mem: hshm,
233            scratch_mem: hscratch,
234            layout: self.layout,
235            entrypoint: self.entrypoint,
236            mapped_rgns: self.mapped_rgns,
237            abort_buffer: self.abort_buffer,
238            snapshot_count: self.snapshot_count,
239        };
240        let guest_mgr = SandboxMemoryManager {
241            shared_mem: gshm,
242            scratch_mem: gscratch,
243            layout: self.layout,
244            entrypoint: self.entrypoint,
245            mapped_rgns: self.mapped_rgns,
246            abort_buffer: Vec::new(), // Guest doesn't need abort buffer
247            snapshot_count: self.snapshot_count,
248        };
249        host_mgr.update_scratch_bookkeeping()?;
250        Ok((host_mgr, guest_mgr))
251    }
252}
253
254impl SandboxMemoryManager<HostSharedMemory> {
255    #[allow(dead_code)]
256    pub(crate) fn parallel_invoke_slots_gva(&self) -> u64 {
257        self.layout.get_parallel_invoke_slots_gva()
258    }
259
260    pub(crate) fn parallel_invoke_slot_scratch_host_offset(&self, lane: usize) -> usize {
261        self.layout
262            .get_parallel_invoke_slot_scratch_host_offset(lane)
263    }
264
265    pub(crate) fn parallel_invoke_slot_host_ptr(
266        &self,
267        lane: usize,
268    ) -> Result<*mut ParallelInvokeSlot> {
269        let offset = self.parallel_invoke_slot_scratch_host_offset(lane);
270        let len = size_of::<ParallelInvokeSlot>();
271        if offset
272            .checked_add(len)
273            .is_none_or(|end| end > self.scratch_mem.mem_size())
274        {
275            return Err(crate::new_error!(
276                "parallel invoke slot {} is outside scratch memory",
277                lane
278            ));
279        }
280        let ptr = self.scratch_mem.base_ptr().wrapping_add(offset) as *mut ParallelInvokeSlot;
281        debug_assert_eq!(
282            (ptr as usize) % align_of::<ParallelInvokeSlot>(),
283            0,
284            "parallel invoke slot is not ABI-aligned"
285        );
286        Ok(ptr)
287    }
288
289    fn parallel_invoke_slot_field_offset(&self, lane: usize, field_offset: usize) -> usize {
290        self.parallel_invoke_slot_scratch_host_offset(lane) + field_offset
291    }
292
293    pub(crate) fn read_parallel_invoke_status(&self, lane: usize) -> Result<u32> {
294        self.scratch_mem.read::<u32>(
295            self.parallel_invoke_slot_field_offset(lane, offset_of!(ParallelInvokeSlot, status)),
296        )
297    }
298
299    pub(crate) fn write_parallel_invoke_status(&self, lane: usize, status: u32) -> Result<()> {
300        self.scratch_mem.write::<u32>(
301            self.parallel_invoke_slot_field_offset(lane, offset_of!(ParallelInvokeSlot, status)),
302            status,
303        )
304    }
305
306    /// Push raw bytes (e.g. a rkyv-archived `Request` envelope) onto
307    /// the guest's input data ring.
308    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
309    pub(crate) fn write_guest_function_call_raw(&mut self, buffer: &[u8]) -> Result<()> {
310        self.scratch_mem.push_buffer(
311            self.layout.get_input_data_buffer_scratch_host_offset(),
312            self.layout.sandbox_memory_config.get_input_data_size(),
313            buffer,
314        )
315    }
316
317    /// Pop the response bytes (e.g. a rkyv-archived `Response`
318    /// envelope) from the guest's output data ring.
319    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
320    pub(crate) fn read_guest_function_call_result_raw(&mut self) -> Result<Vec<u8>> {
321        self.scratch_mem.try_pop_buffer_raw(
322            self.layout.get_output_data_buffer_scratch_host_offset(),
323            self.layout.sandbox_memory_config.get_output_data_size(),
324        )
325    }
326
327    /// Pop raw bytes from the output ring — used by the host's
328    /// `OutBAction::CallFunction` arm to read the guest's request.
329    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
330    pub(crate) fn read_host_function_call_raw(&mut self) -> Result<Vec<u8>> {
331        self.scratch_mem.try_pop_buffer_raw(
332            self.layout.get_output_data_buffer_scratch_host_offset(),
333            self.layout.sandbox_memory_config.get_output_data_size(),
334        )
335    }
336
337    /// Push raw bytes (response to a guest→host call) onto the
338    /// guest's input data ring.
339    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
340    pub(crate) fn write_host_function_response_raw(&mut self, buffer: &[u8]) -> Result<()> {
341        self.scratch_mem.push_buffer(
342            self.layout.get_input_data_buffer_scratch_host_offset(),
343            self.layout.sandbox_memory_config.get_input_data_size(),
344            buffer,
345        )
346    }
347
348    /// Read guest log data from the `SharedMemory` contained within `self`
349    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
350    pub(crate) fn read_guest_log_data(&mut self) -> Result<GuestLogData> {
351        self.scratch_mem.try_pop_buffer_into::<GuestLogData>(
352            self.layout.get_output_data_buffer_scratch_host_offset(),
353            self.layout.sandbox_memory_config.get_output_data_size(),
354        )
355    }
356
357    pub(crate) fn clear_io_buffers(&mut self) {
358        // Clear the output data buffer
359        loop {
360            let Ok(_) = self.scratch_mem.try_pop_buffer_into::<Vec<u8>>(
361                self.layout.get_output_data_buffer_scratch_host_offset(),
362                self.layout.sandbox_memory_config.get_output_data_size(),
363            ) else {
364                break;
365            };
366        }
367        // Clear the input data buffer
368        loop {
369            let Ok(_) = self.scratch_mem.try_pop_buffer_into::<Vec<u8>>(
370                self.layout.get_input_data_buffer_scratch_host_offset(),
371                self.layout.sandbox_memory_config.get_input_data_size(),
372            ) else {
373                break;
374            };
375        }
376    }
377
378    #[inline]
379    fn update_scratch_bookkeeping_item(&mut self, offset: u64, value: u64) -> Result<()> {
380        let scratch_size = self.scratch_mem.mem_size();
381        let base_offset = scratch_size - offset as usize;
382        self.scratch_mem.write::<u64>(base_offset, value)
383    }
384
385    fn update_scratch_bookkeeping(&mut self) -> Result<()> {
386        use nub_host_common::layout::*;
387        let scratch_size = self.scratch_mem.mem_size();
388        self.update_scratch_bookkeeping_item(SCRATCH_TOP_SIZE_OFFSET, scratch_size as u64)?;
389        self.update_scratch_bookkeeping_item(
390            SCRATCH_TOP_ALLOCATOR_OFFSET,
391            self.layout.get_first_free_scratch_gpa(),
392        )?;
393        // Record the GPA of the snapshot's copy of the page tables.
394        // The copy lives at the tail of the snapshot blob; we copy it
395        // into scratch below so the guest walker can run against
396        // mutable, TLB-fresh tables. The guest reads this GPA during
397        // CoW fault-in to follow the original PTs on the first write
398        // — until the HV can execute directly out of the
399        // snapshot-resident PTs, at which point the whole split goes
400        // away.
401        self.update_scratch_bookkeeping_item(
402            SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET,
403            self.layout.get_pt_base_gpa(),
404        )?;
405        self.update_scratch_bookkeeping_item(
406            SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET,
407            self.snapshot_count,
408        )?;
409
410        // Initialise the guest input and output data buffers in
411        // scratch memory. TODO: remove the need for this.
412        self.scratch_mem.write::<u64>(
413            self.layout.get_input_data_buffer_scratch_host_offset(),
414            SandboxMemoryLayout::STACK_POINTER_SIZE_BYTES,
415        )?;
416        self.scratch_mem.write::<u64>(
417            self.layout.get_output_data_buffer_scratch_host_offset(),
418            SandboxMemoryLayout::STACK_POINTER_SIZE_BYTES,
419        )?;
420
421        let slot_start = self.layout.get_parallel_invoke_slots_scratch_host_offset();
422        let slot_end = slot_start + self.layout.get_parallel_invoke_slots_size();
423        self.scratch_mem.with_exclusivity(|scratch| {
424            scratch.as_mut_slice()[slot_start..slot_end].fill(0);
425        })?;
426
427        // Copy page tables from `shared_mem` into scratch. PT bytes
428        // are appended to the snapshot blob at build time and live
429        // just past the end of the guest-visible KVM slot (see
430        // `Snapshot::new`). Keeping them outside the KVM slot avoids
431        // overlapping with `map_file_cow` regions installed
432        // immediately after the snapshot in the guest PA space.
433        let snapshot_pt_end = self.shared_mem.mem_size();
434        let snapshot_pt_size = self.layout.get_pt_size();
435        let snapshot_pt_start = snapshot_pt_end - snapshot_pt_size;
436        self.scratch_mem.with_exclusivity(|scratch| {
437            let bytes = &self.shared_mem.as_slice()[snapshot_pt_start..snapshot_pt_end];
438            #[allow(clippy::needless_borrow)]
439            scratch.copy_from_slice(&bytes, self.layout.get_pt_base_scratch_offset())
440        })??;
441
442        Ok(())
443    }
444}
nub_host_kvm/mem/mgr.rs

nub_host_kvm/mem/
mgr.rs
