Skip to main content

oci_image/
decode.rs

1use std::fs::File;
2use std::io::{self, BufReader, Read, Seek, SeekFrom, Write};
3use std::path::Path;
4use std::sync::mpsc::{self, Receiver, SyncSender};
5use std::sync::{Arc, Mutex};
6use std::thread;
7
8use crate::error::Context;
9use crate::{Digest, Hasher, Result, UnpackOptions};
10
11/// How a layer was decoded.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Decoder {
14    /// An uncompressed tar.
15    Plain,
16    /// One gzip stream decoded on one core as it is read (zlib-rs): small
17    /// layers (below lgz's serial threshold), one thread, or lgz switched off.
18    GzipStream,
19    /// Gzip decoded on all the caller's threads with `lgz`'s one-pass
20    /// streaming decoder, any gzip (no flush points needed), in order, as it
21    /// is read; every member's CRC32 and size trailer is checked.
22    GzipParallel,
23    /// Zstandard.
24    Zstd,
25}
26
27/// The compression a blob starts with.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum Compression {
30    None,
31    Gzip,
32    Zstd,
33}
34
35impl Compression {
36    /// From the first bytes of a blob.
37    pub fn detect(head: &[u8]) -> Self {
38        if head.starts_with(&[0x1f, 0x8b]) {
39            Self::Gzip
40        } else if head.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
41            Self::Zstd
42        } else {
43            Self::None
44        }
45    }
46}
47
48/// A streaming tar reader over an in-memory blob of any supported compression.
49pub fn decoder(blob: &[u8]) -> Result<Box<dyn Read + Send + '_>> {
50    Ok(match Compression::detect(blob) {
51        Compression::None => Box::new(blob),
52        Compression::Gzip => Box::new(flate2::read::MultiGzDecoder::new(blob)),
53        Compression::Zstd => {
54            Box::new(zstd::stream::read::Decoder::new(blob).ctx(|| "zstd header".into())?)
55        }
56    })
57}
58
59/// The sha256 of a layer's decoded tar (its diff_id), as an open layer
60/// computes it.
61pub(crate) enum TarDigest {
62    /// Not asked for.
63    None,
64    /// Hashed on the tar reader's thread as it reads.
65    Inline(Arc<Mutex<Hasher>>),
66    /// Hashed on its own thread from the decoder's in-order output.
67    Thread(thread::JoinHandle<Digest>),
68}
69
70impl TarDigest {
71    /// The digest, once the tar has been read to its end.
72    pub(crate) fn finish(self) -> Option<Digest> {
73        match self {
74            Self::None => None,
75            Self::Inline(h) => Some(std::mem::take(&mut *h.lock().unwrap()).finish()),
76            Self::Thread(t) => t.join().ok(),
77        }
78    }
79}
80
81/// Open a layer blob on disk as a tar stream. With `digest`, the sha256 of
82/// the decoded tar is computed on the way.
83pub(crate) fn open_layer(
84    path: &Path,
85    opts: &UnpackOptions,
86    digest: bool,
87) -> Result<(Box<dyn Read + Send>, Decoder, TarDigest)> {
88    let what = || format!("opening layer {}", path.display());
89    let mut f = File::open(path).ctx(what)?;
90    let mut head = [0u8; 4];
91    let n = read_up_to(&mut f, &mut head).ctx(what)?;
92    f.seek(SeekFrom::Start(0)).ctx(what)?;
93    // The blob's size: what the manifest says, since the store verified the
94    // blob against its digest.
95    let len = f.metadata().ctx(what)?.len();
96    let (reader, decoder): (Box<dyn Read + Send>, Decoder) = match Compression::detect(&head[..n]) {
97        Compression::None => (
98            Box::new(BufReader::with_capacity(1 << 20, f)),
99            Decoder::Plain,
100        ),
101        Compression::Zstd => (
102            Box::new(zstd::stream::read::Decoder::new(f).ctx(what)?),
103            Decoder::Zstd,
104        ),
105        Compression::Gzip if opts.parallel_gzip && opts.threads() > 1 => {
106            let threads = opts.threads();
107            // lgz's own rule: below its serial threshold it takes streaming
108            // zlib-rs itself. One thread decodes below, on the calling thread.
109            let decoder = if len >= lgz::stream::SERIAL_BELOW as u64 {
110                Decoder::GzipParallel
111            } else {
112                Decoder::GzipStream
113            };
114            let (reader, hash) = lgz_stream(f, len, threads, digest);
115            let d = match hash {
116                Some(t) => TarDigest::Thread(t),
117                None => TarDigest::None,
118            };
119            return Ok((Box::new(reader), decoder, d));
120        }
121        Compression::Gzip => (
122            Box::new(flate2::read::MultiGzDecoder::new(BufReader::with_capacity(
123                1 << 20,
124                f,
125            ))),
126            Decoder::GzipStream,
127        ),
128    };
129    if !digest {
130        return Ok((reader, decoder, TarDigest::None));
131    }
132    let hasher = Arc::new(Mutex::new(Hasher::default()));
133    let reader = HashReader {
134        inner: reader,
135        hasher: hasher.clone(),
136    };
137    Ok((Box::new(reader), decoder, TarDigest::Inline(hasher)))
138}
139
140struct HashReader<R> {
141    inner: R,
142    hasher: Arc<Mutex<Hasher>>,
143}
144
145impl<R: Read> Read for HashReader<R> {
146    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
147        let n = self.inner.read(buf)?;
148        self.hasher.lock().unwrap().update(&buf[..n]);
149        Ok(n)
150    }
151}
152
153/// Decoded bytes handed on at a time.
154const PIECE: usize = 1 << 20;
155/// Pieces waiting for the tar reader (and for the hash thread).
156const PIECES: usize = 8;
157
158type Piece = io::Result<Arc<Vec<u8>>>;
159
160/// Decode the gzip `file` (`len` bytes) with `lgz::stream::decode_sized` on
161/// `threads` workers, on a thread of its own, and read the tar from its
162/// in-order output. With `digest`, one more thread hashes the same output.
163/// A decode error (bad data, a CRC32 or size trailer that does not match)
164/// reaches the reader before its end of stream.
165fn lgz_stream(
166    file: File,
167    len: u64,
168    threads: usize,
169    digest: bool,
170) -> (PieceReader, Option<thread::JoinHandle<Digest>>) {
171    let (tx, rx) = mpsc::sync_channel::<Piece>(PIECES);
172    let (hash_tx, hash) = if digest {
173        let (htx, hrx) = mpsc::sync_channel::<Arc<Vec<u8>>>(PIECES);
174        // A32-14, stays a thread: the digest of one layer's decoded stream, beside its decode;
175        // the layer's reader joins it (TarDigest::Thread) once the stream ends.
176        let t = thread::spawn(move || {
177            let mut h = Hasher::default();
178            for piece in hrx {
179                h.update(&piece);
180            }
181            h.finish()
182        });
183        (Some(htx), Some(t))
184    } else {
185        (None, None)
186    };
187    // A32-14, stays a thread: the lgz decode of one layer, feeding the tar reader in order; it
188    // ends with the stream, or at once when the reader is gone (BrokenPipe on the next send).
189    thread::spawn(move || {
190        let mut out = Tee {
191            buf: Vec::with_capacity(PIECE),
192            tx: tx.clone(),
193            hash: hash_tx,
194        };
195        let r = lgz::stream::decode_sized(file, &mut out, threads, Some(len));
196        let end = match r {
197            Ok(_) => out.flush(),
198            Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e.to_string())),
199        };
200        if let Err(e) = end {
201            // The hash thread must not finish as if the stream were whole:
202            // drop its sender only after the reader has the error.
203            let _ = tx.send(Err(e));
204        }
205    });
206    (
207        PieceReader {
208            rx,
209            cur: Arc::new(Vec::new()),
210            pos: 0,
211        },
212        hash,
213    )
214}
215
216/// The decoder's output: copied into pieces and sent, in order, to the tar
217/// reader and the hash thread.
218struct Tee {
219    buf: Vec<u8>,
220    tx: SyncSender<Piece>,
221    hash: Option<SyncSender<Arc<Vec<u8>>>>,
222}
223
224impl Tee {
225    fn send(&mut self) -> io::Result<()> {
226        if self.buf.is_empty() {
227            return Ok(());
228        }
229        let piece = Arc::new(std::mem::replace(&mut self.buf, Vec::with_capacity(PIECE)));
230        let gone = || io::Error::new(io::ErrorKind::BrokenPipe, "the layer reader stopped");
231        if let Some(h) = &self.hash {
232            h.send(piece.clone()).map_err(|_| gone())?;
233        }
234        self.tx.send(Ok(piece)).map_err(|_| gone())
235    }
236}
237
238impl Write for Tee {
239    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
240        let n = b.len().min(PIECE - self.buf.len());
241        self.buf.extend_from_slice(&b[..n]);
242        if self.buf.len() == PIECE {
243            self.send()?;
244        }
245        Ok(n)
246    }
247    fn flush(&mut self) -> io::Result<()> {
248        self.send()
249    }
250}
251
252/// The tar side of [`lgz_stream`].
253pub(crate) struct PieceReader {
254    rx: Receiver<Piece>,
255    cur: Arc<Vec<u8>>,
256    pos: usize,
257}
258
259impl Read for PieceReader {
260    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
261        while self.pos == self.cur.len() {
262            match self.rx.recv() {
263                Ok(piece) => {
264                    self.cur = piece?;
265                    self.pos = 0;
266                }
267                Err(_) => return Ok(0),
268            }
269        }
270        let n = out.len().min(self.cur.len() - self.pos);
271        out[..n].copy_from_slice(&self.cur[self.pos..self.pos + n]);
272        self.pos += n;
273        Ok(n)
274    }
275}
276
277fn read_up_to(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
278    let mut n = 0;
279    while n < buf.len() {
280        match r.read(&mut buf[n..])? {
281            0 => break,
282            k => n += k,
283        }
284    }
285    Ok(n)
286}