From b6ddd8552582dc87f91be9acfd82fb910bdfcd9b Mon Sep 17 00:00:00 2001 From: Valentin Haudiquet Date: Mon, 21 Sep 2026 00:06:20 +0200 Subject: [PATCH] feat(rootfs): cache-aware rootfs preparation with persist hook Add ecr::rootfs with the full rootfs lifecycle behind the library: - RootfsCache::prepare resolves an image reference, downloads through the tarball cache (with the OCI :latest digest freshness check) and extracts into a scratch directory tracked by PreparedRootfs. - PreparedRootfs::persist packs the current rootfs back into its cache entry (compressed to match the entry's extension, symlinks and permissions preserved) and marks it with a .provisioned sidecar. - RootfsCache::prepare_provisioned composes both into the hot-cell flow: provision once, and every later call sharing the cache skips the download and the provisioning step. extract: an "oci-" cache entry without a layers.manifest is a persisted provisioned rootfs; extract it as a plain archive. The CLI now drives prepare and drops its inline cache/orchestration code and the dirs/tempfile dependencies. The binfmt check moves ahead of the download so foreign-arch runs fail before pulling an image. --- Cargo.lock | 2 - crates/ecr-cli/Cargo.toml | 5 - crates/ecr-cli/src/main.rs | 162 +--------- crates/ecr/src/extract.rs | 21 +- crates/ecr/src/lib.rs | 9 +- crates/ecr/src/rootfs.rs | 612 +++++++++++++++++++++++++++++++++++++ 6 files changed, 653 insertions(+), 158 deletions(-) create mode 100644 crates/ecr/src/rootfs.rs diff --git a/Cargo.lock b/Cargo.lock index 5d914f2..2e0ce8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -328,9 +328,7 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", - "dirs", "ecr", - "tempfile", ] [[package]] diff --git a/crates/ecr-cli/Cargo.toml b/crates/ecr-cli/Cargo.toml index b7279cb..4295632 100644 --- a/crates/ecr-cli/Cargo.toml +++ b/crates/ecr-cli/Cargo.toml @@ -21,8 +21,3 @@ clap = { version = "4", features = ["derive", "env"] } # Error handling anyhow = "1" - -# Interim: used by the CLI's inline cache/extraction orchestration until it -# moves behind the library's rootfs API -dirs = "6" -tempfile = "3" diff --git a/crates/ecr-cli/src/main.rs b/crates/ecr-cli/src/main.rs index 30df7a1..0cb49c7 100644 --- a/crates/ecr-cli/src/main.rs +++ b/crates/ecr-cli/src/main.rs @@ -8,14 +8,9 @@ use clap::Parser; use cli::Args; use ecr::chroot; use ecr::config::Config; -use ecr::distro::{ - map_arch, parse_image_ref, resolve_distro_url, resolve_distro_version, Distro, ImageSource, -}; -use ecr::download::{digest_sidecar, download_image, fetch_oci_digest}; use ecr::exec::ExecOptions; -use ecr::extract::extract_tarball; use ecr::mount::BindTarget; -use ecr::{kernel, qemu, qemu_vm, utils, veprintln, verbose}; +use ecr::{kernel, qemu, qemu_vm, utils, veprintln, verbose, PrepareRequest, RootfsCache}; fn main() -> Result<()> { let args = Args::parse(); @@ -27,91 +22,23 @@ fn main() -> Result<()> { let config = Config::load()?; // Get architecture - let host_arch = get_host_arch(); + let host_arch = utils::get_host_arch().debian_name().to_string(); let arch = args.arch.clone().unwrap_or_else(|| host_arch.clone()); - // Parse image reference - let image_source = parse_image_ref(&args.distro, &arch)?; - - // For DirectTarball, resolve floating aliases ("latest", "lts") to a concrete - // version string *before* computing the cache key. This ensures we cache as - // e.g. "ubuntu-noble-amd64" rather than "ubuntu-latest-amd64", so a future - // release automatically gets its own cache entry. - let image_source = match image_source { - ImageSource::DirectTarball { distro, version } => { - let resolved = resolve_distro_version(&distro, version.as_deref(), &arch)?; - ImageSource::DirectTarball { - distro, - version: Some(resolved), - } - } - other => other, - }; - - // Determine cache directory and filename - let cache_dir = dirs::cache_dir() - .expect("Could not determine cache directory") - .join("ecr"); - let cache_filename = generate_cache_filename(&image_source, &arch); - let cache_path = cache_dir.join(&cache_filename); - - // OCI images with a floating tag (":latest") need a freshness check: - // fetch the current manifest digest from the registry and compare it - // against the digest stored from the last download. Only re-pull when - // the digest has actually changed. On a network error we fall back to - // the cached image with a warning rather than hard-failing. - let oci_digest_changed = if cache_path.exists() { - if let ImageSource::OciImage { - registry, - repository, - tag, - .. - } = &image_source - { - if tag == "latest" { - match fetch_oci_digest(registry, repository, tag) { - Ok(current) => { - let stored = std::fs::read_to_string(digest_sidecar(&cache_path)).ok(); - stored.as_deref() != Some(current.trim()) - } - Err(e) => { - eprintln!( - "Warning: could not check image freshness ({}); using cache", - e - ); - false - } - } - } else { - false // pinned tags are assumed immutable - } - } else { - false - } - } else { - false // cache absent — download triggered by !cache_path.exists() below - }; - - // Download if not cached, --no-cache, or the remote digest has moved - if args.no_cache || !cache_path.exists() || oci_digest_changed { - std::fs::create_dir_all(&cache_dir)?; - download_image(&image_source, &cache_path, &arch)?; - } else { - veprintln!("Using cached tarball: {}", cache_path.display()); - } - - // Check QEMU if foreign architecture (for namespace mode) - // For VM mode, we don't need binfmt_misc since we're using system emulation + // Check QEMU binfmt early for foreign architectures (namespace mode only; + // VM mode uses system emulation). Failing here avoids the download. if args.kernel.is_none() && arch != host_arch { qemu::check_binfmt(&arch)?; } - // Create temp directory for extraction - let temp_dir = tempfile::tempdir()?; - let rootfs = temp_dir.path().to_path_buf(); - - veprintln!("Extracting to: {}", rootfs.display()); - extract_tarball(&cache_path, &rootfs)?; + // Resolve, download (through the cache) and extract the image + let cache = RootfsCache::new(RootfsCache::default_dir()?); + let request = PrepareRequest { + image: args.distro.clone(), + arch: args.arch.clone().unwrap_or_default(), + no_cache: args.no_cache, + }; + let rootfs = cache.prepare(&request)?; // Branch based on --kernel flag // Option> (require_equals: the value must use =PATH syntax @@ -136,7 +63,7 @@ fn main() -> Result<()> { } None => { veprintln!("QEMU mode: downloading default kernel..."); - kernel::get_default_kernel(&cache_dir, &arch)? + kernel::get_default_kernel(cache.dir(), &arch)? } }; @@ -148,7 +75,7 @@ fn main() -> Result<()> { let result = qemu_vm::launch_qemu(qemu_vm::QemuConfig { kernel_path, - rootfs_path: rootfs, + rootfs_path: rootfs.dir().to_path_buf(), memory: args.memory.clone(), arch: arch.clone(), command, @@ -162,10 +89,10 @@ fn main() -> Result<()> { result } else { // Namespace/chroot mode - let exit_code = namespace_mode(&args, &rootfs, &config, &arch)?; + let exit_code = namespace_mode(&args, rootfs.dir(), &config, &arch)?; // Propagate the command's exit code, cleaning up the extracted rootfs // first: process::exit does not run destructors. - drop(temp_dir); + drop(rootfs); if exit_code != 0 { std::process::exit(exit_code); } @@ -249,60 +176,3 @@ fn cli_mount_target(source: &Path, parent: &str) -> Result { .ok_or_else(|| anyhow::anyhow!("Invalid bind path: {}", source.display()))?; Ok(Path::new("/").join(parent).join(basename)) } - -/// Generate a cache filename based on the image source -fn generate_cache_filename(source: &ImageSource, arch: &str) -> String { - match source { - ImageSource::DirectTarball { distro, version } => { - let distro_name = match distro { - Distro::Ubuntu => "ubuntu", - Distro::Alpine => "alpine", - }; - let distro_arch = map_arch(*distro, arch); - // Get extension from URL - let url = resolve_distro_url(distro, version.as_deref(), arch).unwrap_or_default(); - let ext = get_tarball_extension(&url); - format!( - "{}-{}-{}.{}", - distro_name, - version.as_deref().unwrap_or("latest"), - distro_arch, - ext - ) - } - ImageSource::OciImage { - registry, - repository, - tag, - architecture, - } => { - // Sanitize for filename - let safe_registry = registry.replace(['.', ':'], "_"); - let safe_repo = repository.replace(['/', ':'], "_"); - format!( - "oci-{}-{}-{}-{}.tar.gz", - safe_registry, safe_repo, tag, architecture - ) - } - } -} - -fn get_tarball_extension(url: &str) -> &str { - // Extract extension from URL (e.g., .tar.gz, .tar.xz, .tar.zst) - if url.ends_with(".tar.zst") { - "tar.zst" - } else if url.ends_with(".tar.xz") { - "tar.xz" - } else if url.ends_with(".tar.gz") { - "tar.gz" - } else if url.ends_with(".tar.bz2") { - "tar.bz2" - } else { - "tar.gz" // default - } -} - -fn get_host_arch() -> String { - // Use the consolidated architecture detection from utils - utils::get_host_arch().debian_name().to_string() -} diff --git a/crates/ecr/src/extract.rs b/crates/ecr/src/extract.rs index ffeef67..612dc21 100644 --- a/crates/ecr/src/extract.rs +++ b/crates/ecr/src/extract.rs @@ -18,6 +18,14 @@ pub fn extract_tarball(tarball: &Path, dest: &Path) -> Result<()> { return extract_multi_layer_oci(tarball, dest); } + extract_single_archive(tarball, dest) +} + +/// Extract a plain rootfs tarball, detecting the compression format from +/// the filename with a magic-byte fallback. +fn extract_single_archive(tarball: &Path, dest: &Path) -> Result<()> { + let filename = tarball.file_name().and_then(|n| n.to_str()).unwrap_or(""); + let file = File::open(tarball) .with_context(|| format!("Failed to open tarball: {}", tarball.display()))?; @@ -100,9 +108,18 @@ fn extract_multi_layer_oci(tarball: &Path, dest: &Path) -> Result<()> { .context("Failed to unpack OCI bundle")?; } + let manifest_path = temp_dir.path().join("layers.manifest"); + if !manifest_path.exists() { + // Not our layer-bundle layout: a provisioned rootfs persisted over + // an OCI cache entry (the dispatch above keys on the "oci-" filename + // prefix). The outer archive is a plain rootfs tarball. + veprintln!("No layers.manifest in bundle; extracting as plain rootfs tarball"); + return extract_single_archive(tarball, dest); + } + // Read the layers manifest - let manifest = std::fs::read_to_string(temp_dir.path().join("layers.manifest")) - .context("Failed to read layers.manifest")?; + let manifest = + std::fs::read_to_string(&manifest_path).context("Failed to read layers.manifest")?; // Apply each layer in order with full whiteout handling. for layer_name in manifest.lines() { diff --git a/crates/ecr/src/lib.rs b/crates/ecr/src/lib.rs index 9a9e1ab..0292d66 100644 --- a/crates/ecr/src/lib.rs +++ b/crates/ecr/src/lib.rs @@ -6,9 +6,10 @@ //! unprivileged user/PID/mount/UTS namespaces — or boots it in a QEMU VM //! (see the `qemu_vm` module). //! -//! The main entry point is [`exec`]: it runs a command inside a prepared -//! rootfs with a caller-composed environment, explicit bind targets and a -//! target architecture. +//! The main entry points are [`rootfs`] (cache-aware rootfs preparation and +//! persistence) and [`exec`] (running a command inside a prepared rootfs +//! with a caller-composed environment, explicit bind targets and a target +//! architecture). pub mod chroot; pub mod config; @@ -21,6 +22,7 @@ pub mod mount; pub mod namespace; pub mod qemu; pub mod qemu_vm; +pub mod rootfs; pub mod utils; pub mod verbose; @@ -36,3 +38,4 @@ macro_rules! veprintln { pub use exec::{exec, ExecOptions}; pub use mount::BindTarget; +pub use rootfs::{PrepareRequest, PreparedRootfs, RootfsCache}; diff --git a/crates/ecr/src/rootfs.rs b/crates/ecr/src/rootfs.rs new file mode 100644 index 0000000..bc2fdb4 --- /dev/null +++ b/crates/ecr/src/rootfs.rs @@ -0,0 +1,612 @@ +//! Cache-aware rootfs lifecycle: resolve an image reference, download +//! through the tarball cache, extract into a scratch directory, and +//! optionally persist a provisioned rootfs back into the cache. +//! +//! The hot-cell flow is [`RootfsCache::prepare_provisioned`]: provision +//! once (package installs, user setup, …), and every subsequent call — +//! in this or another process sharing the cache — hits the persisted +//! entry and skips both the download and the provisioning step: +//! +//! ```no_run +//! use ecr::{chroot, ExecOptions, PrepareRequest, RootfsCache}; +//! +//! # fn main() -> anyhow::Result<()> { +//! let cache = RootfsCache::new(RootfsCache::default_dir()?); +//! let req = PrepareRequest::new("alpine:3.23"); +//! +//! let rootfs = cache.prepare_provisioned(&req, |rootfs| { +//! let code = rootfs.exec(&ExecOptions { +//! command: vec!["apk".into(), "add".into(), "curl".into()], +//! env: chroot::default_env(rootfs.dir()), +//! ..ExecOptions::default() +//! })?; +//! if code != 0 { +//! anyhow::bail!("provisioning command failed with exit code {}", code); +//! } +//! Ok(()) +//! })?; +//! +//! // Amortized run: the provisioned entry is hit directly +//! let code = rootfs.exec(&ExecOptions { +//! command: vec!["curl".into(), "--version".into()], +//! env: chroot::default_env(rootfs.dir()), +//! ..ExecOptions::default() +//! })?; +//! # let _ = code; Ok(()) +//! # } +//! ``` + +use std::path::{Path, PathBuf}; + +use anyhow::{anyhow, Context, Result}; +use tempfile::TempDir; + +use crate::distro::{ + map_arch, parse_image_ref, resolve_distro_url, resolve_distro_version, Distro, ImageSource, +}; +use crate::download::{digest_sidecar, download_image, fetch_oci_digest}; +use crate::exec::ExecOptions; +use crate::extract::extract_tarball; +use crate::veprintln; + +/// Sidecar suffix marking a cache entry that holds a provisioned rootfs. +const PROVISIONED_SUFFIX: &str = "provisioned"; + +/// Sidecar path for a cache entry: `foo.tar.gz` → `foo.tar.gz.`. +fn sidecar(path: &Path, suffix: &str) -> PathBuf { + let mut s = path.as_os_str().to_owned(); + s.push("."); + s.push(suffix); + PathBuf::from(s) +} + +/// The image tarball cache (default location: `~/.cache/ecr`). +#[derive(Debug, Clone)] +pub struct RootfsCache { + dir: PathBuf, +} + +impl RootfsCache { + pub fn new(dir: impl Into) -> Self { + RootfsCache { dir: dir.into() } + } + + /// Default cache directory: `~/.cache/ecr` (or the platform equivalent). + pub fn default_dir() -> Result { + dirs::cache_dir() + .map(|p| p.join("ecr")) + .ok_or_else(|| anyhow!("Could not determine cache directory")) + } + + pub fn dir(&self) -> &Path { + &self.dir + } + + /// Cache entry (tarball path) for an image source and architecture. + pub fn tarball_path(&self, source: &ImageSource, arch: &str) -> PathBuf { + self.dir.join(generate_cache_filename(source, arch)) + } + + /// Resolve `req`, download through the cache when needed, and extract + /// into a fresh scratch directory (removed when the returned + /// [`PreparedRootfs`] is dropped). + pub fn prepare(&self, req: &PrepareRequest) -> Result { + let arch = req.resolved_arch(); + let source = resolve_source(&req.image, &arch)?; + let cache_path = self.tarball_path(&source, &arch); + + // OCI images with a floating tag (":latest") need a freshness check: + // fetch the current manifest digest from the registry and compare it + // against the digest stored from the last download. Only re-pull when + // the digest has actually changed. On a network error we fall back to + // the cached image with a warning rather than hard-failing. + let oci_digest_changed = if cache_path.exists() { + if let ImageSource::OciImage { + registry, + repository, + tag, + .. + } = &source + { + if tag == "latest" { + match fetch_oci_digest(registry, repository, tag) { + Ok(current) => { + let stored = std::fs::read_to_string(digest_sidecar(&cache_path)).ok(); + stored.as_deref() != Some(current.trim()) + } + Err(e) => { + eprintln!( + "Warning: could not check image freshness ({}); using cache", + e + ); + false + } + } + } else { + false // pinned tags are assumed immutable + } + } else { + false + } + } else { + false // cache absent — download triggered by !cache_path.exists() below + }; + + // Download when not cached, explicitly requested, or the remote digest + // has moved + if req.no_cache || !cache_path.exists() || oci_digest_changed { + std::fs::create_dir_all(&self.dir).with_context(|| { + format!("Failed to create cache directory {}", self.dir.display()) + })?; + download_image(&source, &cache_path, &arch)?; + } else { + veprintln!("Using cached tarball: {}", cache_path.display()); + } + + extract_prepared(cache_path, arch) + } + + /// Prepare a provisioned rootfs with cache amortization: when the cache + /// entry is already marked provisioned it is extracted as-is; otherwise + /// the rootfs is prepared from the pristine image, `provision` runs + /// once, and the result is persisted into the cache so subsequent calls + /// (including from other processes sharing this cache) skip both the + /// download and the provisioning step. + /// + /// A provisioned entry no longer tracks the upstream image: for + /// `:latest` OCI images the freshness check is skipped on a marker hit. + /// Re-provision by removing the entry and its `.provisioned` sidecar, or + /// by passing `no_cache: true` (which re-downloads, re-provisions and + /// re-persists). + pub fn prepare_provisioned( + &self, + req: &PrepareRequest, + provision: F, + ) -> Result + where + F: FnOnce(&PreparedRootfs) -> Result<()>, + { + let arch = req.resolved_arch(); + let source = resolve_source(&req.image, &arch)?; + let cache_path = self.tarball_path(&source, &arch); + let marker = sidecar(&cache_path, PROVISIONED_SUFFIX); + + if !req.no_cache && cache_path.exists() && marker.exists() { + veprintln!("Using provisioned cache entry: {}", cache_path.display()); + return extract_prepared(cache_path, arch); + } + + let rootfs = self.prepare(req)?; + provision(&rootfs)?; + rootfs.persist().with_context(|| { + format!( + "Failed to persist provisioned rootfs to {}", + cache_path.display() + ) + })?; + Ok(rootfs) + } +} + +/// Request for [`RootfsCache::prepare`] / [`RootfsCache::prepare_provisioned`]. +#[derive(Debug, Clone)] +pub struct PrepareRequest { + /// Image reference: a distro name ("ubuntu", "ubuntu:noble", + /// "alpine:3.23") or an OCI reference ("debian", + /// "ghcr.io/org/image:tag", "docker://…"). + pub image: String, + + /// Target architecture (e.g. "amd64", "arm64"). Empty selects the host + /// architecture. + pub arch: String, + + /// Ignore the cache and download a fresh image. + pub no_cache: bool, +} + +impl PrepareRequest { + pub fn new(image: impl Into) -> Self { + PrepareRequest { + image: image.into(), + arch: String::new(), + no_cache: false, + } + } + + /// The effective architecture: the request value, or the host arch. + pub fn resolved_arch(&self) -> String { + if self.arch.is_empty() { + crate::utils::get_host_arch().debian_name().to_string() + } else { + self.arch.clone() + } + } +} + +/// A rootfs extracted into a scratch directory, ready for [`ExecOptions`] / +/// [`crate::exec`] or QEMU. The directory is removed when this value is +/// dropped, unless [`PreparedRootfs::persist`] ran first. +#[derive(Debug)] +pub struct PreparedRootfs { + rootfs: TempDir, + cache_path: PathBuf, + arch: String, +} + +impl PreparedRootfs { + /// The extracted rootfs directory. + pub fn dir(&self) -> &Path { + self.rootfs.path() + } + + /// Architecture this rootfs was prepared for. + pub fn arch(&self) -> &str { + &self.arch + } + + /// The cache entry this rootfs was prepared from. + pub fn cache_entry(&self) -> &Path { + &self.cache_path + } + + /// Run a command inside the rootfs (see [`crate::exec`]). When + /// `opts.arch` is empty, this rootfs's architecture is used. + pub fn exec(&self, opts: &ExecOptions) -> Result { + if opts.arch.is_empty() { + let mut opts = opts.clone(); + opts.arch = self.arch.clone(); + return crate::exec::exec(self.dir(), &opts); + } + crate::exec::exec(self.dir(), opts) + } + + /// Persist the current contents of the rootfs into the cache entry it + /// was prepared from, replacing the pristine image tarball, and mark + /// the entry provisioned. Later preparations for the same image and + /// architecture hit this entry directly; note that plain + /// [`RootfsCache::prepare`] also returns the provisioned contents + /// (indistinguishable from a pristine hit) — only + /// [`RootfsCache::prepare_provisioned`] interprets the marker. + pub fn persist(&self) -> Result<()> { + pack_rootfs(self.dir(), &self.cache_path)?; + std::fs::write(sidecar(&self.cache_path, PROVISIONED_SUFFIX), b"").with_context(|| { + format!( + "Failed to write provisioned marker for {}", + self.cache_path.display() + ) + })?; + veprintln!("Persisted rootfs to cache: {}", self.cache_path.display()); + Ok(()) + } +} + +/// Parse an image reference and resolve floating version aliases ("latest", +/// "lts", "edge", …) to a concrete version, so the cache key is stable +/// (e.g. "ubuntu-noble-amd64", not "ubuntu-latest-amd64") and a future +/// release automatically gets its own cache entry. +fn resolve_source(image: &str, arch: &str) -> Result { + let source = parse_image_ref(image, arch)?; + match source { + ImageSource::DirectTarball { distro, version } => { + let resolved = resolve_distro_version(&distro, version.as_deref(), arch)?; + Ok(ImageSource::DirectTarball { + distro, + version: Some(resolved), + }) + } + other => Ok(other), + } +} + +fn extract_prepared(cache_path: PathBuf, arch: String) -> Result { + let temp = tempfile::tempdir().context("Failed to create temporary directory")?; + let rootfs = temp.path().to_path_buf(); + + veprintln!("Extracting to: {}", rootfs.display()); + extract_tarball(&cache_path, &rootfs)?; + + Ok(PreparedRootfs { + rootfs: temp, + cache_path, + arch, + }) +} + +/// Write a tarball of `rootfs` at `dest`, compressed according to dest's +/// extension (.tar.gz, .tar.xz, .tar.zst, plain .tar), atomically via a +/// `.partial` file. +fn pack_rootfs(rootfs: &Path, dest: &Path) -> Result<()> { + // Packing the tree into a tarball placed inside it would recurse + // forever: the walker keeps seeing the archive grow. + let rootfs_canonical = rootfs + .canonicalize() + .with_context(|| format!("Failed to resolve rootfs directory {}", rootfs.display()))?; + if let Some(parent) = dest.parent() { + if let Ok(parent_canonical) = parent.canonicalize() { + if parent_canonical.starts_with(&rootfs_canonical) { + return Err(anyhow!( + "Cache entry {} must live outside the rootfs", + dest.display() + )); + } + } + } + + let filename = dest.file_name().and_then(|n| n.to_str()).unwrap_or(""); + let partial = dest.with_extension("partial"); + + let file = std::fs::File::create(&partial) + .with_context(|| format!("Failed to create {}", partial.display()))?; + + let result: Result<()> = (|| { + if filename.ends_with(".tar.gz") || filename.ends_with(".tgz") { + let encoder = flate2::write::GzEncoder::new(file, flate2::Compression::default()); + let encoder = write_rootfs_tar(rootfs, encoder)?; + encoder.finish().context("Failed to finalise gzip stream")?; + } else if filename.ends_with(".tar.xz") || filename.ends_with(".txz") { + let encoder = xz2::write::XzEncoder::new(file, 6); + let encoder = write_rootfs_tar(rootfs, encoder)?; + encoder.finish().context("Failed to finalise xz stream")?; + } else if filename.ends_with(".tar.zst") || filename.ends_with(".tar.zstd") { + let encoder = zstd::stream::write::Encoder::new(file, 0) + .context("Failed to create zstd encoder")?; + let encoder = write_rootfs_tar(rootfs, encoder)?; + encoder.finish().context("Failed to finalise zstd stream")?; + } else { + write_rootfs_tar(rootfs, file)?; + } + Ok(()) + })(); + + if result.is_err() { + // Best-effort cleanup; ignore errors (file may not exist if creation failed). + let _ = std::fs::remove_file(&partial); + return result; + } + + std::fs::rename(&partial, dest) + .with_context(|| format!("Failed to move packed rootfs to {}", dest.display())) +} + +/// Write the rootfs tree into a tar archive streamed to `writer`, returning +/// the writer so the caller can finalise the compression stream. +fn write_rootfs_tar(rootfs: &Path, writer: W) -> Result { + let mut builder = tar::Builder::new(writer); + // Store symlinks as symlinks; following them would duplicate the target's + // contents into the archive (and recurse into symlinked directories). + builder.follow_symlinks(false); + + let entries = std::fs::read_dir(rootfs) + .with_context(|| format!("Failed to read {}", rootfs.display()))?; + + for entry in entries { + let entry = + entry.with_context(|| format!("Failed to read entry in {}", rootfs.display()))?; + let path = entry.path(); + let name = entry.file_name(); + + // symlink_metadata: never follow symlinks — a top-level symlink to a + // directory must be stored as a link, not recursed into. + let meta = std::fs::symlink_metadata(&path) + .with_context(|| format!("Failed to stat {}", path.display()))?; + + if meta.is_dir() { + builder + .append_dir_all(&name, &path) + .with_context(|| format!("Failed to archive {}", path.display()))?; + } else { + builder + .append_path_with_name(&path, &name) + .with_context(|| format!("Failed to archive {}", path.display()))?; + } + } + + builder + .into_inner() + .context("Failed to finalise tar archive") +} + +/// Generate a cache filename based on the image source +fn generate_cache_filename(source: &ImageSource, arch: &str) -> String { + match source { + ImageSource::DirectTarball { distro, version } => { + let distro_name = match distro { + Distro::Ubuntu => "ubuntu", + Distro::Alpine => "alpine", + }; + let distro_arch = map_arch(*distro, arch); + // Get extension from URL + let url = resolve_distro_url(distro, version.as_deref(), arch).unwrap_or_default(); + let ext = get_tarball_extension(&url); + format!( + "{}-{}-{}.{}", + distro_name, + version.as_deref().unwrap_or("latest"), + distro_arch, + ext + ) + } + ImageSource::OciImage { + registry, + repository, + tag, + architecture, + } => { + // Sanitize for filename + let safe_registry = registry.replace(['.', ':'], "_"); + let safe_repo = repository.replace(['/', ':'], "_"); + format!( + "oci-{}-{}-{}-{}.tar.gz", + safe_registry, safe_repo, tag, architecture + ) + } + } +} + +fn get_tarball_extension(url: &str) -> &str { + // Extract extension from URL (e.g., .tar.gz, .tar.xz, .tar.zst) + if url.ends_with(".tar.zst") { + "tar.zst" + } else if url.ends_with(".tar.xz") { + "tar.xz" + } else if url.ends_with(".tar.gz") { + "tar.gz" + } else if url.ends_with(".tar.bz2") { + "tar.bz2" + } else { + "tar.gz" // default + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::os::unix::fs::PermissionsExt; + + fn direct_tarball(distro: Distro, version: &str) -> ImageSource { + ImageSource::DirectTarball { + distro, + version: Some(version.to_string()), + } + } + + #[test] + fn cache_filename_direct_tarball() { + // Explicit codename/version: no network lookups + let cache = RootfsCache::new("/cache"); + assert_eq!( + cache.tarball_path(&direct_tarball(Distro::Ubuntu, "noble"), "amd64"), + PathBuf::from("/cache/ubuntu-noble-amd64.tar.gz") + ); + assert_eq!( + cache.tarball_path(&direct_tarball(Distro::Alpine, "3.23.1"), "amd64"), + PathBuf::from("/cache/alpine-3.23.1-x86_64.tar.gz") + ); + } + + #[test] + fn cache_filename_oci_sanitized() { + let cache = RootfsCache::new("/cache"); + let source = ImageSource::OciImage { + registry: "docker.io".to_string(), + repository: "library/ubuntu".to_string(), + tag: "latest".to_string(), + architecture: "amd64".to_string(), + }; + assert_eq!( + cache.tarball_path(&source, "amd64"), + PathBuf::from("/cache/oci-docker_io-library_ubuntu-latest-amd64.tar.gz") + ); + } + + #[test] + fn resolved_arch_defaults_to_host() { + let req = PrepareRequest::new("alpine"); + assert_eq!( + req.resolved_arch(), + crate::utils::get_host_arch().debian_name() + ); + + let mut req = PrepareRequest::new("alpine"); + req.arch = "arm64".to_string(); + assert_eq!(req.resolved_arch(), "arm64"); + } + + #[test] + fn sidecar_appends_suffix() { + assert_eq!( + sidecar(Path::new("/cache/foo.tar.gz"), PROVISIONED_SUFFIX), + PathBuf::from("/cache/foo.tar.gz.provisioned") + ); + } + + /// Build a minimal rootfs tree with a nested executable and a symlink, + /// pack it, and extract it back through the public extract_tarball path. + fn roundtrip(dest_name: &str) { + let src = tempfile::tempdir().unwrap(); + std::fs::write(src.path().join("hello.txt"), b"hello").unwrap(); + std::fs::create_dir_all(src.path().join("usr/bin")).unwrap(); + std::fs::write(src.path().join("usr/bin/tool"), b"#!/bin/sh\n").unwrap(); + std::fs::set_permissions( + src.path().join("usr/bin/tool"), + std::fs::Permissions::from_mode(0o755), + ) + .unwrap(); + std::os::unix::fs::symlink("hello.txt", src.path().join("link.txt")).unwrap(); + + let cache = tempfile::tempdir().unwrap(); + let dest = cache.path().join(dest_name); + pack_rootfs(src.path(), &dest).unwrap(); + + let out = tempfile::tempdir().unwrap(); + crate::extract::extract_tarball(&dest, out.path()).unwrap(); + + assert_eq!( + std::fs::read(out.path().join("hello.txt")).unwrap(), + b"hello" + ); + let mode = std::fs::metadata(out.path().join("usr/bin/tool")) + .unwrap() + .permissions() + .mode(); + assert_eq!(mode & 0o777, 0o755, "executable bit must survive"); + assert_eq!( + std::fs::read_link(out.path().join("link.txt")).unwrap(), + Path::new("hello.txt") + ); + } + + #[test] + fn persist_roundtrip_gzip() { + roundtrip("rootfs.tar.gz"); + } + + #[test] + fn persist_roundtrip_xz() { + roundtrip("rootfs.tar.xz"); + } + + #[test] + fn persist_roundtrip_zstd() { + roundtrip("rootfs.tar.zst"); + } + + #[test] + fn persist_roundtrip_plain_tar() { + roundtrip("rootfs.tar"); + } + + #[test] + fn persist_writes_entry_and_provisioned_marker() { + let rootfs_dir = tempfile::tempdir().unwrap(); + std::fs::write(rootfs_dir.path().join("file.txt"), b"data").unwrap(); + + let cache_dir = tempfile::tempdir().unwrap(); + let cache_path = cache_dir.path().join("ubuntu-noble-amd64.tar.gz"); + + let prepared = PreparedRootfs { + rootfs: rootfs_dir, + cache_path: cache_path.clone(), + arch: "amd64".to_string(), + }; + prepared.persist().unwrap(); + + assert!(cache_path.exists()); + assert!(sidecar(&cache_path, PROVISIONED_SUFFIX).exists()); + + // The persisted entry extracts back to the original contents + let out = tempfile::tempdir().unwrap(); + crate::extract::extract_tarball(&cache_path, out.path()).unwrap(); + assert_eq!(std::fs::read(out.path().join("file.txt")).unwrap(), b"data"); + } + + #[test] + fn pack_failure_leaves_no_partial_file() { + // A nonexistent rootfs must fail cleanly and clean up the .partial + let cache = tempfile::tempdir().unwrap(); + let dest = cache.path().join("x.tar.gz"); + assert!(pack_rootfs(Path::new("/nonexistent-rootfs"), &dest).is_err()); + assert!(!dest.exists()); + assert!(!dest.with_extension("partial").exists()); + } +}