chapter_tgz/

lib.rs

//! [![github]](https://github.com/dtolnay/chapter-tgz) [![crates-io]](https://crates.io/crates/chapter-tgz) [![docs-rs]](https://docs.rs/chapter-tgz)
//!
//! [github]: https://img.shields.io/badge/github-8da0cb?style=for-the-badge&labelColor=555555&logo=github
//! [crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust
//! [docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs
//!
//! <br>
//!
//! This is a library for creating and consuming specially crafted .tar.gz files
//! with the following properties:
//!
//! 1. **Efficient access to specific predefined points in the tar** ("chapter
//!    boundaries"). A chapter consists of zero or more consecutive tar entries.
//!    This can be used to skip over groups of tar entries in _O(1)_ time
//!    without performing the work of gzip decompression on the intervening
//!    entries.
//!
//! 2. **Parallel decompression:** different chapters of the same tgz can be
//!    read simultaneously by different threads. Extracting a later entry is not
//!    stalled on processing all previous entries as in a conventional tgz file.
//!
//! 3. **Perfectly compatible with existing readers** that do not know about
//!    chapter information. All existing software will be able to read these
//!    files as ordinary tgz files. Chapter information is encoded in the form
//!    of valid empty gzip blocks with peculiar Huffman code alphabets.
//!
//! 4. **Perfectly compatible with existing writers** that do not embed chapter
//!    information. Tgz files without chapter information are handled as if
//!    there was a single chapter encompassing all of their entries.
//!
//! Refer to example code on [`TgzReader`] that demonstrates leveraging chapter
//! information to skip over entries or process chapters on a thread pool.

#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::doc_markdown,
    clippy::elidable_lifetime_names,
    clippy::missing_errors_doc,
    clippy::missing_panics_doc,
    clippy::must_use_candidate,
    clippy::unreadable_literal
)]

mod count;
mod decode;
mod encode;
mod independent;
pub mod io;
mod refmut;

