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 const ZSTD_OFF: &str = "a zstd layer: this build of oci-image has no feature `zstd`";
50
51#[cfg(feature = "zstd")]
53fn zstd_reader<'a, R: Read + Send + 'a>(
54 r: R,
55 what: impl FnOnce() -> String,
56) -> Result<Box<dyn Read + Send + 'a>> {
57 Ok(Box::new(zstd::stream::read::Decoder::new(r).ctx(what)?))
58}
59
60#[cfg(not(feature = "zstd"))]
62fn zstd_reader<'a, R: Read + Send + 'a>(
63 _r: R,
64 _what: impl FnOnce() -> String,
65) -> Result<Box<dyn Read + Send + 'a>> {
66 Err(crate::Error::Invalid(ZSTD_OFF.into()))
67}
68
69pub fn decoder(blob: &[u8]) -> Result<Box<dyn Read + Send + '_>> {
71 Ok(match Compression::detect(blob) {
72 Compression::None => Box::new(blob),
73 Compression::Gzip => Box::new(flate2::read::MultiGzDecoder::new(blob)),
74 Compression::Zstd => zstd_reader(blob, || "zstd header".into())?,
75 })
76}
77
78pub(crate) enum TarDigest {
81 None,
83 Inline(Arc<Mutex<Hasher>>),
85 Thread(thread::JoinHandle<Digest>),
87}
88
89impl TarDigest {
90 pub(crate) fn finish(self) -> Option<Digest> {
92 match self {
93 Self::None => None,
94 Self::Inline(h) => Some(std::mem::take(&mut *h.lock().unwrap()).finish()),
95 Self::Thread(t) => t.join().ok(),
96 }
97 }
98}
99
100pub(crate) fn open_layer(
103 path: &Path,
104 opts: &UnpackOptions,
105 digest: bool,
106) -> Result<(Box<dyn Read + Send>, Decoder, TarDigest)> {
107 let what = || format!("opening layer {}", path.display());
108 let mut f = File::open(path).ctx(what)?;
109 let mut head = [0u8; 4];
110 let n = read_up_to(&mut f, &mut head).ctx(what)?;
111 f.seek(SeekFrom::Start(0)).ctx(what)?;
112 let len = f.metadata().ctx(what)?.len();
115 let (reader, decoder): (Box<dyn Read + Send>, Decoder) = match Compression::detect(&head[..n]) {
116 Compression::None => (
117 Box::new(BufReader::with_capacity(1 << 20, f)),
118 Decoder::Plain,
119 ),
120 Compression::Zstd => (zstd_reader(f, what)?, Decoder::Zstd),
121 Compression::Gzip if opts.parallel_gzip && opts.threads() > 1 => {
122 let threads = opts.threads();
123 let decoder = if len >= lgz::stream::SERIAL_BELOW as u64 {
126 Decoder::GzipParallel
127 } else {
128 Decoder::GzipStream
129 };
130 let (reader, hash) = lgz_stream(f, len, threads, digest);
131 let d = match hash {
132 Some(t) => TarDigest::Thread(t),
133 None => TarDigest::None,
134 };
135 return Ok((Box::new(reader), decoder, d));
136 }
137 Compression::Gzip => (
138 Box::new(flate2::read::MultiGzDecoder::new(BufReader::with_capacity(
139 1 << 20,
140 f,
141 ))),
142 Decoder::GzipStream,
143 ),
144 };
145 if !digest {
146 return Ok((reader, decoder, TarDigest::None));
147 }
148 let hasher = Arc::new(Mutex::new(Hasher::default()));
149 let reader = HashReader {
150 inner: reader,
151 hasher: hasher.clone(),
152 };
153 Ok((Box::new(reader), decoder, TarDigest::Inline(hasher)))
154}
155
156struct HashReader<R> {
157 inner: R,
158 hasher: Arc<Mutex<Hasher>>,
159}
160
161impl<R: Read> Read for HashReader<R> {
162 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
163 let n = self.inner.read(buf)?;
164 self.hasher.lock().unwrap().update(&buf[..n]);
165 Ok(n)
166 }
167}
168
169const PIECE: usize = 1 << 20;
171const PIECES: usize = 8;
173
174type Piece = io::Result<Arc<Vec<u8>>>;
175
176fn lgz_stream(
182 file: File,
183 len: u64,
184 threads: usize,
185 digest: bool,
186) -> (PieceReader, Option<thread::JoinHandle<Digest>>) {
187 let (tx, rx) = mpsc::sync_channel::<Piece>(PIECES);
188 let (hash_tx, hash) = if digest {
189 let (htx, hrx) = mpsc::sync_channel::<Arc<Vec<u8>>>(PIECES);
190 let t = thread::spawn(move || {
193 let mut h = Hasher::default();
194 for piece in hrx {
195 h.update(&piece);
196 }
197 h.finish()
198 });
199 (Some(htx), Some(t))
200 } else {
201 (None, None)
202 };
203 thread::spawn(move || {
206 let mut out = Tee {
207 buf: Vec::with_capacity(PIECE),
208 tx: tx.clone(),
209 hash: hash_tx,
210 };
211 let r = lgz::stream::decode_sized(file, &mut out, threads, Some(len));
212 let end = match r {
213 Ok(_) => out.flush(),
214 Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e.to_string())),
215 };
216 if let Err(e) = end {
217 let _ = tx.send(Err(e));
220 }
221 });
222 (
223 PieceReader {
224 rx,
225 cur: Arc::new(Vec::new()),
226 pos: 0,
227 },
228 hash,
229 )
230}
231
232struct Tee {
235 buf: Vec<u8>,
236 tx: SyncSender<Piece>,
237 hash: Option<SyncSender<Arc<Vec<u8>>>>,
238}
239
240impl Tee {
241 fn send(&mut self) -> io::Result<()> {
242 if self.buf.is_empty() {
243 return Ok(());
244 }
245 let piece = Arc::new(std::mem::replace(&mut self.buf, Vec::with_capacity(PIECE)));
246 let gone = || io::Error::new(io::ErrorKind::BrokenPipe, "the layer reader stopped");
247 if let Some(h) = &self.hash {
248 h.send(piece.clone()).map_err(|_| gone())?;
249 }
250 self.tx.send(Ok(piece)).map_err(|_| gone())
251 }
252}
253
254impl Write for Tee {
255 fn write(&mut self, b: &[u8]) -> io::Result<usize> {
256 let n = b.len().min(PIECE - self.buf.len());
257 self.buf.extend_from_slice(&b[..n]);
258 if self.buf.len() == PIECE {
259 self.send()?;
260 }
261 Ok(n)
262 }
263 fn flush(&mut self) -> io::Result<()> {
264 self.send()
265 }
266}
267
268pub(crate) struct PieceReader {
270 rx: Receiver<Piece>,
271 cur: Arc<Vec<u8>>,
272 pos: usize,
273}
274
275impl Read for PieceReader {
276 fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
277 while self.pos == self.cur.len() {
278 match self.rx.recv() {
279 Ok(piece) => {
280 self.cur = piece?;
281 self.pos = 0;
282 }
283 Err(_) => return Ok(0),
284 }
285 }
286 let n = out.len().min(self.cur.len() - self.pos);
287 out[..n].copy_from_slice(&self.cur[self.pos..self.pos + n]);
288 self.pos += n;
289 Ok(n)
290 }
291}
292
293fn read_up_to(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
294 let mut n = 0;
295 while n < buf.len() {
296 match r.read(&mut buf[n..])? {
297 0 => break,
298 k => n += k,
299 }
300 }
301 Ok(n)
302}