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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Decoder {
14 Plain,
16 GzipStream,
19 GzipParallel,
23 Zstd,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum Compression {
30 None,
31 Gzip,
32 Zstd,
33}
34
35impl Compression {
36 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
48pub 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
59pub(crate) enum TarDigest {
62 None,
64 Inline(Arc<Mutex<Hasher>>),
66 Thread(thread::JoinHandle<Digest>),
68}
69
70impl TarDigest {
71 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
81pub(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 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 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
153const PIECE: usize = 1 << 20;
155const PIECES: usize = 8;
157
158type Piece = io::Result<Arc<Vec<u8>>>;
159
160fn 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 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 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 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
216struct 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
252pub(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}