use crate::count::Count;
use crate::decode::{decode_from_end, decode_from_start};
use crate::encode::encode;
use crate::io::{ChapterReader, ChapterReaderImpl, ChapterWriter};
use crate::refmut::RefMut;
use flate2::Crc;
use flate2::read::{DeflateDecoder, GzDecoder};
use flate2::write::DeflateEncoder;
use std::io::{Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::mem;

pub extern crate tar;

pub use crate::independent::IndependentRead;
#[doc(no_inline)]
pub use flate2::Compression;

/// Compressor for producing .tar.gz files with embedded chapter information.
///
/// Most of the API surface area is in the [`tar`] crate.
///
/// # Example
///
/// This example creates a tgz with several large files containing random data,
/// each in its own chapter.
///
/// ```no_run
/// use chapter_tgz::{Compression, TgzWriter};
/// use rand::RngReader;
/// use rand::rngs::SmallRng;
/// use std::fs;
/// use std::io::{self, Read as _};
///
/// fn main() -> io::Result<()> {
///     let mut tgz = TgzWriter::new(Vec::new(), Compression::fast());
///     let mut rng: SmallRng = rand::make_rng();
///     for i in 0..20 {
///         let mut chapter = tgz.create_chapter();
///         let mut header = tar::Header::new_gnu();
///         header.set_size(100_000_000);
///         let path = format!("random/{i}");
///         let data = RngReader(&mut rng).take(100_000_000);
///         chapter.append_data(&mut header, path, data)?;
///     }
///     let compressed = tgz.into_inner()?;
///     fs::write("example.tar.gz", compressed)?;
///     Ok(())
/// }
/// ```
pub struct TgzWriter<W: Write> {
    state: TgzWriterState<W>,
    crc: Crc,
    level: Compression,
    chapters: u32,
    boundaries_written: u32,
    last_boundary: u64,
}

enum TgzWriterState<W: Write> {
    Empty(W),
    Chapter(DeflateEncoder<Count<W>>),
    Finished(W),
    Failed(ErrorKind),
}

/// Decompressor for reading .tar.gz files containing embedded chapter
/// information.
///
/// Tgz files without chapter information will simply be observed to have a
/// single chapter encompassing their entire contents.
pub struct TgzReader<R> {
    data: R,
    boundaries: Vec<u64>,
    chapter: u32,
}

static GZIP_HEADER: [u8; 10] = [
    0x1F, // ID1 (magic number)
    0x8B, // ID2
    8,    // CM (compression method = Deflate)
    0,    // FLG (flags)
    0,    // MTIME 1/4
    0,    // MTIME 2/4
    0,    // MTIME 3/4
    0,    // MTIME 4/4
    0,    // XFL (extra flags = none)
    255,  // OS (filesystem = unknown)
];

static TAR_TERMINATION_SECTIONS: [u8; 1024] = [0; 1024];

impl<W> TgzWriter<W>
where
    W: Write,
{
    pub fn new(out: W, level: Compression) -> Self {
        TgzWriter {
            state: TgzWriterState::Empty(out),
            crc: Crc::new(),
            level,
            chapters: 0,
            boundaries_written: 0,
            last_boundary: 0,
        }
    }

    /// Begin writing the next chapters.
    ///
    /// A chapter may contain zero or more tar entries.
    ///
    /// The total number of chapters is limited to u32::MAX.
    ///
    /// # Panics
    ///
    /// Panics if 4,294,967,295 previous chapters have already been created in
    /// this tgz.
    #[track_caller]
    pub fn create_chapter(&mut self) -> tar::Builder<ChapterWriter<'_, W>> {
        if let TgzWriterState::Finished(_) = self.state {
            panic!("called create_chapter on a TgzWriter that is already finished");
        }

        let index = self.chapters;
        self.chapters = self.chapters.strict_add(1);

        tar::Builder::new(ChapterWriter {
            tgz: self,
            index,
            deferred_termination_sections: false,
        })
    }

    /// Number of chapters created so far.
    pub fn chapters(&self) -> u32 {
        self.chapters
    }

    /// Flush the last chapter information to the underlying writer.
    ///
    /// This is done automatically by `Drop`, but calling it explicitly allows
    /// any error to be handled. `Drop` would ignore the error.
    ///
    /// # Panics
    ///
    /// Panics if `finish` has already been called on this tgz.
    #[track_caller]
    pub fn finish(&mut self) -> Result<()> {
        if let TgzWriterState::Finished(_) = self.state {
            panic!("TgzWriter got finished twice");
        }

        self.do_finish()?;
        Ok(())
    }

    /// Flush the last chapter information to the underlying writer and return
    /// the writer object.
    ///
    /// It is not necessary to call `finish` before `into_inner`.
    pub fn into_inner(mut self) -> Result<W> {
        self.do_finish()?;
        let TgzWriterState::Finished(writer) =
            mem::replace(&mut self.state, TgzWriterState::Failed(ErrorKind::Other))
        else {
            unreachable!()
        };
        Ok(writer)
    }

    fn write_boundary(&mut self) -> Result<()> {
        match mem::replace(&mut self.state, TgzWriterState::Failed(ErrorKind::Other)) {
            TgzWriterState::Empty(mut writer) => {
                writer.write_all(&GZIP_HEADER)?;
                self.boundaries_written += 1;
                self.state =
                    TgzWriterState::Chapter(DeflateEncoder::new(Count::new(writer), self.level));
                Ok(())
            }
            TgzWriterState::Chapter(out) => {
                let mut writer = out.flush_finish()?;
                let bfinal = false;
                let payload = writer.position().strict_sub(self.last_boundary);
                self.last_boundary = writer.position();
                let boundary = encode(bfinal, payload);
                writer.write_all(boundary.as_slice())?;
                self.boundaries_written += 1;
                self.state = TgzWriterState::Chapter(DeflateEncoder::new(writer, self.level));
                Ok(())
            }
            TgzWriterState::Finished(_) => unreachable!(),
            TgzWriterState::Failed(kind) => Err(Error::new(kind, "TgzWriter error")),
        }
    }

    fn do_finish(&mut self) -> Result<()> {
        while self.boundaries_written == 0 || self.boundaries_written < self.chapters {
            self.write_boundary()?;
        }

        match mem::replace(&mut self.state, TgzWriterState::Failed(ErrorKind::Other)) {
            TgzWriterState::Empty(_) => unreachable!(),
            TgzWriterState::Chapter(mut out) => {
                out.write_all(&TAR_TERMINATION_SECTIONS)?;
                self.crc.update(&TAR_TERMINATION_SECTIONS);
                let writer = out.flush_finish()?;
                let bfinal = true;
                let mut payload = writer.position().strict_sub(self.last_boundary);
                if self.chapters == 0 {
                    payload += GZIP_HEADER.len() as u64;
                }
                let boundary = encode(bfinal, payload);
                let mut writer = writer.into_inner();
                writer.write_all(boundary.as_slice())?;
                writer.write_all(&self.crc.sum().to_le_bytes())?;
                writer.write_all(&self.crc.amount().to_le_bytes())?;
                writer.flush()?;
                self.state = TgzWriterState::Finished(writer);
                Ok(())
            }
            TgzWriterState::Finished(writer) => {
                self.state = TgzWriterState::Finished(writer);
                Ok(())
            }
            TgzWriterState::Failed(kind) => {
                self.state = TgzWriterState::Failed(kind);
                Err(Error::new(kind, "TgzWriter error"))
            }
        }
    }
}

