tsaur_core/
container.rs

1//! Container awareness for ZIP-based formats (DOCX/XLSX/PPTX/EPUB/JAR/ODT/plain ZIP).
2//!
3//! The container is parsed with a small, strict ZIP framing parser that keeps the *exact* bytes
4//! of every local header, central-directory record, data descriptor and the end-of-central-directory
5//! block. Member payloads are handled by kind:
6//!
7//! * `stored` — the payload is the member's plaintext;
8//! * `deflate` — the raw deflate stream is inverted with **preflate** (Microsoft, Apache-2.0),
9//!   which yields the plaintext plus a small "corrections" blob that recreates the original stream
10//!   bit-exact whatever encoder produced it (zlib, zlib-ng, miniz, .NET, ...);
11//! * `raw` — anything else (unknown methods, encrypted members, streams preflate cannot model):
12//!   the compressed bytes themselves are kept.
13//!
14//! The plaintext/raw payloads are what T-saur chunks, deduplicates and compresses. Rebuilding is
15//! deterministic and is verified against the original hash at pack time; if the framing has gaps,
16//! ZIP64 records or anything else we do not model, the whole file is stored `raw` instead.
17
18use crate::error::{Error, Result};
19use crate::manifest::{StreamRecipe, StreamSegment, ZipMemberMeta, ZipRecipe};
20use preflate_rs::PreflateConfig;
21use serde_bytes::ByteBuf;
22
23pub const ZIP_EXTENSIONS: &[&str] = &["zip", "docx", "xlsx", "pptx", "epub", "jar", "odt", "ods", "odp", "xpi", "whl"];
24
25pub const KIND_STORED: u8 = 0;
26pub const KIND_DEFLATE: u8 = 1;
27pub const KIND_RAW: u8 = 2;
28
29const SIG_LOCAL: u32 = 0x0403_4b50;
30const SIG_CENTRAL: u32 = 0x0201_4b50;
31const SIG_EOCD: u32 = 0x0605_4b50;
32const SIG_ZIP64_LOCATOR: u32 = 0x0706_4b50;
33const SIG_DESCRIPTOR: u32 = 0x0807_4b50;
34
35/// Deflate streams above this plaintext size are kept raw (bounded memory, anti-bomb).
36pub const PREFLATE_PLAINTEXT_LIMIT: usize = 256 << 20;
37
38pub fn is_zip_like(name: &str) -> bool {
39    let ext = name.rsplit('.').next().unwrap_or("").to_ascii_lowercase();
40    ZIP_EXTENSIONS.contains(&ext.as_str())
41}
42
43fn u16_at(b: &[u8], i: usize) -> u16 {
44    u16::from_le_bytes([b[i], b[i + 1]])
45}
46
47fn u32_at(b: &[u8], i: usize) -> u32 {
48    u32::from_le_bytes([b[i], b[i + 1], b[i + 2], b[i + 3]])
49}
50
51struct CentralRecord {
52    bytes_start: usize,
53    bytes_end: usize,
54    flags: u16,
55    method: u16,
56    csize: u32,
57    usize_: u32,
58    local_off: u32,
59    name: String,
60}
61
62/// Parse the ZIP framing and invert the member payloads. Returns None for anything we do not
63/// model exactly (the caller then stores the file raw).
64pub fn explode(data: &[u8]) -> Option<(ZipRecipe, Vec<Vec<u8>>)> {
65    if data.len() < 22 {
66        return None;
67    }
68    // 1. End of central directory (search backwards, max comment 65535).
69    let min_start = data.len().saturating_sub(22 + 65_535);
70    let mut eocd_off = None;
71    let mut i = data.len() - 22;
72    loop {
73        if u32_at(data, i) == SIG_EOCD {
74            let clen = u16_at(data, i + 20) as usize;
75            if i + 22 + clen == data.len() {
76                eocd_off = Some(i);
77                break;
78            }
79        }
80        if i == min_start {
81            break;
82        }
83        i -= 1;
84    }
85    let eocd_off = eocd_off?;
86    if eocd_off >= 20 && u32_at(data, eocd_off - 20) == SIG_ZIP64_LOCATOR {
87        return None; // ZIP64 is not modelled yet
88    }
89    let disk = u16_at(data, eocd_off + 4);
90    let cd_disk = u16_at(data, eocd_off + 6);
91    let entries = u16_at(data, eocd_off + 10) as usize;
92    let cd_size = u32_at(data, eocd_off + 12) as usize;
93    let cd_off = u32_at(data, eocd_off + 16) as usize;
94    if disk != 0 || cd_disk != 0 || entries == 0xFFFF || cd_off + cd_size != eocd_off {
95        return None;
96    }
97
98    // 2. Central directory records (exact bytes kept).
99    let mut central = Vec::with_capacity(entries);
100    let mut p = cd_off;
101    for _ in 0..entries {
102        if p + 46 > eocd_off || u32_at(data, p) != SIG_CENTRAL {
103            return None;
104        }
105        let flags = u16_at(data, p + 8);
106        let method = u16_at(data, p + 10);
107        let csize = u32_at(data, p + 20);
108        let usize_ = u32_at(data, p + 24);
109        let nlen = u16_at(data, p + 28) as usize;
110        let elen = u16_at(data, p + 30) as usize;
111        let clen = u16_at(data, p + 32) as usize;
112        let local_off = u32_at(data, p + 42);
113        let end = p + 46 + nlen + elen + clen;
114        if end > eocd_off || csize == u32::MAX || usize_ == u32::MAX || local_off == u32::MAX {
115            return None;
116        }
117        let name = String::from_utf8_lossy(&data[p + 46..p + 46 + nlen]).to_string();
118        central.push(CentralRecord { bytes_start: p, bytes_end: end, flags, method, csize, usize_, local_off, name });
119        p = end;
120    }
121    if p != eocd_off {
122        return None;
123    }
124
125    // 3. Local headers + payloads, in file order; the layout must be contiguous.
126    let mut order: Vec<usize> = (0..central.len()).collect();
127    order.sort_by_key(|&k| central[k].local_off);
128    let mut members: Vec<Option<ZipMemberMeta>> = (0..central.len()).map(|_| None).collect();
129    let mut parts: Vec<Option<Vec<u8>>> = (0..central.len()).map(|_| None).collect();
130    let mut cursor = 0usize;
131    let config = PreflateConfig { plain_text_limit: PREFLATE_PLAINTEXT_LIMIT, verify_compression: true, ..PreflateConfig::default() };
132    for &k in &order {
133        let c = &central[k];
134        let lo = c.local_off as usize;
135        if lo != cursor || lo + 30 > cd_off || u32_at(data, lo) != SIG_LOCAL {
136            return None;
137        }
138        let nlen = u16_at(data, lo + 26) as usize;
139        let elen = u16_at(data, lo + 28) as usize;
140        let data_start = lo + 30 + nlen + elen;
141        let data_end = data_start + c.csize as usize;
142        if data_end > cd_off {
143            return None;
144        }
145        let payload = &data[data_start..data_end];
146        // optional data descriptor (general purpose bit 3)
147        let mut desc_end = data_end;
148        if c.flags & 0x0008 != 0 {
149            if data_end + 16 <= cd_off && u32_at(data, data_end) == SIG_DESCRIPTOR {
150                desc_end = data_end + 16;
151            } else if data_end + 12 <= cd_off {
152                desc_end = data_end + 12;
153            } else {
154                return None;
155            }
156        }
157        let encrypted = c.flags & 0x0001 != 0;
158        let (kind, plain, corrections) = if encrypted {
159            (KIND_RAW, payload.to_vec(), Vec::new())
160        } else {
161            match c.method {
162                0 => (KIND_STORED, payload.to_vec(), Vec::new()),
163                8 => match preflate_rs::preflate_whole_deflate_stream(payload, &config) {
164                    Ok((res, text)) if res.compressed_size == payload.len() && text.text().len() == c.usize_ as usize => (KIND_DEFLATE, text.text().to_vec(), res.corrections),
165                    _ => (KIND_RAW, payload.to_vec(), Vec::new()),
166                },
167                _ => (KIND_RAW, payload.to_vec(), Vec::new()),
168            }
169        };
170        members[k] = Some(ZipMemberMeta {
171            name: c.name.clone(),
172            kind,
173            local_header: ByteBuf::from(data[lo..data_start].to_vec()),
174            descriptor: ByteBuf::from(data[data_end..desc_end].to_vec()),
175            central: ByteBuf::from(data[c.bytes_start..c.bytes_end].to_vec()),
176            corrections: ByteBuf::from(corrections),
177            size: plain.len() as u64,
178            chunks: Vec::new(),
179        });
180        parts[k] = Some(plain);
181        cursor = desc_end;
182    }
183    if cursor != cd_off {
184        return None; // gap between the last member and the central directory
185    }
186    let members: Vec<ZipMemberMeta> = members.into_iter().map(|m| m.unwrap()).collect();
187    let parts: Vec<Vec<u8>> = parts.into_iter().map(|p| p.unwrap()).collect();
188    Some((ZipRecipe { local_order: order.iter().map(|&k| k as u32).collect(), eocd: ByteBuf::from(data[eocd_off..].to_vec()), members }, parts))
189}
190
191/// Rebuild the container from the recipe and the member payloads (same order as `recipe.members`).
192pub fn rebuild(recipe: &ZipRecipe, parts: &[Vec<u8>]) -> Result<Vec<u8>> {
193    if recipe.members.len() != parts.len() || recipe.local_order.len() != parts.len() {
194        return Err(Error::Corrupt("zip recipe/member count mismatch".into()));
195    }
196    let mut out = Vec::new();
197    for &k in &recipe.local_order {
198        let m = recipe.members.get(k as usize).ok_or_else(|| Error::Corrupt("zip local order out of range".into()))?;
199        out.extend_from_slice(&m.local_header);
200        match m.kind {
201            KIND_STORED | KIND_RAW => out.extend_from_slice(&parts[k as usize]),
202            KIND_DEFLATE => {
203                let stream = preflate_rs::recreate_whole_deflate_stream(&parts[k as usize], &m.corrections).map_err(|e| Error::Corrupt(format!("preflate recreate {}: {e}", m.name)))?;
204                out.extend_from_slice(&stream);
205            }
206            other => return Err(Error::Corrupt(format!("unknown zip member kind {other}"))),
207        }
208        out.extend_from_slice(&m.descriptor);
209    }
210    for m in &recipe.members {
211        out.extend_from_slice(&m.central);
212    }
213    out.extend_from_slice(&recipe.eocd);
214    Ok(out)
215}
216
217/// Explode and prove that the rebuild reproduces the exact input bytes.
218pub fn try_bitexact(data: &[u8]) -> Option<(ZipRecipe, Vec<Vec<u8>>)> {
219    let (recipe, parts) = explode(data)?;
220    match rebuild(&recipe, &parts) {
221        Ok(out) if out == data => Some((recipe, parts)),
222        _ => None,
223    }
224}
225
226// ----------------------------------------------------------------------------------------------
227// Embedded zlib streams (PDF FlateDecode objects). Precomp-style: every `stream ... endstream`
228// payload that starts with a valid zlib header is inverted with preflate; the Adler-32 trailer
229// and the 2-byte header are kept so the rebuild is bit-exact. Everything else is literal.
230// ----------------------------------------------------------------------------------------------
231
232pub const STREAM_EXTENSIONS: &[&str] = &["pdf", "jpg", "jpeg", "jpe"];
233
234pub fn is_stream_container(name: &str) -> bool {
235    let ext = name.rsplit('.').next().unwrap_or("").to_ascii_lowercase();
236    STREAM_EXTENSIONS.contains(&ext.as_str())
237}
238
239pub const SEG_LITERAL: u8 = 0;
240pub const SEG_ZLIB: u8 = 1;
241/// JPEG payload recompressed losslessly with Lepton (the chunked payload is the Lepton stream).
242pub const SEG_JPEG: u8 = 2;
243
244/// JPEG streams above this size are kept literal (bounded memory/time).
245pub const LEPTON_MAX_INPUT: usize = 64 << 20;
246
247#[cfg(feature = "lepton")]
248fn lepton_features() -> lepton_jpeg::EnabledFeatures {
249    lepton_jpeg::EnabledFeatures::compat_lepton_vector_write()
250}
251
252/// Lepton-encode a JPEG; the library verifies the round trip itself. Single-threaded so the
253/// output is deterministic whatever the machine.
254#[cfg(feature = "lepton")]
255pub fn jpeg_to_lepton(jpeg: &[u8]) -> Option<Vec<u8>> {
256    if jpeg.len() < 4 || jpeg.len() > LEPTON_MAX_INPUT || jpeg[..3] != [0xFF, 0xD8, 0xFF] {
257        return None;
258    }
259    let pool = lepton_jpeg::SingleThreadPool {};
260    match lepton_jpeg::catch_unwind_result(|| lepton_jpeg::encode_lepton_verify(jpeg, &lepton_features(), &pool)) {
261        Ok((out, _metrics)) if out.len() < jpeg.len() => Some(out),
262        _ => None,
263    }
264}
265
266#[cfg(feature = "lepton")]
267pub fn lepton_to_jpeg(lepton: &[u8], expected_len: u64) -> Result<Vec<u8>> {
268    let pool = lepton_jpeg::SingleThreadPool {};
269    let mut out = Vec::with_capacity(expected_len as usize);
270    let mut reader = std::io::Cursor::new(lepton);
271    lepton_jpeg::catch_unwind_result(|| lepton_jpeg::decode_lepton(&mut reader, &mut out, &lepton_features(), &pool)).map_err(|e| Error::Corrupt(format!("lepton decode: {e}")))?;
272    Ok(out)
273}
274
275/// Without the `lepton` feature JPEGs are stored as they are (still bit-exact, just not smaller).
276#[cfg(not(feature = "lepton"))]
277pub fn jpeg_to_lepton(_jpeg: &[u8]) -> Option<Vec<u8>> {
278    None
279}
280
281#[cfg(not(feature = "lepton"))]
282pub fn lepton_to_jpeg(_lepton: &[u8], _expected_len: u64) -> Result<Vec<u8>> {
283    Err(Error::Invalid("this build has no Lepton support (Cargo feature `lepton`): JPEG segments of this archive cannot be decoded".into()))
284}
285
286fn adler32(data: &[u8]) -> u32 {
287    let (mut a, mut b) = (1u32, 0u32);
288    for chunk in data.chunks(5552) {
289        for &x in chunk {
290            a += x as u32;
291            b += a;
292        }
293        a %= 65521;
294        b %= 65521;
295    }
296    (b << 16) | a
297}
298
299fn find(hay: &[u8], needle: &[u8], from: usize) -> Option<usize> {
300    if needle.is_empty() || from >= hay.len() {
301        return None;
302    }
303    hay[from..].windows(needle.len()).position(|w| w == needle).map(|p| p + from)
304}
305
306fn literal_segment(bytes: &[u8]) -> StreamSegment {
307    StreamSegment { kind: SEG_LITERAL, header: ByteBuf::new(), trailer: ByteBuf::new(), corrections: ByteBuf::new(), size: bytes.len() as u64, chunks: Vec::new() }
308}
309
310/// Split a PDF (or similar) into literal segments, inverted zlib streams and, when `jpeg` is
311/// set and the build has Lepton, recompressed JPEG streams. A file that is itself a JPEG becomes
312/// a single JPEG segment (or, with `jpeg` off, no recipe at all: it is stored as it is, which
313/// every build can read).
314pub fn explode_streams(data: &[u8], jpeg: bool) -> Option<(StreamRecipe, Vec<Vec<u8>>)> {
315    if data.len() >= 4 && data[..3] == [0xFF, 0xD8, 0xFF] {
316        if !jpeg {
317            return None;
318        }
319        let lepton = jpeg_to_lepton(data)?;
320        let seg = StreamSegment {
321            kind: SEG_JPEG,
322            header: ByteBuf::new(),
323            trailer: ByteBuf::from((data.len() as u64).to_le_bytes().to_vec()),
324            corrections: ByteBuf::new(),
325            size: lepton.len() as u64,
326            chunks: Vec::new(),
327        };
328        return Some((StreamRecipe { segments: vec![seg] }, vec![lepton]));
329    }
330    let config = PreflateConfig { plain_text_limit: PREFLATE_PLAINTEXT_LIMIT, verify_compression: true, ..PreflateConfig::default() };
331    let mut segments = Vec::new();
332    let mut parts = Vec::new();
333    let mut lit_start = 0usize;
334    let mut pos = 0usize;
335    let mut inverted = 0usize;
336    while let Some(s) = find(data, b"stream", pos) {
337        // must be the keyword `stream` followed by an end-of-line, not `endstream`
338        let mut data_start = s + 6;
339        if s >= 3 && &data[s - 3..s] == b"end" {
340            pos = s + 6;
341            continue;
342        }
343        if data.get(data_start) == Some(&b'\r') {
344            data_start += 1;
345        }
346        if data.get(data_start) == Some(&b'\n') {
347            data_start += 1;
348        } else {
349            pos = s + 6;
350            continue;
351        }
352        let Some(end) = find(data, b"endstream", data_start) else { break };
353        let payload = &data[data_start..end];
354        // the payload usually ends with an EOL before `endstream`; the JPEG/zlib detectors compute the exact length
355        let mut consumed: Option<(usize, StreamSegment, Vec<u8>)> = None;
356        if payload.len() > 6 && payload[0] & 0x0F == 8 && ((payload[0] as u16) << 8 | payload[1] as u16).is_multiple_of(31) && payload[1] & 0x20 == 0 {
357            if let Ok((res, text)) = preflate_rs::preflate_whole_deflate_stream(&payload[2..], &config) {
358                let zlen = 2 + res.compressed_size + 4;
359                if zlen <= payload.len() {
360                    let trailer = &payload[2 + res.compressed_size..zlen];
361                    if u32::from_be_bytes([trailer[0], trailer[1], trailer[2], trailer[3]]) == adler32(text.text()) {
362                        let seg = StreamSegment {
363                            kind: SEG_ZLIB,
364                            header: ByteBuf::from(payload[..2].to_vec()),
365                            trailer: ByteBuf::from(payload[zlen - 4..zlen].to_vec()),
366                            corrections: ByteBuf::from(res.corrections),
367                            size: text.text().len() as u64,
368                            chunks: Vec::new(),
369                        };
370                        consumed = Some((zlen, seg, text.text().to_vec()));
371                    }
372                }
373            }
374        } else if jpeg && payload.len() > 4 && payload[..3] == [0xFF, 0xD8, 0xFF] {
375            // JPEG (DCTDecode): trim trailing EOL bytes after the EOI marker, then Lepton-encode
376            let mut jlen = payload.len();
377            while jlen > 2 && (payload[jlen - 1] == b'\n' || payload[jlen - 1] == b'\r') {
378                jlen -= 1;
379            }
380            if payload[jlen - 2..jlen] == [0xFF, 0xD9] {
381                if let Some(lepton) = jpeg_to_lepton(&payload[..jlen]) {
382                    let seg = StreamSegment {
383                        kind: SEG_JPEG,
384                        header: ByteBuf::new(),
385                        trailer: ByteBuf::from((jlen as u64).to_le_bytes().to_vec()),
386                        corrections: ByteBuf::new(),
387                        size: lepton.len() as u64,
388                        chunks: Vec::new(),
389                    };
390                    consumed = Some((jlen, seg, lepton));
391                }
392            }
393        }
394        if let Some((zlen, seg, plain)) = consumed {
395            if data_start > lit_start {
396                let lit = data[lit_start..data_start].to_vec();
397                segments.push(literal_segment(&lit));
398                parts.push(lit);
399            }
400            segments.push(seg);
401            parts.push(plain);
402            lit_start = data_start + zlen;
403            inverted += 1;
404        }
405        pos = end + 9;
406    }
407    if inverted == 0 {
408        return None;
409    }
410    if data.len() > lit_start {
411        let lit = data[lit_start..].to_vec();
412        segments.push(literal_segment(&lit));
413        parts.push(lit);
414    }
415    Some((StreamRecipe { segments }, parts))
416}
417
418pub fn rebuild_streams(recipe: &StreamRecipe, parts: &[Vec<u8>]) -> Result<Vec<u8>> {
419    if recipe.segments.len() != parts.len() {
420        return Err(Error::Corrupt("stream recipe/segment count mismatch".into()));
421    }
422    let mut out = Vec::new();
423    for (seg, part) in recipe.segments.iter().zip(parts) {
424        match seg.kind {
425            SEG_LITERAL => out.extend_from_slice(part),
426            SEG_ZLIB => {
427                out.extend_from_slice(&seg.header);
428                let stream = preflate_rs::recreate_whole_deflate_stream(part, &seg.corrections).map_err(|e| Error::Corrupt(format!("preflate recreate stream: {e}")))?;
429                out.extend_from_slice(&stream);
430                out.extend_from_slice(&seg.trailer);
431            }
432            SEG_JPEG => {
433                let expected = if seg.trailer.len() == 8 { u64::from_le_bytes(seg.trailer[..8].try_into().unwrap()) } else { 0 };
434                let jpeg = lepton_to_jpeg(part, expected)?;
435                if expected != 0 && jpeg.len() as u64 != expected {
436                    return Err(Error::Corrupt("lepton decode: unexpected JPEG length".into()));
437                }
438                out.extend_from_slice(&jpeg);
439            }
440            other => return Err(Error::Corrupt(format!("unknown stream segment kind {other}"))),
441        }
442    }
443    Ok(out)
444}
445
446pub fn try_bitexact_streams(data: &[u8], jpeg: bool) -> Option<(StreamRecipe, Vec<Vec<u8>>)> {
447    let (recipe, parts) = explode_streams(data, jpeg)?;
448    match rebuild_streams(&recipe, &parts) {
449        Ok(out) if out == data => Some((recipe, parts)),
450        _ => None,
451    }
452}
453
454/// Statistics helper for reports.
455pub fn kinds(recipe: &ZipRecipe) -> (usize, usize, usize) {
456    let mut s = (0, 0, 0);
457    for m in &recipe.members {
458        match m.kind {
459            KIND_STORED => s.0 += 1,
460            KIND_DEFLATE => s.1 += 1,
461            _ => s.2 += 1,
462        }
463    }
464    s
465}
tsaur_core/container.rs

tsaur_core/
container.rs
