use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
use super::capabilities::CapabilityFingerprint;
use super::chunker;
use super::container;
use super::objects::{
self, CommitId, CommitObject, Encoded, FileContent, KIND_CAPS, KIND_CHUNK, KIND_COMMIT,
KIND_FILE, KIND_SHELL, KIND_TREE, ObjectId, TreeEntry, TreeEntryKind,
};
use super::{SNAPSHOT_VERSION, Snapshot, SnapshotOptions};
use crate::interpreter::{ShellState, ShellStateOptions};
const MAX_ANCESTRY: usize = 1_000_000;
const MAX_CHECKOUT_BYTES: u64 = 4 * 1024 * 1024 * 1024;
struct CheckoutBudget {
remaining: u64,
}
impl CheckoutBudget {
fn new() -> Self {
Self {
remaining: MAX_CHECKOUT_BYTES,
}
}
fn spend(&mut self, bytes: u64) -> crate::Result<()> {
self.remaining = self.remaining.checked_sub(bytes).ok_or_else(|| {
crate::Error::Internal(format!(
"snapshot checkout exceeds the {MAX_CHECKOUT_BYTES} byte materialization limit"
))
})?;
Ok(())
}
}
pub trait ObjectSource {
fn get_object(&self, id: &ObjectId) -> Option<&[u8]>;
}
impl ObjectSource for HashMap<ObjectId, Vec<u8>> {
fn get_object(&self, id: &ObjectId) -> Option<&[u8]> {
self.get(id).map(Vec::as_slice)
}
}
impl ObjectSource for BTreeMap<ObjectId, Vec<u8>> {
fn get_object(&self, id: &ObjectId) -> Option<&[u8]> {
self.get(id).map(Vec::as_slice)
}
}
impl<T: ObjectSource + ?Sized> ObjectSource for &T {
fn get_object(&self, id: &ObjectId) -> Option<&[u8]> {
(**self).get_object(id)
}
}
#[derive(Debug, Clone, Default)]
pub struct CommitOptions {
pub(crate) parents: Vec<CommitId>,
pub(crate) meta: BTreeMap<String, String>,
pub(crate) snapshot: SnapshotOptions,
pub(crate) have: HashSet<ObjectId>,
}
impl CommitOptions {
pub fn new() -> Self {
Self::default()
}
pub fn parent(mut self, parent: CommitId) -> Self {
self.parents.push(parent);
self
}
pub fn meta(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.meta.insert(key.into(), value.into());
self
}
pub fn have<'a, I: IntoIterator<Item = &'a ObjectId>>(mut self, ids: I) -> Self {
self.have.extend(ids.into_iter().copied());
self
}
pub fn exclude_filesystem(mut self, exclude: bool) -> Self {
self.snapshot.exclude_filesystem = exclude;
self
}
pub fn exclude_functions(mut self, exclude: bool) -> Self {
self.snapshot.exclude_functions = exclude;
self
}
}
#[derive(Debug, Clone)]
pub struct PackedCommit {
pub(crate) id: CommitId,
pub(crate) objects: BTreeMap<ObjectId, Vec<u8>>,
pub(crate) self_contained: bool,
}
impl PackedCommit {
pub fn id(&self) -> CommitId {
self.id
}
pub fn objects(&self) -> impl Iterator<Item = (ObjectId, &[u8])> {
self.objects.iter().map(|(id, blob)| (*id, blob.as_slice()))
}
pub fn into_objects(self) -> impl Iterator<Item = (ObjectId, Vec<u8>)> {
self.objects.into_iter()
}
pub fn object_count(&self) -> usize {
self.objects.len()
}
pub fn stored_bytes(&self) -> usize {
self.objects.values().map(Vec::len).sum()
}
pub fn is_self_contained(&self) -> bool {
self.self_contained
}
pub fn to_bytes(&self) -> crate::Result<Vec<u8>> {
self.pack(None)
}
pub fn to_bytes_keyed(&self, key: &[u8]) -> crate::Result<Vec<u8>> {
self.pack(Some(key))
}
fn pack(&self, key: Option<&[u8]>) -> crate::Result<Vec<u8>> {
if !self.self_contained {
return Err(crate::Error::Internal(
"cannot pack an incremental commit: it omits objects the store already holds"
.to_string(),
));
}
let body = container::encode(self.id, &self.objects);
Ok(super::seal(&body, key))
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SnapshotDiff {
pub files_added: Vec<String>,
pub files_modified: Vec<String>,
pub files_removed: Vec<String>,
pub shell_changed: bool,
}
impl SnapshotDiff {
pub fn is_empty(&self) -> bool {
self.files_added.is_empty()
&& self.files_modified.is_empty()
&& self.files_removed.is_empty()
&& !self.shell_changed
}
}
pub struct SnapshotGraph;
impl SnapshotGraph {
pub fn read_commit(root: CommitId, source: &impl ObjectSource) -> crate::Result<CommitObject> {
let payload = load(root, KIND_COMMIT, source)?;
objects::decode_commit(&payload)
}
pub fn parents(root: CommitId, source: &impl ObjectSource) -> crate::Result<Vec<CommitId>> {
Ok(Self::read_commit(root, source)?.parents)
}
pub fn meta(
root: CommitId,
source: &impl ObjectSource,
) -> crate::Result<BTreeMap<String, String>> {
Ok(Self::read_commit(root, source)?.meta)
}
pub fn capabilities(
root: CommitId,
source: &impl ObjectSource,
) -> crate::Result<CapabilityFingerprint> {
let commit = Self::read_commit(root, source)?;
let payload = load(commit.caps, KIND_CAPS, source)?;
objects::from_canonical_json(&payload)
}
pub fn ancestry(
root: CommitId,
source: &impl ObjectSource,
limit: usize,
) -> crate::Result<Vec<CommitId>> {
let limit = limit.min(MAX_ANCESTRY);
let mut seen = HashSet::new();
let mut queue = VecDeque::from([root]);
let mut out = Vec::new();
while let Some(id) = queue.pop_front() {
if out.len() >= limit || !seen.insert(id) {
continue;
}
if source.get_object(&id).is_none() {
continue;
}
let commit = Self::read_commit(id, source)?;
out.push(id);
queue.extend(commit.parents);
}
Ok(out)
}
pub fn plan_checkout(
root: CommitId,
source: &impl ObjectSource,
) -> crate::Result<Vec<ObjectId>> {
let mut missing = BTreeSet::new();
walk(root, source, &mut missing, &mut BTreeSet::new())?;
Ok(missing.into_iter().collect())
}
pub fn reachable(root: CommitId, source: &impl ObjectSource) -> crate::Result<Vec<ObjectId>> {
let mut missing = BTreeSet::new();
let mut present = BTreeSet::new();
walk(root, source, &mut missing, &mut present)?;
if let Some(id) = missing.iter().next() {
return Err(crate::Error::Internal(format!(
"snapshot object {id} is missing from the store"
)));
}
Ok(present.into_iter().collect())
}
pub fn diff(
a: CommitId,
b: CommitId,
source: &impl ObjectSource,
) -> crate::Result<SnapshotDiff> {
let commit_a = Self::read_commit(a, source)?;
let commit_b = Self::read_commit(b, source)?;
let tree_a = read_tree(commit_a.tree, source)?;
let tree_b = read_tree(commit_b.tree, source)?;
let map_a: BTreeMap<&str, &TreeEntry> =
tree_a.iter().map(|e| (e.path.as_str(), e)).collect();
let map_b: BTreeMap<&str, &TreeEntry> =
tree_b.iter().map(|e| (e.path.as_str(), e)).collect();
let mut diff = SnapshotDiff {
shell_changed: commit_a.shell != commit_b.shell,
..Default::default()
};
for (path, entry) in &map_b {
match map_a.get(path) {
None => diff.files_added.push((*path).to_string()),
Some(before) if before.mode != entry.mode || before.kind != entry.kind => {
diff.files_modified.push((*path).to_string());
}
Some(_) => {}
}
}
for path in map_a.keys() {
if !map_b.contains_key(path) {
diff.files_removed.push((*path).to_string());
}
}
Ok(diff)
}
pub(crate) fn materialize(
root: CommitId,
source: &impl ObjectSource,
) -> crate::Result<(Snapshot, CapabilityFingerprint)> {
let commit = Self::read_commit(root, source)?;
let caps_payload = load(commit.caps, KIND_CAPS, source)?;
let caps: CapabilityFingerprint = objects::from_canonical_json(&caps_payload)?;
let shell_payload = load(commit.shell, KIND_SHELL, source)?;
let shell: ShellState = objects::from_canonical_json(&shell_payload)?;
let vfs = match commit.tree {
None => None,
Some(tree_id) => {
let entries = objects::decode_tree(&load(tree_id, KIND_TREE, source)?)?;
let mut budget = CheckoutBudget::new();
Some(objects::tree_to_vfs(&entries, |file_id| {
resolve_file(file_id, source, &mut budget)
})?)
}
};
Ok((
Snapshot {
version: SNAPSHOT_VERSION,
shell,
vfs,
session_commands: commit.session_commands,
session_exec_calls: commit.session_exec_calls,
},
caps,
))
}
}
fn load(id: ObjectId, kind: u8, source: &impl ObjectSource) -> crate::Result<Vec<u8>> {
let blob = source
.get_object(&id)
.ok_or_else(|| crate::Error::Internal(format!("snapshot object {id} not found")))?;
let raw = container::from_storage_blob(blob)?;
let object = Encoded::from_storage(id, &raw)?;
object.expect_kind(kind)?;
Ok(object.payload)
}
fn read_tree(tree: Option<ObjectId>, source: &impl ObjectSource) -> crate::Result<Vec<TreeEntry>> {
match tree {
None => Ok(Vec::new()),
Some(id) => objects::decode_tree(&load(id, KIND_TREE, source)?),
}
}
fn resolve_file(
file_id: ObjectId,
source: &impl ObjectSource,
budget: &mut CheckoutBudget,
) -> crate::Result<Vec<u8>> {
match objects::decode_file(&load(file_id, KIND_FILE, source)?)? {
FileContent::Inline(bytes) => {
budget.spend(bytes.len() as u64)?;
Ok(bytes)
}
FileContent::Chunked { size, chunks } => {
budget.spend(size)?;
let mut out =
Vec::with_capacity(usize::try_from(size).unwrap_or_default().min(1 << 24));
for chunk_id in chunks {
let chunk = load(chunk_id, KIND_CHUNK, source)?;
if chunk.len() > chunker::MAX_CHUNK {
return Err(crate::Error::Internal(format!(
"snapshot chunk {chunk_id} is {} bytes, above the {} byte maximum",
chunk.len(),
chunker::MAX_CHUNK
)));
}
if out.len() as u64 + chunk.len() as u64 > size {
return Err(crate::Error::Internal(format!(
"snapshot file {file_id} has chunks exceeding its declared {size} bytes"
)));
}
out.extend_from_slice(&chunk);
}
if out.len() as u64 != size {
return Err(crate::Error::Internal(format!(
"snapshot file {file_id} reassembled to {} bytes but declares {size}",
out.len()
)));
}
Ok(out)
}
}
}
fn walk(
root: CommitId,
source: &impl ObjectSource,
missing: &mut BTreeSet<ObjectId>,
present: &mut BTreeSet<ObjectId>,
) -> crate::Result<()> {
let mut queue: VecDeque<(ObjectId, u8)> = VecDeque::from([(root, KIND_COMMIT)]);
let mut seen: HashSet<ObjectId> = HashSet::new();
while let Some((id, kind)) = queue.pop_front() {
if !seen.insert(id) {
continue;
}
let Some(blob) = source.get_object(&id) else {
missing.insert(id);
continue;
};
let raw = container::from_storage_blob(blob)?;
let object = Encoded::from_storage(id, &raw)?;
object.expect_kind(kind)?;
present.insert(id);
match kind {
KIND_COMMIT => {
let commit = objects::decode_commit(&object.payload)?;
queue.push_back((commit.shell, KIND_SHELL));
queue.push_back((commit.caps, KIND_CAPS));
if let Some(tree) = commit.tree {
queue.push_back((tree, KIND_TREE));
}
}
KIND_TREE => {
for entry in objects::decode_tree(&object.payload)? {
if let TreeEntryKind::File(file_id) = entry.kind {
queue.push_back((file_id, KIND_FILE));
}
}
}
KIND_FILE => {
if let FileContent::Chunked { chunks, .. } = objects::decode_file(&object.payload)?
{
for chunk in chunks {
queue.push_back((chunk, KIND_CHUNK));
}
}
}
_ => {}
}
}
Ok(())
}
impl crate::Bash {
pub(crate) fn commit(&self, options: CommitOptions) -> crate::Result<PackedCommit> {
let shell = self
.interpreter
.shell_state_with_options(ShellStateOptions {
include_functions: !options.snapshot.exclude_functions,
});
let vfs = if options.snapshot.exclude_filesystem {
None
} else {
self.fs().vfs_snapshot()
};
let mut emitted: Vec<Encoded> = Vec::new();
let tree_id = match &vfs {
None => None,
Some(vfs) => {
let (tree, content) = objects::build_tree(vfs)?;
let id = tree.id;
emitted.push(tree);
emitted.extend(content);
Some(id)
}
};
let shell_obj = Encoded::new(KIND_SHELL, objects::canonical_json(&shell)?);
let caps_obj = Encoded::new(
KIND_CAPS,
objects::canonical_json(&CapabilityFingerprint::capture(self))?,
);
let shell_id = shell_obj.id;
let caps_id = caps_obj.id;
emitted.push(shell_obj);
emitted.push(caps_obj);
let counters = self.interpreter.counters();
let commit_obj = objects::encode_commit(&CommitObject {
parents: options.parents.clone(),
tree: tree_id,
shell: shell_id,
caps: caps_id,
session_commands: counters.session_commands,
session_exec_calls: counters.session_exec_calls,
meta: options.meta.clone(),
})?;
let commit_id = commit_obj.id;
emitted.push(commit_obj);
let mut store = BTreeMap::new();
for object in emitted {
if object.id != commit_id && options.have.contains(&object.id) {
continue;
}
store.entry(object.id).or_insert_with(|| {
container::storage_blob_from_parts(object.kind, &object.payload)
});
}
Ok(PackedCommit {
id: commit_id,
objects: store,
self_contained: options.have.is_empty(),
})
}
pub(crate) fn checkout(
&mut self,
root: CommitId,
source: &impl ObjectSource,
policy: super::CheckoutPolicy,
) -> crate::Result<()> {
let (snapshot, caps) = SnapshotGraph::materialize(root, source)?;
self.apply_checked(&snapshot, &caps, policy)
}
pub(crate) fn apply_checked(
&mut self,
snapshot: &Snapshot,
caps: &CapabilityFingerprint,
policy: super::CheckoutPolicy,
) -> crate::Result<()> {
let live = CapabilityFingerprint::capture(self);
policy.check(&caps.compare(&live))?;
if let Some(vfs) = &snapshot.vfs {
super::capabilities::check_state_evidence(vfs)?;
}
self.restore_snapshot_inner(snapshot)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::snapshot::objects::{encode_file, encode_tree};
fn store_object(store: &mut HashMap<ObjectId, Vec<u8>>, object: &Encoded) -> ObjectId {
store.insert(
object.id,
container::storage_blob_from_parts(object.kind, &object.payload),
);
object.id
}
fn chunk_bomb(
store: &mut HashMap<ObjectId, Vec<u8>>,
repeats: usize,
declared_size: u64,
) -> ObjectId {
let chunk = Encoded::new(KIND_CHUNK, vec![b'A'; chunker::MAX_CHUNK]);
let chunk_id = store_object(store, &chunk);
let manifest = encode_file(&FileContent::Chunked {
size: declared_size,
chunks: vec![chunk_id; repeats],
});
store_object(store, &manifest)
}
#[test]
fn a_repeated_chunk_cannot_grow_a_file_past_its_declared_size() {
let mut store = HashMap::new();
let file_id = chunk_bomb(&mut store, 64, 1024);
let mut budget = CheckoutBudget::new();
let err = resolve_file(file_id, &store, &mut budget).unwrap_err();
assert!(
err.to_string().contains("declared"),
"expected a declared-size rejection, got: {err}"
);
}
#[test]
fn an_absurd_declared_size_is_refused_before_allocating() {
let mut store = HashMap::new();
let file_id = chunk_bomb(&mut store, 4, u64::MAX / 2);
let mut budget = CheckoutBudget::new();
let err = resolve_file(file_id, &store, &mut budget).unwrap_err();
assert!(
err.to_string().contains("materialization limit"),
"expected a budget rejection, got: {err}"
);
}
#[test]
fn the_budget_is_shared_across_every_file_in_a_tree() {
let mut budget = CheckoutBudget::new();
assert!(budget.spend(MAX_CHECKOUT_BYTES / 2).is_ok());
assert!(budget.spend(MAX_CHECKOUT_BYTES / 2).is_ok());
assert!(
budget.spend(1).is_err(),
"the budget must be cumulative, not per-file"
);
}
#[test]
fn an_oversized_chunk_is_rejected() {
let mut store = HashMap::new();
let chunk = Encoded::new(KIND_CHUNK, vec![b'B'; chunker::MAX_CHUNK + 1]);
let chunk_id = store_object(&mut store, &chunk);
let manifest = encode_file(&FileContent::Chunked {
size: chunker::MAX_CHUNK as u64 + 1,
chunks: vec![chunk_id],
});
let file_id = store_object(&mut store, &manifest);
let mut budget = CheckoutBudget::new();
let err = resolve_file(file_id, &store, &mut budget).unwrap_err();
assert!(
err.to_string().contains("maximum"),
"expected an oversized-chunk rejection, got: {err}"
);
}
#[test]
fn a_well_formed_chunked_file_still_resolves() {
let mut store = HashMap::new();
let a = Encoded::new(KIND_CHUNK, vec![b'x'; 1000]);
let b = Encoded::new(KIND_CHUNK, vec![b'y'; 500]);
let a_id = store_object(&mut store, &a);
let b_id = store_object(&mut store, &b);
let manifest = encode_file(&FileContent::Chunked {
size: 1500,
chunks: vec![a_id, b_id],
});
let file_id = store_object(&mut store, &manifest);
let mut budget = CheckoutBudget::new();
let out = resolve_file(file_id, &store, &mut budget).unwrap();
assert_eq!(out.len(), 1500);
assert_eq!(&out[..1000], &vec![b'x'; 1000][..]);
assert_eq!(&out[1000..], &vec![b'y'; 500][..]);
}
#[test]
fn a_tree_of_many_bomb_files_is_bounded_in_aggregate() {
let mut store = HashMap::new();
let chunk = Encoded::new(KIND_CHUNK, vec![b'C'; chunker::MAX_CHUNK]);
let chunk_id = store_object(&mut store, &chunk);
let mut entries = Vec::new();
for i in 0..64 {
let manifest = encode_file(&FileContent::Chunked {
size: MAX_CHECKOUT_BYTES / 8,
chunks: vec![chunk_id],
});
let file_id = store_object(&mut store, &manifest);
entries.push(TreeEntry {
path: format!("/f{i}"),
mode: 0o644,
kind: TreeEntryKind::File(file_id),
});
}
let tree = encode_tree(&entries);
let entries = objects::decode_tree(&tree.payload).unwrap();
let mut budget = CheckoutBudget::new();
let result = objects::tree_to_vfs(&entries, |file_id| {
resolve_file(file_id, &store, &mut budget)
});
assert!(
result.is_err(),
"a tree of oversized files must exhaust the shared budget"
);
}
}