Skip to main content

lumen/chunks/
mod.rs

1//! One codec per chunk type ([`spec/03-chunks.md`], [`spec/09-layer-data.md`],
2//! [`spec/10-scene-chunks.md`]).
3//!
4//! Each codec deals in *plaintext payload bytes*: decryption and the zstd layer
5//! are the caller's business, because whether a chunk carries a zstd frame is a
6//! property of its type (section 6.3) while whether it is sealed is a property
7//! of its descriptor.
8
9pub mod extd;
10pub mod head;
11pub mod json_chunks;
12pub mod layr;
13pub mod lhas;
14pub mod ltbl;
15pub mod preview;
16pub mod voxl;
17pub mod zdic;
18
19pub use extd::Extension;
20pub use head::Head;
21pub use lhas::LayerHashes;
22pub use ltbl::{LayerEntry, LayerTable};
23pub use preview::{Preview, PreviewRole};
24pub use zdic::ZstdDictionary;
25
26use crate::check::Check;
27use crate::error::{Error, Result};
28use zstd::bulk::{Compressor, Decompressor};
29
30/// Compress a chunk payload into a zstd frame, optionally against a dictionary.
31pub fn compress(payload: &[u8], level: i32, dictionary: Option<&[u8]>) -> Result<Vec<u8>> {
32    let mut compressor = match dictionary {
33        Some(dictionary) => Compressor::with_dictionary(level, dictionary),
34        None => Compressor::new(level),
35    }
36    .map_err(|e| {
37        Error::new(
38            Check::LayrFrameDecompressedSize,
39            format!("could not start the zstd compressor: {e}"),
40        )
41    })?;
42    compressor.compress(payload).map_err(|e| {
43        Error::new(
44            Check::LayrFrameDecompressedSize,
45            format!("zstd compression failed: {e}"),
46        )
47    })
48}
49
50/// Decompress a zstd frame, checking the result is exactly `uncompressed_size`.
51///
52/// The size is bounded by what the frame's own bytes can justify, which is a
53/// coarse guard for a payload whose bound is unknown until the chunk around it is
54/// parsed. A `LAYR` frame has a better bound - the slices that point into it -
55/// so it goes through [`decompress_frame`] instead (section 11.3).
56pub fn decompress(
57    frame: &[u8],
58    uncompressed_size: u64,
59    dictionary: Option<&[u8]>,
60) -> Result<Vec<u8>> {
61    let bound = (frame.len() as u64).saturating_mul(MAX_ZSTD_RATIO as u64);
62    decompress_within(frame, uncompressed_size, dictionary, bound)
63}
64
65/// Decompress a zstd frame, refusing a declared size above `bound`.
66pub fn decompress_within(
67    frame: &[u8],
68    uncompressed_size: u64,
69    dictionary: Option<&[u8]>,
70    bound: u64,
71) -> Result<Vec<u8>> {
72    // Refuse to reserve memory the file's own structure cannot justify before the
73    // decompressor allocates for the declared size (section 11.3).
74    check_allocation(uncompressed_size, bound, frame.len())?;
75    let capacity = usize::try_from(uncompressed_size).map_err(|_| {
76        Error::new(
77            Check::LayrFrameDecompressedSize,
78            format!("declared uncompressed size {uncompressed_size} is not addressable"),
79        )
80    })?;
81    let mut decompressor = match dictionary {
82        Some(dictionary) => Decompressor::with_dictionary(dictionary),
83        None => Decompressor::new(),
84    }
85    .map_err(|e| {
86        Error::new(
87            Check::LayrFrameDecompressedSize,
88            format!("could not start the zstd decompressor: {e}"),
89        )
90    })?;
91    let decoded = decompressor.decompress(frame, capacity).map_err(|e| {
92        Error::new(
93            Check::LayrFrameDecompressedSize,
94            format!("zstd decompression failed: {e}"),
95        )
96    })?;
97    if decoded.len() as u64 != uncompressed_size {
98        return Err(Error::new(
99            Check::LayrFrameDecompressedSize,
100            format!(
101                "the frame expands to {} bytes, not the declared {uncompressed_size}",
102                decoded.len()
103            ),
104        ));
105    }
106    Ok(decoded)
107}
108
109/// The dictionary ID a frame reports, or `0` when the frame used no dictionary.
110pub fn frame_dict_id(frame: &[u8]) -> Result<u32> {
111    Ok(zstd::zstd_safe::get_dict_id_from_frame(frame)
112        .map(|id| id.get())
113        .unwrap_or(0))
114}
115
116/// The decompressed length a frame's own header declares.
117///
118/// A `LAYR` frame MUST carry its content size: the chunk descriptor describes
119/// the container, not the frame's output, so this header field is the only place
120/// a reader can learn how large that output is before it allocates.
121pub fn frame_content_size(frame: &[u8]) -> Result<u64> {
122    match zstd::zstd_safe::get_frame_content_size(frame) {
123        Ok(Some(size)) => Ok(size),
124        Ok(None) => Err(Error::new(
125            Check::LayrContentSizePresent,
126            "the frame carries no content size, so its output cannot be sized",
127        )),
128        Err(e) => Err(Error::new(
129            Check::LayrContentSizePresent,
130            format!("the frame's content size is unreadable: {e}"),
131        )),
132    }
133}
134
135/// Decompress a `LAYR` frame, sizing its output from the frame's own header and
136/// bounding it by `bound`.
137///
138/// The bound is the caller's, because a `LAYR` frame's plausible output is
139/// derived from the file rather than from the frame: it is
140/// [`allocation_bound`] over the slices the layer table points into this chunk.
141/// A ratio guard cannot stand in for it - an encoder may write a block whose
142/// layers repeat, and zstd then compresses far past any fixed ratio, which is a
143/// conforming file rather than a hostile one.
144pub fn decompress_frame(frame: &[u8], dictionary: Option<&[u8]>, bound: u64) -> Result<Vec<u8>> {
145    let size = frame_content_size(frame)?;
146    decompress_within(frame, size, dictionary, bound)
147}
148
149/// The dictionary ID a raw dictionary reports.
150pub fn dictionary_id(dictionary: &[u8]) -> u32 {
151    zstd::zstd_safe::get_dict_id_from_dict(dictionary)
152        .map(|id| id.get())
153        .unwrap_or(0)
154}
155
156/// Train a zstd dictionary from layer samples (section 6.2).
157pub fn train_dictionary(samples: &[&[u8]], max_size: usize) -> Result<Vec<u8>> {
158    zstd::dict::from_samples(samples, max_size).map_err(|e| {
159        Error::new(
160            Check::ZdicDictSize,
161            format!("zstd dictionary training failed: {e}"),
162        )
163    })
164}
165
166/// The decompressed-size bound a `LAYR` chunk's frame may not exceed.
167///
168/// Grayscale REE costs at most about 5 bytes per pixel plus framing, which is
169/// what a reader checks a frame's declared content size against before
170/// allocating, over the layers whose slices the chunk holds (section 11.3).
171pub fn allocation_bound(layers: u64, total_pixels: u32) -> u64 {
172    layers.saturating_mul(u64::from(total_pixels).saturating_mul(5).saturating_add(64))
173}
174
175/// Reject a declared size above the bound the file's own structure allows.
176pub fn check_allocation(uncompressed_size: u64, bound: u64, payload_len: usize) -> Result<()> {
177    if uncompressed_size > bound {
178        return Err(Error::new(
179            Check::FrameAllocationBound,
180            format!(
181                "declared uncompressed size {uncompressed_size} exceeds the {bound} bytes the \
182                 structure allows for {payload_len} stored bytes"
183            ),
184        ));
185    }
186    Ok(())
187}
188
189/// The largest expansion zstd can produce for a chunk payload, used as a coarse
190/// guard before a bound specific to the chunk is available.
191pub const MAX_ZSTD_RATIO: usize = 4096;
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    use std::io::Write;
197
198    #[test]
199    fn zstd_round_trips_with_and_without_a_dictionary() {
200        let payload: Vec<u8> = (0..4096u32).map(|i| (i % 17) as u8).collect();
201
202        let framed = compress(&payload, 3, None).unwrap();
203        assert_eq!(frame_dict_id(&framed).unwrap(), 0);
204        assert_eq!(
205            decompress(&framed, payload.len() as u64, None).unwrap(),
206            payload
207        );
208
209        // Training needs a corpus of similar samples.
210        let samples: Vec<Vec<u8>> = (0..256u32)
211            .map(|s| (0..256u32).map(|i| ((i + s) % 23) as u8).collect())
212            .collect();
213        let refs: Vec<&[u8]> = samples.iter().map(|s| s.as_slice()).collect();
214        let dictionary = train_dictionary(&refs, 512).unwrap();
215        let dict_id = dictionary_id(&dictionary);
216        assert_ne!(dict_id, 0, "a trained dictionary reports its ID");
217
218        let framed = compress(&payload, 3, Some(&dictionary)).unwrap();
219        assert_eq!(frame_dict_id(&framed).unwrap(), dict_id);
220        assert_eq!(
221            decompress(&framed, payload.len() as u64, Some(&dictionary)).unwrap(),
222            payload
223        );
224
225        // A frame that does not expand to the declared size is rejected.
226        let err = decompress(&framed, payload.len() as u64 + 1, Some(&dictionary)).unwrap_err();
227        assert_eq!(err.check(), Check::LayrFrameDecompressedSize);
228    }
229
230    #[test]
231    fn a_writer_frame_declares_the_size_it_expands_to() {
232        let payload: Vec<u8> = (0..4096u32).map(|i| (i % 17) as u8).collect();
233        let framed = compress(&payload, 3, None).unwrap();
234        assert_eq!(frame_content_size(&framed).unwrap(), payload.len() as u64);
235        // A bound generous enough for anything this frame could expand to: the
236        // test is about the header, not about the bound.
237        let bound = (framed.len() as u64).saturating_mul(MAX_ZSTD_RATIO as u64);
238        assert_eq!(decompress_frame(&framed, None, bound).unwrap(), payload);
239
240        // A `LAYR` frame is sized from its own header, because the chunk
241        // descriptor describes the container and not the frame's output. A frame
242        // written by a streaming encoder declares no content size, and a reader
243        // refuses one rather than guessing what to allocate.
244        let streamed = {
245            let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 3).unwrap();
246            encoder.write_all(&payload).unwrap();
247            encoder.finish().unwrap()
248        };
249        assert_eq!(
250            frame_content_size(&streamed).unwrap_err().check(),
251            Check::LayrContentSizePresent
252        );
253    }
254
255    #[test]
256    fn check_allocation_bounds_the_declared_size() {
257        let bound = (MAX_ZSTD_RATIO - 1) as u64;
258        assert!(check_allocation(bound, bound, 1).is_ok());
259        assert_eq!(
260            check_allocation(bound + 1, bound, 1).unwrap_err().check(),
261            Check::FrameAllocationBound
262        );
263        assert!(allocation_bound(2, 1000) > 0);
264    }
265
266    /// A frame may compress far past the ratio guard and still be conforming: a
267    /// block whose layers repeat does exactly that. What bounds it is the slices
268    /// the layer table points into its chunk, so a `LAYR` frame is decompressed
269    /// against that bound instead of against the ratio.
270    #[test]
271    fn a_layr_frame_is_bounded_by_its_slices_not_by_a_ratio() {
272        let payload = vec![0u8; 1 << 20];
273        let frame = compress(&payload, 3, None).unwrap();
274        let ratio = payload.len() as f64 / frame.len() as f64;
275        assert!(
276            ratio > MAX_ZSTD_RATIO as f64,
277            "this test needs a frame past the ratio guard, got {ratio:.0}:1"
278        );
279
280        // The ratio guard refuses it, which is why it cannot be the LAYR rule.
281        assert_eq!(
282            decompress_frame(&frame, None, (frame.len() as u64) * MAX_ZSTD_RATIO as u64)
283                .unwrap_err()
284                .check(),
285            Check::FrameAllocationBound
286        );
287
288        // A bound from the file's own structure accepts it, and the output is
289        // exactly the declared size.
290        let total_pixels = 1 << 20;
291        let bound = allocation_bound(1, total_pixels);
292        assert!(bound >= payload.len() as u64);
293        assert_eq!(decompress_frame(&frame, None, bound).unwrap(), payload);
294    }
295}