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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub struct Options {
24 pub limits: Limits,
26 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
41pub 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 pub fn new() -> Self {
66 Self::with_options(Options::default())
67 }
68
69 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 pub fn load_dictionary_bundle(&mut self, bytes: &[u8]) -> Result<()> {
93 self.dict_store.load_fat_bundle(bytes)
94 }
95
96 pub fn load_fat_bundle(&mut self, bytes: &[u8]) -> Result<()> {
100 self.load_dictionary_bundle(bytes)
101 }
102
103 pub fn clear_dictionary_bundles(&mut self) {
105 self.dict_store.clear();
106 }
107
108 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 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 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
226pub fn decode(input: &[u8]) -> Result<Vec<u8>> {
228 Decoder::new().decode(input)
229}
230
231pub fn decode_with_options(input: &[u8], options: Options) -> Result<Vec<u8>> {
233 Decoder::with_options(options).decode(input)
234}
235
236pub fn decode_into(input: &[u8], dst: &mut Vec<u8>) -> Result<usize> {
241 Decoder::new().decode_into(input, dst)
242}
243
244pub 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
253pub fn inspect(input: &[u8]) -> Result<FrameInfo> {
255 Decoder::default().inspect(input)
256}
257
258pub 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}