oci-image 0.3.0

use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

use crate::digest::Hasher;
use crate::error::Context;
use crate::{apply_layer_file, Digest, Error, Result, UnpackOptions};

/// A content-addressed store of layer blobs, unpacked layers and images:
///
/// ```text
/// <root>/blobs/sha256/<hex>   verified blobs, one copy each
/// <root>/layers/<diff_id>/    one layer unpacked once, in overlayfs format
///                             (whiteouts as 0/0 devices, opaque xattrs);
///                             shared by every image that has it
/// <root>/images/<id>.json     an image: its layers, bottom first; the
///                             refcount of a layer is the images naming it
/// <root>/rootfs/<chain>/      optional flattened rootfs ([`LayerStore::rootfs`])
///                             for tools that want one tree; not used by
///                             overlays and not collected by gc
/// <root>/tmp/                 downloads and unpacks in progress
/// ```
///
/// A container rootfs is an [`crate::Overlay`] over
/// [`LayerStore::image_layers`]; an updated image unpacks only its new layers.
///
/// Everything lands by rename, so a crash leaves only `tmp/` debris.
#[derive(Debug, Clone)]
pub struct LayerStore {
    root: PathBuf,
}

impl LayerStore {
    /// Open (and create) a store at `root`.
    pub fn open(root: &Path) -> Result<Self> {
        for d in ["blobs/sha256", "layers", "images", "rootfs", "tmp"] {
            let p = root.join(d);
            fs::create_dir_all(&p).ctx(|| format!("creating {}", p.display()))?;
        }
        Ok(Self {
            root: root.to_path_buf(),
        })
    }

    pub fn root(&self) -> &Path {
        &self.root
    }

    pub fn blob_path(&self, digest: &Digest) -> PathBuf {
        self.root.join("blobs/sha256").join(digest.hex())
    }

    pub fn has_blob(&self, digest: &Digest) -> bool {
        self.blob_path(digest).is_file()
    }

    /// Store everything `reader` yields as `digest`; refused (and nothing
    /// kept) if the bytes hash to something else.
    pub fn put_blob(&self, digest: &Digest, mut reader: impl Read) -> Result<PathBuf> {
        let mut w = self.blob_writer()?;
        let mut buf = vec![0u8; 1 << 16];
        loop {
            let n = reader.read(&mut buf).ctx(|| format!("reading {digest}"))?;
            if n == 0 {
                break;
            }
            w.write(&buf[..n])?;
        }
        w.commit(digest)
    }

    /// A blob that arrives in chunks (a registry download).
    pub fn blob_writer(&self) -> Result<BlobWriter> {
        let path = self.tmp_path("blob");
        let file = File::create(&path).ctx(|| format!("creating {}", path.display()))?;
        Ok(BlobWriter {
            file: Some(io::BufWriter::with_capacity(1 << 20, file)),
            path,
            hasher: Hasher::default(),
            dest: self.root.join("blobs/sha256"),
        })
    }

    /// The rootfs for `layers` (bottom first), unpacked on first use.
    pub fn rootfs(&self, layers: &[Digest], opts: &UnpackOptions) -> Result<PathBuf> {
        let dir = self.rootfs_path(layers);
        if dir.is_dir() {
            return Ok(dir);
        }
        if let Some(missing) = layers.iter().find(|d| !self.has_blob(d)) {
            return Err(Error::MissingBlob(missing.clone()));
        }
        let tmp = match self.unpack(layers, opts) {
            // A gzip layer the parallel decoder refused (or decoded to bytes
            // that fail the trailer check): start over with the streaming
            // decoder, which is authoritative.
            Err(_) if opts.parallel_gzip => self.unpack(
                layers,
                &UnpackOptions {
                    parallel_gzip: false,
                    ..opts.clone()
                },
            )?,
            r => r?,
        };
        if fs::rename(&tmp, &dir).is_err() {
            // Another unpack of the same chain won the race.
            let _ = fs::remove_dir_all(&tmp);
            if !dir.is_dir() {
                return Err(Error::Invalid(format!("could not place {}", dir.display())));
            }
        }
        Ok(dir)
    }

    /// Apply `layers` onto a fresh directory under `tmp/`.
    fn unpack(&self, layers: &[Digest], opts: &UnpackOptions) -> Result<PathBuf> {
        let tmp = self.tmp_path("rootfs");
        fs::create_dir_all(&tmp).ctx(|| format!("creating {}", tmp.display()))?;
        let unpacked = layers
            .iter()
            .try_for_each(|d| apply_layer_file(&self.blob_path(d), &tmp, opts).map(drop));
        match unpacked {
            Ok(()) => Ok(tmp),
            Err(e) => {
                let _ = fs::remove_dir_all(&tmp);
                Err(e)
            }
        }
    }

    /// Where the rootfs for `layers` lives (whether or not it exists yet).
    pub fn rootfs_path(&self, layers: &[Digest]) -> PathBuf {
        self.root.join("rootfs").join(chain_id(layers).hex())
    }

