Skip to main content

mc_repack_core/entry/
zip.rs

1use std::io::{BufReader, BufWriter, Cursor, Read, Seek, Write};
2use bytes::Bytes;
3use zip::{write::{FileOptions, SimpleFileOptions}, CompressionMethod, ZipArchive, ZipWriter};
4
5use crate::Result_;
6use super::{EntryReader, EntrySaver, ReadEntry, SavingEntry};
7
8/// An entry reader implementation for ZIP archive. It reads its contents from a provided reader (with seeking).
9pub struct ZipEntryReader<R: Read + Seek> {
10    za: ZipArchive<R>,
11    cur: usize
12}
13impl <R: Read + Seek> ZipEntryReader<R> {
14    /// Creates an entry reader with a specified reader.
15    pub fn new(r: R) -> Result_<Self> {
16        Ok(Self { za: ZipArchive::new(r)?, cur: 0 })
17    }
18}
19impl <R: Read + Seek> ZipEntryReader<BufReader<R>> {
20    /// Creates an entry reader wrapping a specified reader with a [`BufReader`].
21    #[inline]
22    pub fn new_buf(r: R) -> Result_<Self> {
23        Self::new(BufReader::new(r))
24    }
25}
26impl <T: AsRef<[u8]>> ZipEntryReader<Cursor<T>> {
27    /// Creates an entry reader wrapping a specified reader with a [`Cursor`].
28    #[inline]
29    pub fn new_mem(t: T) -> Result_<Self> {
30        Self::new(Cursor::new(t))
31    }
32}
33impl <R: Read + Seek> EntryReader for ZipEntryReader<R> {
34    type RE<'a> = ReadZipFileEntry<'a, R> where R: 'a;
35    fn read_next(&mut self) -> Option<Self::RE<'_>> {
36        let za = &mut self.za;
37        let jfc = za.len();
38        if self.cur >= jfc {
39            None
40        } else {
41            let idx = self.cur;
42            self.cur += 1;
43            Some(ReadZipFileEntry { zip: za, idx })
44        }
45    }
46    #[inline]
47    fn read_len(&self) -> usize {
48        self.za.len()
49    }
50}
51
52/// A read entry of a ZIP archive.
53pub struct ReadZipFileEntry<'a, RS: Read + Seek> {
54    zip: &'a mut ZipArchive<RS>,
55    idx: usize
56}
57impl <'a, RS: Read + Seek> ReadEntry for ReadZipFileEntry<'a, RS> {
58    fn meta(&self) -> (Option<bool>, Box<str>) {
59        self.zip.name_for_index(self.idx).map_or_else(
60            || (None, "".into()),
61            |n| (Some(n.ends_with('/')), n.into())
62        )
63    }
64    fn data(self) -> crate::Result_<Bytes> {
65        let mut obuf = Vec::new();
66        let mut jf = self.zip.by_index(self.idx)?;
67        obuf.reserve_exact(jf.size() as usize);
68        jf.read_to_end(&mut obuf)?;
69        Ok(obuf.into())
70    }
71}
72
73#[cfg(feature = "zip-zopfli")]
74const MAX_LEVEL: i64 = 24;
75
76#[cfg(not(feature = "zip-zopfli"))]
77const MAX_LEVEL: i64 = 9;
78
79/// An entry saver implementation for ZIP archive. It writes entries to it using a provided writer.
80pub struct ZipEntrySaver<W: Write + Seek> {
81    w: ZipWriter<BufWriter<W>>,
82    keep_dirs: bool,
83    opts_deflated: SimpleFileOptions,
84    opts_stored: SimpleFileOptions
85}
86impl <W: Write + Seek> ZipEntrySaver<W> {
87    /// Creates an entry saver with a seekable writer.
88    pub fn new(w: W, keep_dirs: bool) -> Self {
89        Self {
90            w: ZipWriter::new(BufWriter::new(w)),
91            keep_dirs,
92            opts_deflated: FileOptions::default().compression_method(CompressionMethod::Deflated).compression_level(Some(MAX_LEVEL)),
93            opts_stored: FileOptions::default().compression_method(CompressionMethod::Stored),
94        }
95    }
96    /// Creates an entry saver with custom file options for ZIP archive and a seekable writer.
97    pub fn custom(w: W, keep_dirs: bool, opts_stored: SimpleFileOptions, opts_deflated: SimpleFileOptions) -> Self {
98        Self {
99            w: ZipWriter::new(BufWriter::new(w)), keep_dirs, opts_deflated, opts_stored
100        }
101    }
102    /// Creates an entry saver with custom compression level for deflated entries of ZIP archive and a seekable writer.
103    pub fn custom_compress(w: W, keep_dirs: bool, compress: impl Into<i64>) -> Self {
104        Self {
105            w: ZipWriter::new(BufWriter::new(w)),
106            keep_dirs,
107            opts_deflated: FileOptions::default().compression_method(CompressionMethod::Deflated).compression_level(Some(compress.into())),
108            opts_stored: FileOptions::default().compression_method(CompressionMethod::Stored),
109        }
110    }
111}
112impl <W: Write + Seek> EntrySaver for ZipEntrySaver<W> {
113    fn save(&mut self, name: &str, entry: SavingEntry) -> crate::Result_<()> {
114        let z = &mut self.w;
115        match entry {
116            SavingEntry::Directory => {
117                if self.keep_dirs && name != ".cache/" {
118                    z.add_directory(name, self.opts_stored)?;
119                }
120            }
121            SavingEntry::File(data, compress_min) => {
122                z.start_file(name, if compress_check(data, compress_min as usize) {
123                    self.opts_deflated
124                } else {
125                    self.opts_stored
126                })?;
127                z.write_all(data)?;
128            }
129        }
130        Ok(())
131    }
132}
133
134/// Check if data should be compressed. If the compressed size is smaller than original, then the compression should be chosen.
135pub fn compress_check(b: &[u8], compress_min: usize) -> bool {
136    let lb = b.len();
137    if lb > compress_min {
138        if calc_entropy(b) < 7.0 { return true }
139        let mut d = flate2::write::DeflateEncoder::new(std::io::sink(), flate2::Compression::best());
140        if d.write_all(b).and_then(|_| d.try_finish()).is_ok() && d.total_out() as usize + 8 < lb { return true }
141    }
142    false
143}
144
145fn calc_entropy(b: &[u8]) -> f32 {
146    if b.is_empty() { return 0.0; }
147    let mut freq = [0usize; 256];
148    for &b in b { freq[b as usize] += 1; }
149    let total = b.len() as f32;
150    let logt = total.log2();
151    let e = freq.into_iter().filter(|&f| f != 0)
152        .map(|f| -(f as f32) * ((f as f32).log2() - logt))
153        .sum::<f32>() / total;
154    assert!((0.0..=8.0).contains(&e), "Invalid entropy: {}", e);
155    e
156}