oci-image 0.1.0

use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom};
use std::path::Path;

use crate::error::Context;
use crate::pgz::ParallelGz;
use crate::{Result, UnpackOptions};

/// How a layer was decoded.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decoder {
    /// An uncompressed tar.
    Plain,
    /// One gzip stream decoded as it is read (zlib-rs).
    GzipStream,
    /// Gzip with full-flush points (`pigz --independent`, BGZF), cut at the
    /// flushes and decoded in parallel with `lgz`, in order, as it is read;
    /// the CRC32 and size trailer are checked at the end.
    GzipParallel,
    /// Zstandard.
    Zstd,
}

/// The compression a blob starts with.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Compression {
    None,
    Gzip,
    Zstd,
}

impl Compression {
    /// From the first bytes of a blob.
    pub fn detect(head: &[u8]) -> Self {
        if head.starts_with(&[0x1f, 0x8b]) {
            Self::Gzip
        } else if head.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
            Self::Zstd
        } else {
            Self::None
        }
    }
}

/// A streaming tar reader over an in-memory blob of any supported compression.
pub fn decoder(blob: &[u8]) -> Result<Box<dyn Read + Send + '_>> {
    Ok(match Compression::detect(blob) {
        Compression::None => Box::new(blob),
        Compression::Gzip => Box::new(flate2::read::MultiGzDecoder::new(blob)),
        Compression::Zstd => {
            Box::new(zstd::stream::read::Decoder::new(blob).ctx(|| "zstd header".into())?)
        }
    })
}

/// Open a layer blob on disk as a tar stream.
pub(crate) fn open_layer(
    path: &Path,
    opts: &UnpackOptions,
) -> Result<(Box<dyn Read + Send>, Decoder)> {
    let what = || format!("opening layer {}", path.display());
    let mut f = File::open(path).ctx(what)?;
    let mut head = [0u8; 4];
    let n = read_up_to(&mut f, &mut head).ctx(what)?;
    f.seek(SeekFrom::Start(0)).ctx(what)?;
    let len = f.metadata().ctx(what)?.len();
    Ok(match Compression::detect(&head[..n]) {
        Compression::None => (
            Box::new(BufReader::with_capacity(1 << 20, f)),
            Decoder::Plain,
        ),
        Compression::Zstd => (
            Box::new(zstd::stream::read::Decoder::new(f).ctx(what)?),
            Decoder::Zstd,
        ),
        Compression::Gzip
            if opts.parallel_gzip
                && opts.threads() > 1
                && len >= lgz::PARALLEL_THRESHOLD as u64
                && has_full_flush(&mut f).ctx(what)? =>
        {
            (
                Box::new(ParallelGz::new(f, opts.threads())),
                Decoder::GzipParallel,
            )
        }
        Compression::Gzip => (
            Box::new(flate2::read::MultiGzDecoder::new(BufReader::with_capacity(
                1 << 20,
                f,
            ))),
            Decoder::GzipStream,
        ),
    })
}

/// How much of a gzip layer is searched for a full-flush point before it is
/// streamed instead. `pigz --independent` flushes every 128 KiB, BGZF every
/// 64 KiB; one ordinary stream has none at all.
const FLUSH_PROBE: usize = 4 << 20;

/// Whether the first [`FLUSH_PROBE`] bytes hold a validated full flush.
/// Leaves the file at its start.
fn has_full_flush(f: &mut File) -> io::Result<bool> {
    let mut head = vec![0u8; FLUSH_PROBE];
    let n = read_up_to(f, &mut head)?;
    f.seek(SeekFrom::Start(0))?;
    head.truncate(n);
    let Ok(start) = lgz::deflate_offset(&head) else {
        return Ok(false);
    };
    Ok(lgz::deflate_scan::find_next_valid_flush(&head[start..], 0).is_some())
}

fn read_up_to(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
    let mut n = 0;
    while n < buf.len() {
        match r.read(&mut buf[n..])? {
            0 => break,
            k => n += k,
        }
    }
    Ok(n)
}