    fn tmp_path(&self, kind: &str) -> PathBuf {
        static N: AtomicU64 = AtomicU64::new(0);
        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_or(0, |d| d.as_nanos());
        self.root.join("tmp").join(format!(
            "{kind}-{}-{}-{nanos}",
            std::process::id(),
            N.fetch_add(1, Ordering::Relaxed)
        ))
    }
}

/// The id of a stack of layers: sha256 over their digests, one per line.
pub fn chain_id(layers: &[Digest]) -> Digest {
    let mut h = Hasher::default();
    for d in layers {
        h.update(d.to_string().as_bytes());
        h.update(b"\n");
    }
    h.finish()
}

/// A blob being written; hashed as it goes, placed by [`BlobWriter::commit`],
/// removed if dropped uncommitted.
pub struct BlobWriter {
    file: Option<io::BufWriter<File>>,
    path: PathBuf,
    hasher: Hasher,
    dest: PathBuf,
}

impl BlobWriter {
    pub fn write(&mut self, bytes: &[u8]) -> Result<()> {
        self.hasher.update(bytes);
        let f = self.file.as_mut().expect("open until commit");
        f.write_all(bytes)
            .ctx(|| format!("writing {}", self.path.display()))
    }

    /// Verify against `digest` and move into the store.
    pub fn commit(mut self, digest: &Digest) -> Result<PathBuf> {
        let f = self.file.take().expect("open until commit");
        f.into_inner()
            .map_err(|e| e.into_error())
            .ctx(|| format!("writing {}", self.path.display()))?;
        digest.verify(&std::mem::take(&mut self.hasher).finish())?;
        let dest = self.dest.join(digest.hex());
        fs::rename(&self.path, &dest).ctx(|| format!("placing {digest}"))?;
        Ok(dest)
    }
}

impl Drop for BlobWriter {
    fn drop(&mut self) {
        let _ = fs::remove_file(&self.path);
    }
}

/// One layer of an image: its blob (compressed) digest and its diff_id (the
/// digest of the uncompressed tar, from the image config).
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ImageLayer {
    pub blob: Digest,
    pub diff_id: Digest,
}

/// What [`LayerStore::gc`] removed.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GcStats {
    pub layers: usize,
    pub blobs: usize,
}

#[derive(serde::Serialize, serde::Deserialize)]
struct ImageRecord {
    layers: Vec<ImageLayer>,
}

impl LayerStore {
    /// Where the unpacked directory of a layer lives.
    pub fn layer_path(&self, diff_id: &Digest) -> PathBuf {
        self.root.join("layers").join(diff_id.hex())
    }

    pub fn has_layer(&self, diff_id: &Digest) -> bool {
        self.layer_path(diff_id).is_dir()
    }

    /// The layers not yet unpacked: what a pull still has to fetch.
    pub fn missing_layers(&self, layers: &[ImageLayer]) -> Vec<ImageLayer> {
        layers
            .iter()
            .filter(|l| !self.has_layer(&l.diff_id))
            .cloned()
            .collect()
    }

    /// Unpack one layer into its own directory (overlayfs format), once.
    /// The uncompressed stream must hash to `layer.diff_id`.
    pub fn unpack_layer(&self, layer: &ImageLayer, opts: &UnpackOptions) -> Result<PathBuf> {
        let dir = self.layer_path(&layer.diff_id);
        if dir.is_dir() {
            return Ok(dir);
        }
        if !self.has_blob(&layer.blob) {
            return Err(Error::MissingBlob(layer.blob.clone()));
        }
        let attempt = |o: &UnpackOptions| -> Result<PathBuf> {
            let tmp = self.tmp_path("layer");
            fs::create_dir_all(&tmp).ctx(|| format!("creating {}", tmp.display()))?;
            let r = crate::layer::apply_layer_file_digest(&self.blob_path(&layer.blob), &tmp, o)
                .and_then(|(_, d)| layer.diff_id.verify(&d));
            match r {
                Ok(()) => Ok(tmp),
                Err(e) => {
                    let _ = fs::remove_dir_all(&tmp);
                    Err(e)
                }
            }
        };
        let o = UnpackOptions {
            layout: match opts.layout {
                crate::Layout::Flatten => crate::Layout::Overlay(crate::OverlayXattr::Trusted),
                l => l,
            },
            ..opts.clone()
        };
        let tmp = match attempt(&o) {
            Err(e @ Error::DigestMismatch { .. }) => return Err(e),
            Err(_) if o.parallel_gzip => attempt(&UnpackOptions {
                parallel_gzip: false,
                ..o.clone()
            })?,
            r => r?,
        };
        if fs::rename(&tmp, &dir).is_err() {
            let _ = fs::remove_dir_all(&tmp);
            if !dir.is_dir() {
                return Err(Error::Invalid(format!("could not place {}", dir.display())));
            }
        }
        Ok(dir)
    }