impl<W> Drop for TgzWriter<W>
where
    W: Write,
{
    fn drop(&mut self) {
        let _ = self.do_finish();
    }
}

impl<'a, W> Write for ChapterWriter<'a, W>
where
    W: Write,
{
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }

        if self.deferred_termination_sections {
            if self.tgz.boundaries_written > self.index {
                match &mut self.tgz.state {
                    TgzWriterState::Chapter(out) => {
                        if let Err(err) = out.write_all(&TAR_TERMINATION_SECTIONS) {
                            self.tgz.state = TgzWriterState::Failed(err.kind());
                            return Err(err);
                        }
                    }
                    TgzWriterState::Empty(_) | TgzWriterState::Finished(_) => unreachable!(),
                    TgzWriterState::Failed(kind) => {
                        return Err(Error::new(*kind, "TgzWriter error"));
                    }
                }
            }
            self.deferred_termination_sections = false;
        }

        if buf == TAR_TERMINATION_SECTIONS {
            self.deferred_termination_sections = true;
            return Ok(buf.len());
        }

        while self.tgz.boundaries_written <= self.index {
            if let Err(err) = self.tgz.write_boundary() {
                self.tgz.state = TgzWriterState::Failed(err.kind());
                return Err(err);
            }
        }

        let n = match &mut self.tgz.state {
            TgzWriterState::Chapter(out) => out.write(buf)?,
            TgzWriterState::Empty(_) | TgzWriterState::Finished(_) => unreachable!(),
            TgzWriterState::Failed(kind) => {
                return Err(Error::new(*kind, "TgzWriter error"));
            }
        };

        self.tgz.crc.update(&buf[..n]);
        Ok(n)
    }

    fn flush(&mut self) -> Result<()> {
        match &mut self.tgz.state {
            TgzWriterState::Empty(_) => Ok(()),
            TgzWriterState::Chapter(out) => out.flush(),
            TgzWriterState::Finished(_) => unreachable!(),
            TgzWriterState::Failed(kind) => Err(Error::new(*kind, "TgzWriter error")),
        }
    }
}

