a3s_box_runtime/

snapshot.rs

//! VM Snapshot Store — Save and restore VM configuration snapshots.
//!
//! Snapshots are stored as directories under `~/.a3s/snapshots/<id>/`:
//! - `metadata.json` — SnapshotMetadata (config, resources, env, etc.)
//! - `rootfs.json` — versioned payload identity
//! - `rootfs/` — immutable directory payload (legacy-compatible), or
//! - `rootfs-ext4-v1/` — immutable guest-native raw artifact
//!
//! Directory restore creates a shared copy-on-write lower. Guest-native restore
//! creates and verifies a private raw clone, so the restored box has no
//! continued dependency on the snapshot bundle.

use std::cmp::Reverse;
use std::path::{Path, PathBuf};

use a3s_box_core::error::{BoxError, Result};
use a3s_box_core::rootfs_metadata::RootfsMetadataManifest;
#[cfg(test)]
use a3s_box_core::rootfs_metadata::{
    IMAGE_ROOTFS_METADATA_PATH, PREVIOUS_ROOTFS_METADATA_PATH, ROOTFS_METADATA_PATH,
};
use a3s_box_core::snapshot::SnapshotMetadata;
use a3s_box_core::SnapshotStoreBackend;

use crate::file_lock::FileLock;

mod copy;
mod rootfs_bundle;

use copy::{copy_dir_recursive, dir_size, install_rootfs_metadata};
pub use rootfs_bundle::{RestoredSnapshotRootfs, SnapshotRootfsFormat, SNAPSHOT_ROOTFS_SCHEMA};

/// Persistent store for VM snapshots.
pub struct SnapshotStore {
    /// Root directory for all snapshots
    base_dir: PathBuf,
}

impl SnapshotStore {
    pub(crate) fn acquire_exclusive_lock(&self) -> Result<FileLock> {
        FileLock::acquire(&self.base_dir.join(".snapshot-store")).map_err(|error| {
            BoxError::CacheError(format!(
                "Failed to lock snapshot directory {}: {error}",
                self.base_dir.display()
            ))
        })
    }

    /// Create a new snapshot store at the given directory.
    pub fn new(base_dir: &Path) -> Result<Self> {
        std::fs::create_dir_all(base_dir).map_err(|e| {
            BoxError::CacheError(format!(
                "Failed to create snapshot directory {}: {}",
                base_dir.display(),
                e
            ))
        })?;
        let _lock = FileLock::acquire(&base_dir.join(".snapshot-store")).map_err(|e| {
            BoxError::CacheError(format!(
                "Failed to lock snapshot directory {}: {}",
                base_dir.display(),
                e
            ))
        })?;
        // Sweep leftover `.staging-*` dirs from a prior crashed/aborted save
        // (mirrors ImageStore::new). `save` builds the whole snapshot in
        // `.staging-<id>-<pid>-<seq>` and only renames it into `<id>/` after
        // metadata.json; a SIGKILL/OOM/power-loss mid-copy leaves a full
        // rootfs-sized staging dir that list/count/prune never see (they key on
        // metadata.json) and that embeds the dead writer's pid+seq, so no later
        // process can match it — it would leak permanently without this sweep.
        if let Ok(entries) = std::fs::read_dir(base_dir) {
            for entry in entries.flatten() {
                let is_plain_directory = entry
                    .file_type()
                    .is_ok_and(|file_type| file_type.is_dir() && !file_type.is_symlink());
                if is_plain_directory
                    && entry.file_name().to_string_lossy().starts_with(".staging-")
                {
                    let _ = std::fs::remove_dir_all(entry.path());
                }
            }
        }
        Ok(Self {
            base_dir: base_dir.to_path_buf(),
        })
    }

    /// Open the default snapshot store at `~/.a3s/snapshots`.
    pub fn default_path() -> Result<Self> {
        let home = a3s_box_core::dirs_home();
        Self::new(&home.join("snapshots"))
    }

    /// Save a snapshot with the given metadata and rootfs source.
    ///
    /// Copies the rootfs directory into the snapshot bundle.
    /// Returns the updated metadata with `size_bytes` populated.
    pub fn save(
        &self,
        metadata: SnapshotMetadata,
        rootfs_source: &Path,
    ) -> Result<SnapshotMetadata> {
        self.save_inner(metadata, rootfs_source, None)
    }