    /// Record image `id` as `layers` (bottom first), unpacking the layers
    /// not in the store yet. Returns the layer directories, bottom first.
    pub fn add_image(
        &self,
        id: &Digest,
        layers: &[ImageLayer],
        opts: &UnpackOptions,
    ) -> Result<Vec<PathBuf>> {
        let dirs = layers
            .iter()
            .map(|l| self.unpack_layer(l, opts))
            .collect::<Result<Vec<_>>>()?;
        let rec = ImageRecord {
            layers: layers.to_vec(),
        };
        let path = self.image_record(id);
        let tmp = self.tmp_path("image");
        let bytes = serde_json::to_vec_pretty(&rec).map_err(|e| Error::Invalid(e.to_string()))?;
        fs::write(&tmp, bytes).ctx(|| format!("writing {}", tmp.display()))?;
        fs::rename(&tmp, &path).ctx(|| format!("placing {}", path.display()))?;
        Ok(dirs)
    }

    /// The layer directories of a recorded image, bottom first.
    pub fn image_layers(&self, id: &Digest) -> Result<Vec<PathBuf>> {
        Ok(self
            .record(id)?
            .layers
            .iter()
            .map(|l| self.layer_path(&l.diff_id))
            .collect())
    }

    /// Forget image `id`; its layers stay until [`LayerStore::gc`].
    pub fn remove_image(&self, id: &Digest) -> Result<()> {
        let p = self.image_record(id);
        fs::remove_file(&p).ctx(|| format!("removing {}", p.display()))
    }

    /// How many recorded images use the layer.
    pub fn refcount(&self, diff_id: &Digest) -> Result<usize> {
        Ok(self
            .records()?
            .iter()
            .filter(|r| r.layers.iter().any(|l| &l.diff_id == diff_id))
            .count())
    }

    /// Remove layer directories and blobs that no recorded image uses, and
    /// leftovers in `tmp/`. Run it while no pull is in progress.
    pub fn gc(&self) -> Result<GcStats> {
        let recs = self.records()?;
        let diff_ids: std::collections::HashSet<String> = recs
            .iter()
            .flat_map(|r| r.layers.iter().map(|l| l.diff_id.hex().to_string()))
            .collect();
        let blobs: std::collections::HashSet<String> = recs
            .iter()
            .flat_map(|r| r.layers.iter().map(|l| l.blob.hex().to_string()))
            .collect();
        let mut st = GcStats::default();
        for (dir, keep, n) in [
            ("layers", &diff_ids, &mut st.layers),
            ("blobs/sha256", &blobs, &mut st.blobs),
        ] {
            let d = self.root.join(dir);
            for e in fs::read_dir(&d)
                .ctx(|| format!("listing {}", d.display()))?
                .flatten()
            {
                let name = e.file_name().to_string_lossy().into_owned();
                if !keep.contains(&name) {
                    let p = e.path();
                    let r = if p.is_dir() {
                        fs::remove_dir_all(&p)
                    } else {
                        fs::remove_file(&p)
                    };
                    r.ctx(|| format!("removing {}", p.display()))?;
                    *n += 1;
                }
            }
        }
        let t = self.root.join("tmp");
        for e in fs::read_dir(&t)
            .ctx(|| format!("listing {}", t.display()))?
            .flatten()
        {
            let p = e.path();
            let _ = if p.is_dir() {
                fs::remove_dir_all(&p)
            } else {
                fs::remove_file(&p)
            };
        }
        Ok(st)
    }

    /// Whether the store holds what `blob` is: the blob itself, or the layer
    /// directory it unpacks to (its blob pruned by
    /// [`LayerStore::prune_unpacked_blobs`]), as an image record names it. A
    /// pull asks this to skip what it need not download.
    pub fn holds_blob(&self, blob: &Digest) -> bool {
        self.has_blob(blob)
            || self.records().is_ok_and(|rs| {
                rs.iter()
                    .flat_map(|r| &r.layers)
                    .any(|l| &l.blob == blob && self.has_layer(&l.diff_id))
            })
    }

    /// Delete the blobs of layers that are unpacked (to save disk); a later
    /// pull of the same layers needs neither the blob nor a download.
    pub fn prune_unpacked_blobs(&self) -> Result<usize> {
        let mut n = 0;
        for r in self.records()? {
            for l in r.layers {
                if self.has_layer(&l.diff_id) && fs::remove_file(self.blob_path(&l.blob)).is_ok() {
                    n += 1;
                }
            }
        }
        Ok(n)
    }

    fn image_record(&self, id: &Digest) -> PathBuf {
        self.root.join("images").join(format!("{}.json", id.hex()))
    }

    fn record(&self, id: &Digest) -> Result<ImageRecord> {
        let p = self.image_record(id);
        let b = fs::read(&p).ctx(|| format!("reading {}", p.display()))?;
        serde_json::from_slice(&b).map_err(|e| Error::Invalid(format!("{}: {e}", p.display())))
    }

    fn records(&self) -> Result<Vec<ImageRecord>> {
        let d = self.root.join("images");
        let mut out = Vec::new();
        for e in fs::read_dir(&d)
            .ctx(|| format!("listing {}", d.display()))?
            .flatten()
        {
            if let Ok(b) = fs::read(e.path()) {
                if let Ok(r) = serde_json::from_slice(&b) {
                    out.push(r);
                }
            }
        }
        Ok(out)
    }
}