hyperlight_host/mem/
shared_mem_snapshot.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
17use tracing::{Span, instrument};
18
19use super::memory_region::MemoryRegion;
20use super::shared_mem::SharedMemory;
21use crate::Result;
22
23/// A wrapper around a `SharedMemory` reference and a snapshot
24/// of the memory therein
25#[derive(Clone)]
26pub(crate) struct SharedMemorySnapshot {
27    // Unique ID of the sandbox this snapshot was taken from
28    sandbox_id: u64,
29    // Memory of the sandbox at the time this snapshot was taken
30    snapshot: Vec<u8>,
31    /// The memory regions that were mapped when this snapshot was taken (excluding initial sandbox regions)
32    regions: Vec<MemoryRegion>,
33}
34
35impl SharedMemorySnapshot {
36    /// Take a snapshot of the memory in `shared_mem`, then create a new
37    /// instance of `Self` with the snapshot stored therein.
38    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
39    pub(super) fn new<S: SharedMemory>(
40        shared_mem: &mut S,
41        sandbox_id: u64,
42        regions: Vec<MemoryRegion>,
43    ) -> Result<Self> {
44        // TODO: Track dirty pages instead of copying entire memory
45        let snapshot = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
46        Ok(Self {
47            sandbox_id,
48            snapshot,
49            regions,
50        })
51    }
52
53    /// Take another snapshot of the internally-stored `SharedMemory`,
54    /// then store it internally.
55    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
56    #[allow(dead_code)]
57    pub(super) fn replace_snapshot<S: SharedMemory>(&mut self, shared_mem: &mut S) -> Result<()> {
58        self.snapshot = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
59        Ok(())
60    }
61
62    /// Copy the memory from the internally-stored memory snapshot
63    /// into the internally-stored `SharedMemory`.
64    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
65    pub(super) fn restore_from_snapshot<S: SharedMemory>(&self, shared_mem: &mut S) -> Result<()> {
66        shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))??;
67        Ok(())
68    }
69
70    /// The id of the sandbox this snapshot was taken from.
71    pub(crate) fn sandbox_id(&self) -> u64 {
72        self.sandbox_id
73    }
74
75    /// Get the mapped regions from this snapshot
76    pub(crate) fn regions(&self) -> &[MemoryRegion] {
77        &self.regions
78    }
79
80    /// Return the size of the snapshot in bytes.
81    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
82    pub(super) fn mem_size(&self) -> usize {
83        self.snapshot.len()
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use hyperlight_common::mem::PAGE_SIZE_USIZE;
90
91    use crate::mem::shared_mem::ExclusiveSharedMemory;
92
93    #[test]
94    fn restore_replace() {
95        let mut data1 = vec![b'a', b'b', b'c'];
96        data1.resize_with(PAGE_SIZE_USIZE, || 0);
97        let data2 = data1.iter().map(|b| b + 1).collect::<Vec<u8>>();
98        let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
99        gm.copy_from_slice(data1.as_slice(), 0).unwrap();
100        let mut snap = super::SharedMemorySnapshot::new(&mut gm, 0, Vec::new()).unwrap();
101        {
102            // after the first snapshot is taken, make sure gm has the equivalent
103            // of data1
104            assert_eq!(data1, gm.copy_all_to_vec().unwrap());
105        }
106
107        {
108            // modify gm with data2 rather than data1 and restore from
109            // snapshot. we should have the equivalent of data1 again
110            gm.copy_from_slice(data2.as_slice(), 0).unwrap();
111            assert_eq!(data2, gm.copy_all_to_vec().unwrap());
112            snap.restore_from_snapshot(&mut gm).unwrap();
113            assert_eq!(data1, gm.copy_all_to_vec().unwrap());
114        }
115        {
116            // modify gm with data2, then retake the snapshot and restore
117            // from the new snapshot. we should have the equivalent of data2
118            gm.copy_from_slice(data2.as_slice(), 0).unwrap();
119            assert_eq!(data2, gm.copy_all_to_vec().unwrap());
120            snap.replace_snapshot(&mut gm).unwrap();
121            assert_eq!(data2, gm.copy_all_to_vec().unwrap());
122            snap.restore_from_snapshot(&mut gm).unwrap();
123            assert_eq!(data2, gm.copy_all_to_vec().unwrap());
124        }
125    }
126}
hyperlight_host/mem/shared_mem_snapshot.rs

hyperlight_host/mem/
shared_mem_snapshot.rs
