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    /// Copy the memory from the internally-stored memory snapshot
54    /// into the internally-stored `SharedMemory`.
55    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
56    pub(super) fn restore_from_snapshot<S: SharedMemory>(&self, shared_mem: &mut S) -> Result<()> {
57        shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))??;
58        Ok(())
59    }
60
61    /// The id of the sandbox this snapshot was taken from.
62    pub(crate) fn sandbox_id(&self) -> u64 {
63        self.sandbox_id
64    }
65
66    /// Get the mapped regions from this snapshot
67    pub(crate) fn regions(&self) -> &[MemoryRegion] {
68        &self.regions
69    }
70
71    /// Return the size of the snapshot in bytes.
72    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
73    pub(super) fn mem_size(&self) -> usize {
74        self.snapshot.len()
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use hyperlight_common::mem::PAGE_SIZE_USIZE;
81
82    use crate::mem::shared_mem::ExclusiveSharedMemory;
83
84    #[test]
85    fn restore() {
86        // Simplified version of the original test
87        let data1 = vec![b'a'; PAGE_SIZE_USIZE];
88        let data2 = vec![b'b'; PAGE_SIZE_USIZE];
89
90        let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
91        gm.copy_from_slice(&data1, 0).unwrap();
92
93        // Take snapshot of data1
94        let snapshot = super::SharedMemorySnapshot::new(&mut gm, 0, Vec::new()).unwrap();
95
96        // Modify memory to data2
97        gm.copy_from_slice(&data2, 0).unwrap();
98        assert_eq!(gm.as_slice(), &data2[..]);
99
100        // Restore should bring back data1
101        snapshot.restore_from_snapshot(&mut gm).unwrap();
102        assert_eq!(gm.as_slice(), &data1[..]);
103    }
104
105    #[test]
106    fn snapshot_mem_size() {
107        let size = PAGE_SIZE_USIZE * 2;
108        let mut gm = ExclusiveSharedMemory::new(size).unwrap();
109
110        let snapshot = super::SharedMemorySnapshot::new(&mut gm, 0, Vec::new()).unwrap();
111        assert_eq!(snapshot.mem_size(), size);
112    }
113
114    #[test]
115    fn multiple_snapshots_independent() {
116        let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
117
118        // Create first snapshot with pattern A
119        let pattern_a = vec![0xAA; PAGE_SIZE_USIZE];
120        gm.copy_from_slice(&pattern_a, 0).unwrap();
121        let snapshot_a = super::SharedMemorySnapshot::new(&mut gm, 1, Vec::new()).unwrap();
122
123        // Create second snapshot with pattern B
124        let pattern_b = vec![0xBB; PAGE_SIZE_USIZE];
125        gm.copy_from_slice(&pattern_b, 0).unwrap();
126        let snapshot_b = super::SharedMemorySnapshot::new(&mut gm, 2, Vec::new()).unwrap();
127
128        // Clear memory
129        gm.copy_from_slice(&[0; PAGE_SIZE_USIZE], 0).unwrap();
130
131        // Restore snapshot A
132        snapshot_a.restore_from_snapshot(&mut gm).unwrap();
133        assert_eq!(gm.as_slice(), &pattern_a[..]);
134
135        // Restore snapshot B
136        snapshot_b.restore_from_snapshot(&mut gm).unwrap();
137        assert_eq!(gm.as_slice(), &pattern_b[..]);
138    }
139}
hyperlight_host/mem/shared_mem_snapshot.rs

hyperlight_host/mem/
shared_mem_snapshot.rs
