Skip to main content

lumen/chunks/
zdic.rs

1//! The `ZDIC` chunk: the zstd dictionary shared by every `LAYR` frame
2//! ([`spec/09-layer-data.md`] section 4.8).
3
4use crate::check::Check;
5use crate::error::{Error, Result};
6use crate::io::{Reader, Writer};
7
8/// The `ZDIC` chunk header: `zdic_version`, `dict_id`, `dict_size`.
9pub const ZDIC_HEADER_LEN: usize = 12;
10/// zstd's `ZDICT_DICTSIZE_MAX`.
11pub const ZDIC_DICTSIZE_MAX: usize = 112_640;
12
13/// The `ZDIC` chunk.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ZstdDictionary {
16    /// Layout version. `1` for this specification.
17    pub zdic_version: u32,
18    /// The dictionary ID, which every `LAYR` frame must report.
19    pub dict_id: u32,
20    /// Raw dictionary bytes, exactly as `ZDICT_trainFromBuffer` produced them.
21    pub dict_bytes: Vec<u8>,
22}
23
24impl ZstdDictionary {
25    /// Parse the chunk payload.
26    pub fn parse(payload: &[u8]) -> Result<ZstdDictionary> {
27        let mut r = Reader::checked(payload, Check::ZdicDictSize);
28        let zdic_version = r.u32()?;
29        if zdic_version != 1 {
30            return Err(Error::new(
31                Check::ZdicVersion,
32                format!("unsupported zdic_version {zdic_version}"),
33            ));
34        }
35        let dict_id = r.u32()?;
36        let dict_size = r.u32()?;
37        if dict_size as usize > ZDIC_DICTSIZE_MAX {
38            return Err(Error::new(
39                Check::ZdicDictSize,
40                format!("dict_size {dict_size} exceeds zstd's {ZDIC_DICTSIZE_MAX} byte maximum"),
41            ));
42        }
43        let dict_bytes = r.bytes(dict_size as usize)?.to_vec();
44        Ok(ZstdDictionary {
45            zdic_version,
46            dict_id,
47            dict_bytes,
48        })
49    }
50
51    /// Serialize the chunk payload.
52    pub fn to_bytes(&self) -> Vec<u8> {
53        let mut w = Writer::with_capacity(ZDIC_HEADER_LEN + self.dict_bytes.len());
54        w.u32(self.zdic_version);
55        w.u32(self.dict_id);
56        w.u32(self.dict_bytes.len() as u32);
57        w.bytes(&self.dict_bytes);
58        w.into_vec()
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn zdic_round_trips() {
68        let dict = ZstdDictionary {
69            zdic_version: 1,
70            dict_id: 0x1234_5678,
71            dict_bytes: vec![1, 2, 3, 4, 5],
72        };
73        let bytes = dict.to_bytes();
74        assert_eq!(bytes.len(), ZDIC_HEADER_LEN + 5);
75        assert_eq!(ZstdDictionary::parse(&bytes).unwrap(), dict);
76    }
77
78    #[test]
79    fn zdic_rejects_bad_version_size_and_truncation() {
80        let oversize = ZstdDictionary {
81            zdic_version: 1,
82            dict_id: 0,
83            dict_bytes: vec![0; ZDIC_DICTSIZE_MAX + 1],
84        };
85        assert_eq!(
86            ZstdDictionary::parse(&oversize.to_bytes())
87                .unwrap_err()
88                .check(),
89            Check::ZdicDictSize
90        );
91
92        let dict = ZstdDictionary {
93            zdic_version: 1,
94            dict_id: 7,
95            dict_bytes: vec![1, 2, 3],
96        };
97        let mut bytes = dict.to_bytes();
98        bytes[0..4].copy_from_slice(&2u32.to_le_bytes());
99        assert_eq!(
100            ZstdDictionary::parse(&bytes).unwrap_err().check(),
101            Check::ZdicVersion
102        );
103
104        let mut bytes = dict.to_bytes();
105        bytes[8..12].copy_from_slice(&4u32.to_le_bytes()); // declares one byte too many
106        assert_eq!(
107            ZstdDictionary::parse(&bytes).unwrap_err().check(),
108            Check::ZdicDictSize
109        );
110    }
111}