    /// Save a managed Sandbox snapshot with an authoritative terminal rootfs
    /// metadata manifest captured while the source execution is quiesced.
    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
    pub(crate) fn save_managed(
        &self,
        metadata: SnapshotMetadata,
        rootfs_source: &Path,
        rootfs_metadata: &RootfsMetadataManifest,
    ) -> Result<SnapshotMetadata> {
        self.save_inner(metadata, rootfs_source, Some(rootfs_metadata))
    }

    fn save_inner(
        &self,
        mut metadata: SnapshotMetadata,
        rootfs_source: &Path,
        rootfs_metadata: Option<&RootfsMetadataManifest>,
    ) -> Result<SnapshotMetadata> {
        rootfs_bundle::validate_snapshot_id(&metadata.id)?;
        let _lock = FileLock::acquire(&self.base_dir.join(".snapshot-store")).map_err(|e| {
            BoxError::CacheError(format!(
                "Failed to lock snapshot directory {}: {}",
                self.base_dir.display(),
                e
            ))
        })?;
        let snap_dir = self.base_dir.join(&metadata.id);
        match std::fs::symlink_metadata(&snap_dir) {
            Ok(_) => {
                return Err(BoxError::CacheError(format!(
                    "Snapshot '{}' already exists",
                    metadata.id
                )))
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
            Err(error) => return Err(BoxError::IoError(error)),
        }

        // Build the whole snapshot in a staging dir, then atomically rename it to
        // `<id>/` only after metadata.json is written. A crash mid-save then
        // leaves at most a `<id>.staging-*` dir (GC-able), never a partial
        // `<id>/` that get/list ignore (they key on metadata.json) yet that
        // blocks re-create and never prunes.
        let staging_prefix = format!(".staging-{}-{}-", metadata.id, std::process::id());
        let staging = tempfile::Builder::new()
            .prefix(&staging_prefix)
            .tempdir_in(&self.base_dir)
            .map_err(|e| {
                BoxError::CacheError(format!(
                    "Failed to create snapshot staging directory in {}: {}",
                    self.base_dir.display(),
                    e
                ))
            })?;

        // Copy rootfs if source exists
        let rootfs_dest = staging.path().join("rootfs");
        if rootfs_source.exists() {
            copy_dir_recursive(rootfs_source, &rootfs_dest)?;
        } else {
            std::fs::create_dir_all(&rootfs_dest).map_err(|e| {
                BoxError::CacheError(format!("Failed to create snapshot rootfs directory: {}", e))
            })?;
        }
        if let Some(rootfs_metadata) = rootfs_metadata {
            install_rootfs_metadata(&rootfs_dest, rootfs_metadata)?;
        }

        // Calculate size
        metadata.size_bytes = dir_size(&rootfs_dest)?;

        // New snapshots publish an explicit rootfs representation contract.
        // Readers still accept a missing manifest only for the exact legacy
        // directory layout, so existing snapshots remain restorable without
        // letting a future payload be guessed from ambient files.
        rootfs_bundle::write_directory_manifest(staging.path(), &metadata)?;

        // Write metadata into the staging dir.
        let meta_path = staging.path().join("metadata.json");
        let json = serde_json::to_string_pretty(&metadata).map_err(|e| {
            BoxError::SerializationError(format!("Failed to serialize snapshot metadata: {}", e))
        })?;
        std::fs::write(&meta_path, &json).map_err(|e| {
            BoxError::CacheError(format!(
                "Failed to write snapshot metadata {}: {}",
                meta_path.display(),
                e
            ))
        })?;

        // Atomic publish: the snapshot becomes visible (with its metadata) in one
        // step, or not at all.
        std::fs::rename(staging.path(), &snap_dir).map_err(|e| {
            BoxError::CacheError(format!(
                "Failed to publish snapshot {}: {}",
                snap_dir.display(),
                e
            ))
        })?;

        Ok(metadata)
    }

    /// Load snapshot metadata by ID.
    pub fn get(&self, id: &str) -> Result<Option<SnapshotMetadata>> {
        Ok(
            rootfs_bundle::load_snapshot_metadata(&self.base_dir, id)?
                .map(|(_, metadata)| metadata),
        )
    }

    /// Get the rootfs path for a snapshot.
    pub fn rootfs_path(&self, id: &str) -> PathBuf {
        self.base_dir.join(id).join("rootfs")
    }

    /// List all snapshots, sorted by creation time (newest first).
    pub fn list(&self) -> Result<Vec<SnapshotMetadata>> {
        let mut snapshots = Vec::new();

        let entries = match std::fs::read_dir(&self.base_dir) {
            Ok(entries) => entries,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(snapshots),
            Err(e) => {
                return Err(BoxError::CacheError(format!(
                    "Failed to read snapshot directory: {}",
                    e
                )));
            }
        };

        for entry in entries {
            let entry = entry.map_err(|e| {
                BoxError::CacheError(format!("Failed to read snapshot entry: {}", e))
            })?;
            let entry_type = entry.file_type().map_err(BoxError::IoError)?;
            if !entry_type.is_dir() || entry_type.is_symlink() {
                continue;
            }
            let Some(id) = entry.file_name().to_str().map(str::to_owned) else {
                tracing::warn!(
                    path = %entry.path().display(),
                    "Skipping snapshot with a non-UTF-8 store identity"
                );
                continue;
            };
            match self.get(&id) {
                Ok(Some(metadata)) => snapshots.push(metadata),
                Ok(None) => {}
                // Don't silently skip: invalid metadata would make the
                // snapshot invisible to count/total_size/prune while its
                // payload keeps consuming disk. Surface it for diagnosis.
                Err(error) => tracing::warn!(
                    path = %entry.path().display(),
                    %error,
                    "Skipping snapshot with invalid persisted metadata"
                ),
            }
        }

        // Sort newest first
        snapshots.sort_by_key(|snapshot| Reverse(snapshot.created_at));
        Ok(snapshots)
    }

    /// Delete a snapshot by ID.
    pub fn delete(&self, id: &str) -> Result<bool> {
        rootfs_bundle::validate_snapshot_id(id)?;
        let _lock = self.acquire_exclusive_lock()?;
        let rootfs = normalize_snapshot_reference(self.rootfs_path(id));
        if self.referenced_rootfs_paths().contains(&rootfs) {
            return Err(BoxError::StateError(format!(
                "Snapshot '{id}' is still used as a copy-on-write rootfs lower"
            )));
        }
        self.delete_locked(id)
    }

    /// Delete a snapshot even when a restored directory-backed box still
    /// references it. This deliberately unsafe operation exists only for the
    /// CLI's explicit `snapshot rm --force` contract.
    pub fn delete_force(&self, id: &str) -> Result<bool> {
        rootfs_bundle::validate_snapshot_id(id)?;
        let _lock = self.acquire_exclusive_lock()?;
        self.delete_locked(id)
    }

    /// Delete while the caller holds [`Self::acquire_exclusive_lock`].
    ///
    /// This is crate-visible so managed execution reservation can validate and
    /// persist a Snapshot reference under the same lock used by deletion.
    pub(crate) fn delete_locked(&self, id: &str) -> Result<bool> {
        rootfs_bundle::validate_snapshot_id(id)?;
        let snap_dir = self.base_dir.join(id);
        let metadata = match std::fs::symlink_metadata(&snap_dir) {
            Ok(metadata) => metadata,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
            Err(error) => return Err(BoxError::IoError(error)),
        };
        if !metadata.is_dir() || metadata.file_type().is_symlink() {
            return Err(BoxError::CacheError(format!(
                "Snapshot bundle is not a plain directory: {}",
                snap_dir.display()
            )));
        }

        std::fs::remove_dir_all(&snap_dir).map_err(|e| {
            BoxError::CacheError(format!("Failed to delete snapshot {}: {}", id, e))
        })?;
        Ok(true)
    }

    /// Count the number of snapshots.
    pub fn count(&self) -> Result<usize> {
        Ok(self.list()?.len())
    }

    /// Calculate total size of all snapshots in bytes.
    pub fn total_size(&self) -> Result<u64> {
        Ok(self.list()?.iter().map(|s| s.size_bytes).sum())
    }

    /// Prune old snapshots to stay within limits.
    ///
    /// Removes oldest snapshots first until both `max_count` and `max_bytes`
    /// constraints are satisfied. A value of 0 means unlimited.
    pub fn prune(&self, max_count: usize, max_bytes: u64) -> Result<Vec<String>> {
        // A snapshot a restored box shares as its copy-on-write overlay lower must
        // never be evicted — deleting a live lower breaks the box (ESTALE) or stops
        // it from re-starting. Read which are in use from the boxes' `.snapshot-lower`
        // markers and skip them, evicting the oldest *evictable* snapshot instead.
        let protected = self.referenced_rootfs_paths();
        let mut snapshots = self.list()?; // newest-first
        let mut removed = Vec::new();

        loop {
            let over_count = max_count > 0 && snapshots.len() > max_count;
            let total: u64 = snapshots.iter().map(|s| s.size_bytes).sum();
            let over_size = max_bytes > 0 && total > max_bytes;
            if !over_count && !over_size {
                break;
            }

            // Oldest (last) snapshot that is not an in-use CoW lower.
            let idx = snapshots.iter().rposition(|s| {
                !protected.contains(&normalize_snapshot_reference(self.rootfs_path(&s.id)))
            });
            match idx {
                Some(i) => {
                    let snap = snapshots.remove(i);
                    self.delete(&snap.id)?;
                    removed.push(snap.id);
                }
                // Everything left is protected; can't prune further without
                // breaking a live box, so stop (caller stays over the cap).
                None => {
                    tracing::warn!(
                        in_use = snapshots.len(),
                        "snapshot prune kept {} in-use snapshot(s) (each referenced as a \
                         copy-on-write overlay lower); requested limit not fully met",
                        snapshots.len()
                    );
                    break;
                }
            }
        }

        Ok(removed)
    }

    /// Rootfs paths currently referenced by a restored box as its CoW overlay lower
    /// (`<box_dir>/.snapshot-lower`, written by `snapshot restore`). These snapshots
    /// must not be pruned. Boxes live next to snapshots under the a3s home
    /// (`base_dir` = `<home>/snapshots`).
    fn referenced_rootfs_paths(&self) -> std::collections::HashSet<PathBuf> {
        let mut set = std::collections::HashSet::new();
        let boxes = match self.base_dir.parent() {
            Some(home) => home.join("boxes"),
            None => return set,
        };
        if let Ok(entries) = std::fs::read_dir(&boxes) {
            for entry in entries.flatten() {
                if let Ok(content) = std::fs::read_to_string(entry.path().join(".snapshot-lower")) {
                    set.insert(normalize_snapshot_reference(PathBuf::from(content.trim())));
                }
            }
        }
        set
    }
}

fn normalize_snapshot_reference(path: PathBuf) -> PathBuf {
    path.canonicalize().unwrap_or(path)
}

impl SnapshotStoreBackend for SnapshotStore {
    fn save(&self, metadata: SnapshotMetadata, rootfs_source: &Path) -> Result<SnapshotMetadata> {
        self.save(metadata, rootfs_source)
    }

