agentos_native_sidecar_core/
layers.rs

1use crate::SidecarCoreError;
2use agentos_kernel::root_fs::{
3    FilesystemEntry, FilesystemEntryKind, RootFileSystem,
4    RootFilesystemDescriptor as KernelRootFilesystemDescriptor,
5    RootFilesystemMode as KernelRootFilesystemMode, RootFilesystemSnapshot,
6};
7use std::collections::{BTreeMap, BTreeSet};
8
9pub const MAX_VM_LAYERS: usize = 256;
10
11#[derive(Debug)]
12pub struct VmLayerStore {
13    next_layer_id: u64,
14    layers: BTreeMap<String, VmLayer>,
15}
16
17impl Default for VmLayerStore {
18    fn default() -> Self {
19        Self {
20            next_layer_id: 1,
21            layers: BTreeMap::new(),
22        }
23    }
24}
25
26#[derive(Debug)]
27enum VmLayer {
28    Writable(RootFileSystem),
29    Snapshot(RootFilesystemSnapshot),
30    Overlay(VmOverlayLayer),
31}
32
33#[derive(Debug, Clone)]
34struct VmOverlayLayer {
35    mode: KernelRootFilesystemMode,
36    upper_layer_id: Option<String>,
37    lower_layer_ids: Vec<String>,
38}
39
40impl VmLayerStore {
41    pub fn layer_count(&self) -> usize {
42        self.layers.len()
43    }
44
45    pub fn create_writable_layer(&mut self) -> Result<String, SidecarCoreError> {
46        self.ensure_layer_capacity()?;
47        let filesystem = new_writable_layer()?;
48        let layer_id = self.allocate_layer_id()?;
49        self.layers
50            .insert(layer_id.clone(), VmLayer::Writable(filesystem));
51        Ok(layer_id)
52    }
53
54    pub fn seal_layer(&mut self, layer_id: &str) -> Result<String, SidecarCoreError> {
55        let snapshot = match self.layers.get_mut(layer_id) {
56            Some(VmLayer::Writable(filesystem)) => filesystem
57                .snapshot()
58                .map_err(|error| SidecarCoreError::new(format!("seal layer: {error}")))?,
59            Some(VmLayer::Snapshot(_)) | Some(VmLayer::Overlay(_)) => {
60                return Err(SidecarCoreError::new(format!(
61                    "layer {layer_id} is not writable"
62                )));
63            }
64            None => return Err(SidecarCoreError::new(format!("unknown layer: {layer_id}"))),
65        };
66        let sealed_layer_id = self.allocate_layer_id()?;
67        match self
68            .layers
69            .remove(layer_id)
70            .expect("layer should still exist after snapshot")
71        {
72            VmLayer::Writable(_) => {}
73            VmLayer::Snapshot(_) | VmLayer::Overlay(_) => {
74                return Err(SidecarCoreError::new(format!(
75                    "layer {layer_id} is not writable"
76                )));
77            }
78        }
79        self.layers
80            .insert(sealed_layer_id.clone(), VmLayer::Snapshot(snapshot));
81        Ok(sealed_layer_id)
82    }
83
84    pub fn import_snapshot(
85        &mut self,
86        snapshot: RootFilesystemSnapshot,
87    ) -> Result<String, SidecarCoreError> {
88        self.ensure_layer_capacity()?;
89        let layer_id = self.allocate_layer_id()?;
90        self.layers
91            .insert(layer_id.clone(), VmLayer::Snapshot(snapshot));
92        Ok(layer_id)
93    }
94
95    pub fn export_snapshot(
96        &mut self,
97        layer_id: &str,
98    ) -> Result<RootFilesystemSnapshot, SidecarCoreError> {
99        materialize_vm_layer_snapshot(self, layer_id)
100    }
101
102    pub fn create_overlay_layer(
103        &mut self,
104        mode: KernelRootFilesystemMode,
105        upper_layer_id: Option<String>,
106        lower_layer_ids: Vec<String>,
107    ) -> Result<String, SidecarCoreError> {
108        self.ensure_layer_capacity()?;
109        for layer_id in &lower_layer_ids {
110            if !self.layers.contains_key(layer_id) {
111                return Err(SidecarCoreError::new(format!(
112                    "unknown lower layer: {layer_id}"
113                )));
114            }
115        }
116        if let Some(layer_id) = upper_layer_id.as_ref() {
117            if !self.layers.contains_key(layer_id) {
118                return Err(SidecarCoreError::new(format!(
119                    "unknown upper layer: {layer_id}"
120                )));
121            }
122        }
123
124        let layer_id = self.allocate_layer_id()?;
125        self.layers.insert(
126            layer_id.clone(),
127            VmLayer::Overlay(VmOverlayLayer {
128                mode,
129                upper_layer_id,
130                lower_layer_ids,
131            }),
132        );
133        Ok(layer_id)
134    }
135
136    fn ensure_layer_capacity(&self) -> Result<(), SidecarCoreError> {
137        if self.layers.len() >= MAX_VM_LAYERS {
138            return Err(SidecarCoreError::new(format!(
139                "VM layer limit exceeded: limit is {MAX_VM_LAYERS}"
140            )));
141        }
142        Ok(())
143    }
144
145    fn allocate_layer_id(&mut self) -> Result<String, SidecarCoreError> {
146        let layer_id = format!("layer-{}", self.next_layer_id);
147        self.next_layer_id = self
148            .next_layer_id
149            .checked_add(1)
150            .ok_or_else(|| SidecarCoreError::new("VM layer id overflow"))?;
151        Ok(layer_id)
152    }
153}
154
155fn new_writable_layer() -> Result<RootFileSystem, SidecarCoreError> {
156    RootFileSystem::from_descriptor(KernelRootFilesystemDescriptor {
157        mode: KernelRootFilesystemMode::Ephemeral,
158        disable_default_base_layer: true,
159        lowers: Vec::new(),
160        bootstrap_entries: Vec::new(),
161    })
162    .map_err(|error| SidecarCoreError::new(format!("create writable layer: {error}")))
163}
164
165fn materialize_vm_layer_snapshot(
166    layers: &mut VmLayerStore,
167    layer_id: &str,
168) -> Result<RootFilesystemSnapshot, SidecarCoreError> {
169    materialize_vm_layer_snapshot_inner(layers, layer_id, &mut BTreeSet::new())
170}
171
172fn materialize_vm_layer_snapshot_inner(
173    layers: &mut VmLayerStore,
174    layer_id: &str,
175    active: &mut BTreeSet<String>,
176) -> Result<RootFilesystemSnapshot, SidecarCoreError> {
177    if !active.insert(layer_id.to_owned()) {
178        return Err(SidecarCoreError::new(format!(
179            "layer graph cycle detected at {layer_id}"
180        )));
181    }
182
183    let result = if let Some(VmLayer::Snapshot(snapshot)) = layers.layers.get(layer_id) {
184        Ok(snapshot.clone())
185    } else if let Some(VmLayer::Overlay(overlay)) = layers.layers.get(layer_id) {
186        let overlay = overlay.clone();
187        let lowers = overlay
188            .lower_layer_ids
189            .iter()
190            .map(|lower_id| materialize_vm_layer_snapshot_inner(layers, lower_id, active))
191            .collect::<Result<Vec<_>, _>>()?;
192        let bootstrap_entries = match overlay.upper_layer_id.as_deref() {
193            Some(upper_layer_id) => dedupe_overlay_bootstrap_entries(
194                &lowers,
195                materialize_vm_layer_snapshot_inner(layers, upper_layer_id, active)?.entries,
196            ),
197            None => Vec::new(),
198        };
199        let mut root = RootFileSystem::from_descriptor(KernelRootFilesystemDescriptor {
200            mode: overlay.mode,
201            disable_default_base_layer: true,
202            lowers,
203            bootstrap_entries,
204        })
205        .map_err(|error| SidecarCoreError::new(format!("materialize overlay layer: {error}")))?;
206        root.snapshot()
207            .map_err(|error| SidecarCoreError::new(format!("snapshot overlay layer: {error}")))
208    } else if let Some(VmLayer::Writable(filesystem)) = layers.layers.get_mut(layer_id) {
209        filesystem
210            .snapshot()
211            .map_err(|error| SidecarCoreError::new(format!("snapshot writable layer: {error}")))
212    } else {
213        Err(SidecarCoreError::new(format!("unknown layer: {layer_id}")))
214    };
215
216    active.remove(layer_id);
217    result
218}
219
220fn dedupe_overlay_bootstrap_entries(
221    lowers: &[RootFilesystemSnapshot],
222    upper_entries: Vec<FilesystemEntry>,
223) -> Vec<FilesystemEntry> {
224    let mut lower_paths = lowers
225        .iter()
226        .flat_map(|snapshot| snapshot.entries.iter().map(|entry| entry.path.clone()))
227        .collect::<BTreeSet<_>>();
228
229    upper_entries
230        .into_iter()
231        .filter(|entry| {
232            if lower_paths.contains(&entry.path)
233                && matches!(entry.kind, FilesystemEntryKind::Directory)
234            {
235                return false;
236            }
237            lower_paths.insert(entry.path.clone());
238            true
239        })
240        .collect()
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246
247    #[test]
248    fn imports_exports_and_overlays_snapshots() {
249        let mut layers = VmLayerStore::default();
250        let lower = layers
251            .import_snapshot(RootFilesystemSnapshot {
252                entries: vec![FilesystemEntry::file("/value.txt", b"lower".to_vec())],
253            })
254            .expect("import lower");
255        let upper = layers
256            .import_snapshot(RootFilesystemSnapshot {
257                entries: vec![FilesystemEntry::file("/extra.txt", b"upper".to_vec())],
258            })
259            .expect("import upper");
260        let overlay = layers
261            .create_overlay_layer(
262                KernelRootFilesystemMode::Ephemeral,
263                Some(upper),
264                vec![lower],
265            )
266            .expect("create overlay");
267
268        let snapshot = layers.export_snapshot(&overlay).expect("export overlay");
269        assert!(snapshot
270            .entries
271            .iter()
272            .any(|entry| entry.path == "/value.txt"));
273        assert!(snapshot
274            .entries
275            .iter()
276            .any(|entry| entry.path == "/extra.txt"));
277    }
278
279    #[test]
280    fn sealing_non_writable_layer_is_rejected() {
281        let mut layers = VmLayerStore::default();
282        let layer = layers
283            .import_snapshot(RootFilesystemSnapshot {
284                entries: Vec::new(),
285            })
286            .expect("import snapshot");
287
288        let error = layers
289            .seal_layer(&layer)
290            .expect_err("snapshot layer should not seal");
291        assert!(error.to_string().contains("is not writable"));
292    }
293}
agentos_native_sidecar_core/layers.rs

agentos_native_sidecar_core/
layers.rs
