oci-image 0.3.0

use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom, Write};
use std::path::Path;
use std::sync::mpsc::{self, Receiver, SyncSender};
use std::sync::{Arc, Mutex};
use std::thread;

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

/// How a layer was decoded.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decoder {
    /// An uncompressed tar.
    Plain,
    /// One gzip stream decoded on one core as it is read (zlib-rs): small
    /// layers (below lgz's serial threshold), one thread, or lgz switched off.
    GzipStream,
    /// Gzip decoded on all the caller's threads with `lgz`'s one-pass
    /// streaming decoder, any gzip (no flush points needed), in order, as it
    /// is read; every member's CRC32 and size trailer is checked.
    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())?)
        }
    })
}

/// The sha256 of a layer's decoded tar (its diff_id), as an open layer
/// computes it.
pub(crate) enum TarDigest {
    /// Not asked for.
    None,
    /// Hashed on the tar reader's thread as it reads.
    Inline(Arc<Mutex<Hasher>>),
    /// Hashed on its own thread from the decoder's in-order output.
    Thread(thread::JoinHandle<Digest>),
}

impl TarDigest {
    /// The digest, once the tar has been read to its end.
    pub(crate) fn finish(self) -> Option<Digest> {
        match self {
            Self::None => None,
            Self::Inline(h) => Some(std::mem::take(&mut *h.lock().unwrap()).finish()),
            Self::Thread(t) => t.join().ok(),
        }
    }
}

/// Open a layer blob on disk as a tar stream. With `digest`, the sha256 of
/// the decoded tar is computed on the way.
pub(crate) fn open_layer(
    path: &Path,
    opts: &UnpackOptions,
    digest: bool,
) -> Result<(Box<dyn Read + Send>, Decoder, TarDigest)> {
    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)?;
    // The blob's size: what the manifest says, since the store verified the
    // blob against its digest.
    let len = f.metadata().ctx(what)?.len();
    let (reader, decoder): (Box<dyn Read + Send>, Decoder) = 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 => {
            let threads = opts.threads();
            // lgz's own rule: below its serial threshold it takes streaming
            // zlib-rs itself. One thread decodes below, on the calling thread.
            let decoder = if len >= lgz::stream::SERIAL_BELOW as u64 {
                Decoder::GzipParallel
            } else {
                Decoder::GzipStream
            };
            let (reader, hash) = lgz_stream(f, len, threads, digest);
            let d = match hash {
                Some(t) => TarDigest::Thread(t),
                None => TarDigest::None,
            };
            return Ok((Box::new(reader), decoder, d));
        }
        Compression::Gzip => (
            Box::new(flate2::read::MultiGzDecoder::new(BufReader::with_capacity(
                1 << 20,
                f,
            ))),
            Decoder::GzipStream,
        ),
    };
    if !digest {
        return Ok((reader, decoder, TarDigest::None));
    }
    let hasher = Arc::new(Mutex::new(Hasher::default()));
    let reader = HashReader {
        inner: reader,
        hasher: hasher.clone(),
    };
    Ok((Box::new(reader), decoder, TarDigest::Inline(hasher)))
}

struct HashReader<R> {
    inner: R,
    hasher: Arc<Mutex<Hasher>>,
}

impl<R: Read> Read for HashReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let n = self.inner.read(buf)?;
        self.hasher.lock().unwrap().update(&buf[..n]);
        Ok(n)
    }
}

/// Decoded bytes handed on at a time.
const PIECE: usize = 1 << 20;
/// Pieces waiting for the tar reader (and for the hash thread).
const PIECES: usize = 8;

type Piece = io::Result<Arc<Vec<u8>>>;

/// Decode the gzip `file` (`len` bytes) with `lgz::stream::decode_sized` on
/// `threads` workers, on a thread of its own, and read the tar from its
/// in-order output. With `digest`, one more thread hashes the same output.
/// A decode error (bad data, a CRC32 or size trailer that does not match)
/// reaches the reader before its end of stream.
fn lgz_stream(
    file: File,
    len: u64,
    threads: usize,
    digest: bool,
) -> (PieceReader, Option<thread::JoinHandle<Digest>>) {
    let (tx, rx) = mpsc::sync_channel::<Piece>(PIECES);
    let (hash_tx, hash) = if digest {
        let (htx, hrx) = mpsc::sync_channel::<Arc<Vec<u8>>>(PIECES);
        // A32-14, stays a thread: the digest of one layer's decoded stream, beside its decode;
        // the layer's reader joins it (TarDigest::Thread) once the stream ends.
        let t = thread::spawn(move || {
            let mut h = Hasher::default();
            for piece in hrx {
                h.update(&piece);
            }
            h.finish()
        });
        (Some(htx), Some(t))
    } else {
        (None, None)
    };
    // A32-14, stays a thread: the lgz decode of one layer, feeding the tar reader in order; it
    // ends with the stream, or at once when the reader is gone (BrokenPipe on the next send).
    thread::spawn(move || {
        let mut out = Tee {
            buf: Vec::with_capacity(PIECE),
            tx: tx.clone(),
            hash: hash_tx,
        };
        let r = lgz::stream::decode_sized(file, &mut out, threads, Some(len));
        let end = match r {
            Ok(_) => out.flush(),
            Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e.to_string())),
        };
        if let Err(e) = end {
            // The hash thread must not finish as if the stream were whole:
            // drop its sender only after the reader has the error.
            let _ = tx.send(Err(e));
        }
    });
    (
        PieceReader {
            rx,
            cur: Arc::new(Vec::new()),
            pos: 0,
        },
        hash,
    )
}

/// The decoder's output: copied into pieces and sent, in order, to the tar
/// reader and the hash thread.
struct Tee {
    buf: Vec<u8>,
    tx: SyncSender<Piece>,
    hash: Option<SyncSender<Arc<Vec<u8>>>>,
}

impl Tee {
    fn send(&mut self) -> io::Result<()> {
        if self.buf.is_empty() {
            return Ok(());
        }
        let piece = Arc::new(std::mem::replace(&mut self.buf, Vec::with_capacity(PIECE)));
        let gone = || io::Error::new(io::ErrorKind::BrokenPipe, "the layer reader stopped");
        if let Some(h) = &self.hash {
            h.send(piece.clone()).map_err(|_| gone())?;
        }
        self.tx.send(Ok(piece)).map_err(|_| gone())
    }
}

impl Write for Tee {
    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
        let n = b.len().min(PIECE - self.buf.len());
        self.buf.extend_from_slice(&b[..n]);
        if self.buf.len() == PIECE {
            self.send()?;
        }
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.send()
    }
}

/// The tar side of [`lgz_stream`].
pub(crate) struct PieceReader {
    rx: Receiver<Piece>,
    cur: Arc<Vec<u8>>,
    pos: usize,
}

impl Read for PieceReader {
    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
        while self.pos == self.cur.len() {
            match self.rx.recv() {
                Ok(piece) => {
                    self.cur = piece?;
                    self.pos = 0;
                }
                Err(_) => return Ok(0),
            }
        }
        let n = out.len().min(self.cur.len() - self.pos);
        out[..n].copy_from_slice(&self.cur[self.pos..self.pos + n]);
        self.pos += n;
        Ok(n)
    }
}

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)
}