    fn get(&self, id: &str) -> Result<Option<SnapshotMetadata>> {
        self.get(id)
    }

    fn list(&self) -> Result<Vec<SnapshotMetadata>> {
        self.list()
    }

    fn delete(&self, id: &str) -> Result<bool> {
        self.delete(id)
    }

    fn count(&self) -> Result<usize> {
        self.count()
    }

    fn total_size(&self) -> Result<u64> {
        self.total_size()
    }

    fn prune(&self, max_count: usize, max_bytes: u64) -> Result<Vec<String>> {
        self.prune(max_count, max_bytes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn make_metadata(id: &str, name: &str) -> SnapshotMetadata {
        SnapshotMetadata::new(
            id.to_string(),
            name.to_string(),
            "box-source".to_string(),
            "alpine:latest".to_string(),
        )
    }

    fn make_rootfs(tmp: &TempDir) -> PathBuf {
        let rootfs = tmp.path().join("rootfs");
        std::fs::create_dir_all(&rootfs).unwrap();
        std::fs::write(rootfs.join("bin.sh"), "#!/bin/sh\necho hello").unwrap();
        std::fs::create_dir_all(rootfs.join("etc")).unwrap();
        std::fs::write(rootfs.join("etc/config"), "key=value").unwrap();
        rootfs
    }

    #[test]
    fn test_snapshot_store_new() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        assert!(store.base_dir.exists());
    }

    #[test]
    fn test_snapshot_save_and_get() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        let meta = make_metadata("snap-1", "first");
        let saved = store.save(meta, &rootfs).unwrap();

        assert_eq!(saved.id, "snap-1");
        assert!(saved.size_bytes > 0);

        let loaded = store.get("snap-1").unwrap().unwrap();
        assert_eq!(loaded.id, "snap-1");
        assert_eq!(loaded.name, "first");
        assert_eq!(loaded.image, "alpine:latest");
        assert_eq!(loaded.size_bytes, saved.size_bytes);
    }