impl<R> TgzReader<R>
where
    R: Read + Seek,
{
    /// Parse chapter information from a tgz.
    ///
    /// Time complexity: **O(_c_)** where _c_ is the number of chapters. In
    /// particular, the runtime is not sensitive to the number of entries within
    /// a chapter, the compressed size of chapter contents, or the uncompressed
    /// size of chapter contents.
    pub fn open(mut read: R) -> Result<Self> {
        let start = read.stream_position()?;

        let mut prefix_buf = [0u8; GZIP_HEADER.len()];
        let mut payload_buf = [0u8; encode::MAX_BYTES_FINAL];
        let mut boundaries = Vec::new();

        if read.read_exact(&mut prefix_buf).is_ok()
            && prefix_buf == GZIP_HEADER
            && let Ok(seek) = read.seek(SeekFrom::End(-(encode::MAX_BYTES_FINAL as i64 + 4)))
            && read.read_exact(&mut payload_buf).is_ok()
            && let Some((payload_start, mut payload)) = decode_from_end(&payload_buf)
            && {
                let mut boundary = seek.strict_add(payload_start as u64);
                boundaries.push(boundary);
                loop {
                    let Some(prev) = boundary.checked_sub(payload) else {
                        break false;
                    };
                    if prev == start {
                        break true;
                    }
                    boundaries.push(prev);
                    if prev == start.strict_add(GZIP_HEADER.len() as u64) {
                        break true;
                    }
                    if read.seek(SeekFrom::Start(prev)).is_ok()
                        && read
                            .read_exact(&mut payload_buf[..encode::MAX_BYTES_NONFINAL])
                            .is_ok()
                        && let Some(prev_payload) = decode_from_start(&payload_buf)
                    {
                        boundary = prev;
                        payload = prev_payload;
                    } else {
                        break false;
                    }
                }
            }
            && u32::try_from(boundaries.len()).is_ok()
        {
            boundaries.reverse();
        } else {
            let end = read.seek(SeekFrom::End(0))?;
            boundaries.clear();
            boundaries.push(0);
            boundaries.push(end.saturating_sub(start));
            read.seek(SeekFrom::Start(start))?;
        }

        Ok(TgzReader {
            data: read,
            boundaries,
            chapter: 0,
        })
    }

    /// Number of chapters in the tgz.
    ///
    /// Time complexity: **O(1)**
    ///
    /// Files without chapter information (i.e. not produced by this crate's
    /// TgzWriter) have 1 chapter.
    pub fn chapters(&self) -> u32 {
        self.boundaries.len() as u32 - 1
    }

    /// Sequentially begin reading a single chapter. Returns `None` if there is
    /// no next chapter.
    ///
    /// Time complexity: **O(1)**
    ///
    /// # Example
    ///
    /// This example uses a `while let` loop to list the contents of every
    /// chapter in a tgz.
    ///
    /// ```no_run
    /// use chapter_tgz::TgzReader;
    /// use std::fs::File;
    /// use std::io;
    ///
    /// fn main() -> io::Result<()> {
    ///     let file = File::open("example.tar.gz")?;
    ///     let mut tgz = TgzReader::open(file)?;
    ///     let mut i = 0;
    ///     while let Some(mut chapter) = tgz.next_chapter() {
    ///         println!("Chapter {i}:");
    ///         for entry in chapter.entries()? {
    ///             let entry = entry?;
    ///             let path = entry.path()?;
    ///             println!("  {}", path.display());
    ///         }
    ///         i += 1;
    ///     }
    ///     Ok(())
    /// }
    /// ```
    pub fn next_chapter(&mut self) -> Option<tar::Archive<ChapterReader<'_, R>>> {
        if self.boundaries[0] == 0 {
            if self.chapter != 0 {
                return None;
            }
            self.chapter = 1;
            Some(tar::Archive::new(ChapterReader {
                imp: ChapterReaderImpl::Dumb(GzDecoder::new(RefMut::Borrowed(&mut self.data))),
            }))
        } else {
            if self.chapter == self.boundaries.len() as u32 - 1 {
                return None;
            }
            let index = self.chapter;
            let begin = self.boundaries[index as usize];
            let end = self.boundaries[index as usize + 1];
            self.chapter += 1;
            Some(tar::Archive::new(ChapterReader {
                imp: ChapterReaderImpl::SeekTo(begin..end, RefMut::Borrowed(&mut self.data)),
            }))
        }
    }

    /// Begin reading a single chapter by index.
    ///
    /// Time complexity: **O(1)**
    ///
    /// As a side effect, this also affects which chapter `next_chapter` will
    /// pick up next. For example if you call `tgz.next_chapter()` followed by
    /// `jump_to_chapter(10)` followed by `tgz.next_chapter()`, the three
    /// chapters you get will be chapters 0, 10, 11.
    ///
    /// Chapters are indexed starting from 0.
    ///
    /// # Panics
    ///
    /// Requires `i < self.chapters()`.
    #[track_caller]
    pub fn jump_to_chapter(&mut self, i: u32) -> tar::Archive<ChapterReader<'_, R>> {
        assert!(i < self.chapters());
        self.chapter = i;
        self.next_chapter().unwrap()
    }

    /// Begin reading a single chapter in a way that can be parallelized with
    /// other reads from the same tgz.
    ///
    /// Time complexity: **O(1)**
    ///
    /// This does not affect the next chapter read by `next_chapter`, which will
    /// continue with whichever would have been the next chapter if
    /// `independent_read_chapter` had not been called.
    ///
    /// Note the `IndependentRead` trait bound, which enforces that the
    /// underlying reader `R` given to `TgzReader::open` can be cloned, and that
    /// the original and clone will each have a position that can read and seek
    /// independently of the other. This is not the case, for example, for
    /// `File`. While [`File::try_clone`] exists, it is documented that _"Reads,
    /// writes, and seeks will affect both `File` instances simultaneously."_
    ///
    /// [`File::try_clone`]: std::fs::File::try_clone
    ///
    /// # Panics
    ///
    /// Requires `i < self.chapters()`.
    ///
    /// # Example
    ///
    /// This example demonstrates using `independent_read_chapter` to decompress
    /// all chapters in parallel.
    ///
    /// A conveniently runnable form of this example is provided in this crate's
    /// *examples/* directory, also with a progress bar that shows overall
    /// progress.
    ///
    /// ```no_run
    /// use chapter_tgz::TgzReader;
    /// use std::fs;
    /// use std::io::{self, Cursor};
    /// use std::thread;
    ///
    /// fn main() -> io::Result<()> {
    ///     let data = fs::read("example.tar.gz")?;
    ///     let tgz = TgzReader::open(Cursor::new(data))?;
    ///
    ///     let n = tgz.chapters();
    ///     let mut chapters = Vec::with_capacity(n as usize);
    ///     for i in 0..n {
    ///         let mut chapter = tgz.independent_read_chapter(i)?;
    ///         if i % 2 == 0 {
    ///             // Option 1: We can directly enqueue chapters into the thread pool.
    ///             chapters.push(chapter);
    ///         } else if let Some(first_entry) = chapter.entries()?.next()
    ///             && let Some(file_name) = first_entry?.path()?.file_name()
    ///             && let Some(file_name_str) = file_name.to_str()
    ///             && !file_name_str.starts_with("__")
    ///         {
    ///             // Option 2: We can examine chapter entries to decide whether to
    ///             // process or skip a chapter. Reading the first entry's tar header
    ///             // (path, size, PAX extensions) from a chapter is fast.
    ///             chapters.push(tgz.independent_read_chapter(i)?);
    ///         } else {
    ///             // Option 3: Also fine to pass i and call independent_read_chapter
    ///             // on the other thread.
    ///         }
    ///     }
    ///
    ///     thread::scope(|scope| {
    ///         for mut chapter in chapters {
    ///             scope.spawn(move || {
    ///                 if let Err(err) = (|| -> io::Result<()> {
    ///                     for _entry in chapter.entries()? {}
    ///                     Ok(())
    ///                 })() {
    ///                     eprintln!("Error: {err}");
    ///                 }
    ///             });
    ///         }
    ///     });
    ///
    ///     Ok(())
    /// }
    /// ```
    #[track_caller]
    pub fn independent_read_chapter<'a>(&self, i: u32) -> Result<tar::Archive<ChapterReader<'a, R>>>
    where
        R: IndependentRead + 'a,
    {
        assert!(i < self.chapters());
        let data = self.data.independent_clone()?;
        if self.boundaries[0] == 0 {
            Ok(tar::Archive::new(ChapterReader {
                imp: ChapterReaderImpl::Dumb(GzDecoder::new(RefMut::Owned(Box::new(data)))),
            }))
        } else {
            let begin = self.boundaries[i as usize];
            let end = self.boundaries[i as usize + 1];
            Ok(tar::Archive::new(ChapterReader {
                imp: ChapterReaderImpl::SeekTo(begin..end, RefMut::Owned(Box::new(data))),
            }))
        }
    }

    /// Approximate compressed size of a chapter: the distance in bytes between
    /// the start and end of the chapter.
    ///
    /// Time complexity: **O(1)**
    ///
    /// This is primarily intended to enable progress reporting through the use
    /// of a `Read` implementation that monitors reads performed against the
    /// underlying compressed archive. Outside of this use case, the returned
    /// quantity is not guaranteed to be meaningful.
    ///
    /// # Example
    ///
    /// This example demonstrates using `compressed_size_of_chapter` to size an
    /// [Indicatif] progress bar that increments as bytes are read.
    ///
    /// [Indicatif]: https://crates.io/crates/indicatif
    ///
    /// The `ProgressRead` type is the one shown in the documentation of
    /// [`IndependentRead`].
    ///
    /// A conveniently runnable multithreaded form of this example is provided
    /// in this crate's *examples/* directory.
    ///
    /// ```no_run
    /// use chapter_tgz::TgzReader;
    /// use indicatif::ProgressBar;
    /// use std::fs::File;
    /// use std::io::{self, Cursor};
    ///
    /// fn main() -> io::Result<()> {
    ///     let file = File::open("example.tar.gz")?;
    ///     let pb = ProgressBar::no_length();
    ///     let mut tgz = TgzReader::open(ProgressRead {
    ///         inner: file,
    ///         progress: &pb,
    ///     })?;
    ///
    ///     let n = tgz.chapters();
    ///     let mut total_compressed_size = 0;
    ///     for i in 0..n {
    ///         total_compressed_size += tgz.compressed_size_of_chapter(i);
    ///     }
    ///
    ///     pb.reset();
    ///     pb.set_length(total_compressed_size);
    ///
    ///     for i in 0..n {
    ///         for _entry in tgz.jump_to_chapter(i).entries()? {
    ///             // ...
    ///         }
    ///     }
    ///     Ok(())
    /// }
    /// #
    /// # use chapter_tgz::IndependentRead;
    /// # use std::io::{Read, Seek, SeekFrom};
    /// #
    /// # struct ProgressRead<'a, R> {
    /// #     inner: R,
    /// #     progress: &'a ProgressBar,
    /// # }
    /// #
    /// # impl<R: Read> Read for ProgressRead<'_, R> {
    /// #     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
    /// #         let n = self.inner.read(buf)?;
    /// #         self.progress.inc(n as u64);
    /// #         Ok(n)
    /// #     }
    /// # }
    /// #
    /// # impl<R: Seek> Seek for ProgressRead<'_, R> {
    /// #     fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
    /// #         self.inner.seek(pos)
    /// #     }
    /// # }
    /// #
    /// # impl<R: IndependentRead> IndependentRead for ProgressRead<'_, R> {
    /// #     fn independent_clone(&self) -> io::Result<Self> {
    /// #         Ok(ProgressRead {
    /// #             inner: self.inner.independent_clone()?,
    /// #             progress: self.progress,
    /// #         })
    /// #     }
    /// # }
    /// ```
    #[track_caller]
    pub fn compressed_size_of_chapter(&self, i: u32) -> u64 {
        assert!(i < self.chapters());
        let begin = self.boundaries[i as usize];
        let end = self.boundaries[i as usize + 1];
        end - begin
    }
}

impl<'a, R> Read for ChapterReader<'a, R>
where
    R: Read + Seek,
{
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        match &mut self.imp {
            imp @ ChapterReaderImpl::SeekTo(..) => {
                let ChapterReaderImpl::SeekTo(range, mut data) =
                    mem::replace(imp, ChapterReaderImpl::Failed(ErrorKind::Other))
                else {
                    unreachable!()
                };
                data.seek(SeekFrom::Start(range.start))?;
                let mut decoder = DeflateDecoder::new(data.take(range.end - range.start));
                let result = decoder.read(buf);
                *imp = ChapterReaderImpl::Smart(decoder);
                result
            }
            ChapterReaderImpl::Smart(data) => data.read(buf),
            ChapterReaderImpl::Dumb(gz) => gz.read(buf),
            ChapterReaderImpl::Failed(kind) => Err(Error::new(*kind, "TgzReader error")),
        }
    }
}