Skip to main content

ozlrip_decode/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![cfg_attr(feature = "paranoid", forbid(unsafe_code))]
3
4extern crate alloc;
5
6use alloc::vec::Vec;
7use ozlrip_core::{FrameInfo, Limits, Result};
8
9mod dict;
10mod execute;
11mod parse;
12mod standard;
13
14use dict::DictionaryStore;
15
16pub const DEFAULT_PLAN_CACHE_MAX_FRAME_BYTES: usize = 4 * 1024 * 1024;
17
18/// Decoder configuration.
19///
20/// New options should be added here instead of creating new public function
21/// variants.
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub struct Options {
24    /// Defensive parser, graph, allocation, and expansion limits.
25    pub limits: Limits,
26    /// Maximum frame size eligible for reusable decoder plan caching.
27    ///
28    /// Set to `0` to disable plan caching.
29    pub plan_cache_max_frame_bytes: usize,
30}
31
32impl Default for Options {
33    fn default() -> Self {
34        Self {
35            limits: Limits::default(),
36            plan_cache_max_frame_bytes: DEFAULT_PLAN_CACHE_MAX_FRAME_BYTES,
37        }
38    }
39}
40
41/// Reusable OpenZL decoder with scratch buffers and codec state.
42pub struct Decoder {
43    options: Options,
44    scratch: execute::DecodeScratch,
45    dict_store: DictionaryStore,
46    plan_cache: Option<CachedFramePlan>,
47    #[cfg(feature = "zstd")]
48    zstd: zrip::DecompressContext,
49}
50
51struct CachedFramePlan {
52    frame: Vec<u8>,
53    plan: parse::FramePlan,
54    execution: CachedExecutionPlan,
55}
56
57enum CachedExecutionPlan {
58    DirectAppend(execute::DirectAppendChunkPlans),
59    ChunkExecution(execute::ChunkExecutionPlans),
60    Unplanned,
61}
62
63impl Decoder {
64    /// Creates a decoder with [`Options::default`].
65    pub fn new() -> Self {
66        Self::with_options(Options::default())
67    }
68
69    /// Creates a decoder with explicit options.
70    pub fn with_options(options: Options) -> Self {
71        Self {
72            options,
73            scratch: execute::DecodeScratch::new(),
74            dict_store: DictionaryStore::new(),
75            plan_cache: None,
76            #[cfg(feature = "zstd")]
77            zstd: zrip::DecompressContext::new(),
78        }
79    }
80
81    pub fn options(&self) -> Options {
82        self.options
83    }
84
85    pub fn limits(&self) -> Limits {
86        self.options.limits
87    }
88
89    /// Loads an OpenZL dictionary bundle for later dictionary-backed decode.
90    ///
91    /// Currently only zstd dictionary materialization is implemented.
92    pub fn load_dictionary_bundle(&mut self, bytes: &[u8]) -> Result<()> {
93        self.dict_store.load_fat_bundle(bytes)
94    }
95
96    /// Loads a fat OpenZL dictionary bundle for later dictionary-backed decode.
97    ///
98    /// Prefer [`Decoder::load_dictionary_bundle`] for new code.
99    pub fn load_fat_bundle(&mut self, bytes: &[u8]) -> Result<()> {
100        self.load_dictionary_bundle(bytes)
101    }
102
103    /// Clears all loaded dictionary bundles.
104    pub fn clear_dictionary_bundles(&mut self) {
105        self.dict_store.clear();
106    }
107
108    /// Decodes one OpenZL frame and appends the decoded bytes to `dst`.
109    ///
110    /// Returns the number of bytes appended. On error, `dst` is restored to its
111    /// original length.
112    pub fn decode_into(&mut self, input: &[u8], dst: &mut Vec<u8>) -> Result<usize> {
113        #[cfg(feature = "zstd")]
114        if let Some(frame) = parse::parse_single_zstd_frame(input, self.options.limits)? {
115            return execute::decode_single_zstd_frame_with_context(
116                input,
117                frame,
118                dst,
119                self.options.limits,
120                &mut self.zstd,
121            );
122        }
123        if let Some(cached) = self.plan_cache.as_ref()
124            && cached.frame == input
125        {
126            let mut runtime = execute::DecodeRuntime::new(
127                &mut self.scratch,
128                &mut self.dict_store,
129                #[cfg(feature = "zstd")]
130                &mut self.zstd,
131            );
132            match &cached.execution {
133                CachedExecutionPlan::DirectAppend(plans) => {
134                    return execute::decode_plan_with_cached_direct_append_plans(
135                        input,
136                        &cached.plan,
137                        plans,
138                        dst,
139                        self.options.limits,
140                        &mut runtime,
141                    );
142                }
143                CachedExecutionPlan::ChunkExecution(plans) => {
144                    return execute::decode_plan_with_cached_chunk_execution_plans(
145                        input,
146                        &cached.plan,
147                        plans,
148                        dst,
149                        self.options.limits,
150                        &mut runtime,
151                    );
152                }
153                CachedExecutionPlan::Unplanned => {}
154            }
155            return execute::decode_plan_with_context(
156                input,
157                &cached.plan,
158                dst,
159                self.options.limits,
160                &mut runtime,
161            );
162        }
163
164        let plan = parse::parse_frame_plan(input, self.options.limits)?;
165        self.remember_frame_plan(input, &plan);
166        let mut runtime = execute::DecodeRuntime::new(
167            &mut self.scratch,
168            &mut self.dict_store,
169            #[cfg(feature = "zstd")]
170            &mut self.zstd,
171        );
172        execute::decode_plan_with_context(input, &plan, dst, self.options.limits, &mut runtime)
173    }
174
175    fn remember_frame_plan(&mut self, input: &[u8], plan: &parse::FramePlan) {
176        if self.options.plan_cache_max_frame_bytes == 0
177            || input.len() > self.options.plan_cache_max_frame_bytes
178        {
179            self.plan_cache = None;
180            return;
181        }
182
183        let mut frame = Vec::new();
184        if frame.try_reserve_exact(input.len()).is_err() {
185            self.plan_cache = None;
186            return;
187        }
188        frame.extend_from_slice(input);
189        let execution = Self::prepare_cached_execution_plan(plan);
190        self.plan_cache = Some(CachedFramePlan {
191            frame,
192            plan: plan.clone(),
193            execution,
194        });
195    }
196
197    fn prepare_cached_execution_plan(plan: &parse::FramePlan) -> CachedExecutionPlan {
198        if let Ok(Some(plans)) = execute::prepare_direct_append_chunk_plans(plan) {
199            return CachedExecutionPlan::DirectAppend(plans);
200        }
201        match execute::prepare_chunk_execution_plans(plan) {
202            Ok(plans) => CachedExecutionPlan::ChunkExecution(plans),
203            Err(_) => CachedExecutionPlan::Unplanned,
204        }
205    }
206
207    /// Decodes one OpenZL frame into a new `Vec`.
208    pub fn decode(&mut self, input: &[u8]) -> Result<Vec<u8>> {
209        let mut output = Vec::new();
210        self.decode_into(input, &mut output)?;
211        Ok(output)
212    }
213
214    /// Parses and validates frame metadata without executing decode nodes.
215    pub fn inspect(&self, input: &[u8]) -> Result<FrameInfo> {
216        parse::inspect_frame(input, self.options.limits)
217    }
218}
219
220impl Default for Decoder {
221    fn default() -> Self {
222        Self::new()
223    }
224}
225
226/// Decodes one OpenZL frame into a new `Vec` using [`Options::default`].
227pub fn decode(input: &[u8]) -> Result<Vec<u8>> {
228    Decoder::new().decode(input)
229}
230
231/// Decodes one OpenZL frame into a new `Vec` using explicit options.
232pub fn decode_with_options(input: &[u8], options: Options) -> Result<Vec<u8>> {
233    Decoder::with_options(options).decode(input)
234}
235
236/// Decodes one OpenZL frame and appends the decoded bytes to `dst`.
237///
238/// Returns the number of bytes appended. On error, `dst` is restored to its
239/// original length.
240pub fn decode_into(input: &[u8], dst: &mut Vec<u8>) -> Result<usize> {
241    Decoder::new().decode_into(input, dst)
242}
243
244/// Decodes one OpenZL frame into `dst` using explicit options.
245pub fn decode_into_with_options(
246    input: &[u8],
247    dst: &mut Vec<u8>,
248    options: Options,
249) -> Result<usize> {
250    Decoder::with_options(options).decode_into(input, dst)
251}
252
253/// Parses and validates frame metadata using [`Options::default`].
254pub fn inspect(input: &[u8]) -> Result<FrameInfo> {
255    Decoder::default().inspect(input)
256}
257
258/// Parses and validates frame metadata using explicit options.
259pub fn inspect_with_options(input: &[u8], options: Options) -> Result<FrameInfo> {
260    Decoder::with_options(options).inspect(input)
261}
262
263#[cfg(all(test, feature = "zstd"))]
264mod tests {
265    use super::*;
266    use ozlrip_core::ErrorKind;
267
268    const MAGIC_BASE: u32 = 0xd7b1_a5c0;
269    const BUNDLE_INFO_MAGIC: u32 = 0x4942_ccda;
270    const PACKED_DICT_MAGIC: u32 = 0x4944_ccda;
271    const BUNDLE_ID: [u8; 32] = [7; 32];
272    const OTHER_BUNDLE_ID: [u8; 32] = [8; 32];
273    const DICT_ID: [u8; 32] = [9; 32];
274
275    fn magic(version: u32) -> [u8; 4] {
276        (MAGIC_BASE + version).to_le_bytes()
277    }
278
279    fn push_var_u64(out: &mut Vec<u8>, mut value: u64) {
280        while value >= 0x80 {
281            out.push(u8::try_from(value & 0x7f).unwrap() | 0x80);
282            value >>= 7;
283        }
284        out.push(u8::try_from(value).unwrap());
285    }
286
287    #[cfg(feature = "zstd")]
288    fn test_zstd_dict_bytes() -> Vec<u8> {
289        use zrip_core::dict::DICT_MAGIC;
290        use zrip_core::fse::table_builder::serialize_fse_table_description;
291
292        let mut dict = Vec::new();
293        dict.extend_from_slice(&DICT_MAGIC.to_le_bytes());
294        dict.extend_from_slice(&1u32.to_le_bytes());
295        dict.push(128);
296        dict.push(0x10);
297
298        let mut of_dist = vec![0i16; 32];
299        of_dist[0] = 1 << 8;
300        dict.extend_from_slice(&serialize_fse_table_description(&of_dist, 8));
301
302        let mut ml_dist = vec![0i16; 53];
303        ml_dist[0] = 1 << 6;
304        dict.extend_from_slice(&serialize_fse_table_description(&ml_dist, 6));
305
306        let mut ll_dist = vec![0i16; 36];
307        ll_dist[0] = 1 << 6;
308        dict.extend_from_slice(&serialize_fse_table_description(&ll_dist, 6));
309
310        dict.extend_from_slice(&1u32.to_le_bytes());
311        dict.extend_from_slice(&4u32.to_le_bytes());
312        dict.extend_from_slice(&8u32.to_le_bytes());
313        dict.extend_from_slice(b"dictionary-backed OpenZL zstd content");
314        dict
315    }
316
317    #[cfg(feature = "zstd")]
318    fn fat_bundle(bundle_id: [u8; 32], raw_dict: &[u8]) -> Vec<u8> {
319        let mut content = Vec::new();
320        content.extend_from_slice(&1u32.to_le_bytes());
321        content.extend_from_slice(&1i32.to_le_bytes());
322        content.extend_from_slice(raw_dict);
323
324        let mut bundle = Vec::new();
325        bundle.extend_from_slice(&BUNDLE_INFO_MAGIC.to_le_bytes());
326        bundle.extend_from_slice(&bundle_id);
327        bundle.push(1);
328        bundle.extend_from_slice(&1u32.to_le_bytes());
329        bundle.extend_from_slice(&DICT_ID);
330        bundle.extend_from_slice(&PACKED_DICT_MAGIC.to_le_bytes());
331        bundle.extend_from_slice(&DICT_ID);
332        bundle.extend_from_slice(&standard::ZSTD_ID.to_le_bytes());
333        bundle.push(0);
334        bundle.extend_from_slice(&u32::try_from(content.len()).unwrap().to_le_bytes());
335        bundle.extend_from_slice(&content);
336        bundle
337    }
338
339    #[cfg(feature = "zstd")]
340    fn zstd_dict_frame(bundle_id: &[u8], stored: &[u8], decoded_len: usize) -> Vec<u8> {
341        let mut input = Vec::new();
342        input.extend_from_slice(&magic(25));
343        input.push(1 << 3);
344        input.push(u8::try_from(bundle_id.len()).unwrap());
345        input.extend_from_slice(bundle_id);
346        input.push(1);
347        push_var_u64(&mut input, u64::try_from(decoded_len + 1).unwrap());
348        input.push(2);
349        input.push(1);
350        input.push(0);
351        input.push(u8::try_from(standard::ZSTD_ID).unwrap());
352        input.push(0);
353        input.push(0);
354        input.push(0);
355        input.push(1);
356        input.push(0);
357        input.push(0);
358        push_var_u64(&mut input, u64::try_from(stored.len()).unwrap());
359        input.extend_from_slice(stored);
360        input.push(0);
361        input
362    }
363
364    #[cfg(feature = "zstd")]
365    fn push_zstd_block_header(out: &mut Vec<u8>, last: bool, block_type: u32, block_size: usize) {
366        let raw = (u32::try_from(block_size).unwrap() << 3) | (block_type << 1) | u32::from(last);
367        out.extend_from_slice(&raw.to_le_bytes()[..3]);
368    }
369
370    #[cfg(feature = "zstd")]
371    fn zstd_raw_frame_with_dict_id(bytes: &[u8], dict_id: u8) -> Vec<u8> {
372        let mut frame = Vec::new();
373        frame.extend_from_slice(&zrip_core::frame::ZSTD_MAGIC.to_le_bytes());
374        frame.push(0x21);
375        frame.push(dict_id);
376        frame.push(u8::try_from(bytes.len()).unwrap());
377        push_zstd_block_header(&mut frame, true, 0, bytes.len());
378        frame.extend_from_slice(bytes);
379        frame
380    }
381
382    #[cfg(feature = "zstd")]
383    fn zstd_dict_test_frame() -> (Vec<u8>, Vec<u8>, Vec<u8>) {
384        let raw_dict = test_zstd_dict_bytes();
385        let expected = b"dictionary-backed OpenZL zstd frame".to_vec();
386        let compressed = zstd_raw_frame_with_dict_id(&expected, 1);
387        let mut stored = Vec::new();
388        push_var_u64(&mut stored, 1);
389        stored.extend_from_slice(&compressed[4..]);
390        let frame = zstd_dict_frame(&BUNDLE_ID, &stored, expected.len());
391        (frame, fat_bundle(BUNDLE_ID, &raw_dict), expected)
392    }
393
394    #[cfg(feature = "zstd")]
395    #[test]
396    fn decodes_zstd_dictionary_frame_after_loading_bundle() {
397        let (frame, bundle, expected) = zstd_dict_test_frame();
398        let mut decoder = Decoder::new();
399        decoder.load_dictionary_bundle(&bundle).unwrap();
400        let mut output = vec![1, 2];
401
402        let written = decoder.decode_into(&frame, &mut output).unwrap();
403
404        assert_eq!(written, expected.len());
405        assert_eq!(&output[..2], &[1, 2]);
406        assert_eq!(&output[2..], expected);
407    }
408
409    #[cfg(feature = "zstd")]
410    #[test]
411    fn rejects_dictionary_frame_without_loaded_bundle_without_mutating_destination() {
412        let (frame, _bundle, _expected) = zstd_dict_test_frame();
413        let mut decoder = Decoder::new();
414        let mut output = vec![1, 2];
415
416        let err = decoder.decode_into(&frame, &mut output).unwrap_err();
417
418        assert_eq!(err.kind(), ErrorKind::Unsupported);
419        assert_eq!(output, [1, 2]);
420    }
421
422    #[cfg(feature = "zstd")]
423    #[test]
424    fn rejects_dictionary_frame_with_wrong_bundle_without_mutating_destination() {
425        let (frame, _bundle, _expected) = zstd_dict_test_frame();
426        let raw_dict = test_zstd_dict_bytes();
427        let mut decoder = Decoder::new();
428        decoder
429            .load_fat_bundle(&fat_bundle(OTHER_BUNDLE_ID, &raw_dict))
430            .unwrap();
431        let mut output = vec![1, 2];
432
433        let err = decoder.decode_into(&frame, &mut output).unwrap_err();
434
435        assert_eq!(err.kind(), ErrorKind::Unsupported);
436        assert_eq!(output, [1, 2]);
437    }
438
439    #[cfg(feature = "zstd")]
440    #[test]
441    fn clear_dictionary_bundles_drops_loaded_bundle() {
442        let (frame, bundle, _expected) = zstd_dict_test_frame();
443        let mut decoder = Decoder::new();
444        decoder.load_fat_bundle(&bundle).unwrap();
445        decoder.clear_dictionary_bundles();
446        let mut output = vec![1, 2];
447
448        let err = decoder.decode_into(&frame, &mut output).unwrap_err();
449
450        assert_eq!(err.kind(), ErrorKind::Unsupported);
451        assert_eq!(output, [1, 2]);
452    }
453}