    #[test]
    fn test_snapshot_get_nonexistent() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        assert!(store.get("nonexistent").unwrap().is_none());
    }

    #[test]
    fn test_snapshot_save_duplicate_fails() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        let meta = make_metadata("snap-dup", "dup");
        store.save(meta.clone(), &rootfs).unwrap();

        let result = store.save(meta, &rootfs);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("already exists"));
    }

    #[test]
    fn test_snapshot_rootfs_copied() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        let meta = make_metadata("snap-fs", "fs-test");
        store.save(meta, &rootfs).unwrap();

        let snap_rootfs = store.rootfs_path("snap-fs");
        assert!(snap_rootfs.join("bin.sh").exists());
        assert!(snap_rootfs.join("etc/config").exists());
        assert_eq!(
            std::fs::read_to_string(snap_rootfs.join("etc/config")).unwrap(),
            "key=value"
        );
    }

    #[test]
    fn test_snapshot_list_empty() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let list = store.list().unwrap();
        assert!(list.is_empty());
    }

    #[test]
    fn test_snapshot_list_multiple() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        for i in 0..3 {
            let meta = make_metadata(&format!("snap-{}", i), &format!("snap-{}", i));
            store.save(meta, &rootfs).unwrap();
            std::thread::sleep(std::time::Duration::from_millis(10));
        }

        let list = store.list().unwrap();
        assert_eq!(list.len(), 3);
        // Newest first
        assert_eq!(list[0].id, "snap-2");
        assert_eq!(list[2].id, "snap-0");
    }

    #[test]
    fn test_snapshot_delete() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        let meta = make_metadata("snap-del", "delete-me");
        store.save(meta, &rootfs).unwrap();

        assert!(store.delete("snap-del").unwrap());
        assert!(store.get("snap-del").unwrap().is_none());
    }

    #[test]
    fn test_snapshot_delete_nonexistent() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        assert!(!store.delete("nope").unwrap());
    }

    #[test]
    fn snapshot_ids_cannot_escape_the_store() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        for id in ["", ".", "..", "../outside", "nested/snapshot"] {
            let save_error = store
                .save(make_metadata(id, "invalid"), &rootfs)
                .unwrap_err()
                .to_string();
            assert!(save_error.contains("invalid snapshot id"), "{save_error}");
            let get_error = store.get(id).unwrap_err().to_string();
            assert!(get_error.contains("invalid snapshot id"), "{get_error}");
            let delete_error = store.delete_force(id).unwrap_err().to_string();
            assert!(
                delete_error.contains("invalid snapshot id"),
                "{delete_error}"
            );
        }
    }

    #[cfg(unix)]
    #[test]
    fn snapshot_delete_never_follows_a_bundle_symlink() {
        let tmp = TempDir::new().unwrap();
        let snapshots = tmp.path().join("snapshots");
        let store = SnapshotStore::new(&snapshots).unwrap();
        let outside = tmp.path().join("outside");
        std::fs::create_dir(&outside).unwrap();
        std::fs::write(outside.join("keep"), b"external").unwrap();
        std::os::unix::fs::symlink(&outside, snapshots.join("forged")).unwrap();

        let error = store.delete_force("forged").unwrap_err().to_string();

        assert!(error.contains("not a plain directory"), "{error}");
        assert_eq!(std::fs::read(outside.join("keep")).unwrap(), b"external");
        assert!(snapshots.join("forged").symlink_metadata().is_ok());
    }

    #[test]
    fn delete_refuses_a_referenced_directory_snapshot_without_force() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);
        store
            .save(make_metadata("snap-used", "used"), &rootfs)
            .unwrap();
        let box_dir = tmp.path().join("boxes/restored");
        std::fs::create_dir_all(&box_dir).unwrap();
        std::fs::write(
            box_dir.join(".snapshot-lower"),
            store
                .rootfs_path("snap-used")
                .canonicalize()
                .unwrap()
                .to_string_lossy()
                .as_bytes(),
        )
        .unwrap();

        let error = store.delete("snap-used").unwrap_err().to_string();
        assert!(error.contains("still used"), "{error}");
        assert!(store.get("snap-used").unwrap().is_some());
        assert!(store.delete_force("snap-used").unwrap());
    }

    #[test]
    fn test_snapshot_count() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        assert_eq!(store.count().unwrap(), 0);

        store.save(make_metadata("s1", "s1"), &rootfs).unwrap();
        store.save(make_metadata("s2", "s2"), &rootfs).unwrap();
        assert_eq!(store.count().unwrap(), 2);
    }

    #[test]
    fn test_snapshot_total_size() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        store.save(make_metadata("s1", "s1"), &rootfs).unwrap();
        let total = store.total_size().unwrap();
        assert!(total > 0);
    }

    #[test]
    fn test_snapshot_prune_by_count() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        for i in 0..5 {
            let meta = make_metadata(&format!("s{}", i), &format!("s{}", i));
            store.save(meta, &rootfs).unwrap();
            std::thread::sleep(std::time::Duration::from_millis(10));
        }

        let removed = store.prune(3, 0).unwrap();
        assert_eq!(removed.len(), 2);
        assert_eq!(store.count().unwrap(), 3);

        // Oldest should be removed
        assert!(store.get("s0").unwrap().is_none());
        assert!(store.get("s1").unwrap().is_none());
        // Newest should remain
        assert!(store.get("s4").unwrap().is_some());
        assert!(store.get("s3").unwrap().is_some());
        assert!(store.get("s2").unwrap().is_some());
    }

    #[test]
    fn test_snapshot_prune_no_limits() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        for i in 0..3 {
            store
                .save(
                    make_metadata(&format!("s{}", i), &format!("s{}", i)),
                    &rootfs,
                )
                .unwrap();
        }

        let removed = store.prune(0, 0).unwrap();
        assert!(removed.is_empty());
        assert_eq!(store.count().unwrap(), 3);
    }

    #[test]
    fn prune_skips_in_use_snapshots() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);

        for i in 0..5 {
            store
                .save(
                    make_metadata(&format!("s{}", i), &format!("s{}", i)),
                    &rootfs,
                )
                .unwrap();
            std::thread::sleep(std::time::Duration::from_millis(10));
        }

        // Mark the OLDEST (s0) as a box's in-use CoW overlay lower.
        let box_dir = tmp.path().join("boxes").join("box1");
        std::fs::create_dir_all(&box_dir).unwrap();
        std::fs::write(
            box_dir.join(".snapshot-lower"),
            store.rootfs_path("s0").to_string_lossy().as_bytes(),
        )
        .unwrap();

        // keep=3 would normally evict the oldest two (s0, s1); s0 is protected, so
        // it survives and s1 + s2 are evicted instead.
        let removed = store.prune(3, 0).unwrap();
        assert_eq!(removed.len(), 2);
        assert!(store.get("s0").unwrap().is_some(), "in-use s0 must be kept");
        assert!(store.get("s1").unwrap().is_none());
        assert!(store.get("s2").unwrap().is_none());
        assert!(store.get("s3").unwrap().is_some());
        assert!(store.get("s4").unwrap().is_some());
    }

    #[test]
    fn prune_keeps_everything_when_all_in_use() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);
        store.save(make_metadata("a", "a"), &rootfs).unwrap();
        store.save(make_metadata("b", "b"), &rootfs).unwrap();

        // Both snapshots are in use as a box's CoW overlay lower.
        for (i, id) in ["a", "b"].iter().enumerate() {
            let bd = tmp.path().join("boxes").join(format!("bx{i}"));
            std::fs::create_dir_all(&bd).unwrap();
            std::fs::write(
                bd.join(".snapshot-lower"),
                store.rootfs_path(id).to_string_lossy().as_bytes(),
            )
            .unwrap();
        }

        // Even asked to keep only 1, prune evicts nothing — both are protected.
        let removed = store.prune(1, 0).unwrap();
        assert!(removed.is_empty(), "in-use snapshots must never be pruned");
        assert_eq!(store.count().unwrap(), 2);
    }

    #[test]
    fn test_snapshot_save_with_empty_rootfs() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let empty_rootfs = tmp.path().join("nonexistent_rootfs");

        let meta = make_metadata("snap-empty", "empty");
        let saved = store.save(meta, &empty_rootfs).unwrap();
        assert_eq!(saved.id, "snap-empty");

        // Rootfs dir should still be created (empty)
        assert!(store.rootfs_path("snap-empty").exists());
    }

    #[test]
    fn test_dir_size() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path().join("sized");
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("a.txt"), "hello").unwrap();
        std::fs::write(dir.join("b.txt"), "world!").unwrap();

        let size = dir_size(&dir).unwrap();
        assert_eq!(size, 11); // 5 + 6
    }

    #[test]
    fn test_dir_size_nested() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path().join("nested");
        std::fs::create_dir_all(dir.join("sub")).unwrap();
        std::fs::write(dir.join("a.txt"), "abc").unwrap();
        std::fs::write(dir.join("sub/b.txt"), "defgh").unwrap();

        let size = dir_size(&dir).unwrap();
        assert_eq!(size, 8); // 3 + 5
    }

    #[test]
    fn test_dir_size_empty() {
        let tmp = TempDir::new().unwrap();
        let dir = tmp.path().join("empty");
        std::fs::create_dir_all(&dir).unwrap();
        assert_eq!(dir_size(&dir).unwrap(), 0);
    }

    #[test]
    fn test_copy_dir_recursive() {
        let tmp = TempDir::new().unwrap();
        let src = tmp.path().join("src");
        let dst = tmp.path().join("dst");

        std::fs::create_dir_all(src.join("sub")).unwrap();
        std::fs::write(src.join("a.txt"), "hello").unwrap();
        std::fs::write(src.join("sub/b.txt"), "world").unwrap();

        copy_dir_recursive(&src, &dst).unwrap();

        assert!(dst.join("a.txt").exists());
        assert!(dst.join("sub/b.txt").exists());
        assert_eq!(std::fs::read_to_string(dst.join("a.txt")).unwrap(), "hello");
        assert_eq!(
            std::fs::read_to_string(dst.join("sub/b.txt")).unwrap(),
            "world"
        );
    }

    #[cfg(unix)]
    #[test]
    fn snapshot_size_does_not_follow_absolute_symlinks() {
        use std::os::unix::ffi::OsStrExt;

        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = tmp.path().join("rootfs");
        std::fs::create_dir(&rootfs).unwrap();
        let outside = tmp.path().join("outside");
        std::fs::write(&outside, vec![0_u8; 64 * 1024]).unwrap();
        std::os::unix::fs::symlink(&outside, rootfs.join("outside-link")).unwrap();

        let saved = store
            .save(make_metadata("symlink-size", "symlink-size"), &rootfs)
            .unwrap();

        assert_eq!(
            saved.size_bytes,
            outside.as_os_str().as_bytes().len() as u64
        );
    }

    #[cfg(unix)]
    #[test]
    fn snapshot_preserves_hardlinks_and_xattrs() {
        use std::os::unix::fs::MetadataExt;

        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = tmp.path().join("rootfs");
        std::fs::create_dir(&rootfs).unwrap();
        let first = rootfs.join("first");
        let second = rootfs.join("second");
        std::fs::write(&first, b"shared inode").unwrap();
        std::fs::hard_link(&first, &second).unwrap();
        xattr::set(&first, "user.a3s.snapshot", b"preserved").unwrap();

        let saved = store
            .save(make_metadata("hardlink-xattr", "hardlink-xattr"), &rootfs)
            .unwrap();
        let captured = store.rootfs_path(&saved.id);
        let captured_first = std::fs::metadata(captured.join("first")).unwrap();
        let captured_second = std::fs::metadata(captured.join("second")).unwrap();

        assert_eq!(captured_first.dev(), captured_second.dev());
        assert_eq!(captured_first.ino(), captured_second.ino());
        assert_eq!(saved.size_bytes, b"shared inode".len() as u64);
        assert_eq!(
            xattr::get(captured.join("first"), "user.a3s.snapshot").unwrap(),
            Some(b"preserved".to_vec())
        );
    }

    #[cfg(unix)]
    #[test]
    fn snapshot_rejects_fifo_without_leaking_staging() {
        use std::os::unix::ffi::OsStrExt;

        let tmp = TempDir::new().unwrap();
        let snapshots = tmp.path().join("snapshots");
        let store = SnapshotStore::new(&snapshots).unwrap();
        let rootfs = tmp.path().join("rootfs");
        std::fs::create_dir(&rootfs).unwrap();
        let fifo = rootfs.join("blocking-fifo");
        let fifo_path = std::ffi::CString::new(fifo.as_os_str().as_bytes()).unwrap();
        assert_eq!(unsafe { libc::mkfifo(fifo_path.as_ptr(), 0o600) }, 0);

        let error = store
            .save(make_metadata("fifo", "fifo"), &rootfs)
            .unwrap_err();

        assert!(error.to_string().contains("unsupported special file"));
        assert!(store.get("fifo").unwrap().is_none());
        assert!(
            std::fs::read_dir(snapshots)
                .unwrap()
                .flatten()
                .all(|entry| !entry.file_name().to_string_lossy().starts_with(".staging-")),
            "failed Snapshot must remove its staging tree"
        );
    }

    #[test]
    fn managed_snapshot_installs_only_terminal_rootfs_metadata() {
        let tmp = TempDir::new().unwrap();
        let store = SnapshotStore::new(&tmp.path().join("snapshots")).unwrap();
        let rootfs = make_rootfs(&tmp);
        std::fs::write(
            rootfs.join(IMAGE_ROOTFS_METADATA_PATH.trim_start_matches('/')),
            b"stale image metadata",
        )
        .unwrap();
        std::fs::write(
            rootfs.join(PREVIOUS_ROOTFS_METADATA_PATH.trim_start_matches('/')),
            b"stale replay metadata",
        )
        .unwrap();
        let manifest = RootfsMetadataManifest::new(Vec::new());

        store
            .save_managed(
                make_metadata("managed-metadata", "managed-metadata"),
                &rootfs,
                &manifest,
            )
            .unwrap();

        let captured = store.rootfs_path("managed-metadata");
        assert!(!captured
            .join(IMAGE_ROOTFS_METADATA_PATH.trim_start_matches('/'))
            .exists());
        assert!(!captured
            .join(PREVIOUS_ROOTFS_METADATA_PATH.trim_start_matches('/'))
            .exists());
        let stored: RootfsMetadataManifest = serde_json::from_slice(
            &std::fs::read(captured.join(ROOTFS_METADATA_PATH.trim_start_matches('/'))).unwrap(),
        )
        .unwrap();
        assert_eq!(stored, manifest);
    }

    #[test]
    fn new_sweeps_leftover_staging_dirs() {
        let tmp = TempDir::new().unwrap();
        let base = tmp.path().join("snapshots");
        std::fs::create_dir_all(&base).unwrap();

        // A full rootfs-sized staging dir leaked by a crashed save...
        let leaked = base.join(".staging-snap1-12345-0");
        std::fs::create_dir_all(leaked.join("rootfs")).unwrap();
        std::fs::write(leaked.join("rootfs").join("big"), vec![0u8; 4096]).unwrap();
        // ...alongside a real snapshot (keyed by metadata.json).
        let real = base.join("snap1");
        std::fs::create_dir_all(&real).unwrap();
        std::fs::write(real.join("metadata.json"), "{}").unwrap();

        let _store = SnapshotStore::new(&base).unwrap();

        assert!(
            !leaked.exists(),
            "leaked .staging-* dir must be swept on open"
        );
        assert!(real.exists(), "a real snapshot dir must be preserved");
    }
}