1use super::frame;
8use crate::decoding;
9use crate::decoding::dictionary::Dictionary;
10use crate::decoding::errors::FrameDecoderError;
11use crate::decoding::scratch::DecoderScratch;
12use crate::io::{Error, Read, Write};
13use alloc::collections::BTreeMap;
14use alloc::vec::Vec;
15use core::convert::TryInto;
16
17pub const DEFAULT_MAX_WINDOW_SIZE: u64 = 1024 * 1024 * 100;
21
22pub struct FrameDecoder {
76 state: Option<FrameDecoderState>,
77 prepared_scratch: Option<DecoderScratch>,
78 dicts: BTreeMap<u32, Dictionary>,
79}
80
81struct FrameDecoderState {
82 pub frame_header: frame::FrameHeader,
83 decoder_scratch: DecoderScratch,
84 frame_finished: bool,
85 block_counter: usize,
86 bytes_read_counter: u64,
87 decoded_size_counter: u64,
88 check_sum: Option<u32>,
89 using_dict: Option<u32>,
90}
91
92#[derive(Clone, Copy)]
93pub(crate) enum PreparedDictionaryBytes<'a> {
94 Formatted(&'a [u8]),
95 Raw(&'a [u8]),
96}
97
98pub enum BlockDecodingStrategy {
99 All,
100 UptoBlocks(usize),
101 UptoBytes(usize),
102}
103
104impl FrameDecoderState {
105 fn from_header(
106 frame: frame::FrameHeader,
107 header_size: u8,
108 window_size: usize,
109 scratch: Option<DecoderScratch>,
110 ) -> FrameDecoderState {
111 let mut decoder_scratch = scratch.unwrap_or_else(|| DecoderScratch::new(window_size));
112 decoder_scratch.reset(window_size);
113 FrameDecoderState {
114 frame_header: frame,
115 frame_finished: false,
116 block_counter: 0,
117 decoder_scratch,
118 bytes_read_counter: u64::from(header_size),
119 decoded_size_counter: 0,
120 check_sum: None,
121 using_dict: None,
122 }
123 }
124
125 pub fn reset(
126 &mut self,
127 source: impl Read,
128 max_window_size: u64,
129 ) -> Result<(), FrameDecoderError> {
130 let (frame_header, header_size) = frame::read_frame_header(source)?;
131 let window_size = frame_header.window_size()?;
132
133 if window_size > max_window_size {
134 return Err(FrameDecoderError::WindowSizeTooBig {
135 requested: window_size,
136 });
137 }
138
139 self.frame_header = frame_header;
140 self.frame_finished = false;
141 self.block_counter = 0;
142 self.decoder_scratch.reset(window_size as usize);
143 self.bytes_read_counter = u64::from(header_size);
144 self.decoded_size_counter = 0;
145 self.check_sum = None;
146 self.using_dict = None;
147 Ok(())
148 }
149}
150
151fn wipe_dictionary_content(dict: &mut Dictionary) {
155 #[cfg(all(test, feature = "std"))]
156 super::wiping::note_initialized(&dict.dict_content);
157 unsafe {
159 super::wiping::wipe_allocation(dict.dict_content.as_mut_ptr(), dict.dict_content.capacity())
160 };
161}
162
163impl Drop for FrameDecoder {
164 fn drop(&mut self) {
165 for dictionary in self.dicts.values_mut() {
166 wipe_dictionary_content(dictionary);
167 }
168 }
169}
170
171impl Default for FrameDecoder {
172 fn default() -> Self {
173 Self::new()
174 }
175}
176
177impl FrameDecoder {
178 pub fn new() -> FrameDecoder {
182 FrameDecoder {
183 state: None,
184 prepared_scratch: None,
185 dicts: BTreeMap::new(),
186 }
187 }
188
189 pub(crate) fn prepare_workspace(&mut self, max_window_size: usize, max_dictionary_size: usize) -> Result<(), crate::workspace::CapacityError> {
190 Ok({
191 let mut scratch = self
192 .state
193 .take()
194 .map(|state| state.decoder_scratch)
195 .or_else(|| self.prepared_scratch.take())
196 .unwrap_or_else(|| DecoderScratch::new(0));
197 scratch.reset(0);
198 scratch.reserve_workspace(max_window_size, max_dictionary_size)?;
199 self.prepared_scratch = Some(scratch);
200 })
201}
202
203 pub(crate) fn install_prepared_scratch(&mut self, scratch: DecoderScratch) {
204 self.state = None;
205 self.prepared_scratch = Some(scratch);
206 }
207
208 pub fn init(&mut self, source: impl Read) -> Result<(), FrameDecoderError> {
215 self.reset(source)
216 }
217
218 pub fn reset(&mut self, source: impl Read) -> Result<(), FrameDecoderError> {
225 self.reset_with_window_limit(source, DEFAULT_MAX_WINDOW_SIZE)
226 }
227
228 pub(crate) fn reset_with_window_limit(
229 &mut self,
230 source: impl Read,
231 max_window_size: u64,
232 ) -> Result<(), FrameDecoderError> {
233 self.reset_with_window_limit_and_dictionary(source, max_window_size, None)
234 }
235
236 fn reset_with_window_limit_and_dictionary(
237 &mut self,
238 source: impl Read,
239 max_window_size: u64,
240 dictionary: Option<PreparedDictionaryBytes<'_>>,
241 ) -> Result<(), FrameDecoderError> {
242 use FrameDecoderError as err;
243 let state = match &mut self.state {
244 Some(s) => {
245 s.reset(source, max_window_size)?;
246 s
247 }
248 None => {
249 let (frame_header, header_size) = frame::read_frame_header(source)?;
253 let window_size = frame_header.window_size()?;
254 if window_size > max_window_size {
255 return Err(FrameDecoderError::WindowSizeTooBig {
256 requested: window_size,
257 });
258 }
259 self.state = Some(FrameDecoderState::from_header(
260 frame_header,
261 header_size,
262 window_size as usize,
263 self.prepared_scratch.take(),
264 ));
265 self.state.as_mut().unwrap()
266 }
267 };
268 if let Some(dict_id) = state.frame_header.dictionary_id() {
269 if let Some(PreparedDictionaryBytes::Formatted(dictionary)) = dictionary {
270 let provided_id = state.decoder_scratch.init_from_dict_bytes(dictionary)?;
271 if provided_id != dict_id {
272 return Err(err::DictNotProvided { dict_id });
273 }
274 } else if dictionary.is_some() {
275 return Err(err::DictNotProvided { dict_id });
276 } else {
277 let dict = self
278 .dicts
279 .get(&dict_id)
280 .ok_or(err::DictNotProvided { dict_id })?;
281 state.decoder_scratch.init_from_dict(dict)?;
282 }
283 state.using_dict = Some(dict_id);
284 } else if let Some(PreparedDictionaryBytes::Raw(dictionary)) = dictionary {
285 state.decoder_scratch.init_from_raw_dictionary(dictionary)?;
286 }
287 Ok(())
288 }
289
290 pub fn add_dict(&mut self, dict: Dictionary) -> Result<(), FrameDecoderError> {
292 if let Some(mut previous) = self.dicts.insert(dict.id, dict) {
293 wipe_dictionary_content(&mut previous);
294 }
295 Ok(())
296 }
297
298 pub fn force_dict(&mut self, dict_id: u32) -> Result<(), FrameDecoderError> {
299 use FrameDecoderError as err;
300 let Some(state) = self.state.as_mut() else {
301 return Err(err::NotYetInitialized);
302 };
303
304 let dict = self
305 .dicts
306 .get(&dict_id)
307 .ok_or(err::DictNotProvided { dict_id })?;
308 state.decoder_scratch.init_from_dict(dict)?;
309 state.using_dict = Some(dict_id);
310
311 Ok(())
312 }
313
314 pub fn content_size(&self) -> u64 {
316 match &self.state {
317 None => 0,
318 Some(s) => s.frame_header.frame_content_size(),
319 }
320 }
321
322 pub fn get_checksum_from_data(&self) -> Option<u32> {
324 let state = self.state.as_ref()?;
325
326 state.check_sum
327 }
328
329 #[cfg(feature = "hash")]
332 pub fn get_calculated_checksum(&self) -> Option<u32> {
333 use core::hash::Hasher;
334
335 let state = self.state.as_ref()?;
336 let cksum_64bit = state.decoder_scratch.buffer.hash.finish();
337 Some(cksum_64bit as u32)
339 }
340
341 pub fn bytes_read_from_source(&self) -> u64 {
343 let state = match &self.state {
344 None => return 0,
345 Some(s) => s,
346 };
347 state.bytes_read_counter
348 }
349
350 pub fn decoded_size(&self) -> u64 {
352 let state = match &self.state {
353 None => return 0,
354 Some(s) => s,
355 };
356 state.decoded_size_counter
357 }
358
359 pub fn is_finished(&self) -> bool {
363 let state = match &self.state {
364 None => return true,
365 Some(s) => s,
366 };
367 if state.frame_header.descriptor.content_checksum_flag() {
368 state.frame_finished && state.check_sum.is_some()
369 } else {
370 state.frame_finished
371 }
372 }
373
374 pub fn blocks_decoded(&self) -> usize {
376 let state = match &self.state {
377 None => return 0,
378 Some(s) => s,
379 };
380 state.block_counter
381 }
382
383 pub fn decode_blocks(
388 &mut self,
389 mut source: impl Read,
390 strat: BlockDecodingStrategy,
391 ) -> Result<bool, FrameDecoderError> {
392 use FrameDecoderError as err;
393 let state = self.state.as_mut().ok_or(err::NotYetInitialized)?;
394
395 let mut block_dec = decoding::block_decoder::new();
396
397 let buffer_size_before = state.decoder_scratch.buffer.len();
398 let block_counter_before = state.block_counter;
399 loop {
400 vprintln!("################");
401 vprintln!("Next Block: {}", state.block_counter);
402 vprintln!("################");
403 let (block_header, block_header_size) = block_dec
404 .read_block_header(&mut source)
405 .map_err(err::FailedToReadBlockHeader)?;
406 state.bytes_read_counter += u64::from(block_header_size);
407
408 vprintln!();
409 vprintln!(
410 "Found {} block with size: {}, which will be of size: {}",
411 block_header.block_type,
412 block_header.content_size,
413 block_header.decompressed_size
414 );
415
416 let decoded_size_before_block = state.decoder_scratch.buffer.len();
417 let bytes_read_in_block_body = block_dec
418 .decode_block_content(&block_header, &mut state.decoder_scratch, &mut source)
419 .map_err(err::FailedToReadBlockBody)?;
420 state.bytes_read_counter += bytes_read_in_block_body;
421 state.decoded_size_counter +=
422 (state.decoder_scratch.buffer.len() - decoded_size_before_block) as u64;
423
424 state.block_counter += 1;
425
426 vprintln!("Output: {}", state.decoder_scratch.buffer.len());
427
428 if block_header.last_block {
429 state.frame_finished = true;
430 if state.frame_header.descriptor.content_checksum_flag() {
431 let mut chksum = [0u8; 4];
432 source
433 .read_exact(&mut chksum)
434 .map_err(err::FailedToReadChecksum)?;
435 state.bytes_read_counter += 4;
436 let chksum = u32::from_le_bytes(chksum);
437 state.check_sum = Some(chksum);
438 }
439 break;
440 }
441
442 match strat {
443 BlockDecodingStrategy::All => { }
444 BlockDecodingStrategy::UptoBlocks(n) => {
445 if state.block_counter - block_counter_before >= n {
446 break;
447 }
448 }
449 BlockDecodingStrategy::UptoBytes(n) => {
450 if state.decoder_scratch.buffer.len() - buffer_size_before >= n {
451 break;
452 }
453 }
454 }
455 }
456
457 Ok(state.frame_finished)
458 }
459
460 pub fn collect(&mut self) -> Option<Vec<u8>> {
463 let finished = self.is_finished();
464 let state = self.state.as_mut()?;
465 if finished {
466 Some(state.decoder_scratch.buffer.drain())
467 } else {
468 state.decoder_scratch.buffer.drain_to_window_size()
469 }
470 }
471
472 pub fn collect_to_writer(&mut self, w: impl Write) -> Result<usize, Error> {
475 let finished = self.is_finished();
476 let state = match &mut self.state {
477 None => return Ok(0),
478 Some(s) => s,
479 };
480 if finished {
481 state.decoder_scratch.buffer.drain_to_writer(w)
482 } else {
483 state.decoder_scratch.buffer.drain_to_window_size_writer(w)
484 }
485 }
486
487 pub fn can_collect(&self) -> usize {
491 let finished = self.is_finished();
492 let state = match &self.state {
493 None => return 0,
494 Some(s) => s,
495 };
496 if finished {
497 state.decoder_scratch.buffer.can_drain()
498 } else {
499 state
500 .decoder_scratch
501 .buffer
502 .can_drain_to_window_size()
503 .unwrap_or(0)
504 }
505 }
506
507 pub fn decode_from_to(
521 &mut self,
522 source: &[u8],
523 target: &mut [u8],
524 ) -> Result<(usize, usize), FrameDecoderError> {
525 use FrameDecoderError as err;
526 let bytes_read_at_start = match &self.state {
527 Some(s) => s.bytes_read_counter,
528 None => 0,
529 };
530
531 if !self.is_finished() || self.state.is_none() {
532 let mut mt_source = source;
533
534 if self.state.is_none() {
535 self.init(&mut mt_source)?;
536 }
537
538 {
540 let state = match &mut self.state {
541 Some(s) => s,
542 None => panic!("Bug in library"),
543 };
544 let mut block_dec = decoding::block_decoder::new();
545
546 if state.frame_header.descriptor.content_checksum_flag()
547 && state.frame_finished
548 && state.check_sum.is_none()
549 {
550 if mt_source.len() >= 4 {
552 let chksum = mt_source[..4].try_into().expect("optimized away");
553 state.bytes_read_counter += 4;
554 let chksum = u32::from_le_bytes(chksum);
555 state.check_sum = Some(chksum);
556 }
557 return Ok((4, 0));
558 }
559
560 loop {
561 if mt_source.len() < 3 {
563 break;
564 }
565 let (block_header, block_header_size) = block_dec
566 .read_block_header(&mut mt_source)
567 .map_err(err::FailedToReadBlockHeader)?;
568
569 if mt_source.len() < block_header.content_size as usize {
572 break;
573 }
574 state.bytes_read_counter += u64::from(block_header_size);
575
576 let bytes_read_in_block_body = block_dec
577 .decode_block_content(
578 &block_header,
579 &mut state.decoder_scratch,
580 &mut mt_source,
581 )
582 .map_err(err::FailedToReadBlockBody)?;
583 state.bytes_read_counter += bytes_read_in_block_body;
584 state.block_counter += 1;
585
586 if block_header.last_block {
587 state.frame_finished = true;
588 if state.frame_header.descriptor.content_checksum_flag() {
589 if mt_source.len() >= 4 {
591 let chksum = mt_source[..4].try_into().expect("optimized away");
592 state.bytes_read_counter += 4;
593 let chksum = u32::from_le_bytes(chksum);
594 state.check_sum = Some(chksum);
595 }
596 }
597 break;
598 }
599 }
600 }
601 }
602
603 let result_len = self.read(target).map_err(err::FailedToDrainDecodebuffer)?;
604 let bytes_read_at_end = match &mut self.state {
605 Some(s) => s.bytes_read_counter,
606 None => panic!("Bug in library"),
607 };
608 let read_len = bytes_read_at_end - bytes_read_at_start;
609 Ok((read_len as usize, result_len))
610 }
611
612 pub fn decode_all(
623 &mut self,
624 input: &[u8],
625 output: &mut [u8],
626 ) -> Result<usize, FrameDecoderError> {
627 self.decode_all_with_window_limit(input, output, DEFAULT_MAX_WINDOW_SIZE)
628 }
629
630 pub(crate) fn decode_all_with_window_limit(
631 &mut self,
632 input: &[u8],
633 output: &mut [u8],
634 max_window_size: u64,
635 ) -> Result<usize, FrameDecoderError> {
636 self.decode_all_with_window_limit_and_dictionary(input, output, max_window_size, None)
637 }
638
639 pub(crate) fn decode_all_with_window_limit_and_dictionary(
640 &mut self,
641 mut input: &[u8],
642 mut output: &mut [u8],
643 max_window_size: u64,
644 dictionary: Option<PreparedDictionaryBytes<'_>>,
645 ) -> Result<usize, FrameDecoderError> {
646 let mut total_bytes_written = 0;
647 while !input.is_empty() {
648 match self.reset_with_window_limit_and_dictionary(
649 &mut input,
650 max_window_size,
651 dictionary,
652 ) {
653 Ok(_) => {}
654 Err(FrameDecoderError::ReadFrameHeaderError(
655 crate::decoding::errors::ReadFrameHeaderError::SkipFrame { length, .. },
656 )) => {
657 input = input
658 .get(length as usize..)
659 .ok_or(FrameDecoderError::FailedToSkipFrame)?;
660 continue;
661 }
662 Err(e) => return Err(e),
663 };
664 loop {
665 self.decode_blocks(&mut input, BlockDecodingStrategy::UptoBytes(1024 * 1024))?;
666 let bytes_written = self
667 .read(output)
668 .map_err(FrameDecoderError::FailedToDrainDecodebuffer)?;
669 output = &mut output[bytes_written..];
670 total_bytes_written += bytes_written;
671 if self.can_collect() != 0 {
672 return Err(FrameDecoderError::TargetTooSmall);
673 }
674 if self.is_finished() {
675 #[cfg(feature = "hash")]
678 if let Some(expected) = self.get_checksum_from_data() {
679 if self.get_calculated_checksum() != Some(expected) {
680 return Err(FrameDecoderError::ChecksumMismatch);
681 }
682 }
683 break;
684 }
685 }
686 }
687
688 Ok(total_bytes_written)
689 }
690
691 pub fn decode_all_to_vec(
704 &mut self,
705 input: &[u8],
706 output: &mut Vec<u8>,
707 ) -> Result<(), FrameDecoderError> {
708 let len = output.len();
709 let cap = output.capacity();
710 output.resize(cap, 0);
711 match self.decode_all(input, &mut output[len..]) {
712 Ok(bytes_written) => {
713 let new_len = core::cmp::min(len + bytes_written, cap); output.resize(new_len, 0);
715 Ok(())
716 }
717 Err(e) => {
718 output.resize(len, 0);
719 Err(e)
720 }
721 }
722 }
723}
724
725impl Read for FrameDecoder {
728 fn read(&mut self, target: &mut [u8]) -> Result<usize, Error> {
729 let state = match &mut self.state {
730 None => return Ok(0),
731 Some(s) => s,
732 };
733 if state.frame_finished {
734 state.decoder_scratch.buffer.read_all(target)
735 } else {
736 state.decoder_scratch.buffer.read(target)
737 }
738 }
739}