| 1 | //! Framedecoder is the man struct users interact with to decode zstd frames | 
| 2 | //! | 
|---|
| 3 | //! Zstandard compressed data is made of one or more [Frame]s. Each frame is independent and can be | 
|---|
| 4 | //! decompressed independently of other frames. This module contains structures | 
|---|
| 5 | //! and utilities that can be used to decode a frame. | 
|---|
| 6 |  | 
|---|
| 7 | use super::frame; | 
|---|
| 8 | use crate::decoding::dictionary::Dictionary; | 
|---|
| 9 | use crate::decoding::scratch::DecoderScratch; | 
|---|
| 10 | use crate::decoding::{self, dictionary}; | 
|---|
| 11 | use crate::io::{Error, Read, Write}; | 
|---|
| 12 | use alloc::collections::BTreeMap; | 
|---|
| 13 | use alloc::vec::Vec; | 
|---|
| 14 | use core::convert::TryInto; | 
|---|
| 15 | #[ cfg(feature = "std")] | 
|---|
| 16 | use std::error::Error as StdError; | 
|---|
| 17 |  | 
|---|
| 18 | /// This implements a decoder for zstd frames. | 
|---|
| 19 | /// | 
|---|
| 20 | /// This decoder is able to decode frames only partially and gives control | 
|---|
| 21 | /// over how many bytes/blocks will be decoded at a time (so you don't have to decode a 10GB file into memory all at once). | 
|---|
| 22 | /// It reads bytes as needed from a provided source and can be read from to collect partial results. | 
|---|
| 23 | /// | 
|---|
| 24 | /// If you want to just read the whole frame with an `io::Read` without having to deal with manually calling [FrameDecoder::decode_blocks] | 
|---|
| 25 | /// you can use the provided StreamingDecoder with wraps this FrameDecoder | 
|---|
| 26 | /// | 
|---|
| 27 | /// Workflow is as follows: | 
|---|
| 28 | /// ``` | 
|---|
| 29 | /// use ruzstd::frame_decoder::BlockDecodingStrategy; | 
|---|
| 30 | /// | 
|---|
| 31 | /// # #[ cfg(feature = "std")] | 
|---|
| 32 | /// use std::io::{Read, Write}; | 
|---|
| 33 | /// | 
|---|
| 34 | /// // no_std environments can use the crate's own Read traits | 
|---|
| 35 | /// # #[ cfg(not(feature = "std"))] | 
|---|
| 36 | /// use ruzstd::io::{Read, Write}; | 
|---|
| 37 | /// | 
|---|
| 38 | /// fn decode_this(mut file: impl Read) { | 
|---|
| 39 | ///     //Create a new decoder | 
|---|
| 40 | ///     let mut frame_dec = ruzstd::FrameDecoder::new(); | 
|---|
| 41 | ///     let mut result = Vec::new(); | 
|---|
| 42 | /// | 
|---|
| 43 | ///     // Use reset or init to make the decoder ready to decode the frame from the io::Read | 
|---|
| 44 | ///     frame_dec.reset(&mut file).unwrap(); | 
|---|
| 45 | /// | 
|---|
| 46 | ///     // Loop until the frame has been decoded completely | 
|---|
| 47 | ///     while !frame_dec.is_finished() { | 
|---|
| 48 | ///         // decode (roughly) batch_size many bytes | 
|---|
| 49 | ///         frame_dec.decode_blocks(&mut file, BlockDecodingStrategy::UptoBytes(1024)).unwrap(); | 
|---|
| 50 | /// | 
|---|
| 51 | ///         // read from the decoder to collect bytes from the internal buffer | 
|---|
| 52 | ///         let bytes_read = frame_dec.read(result.as_mut_slice()).unwrap(); | 
|---|
| 53 | /// | 
|---|
| 54 | ///         // then do something with it | 
|---|
| 55 | ///         do_something(&result[0..bytes_read]); | 
|---|
| 56 | ///     } | 
|---|
| 57 | /// | 
|---|
| 58 | ///     // handle the last chunk of data | 
|---|
| 59 | ///     while frame_dec.can_collect() > 0 { | 
|---|
| 60 | ///         let x = frame_dec.read(result.as_mut_slice()).unwrap(); | 
|---|
| 61 | /// | 
|---|
| 62 | ///         do_something(&result[0..x]); | 
|---|
| 63 | ///     } | 
|---|
| 64 | /// } | 
|---|
| 65 | /// | 
|---|
| 66 | /// fn do_something(data: &[u8]) { | 
|---|
| 67 | /// # #[ cfg(feature = "std")] | 
|---|
| 68 | ///     std::io::stdout().write_all(data).unwrap(); | 
|---|
| 69 | /// } | 
|---|
| 70 | /// ``` | 
|---|
| 71 | pub struct FrameDecoder { | 
|---|
| 72 | state: Option<FrameDecoderState>, | 
|---|
| 73 | dicts: BTreeMap<u32, Dictionary>, | 
|---|
| 74 | } | 
|---|
| 75 |  | 
|---|
| 76 | struct FrameDecoderState { | 
|---|
| 77 | pub frame: frame::Frame, | 
|---|
| 78 | decoder_scratch: DecoderScratch, | 
|---|
| 79 | frame_finished: bool, | 
|---|
| 80 | block_counter: usize, | 
|---|
| 81 | bytes_read_counter: u64, | 
|---|
| 82 | check_sum: Option<u32>, | 
|---|
| 83 | using_dict: Option<u32>, | 
|---|
| 84 | } | 
|---|
| 85 |  | 
|---|
| 86 | pub enum BlockDecodingStrategy { | 
|---|
| 87 | All, | 
|---|
| 88 | UptoBlocks(usize), | 
|---|
| 89 | UptoBytes(usize), | 
|---|
| 90 | } | 
|---|
| 91 |  | 
|---|
| 92 | #[ derive(Debug)] | 
|---|
| 93 | #[ non_exhaustive] | 
|---|
| 94 | pub enum FrameDecoderError { | 
|---|
| 95 | ReadFrameHeaderError(frame::ReadFrameHeaderError), | 
|---|
| 96 | FrameHeaderError(frame::FrameHeaderError), | 
|---|
| 97 | WindowSizeTooBig { requested: u64 }, | 
|---|
| 98 | DictionaryDecodeError(dictionary::DictionaryDecodeError), | 
|---|
| 99 | FailedToReadBlockHeader(decoding::block_decoder::BlockHeaderReadError), | 
|---|
| 100 | FailedToReadBlockBody(decoding::block_decoder::DecodeBlockContentError), | 
|---|
| 101 | FailedToReadChecksum(Error), | 
|---|
| 102 | NotYetInitialized, | 
|---|
| 103 | FailedToInitialize(frame::FrameHeaderError), | 
|---|
| 104 | FailedToDrainDecodebuffer(Error), | 
|---|
| 105 | FailedToSkipFrame, | 
|---|
| 106 | TargetTooSmall, | 
|---|
| 107 | DictNotProvided { dict_id: u32 }, | 
|---|
| 108 | } | 
|---|
| 109 |  | 
|---|
| 110 | #[ cfg(feature = "std")] | 
|---|
| 111 | impl StdError for FrameDecoderError { | 
|---|
| 112 | fn source(&self) -> Option<&(dyn StdError + 'static)> { | 
|---|
| 113 | match self { | 
|---|
| 114 | FrameDecoderError::ReadFrameHeaderError(source: &ReadFrameHeaderError) => Some(source), | 
|---|
| 115 | FrameDecoderError::FrameHeaderError(source: &FrameHeaderError) => Some(source), | 
|---|
| 116 | FrameDecoderError::DictionaryDecodeError(source: &DictionaryDecodeError) => Some(source), | 
|---|
| 117 | FrameDecoderError::FailedToReadBlockHeader(source: &BlockHeaderReadError) => Some(source), | 
|---|
| 118 | FrameDecoderError::FailedToReadBlockBody(source: &DecodeBlockContentError) => Some(source), | 
|---|
| 119 | FrameDecoderError::FailedToReadChecksum(source: &Error) => Some(source), | 
|---|
| 120 | FrameDecoderError::FailedToInitialize(source: &FrameHeaderError) => Some(source), | 
|---|
| 121 | FrameDecoderError::FailedToDrainDecodebuffer(source: &Error) => Some(source), | 
|---|
| 122 | _ => None, | 
|---|
| 123 | } | 
|---|
| 124 | } | 
|---|
| 125 | } | 
|---|
| 126 |  | 
|---|
| 127 | impl core::fmt::Display for FrameDecoderError { | 
|---|
| 128 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> ::core::fmt::Result { | 
|---|
| 129 | match self { | 
|---|
| 130 | FrameDecoderError::ReadFrameHeaderError(e) => { | 
|---|
| 131 | write!(f, "{:?} ", e) | 
|---|
| 132 | } | 
|---|
| 133 | FrameDecoderError::FrameHeaderError(e) => { | 
|---|
| 134 | write!(f, "{:?} ", e) | 
|---|
| 135 | } | 
|---|
| 136 | FrameDecoderError::WindowSizeTooBig { requested } => { | 
|---|
| 137 | write!( | 
|---|
| 138 | f, | 
|---|
| 139 | "Specified window_size is too big; Requested: {} , Max: {} ", | 
|---|
| 140 | requested, MAX_WINDOW_SIZE, | 
|---|
| 141 | ) | 
|---|
| 142 | } | 
|---|
| 143 | FrameDecoderError::DictionaryDecodeError(e) => { | 
|---|
| 144 | write!(f, "{:?} ", e) | 
|---|
| 145 | } | 
|---|
| 146 | FrameDecoderError::FailedToReadBlockHeader(e) => { | 
|---|
| 147 | write!(f, "Failed to parse/decode block body: {} ", e) | 
|---|
| 148 | } | 
|---|
| 149 | FrameDecoderError::FailedToReadBlockBody(e) => { | 
|---|
| 150 | write!(f, "Failed to parse block header: {} ", e) | 
|---|
| 151 | } | 
|---|
| 152 | FrameDecoderError::FailedToReadChecksum(e) => { | 
|---|
| 153 | write!(f, "Failed to read checksum: {} ", e) | 
|---|
| 154 | } | 
|---|
| 155 | FrameDecoderError::NotYetInitialized => { | 
|---|
| 156 | write!(f, "Decoder must initialized or reset before using it",) | 
|---|
| 157 | } | 
|---|
| 158 | FrameDecoderError::FailedToInitialize(e) => { | 
|---|
| 159 | write!(f, "Decoder encountered error while initializing: {} ", e) | 
|---|
| 160 | } | 
|---|
| 161 | FrameDecoderError::FailedToDrainDecodebuffer(e) => { | 
|---|
| 162 | write!( | 
|---|
| 163 | f, | 
|---|
| 164 | "Decoder encountered error while draining the decodebuffer: {} ", | 
|---|
| 165 | e, | 
|---|
| 166 | ) | 
|---|
| 167 | } | 
|---|
| 168 | FrameDecoderError::FailedToSkipFrame => { | 
|---|
| 169 | write!( | 
|---|
| 170 | f, | 
|---|
| 171 | "Failed to skip bytes for the length given in the frame header" | 
|---|
| 172 | ) | 
|---|
| 173 | } | 
|---|
| 174 | FrameDecoderError::TargetTooSmall => { | 
|---|
| 175 | write!(f, "Target must have at least as many bytes as the contentsize of the frame reports") | 
|---|
| 176 | } | 
|---|
| 177 | FrameDecoderError::DictNotProvided { dict_id } => { | 
|---|
| 178 | write!(f, "Frame header specified dictionary id 0x{:X}  that wasnt provided by add_dict() or reset_with_dict()", dict_id) | 
|---|
| 179 | } | 
|---|
| 180 | } | 
|---|
| 181 | } | 
|---|
| 182 | } | 
|---|
| 183 |  | 
|---|
| 184 | impl From<dictionary::DictionaryDecodeError> for FrameDecoderError { | 
|---|
| 185 | fn from(val: dictionary::DictionaryDecodeError) -> Self { | 
|---|
| 186 | Self::DictionaryDecodeError(val) | 
|---|
| 187 | } | 
|---|
| 188 | } | 
|---|
| 189 |  | 
|---|
| 190 | impl From<decoding::block_decoder::BlockHeaderReadError> for FrameDecoderError { | 
|---|
| 191 | fn from(val: decoding::block_decoder::BlockHeaderReadError) -> Self { | 
|---|
| 192 | Self::FailedToReadBlockHeader(val) | 
|---|
| 193 | } | 
|---|
| 194 | } | 
|---|
| 195 |  | 
|---|
| 196 | impl From<frame::FrameHeaderError> for FrameDecoderError { | 
|---|
| 197 | fn from(val: frame::FrameHeaderError) -> Self { | 
|---|
| 198 | Self::FrameHeaderError(val) | 
|---|
| 199 | } | 
|---|
| 200 | } | 
|---|
| 201 |  | 
|---|
| 202 | impl From<frame::ReadFrameHeaderError> for FrameDecoderError { | 
|---|
| 203 | fn from(val: frame::ReadFrameHeaderError) -> Self { | 
|---|
| 204 | Self::ReadFrameHeaderError(val) | 
|---|
| 205 | } | 
|---|
| 206 | } | 
|---|
| 207 |  | 
|---|
| 208 | const MAX_WINDOW_SIZE: u64 = 1024 * 1024 * 100; | 
|---|
| 209 |  | 
|---|
| 210 | impl FrameDecoderState { | 
|---|
| 211 | pub fn new(source: impl Read) -> Result<FrameDecoderState, FrameDecoderError> { | 
|---|
| 212 | let (frame, header_size) = frame::read_frame_header(source)?; | 
|---|
| 213 | let window_size = frame.header.window_size()?; | 
|---|
| 214 | Ok(FrameDecoderState { | 
|---|
| 215 | frame, | 
|---|
| 216 | frame_finished: false, | 
|---|
| 217 | block_counter: 0, | 
|---|
| 218 | decoder_scratch: DecoderScratch::new(window_size as usize), | 
|---|
| 219 | bytes_read_counter: u64::from(header_size), | 
|---|
| 220 | check_sum: None, | 
|---|
| 221 | using_dict: None, | 
|---|
| 222 | }) | 
|---|
| 223 | } | 
|---|
| 224 |  | 
|---|
| 225 | pub fn reset(&mut self, source: impl Read) -> Result<(), FrameDecoderError> { | 
|---|
| 226 | let (frame, header_size) = frame::read_frame_header(source)?; | 
|---|
| 227 | let window_size = frame.header.window_size()?; | 
|---|
| 228 |  | 
|---|
| 229 | if window_size > MAX_WINDOW_SIZE { | 
|---|
| 230 | return Err(FrameDecoderError::WindowSizeTooBig { | 
|---|
| 231 | requested: window_size, | 
|---|
| 232 | }); | 
|---|
| 233 | } | 
|---|
| 234 |  | 
|---|
| 235 | self.frame = frame; | 
|---|
| 236 | self.frame_finished = false; | 
|---|
| 237 | self.block_counter = 0; | 
|---|
| 238 | self.decoder_scratch.reset(window_size as usize); | 
|---|
| 239 | self.bytes_read_counter = u64::from(header_size); | 
|---|
| 240 | self.check_sum = None; | 
|---|
| 241 | self.using_dict = None; | 
|---|
| 242 | Ok(()) | 
|---|
| 243 | } | 
|---|
| 244 | } | 
|---|
| 245 |  | 
|---|
| 246 | impl Default for FrameDecoder { | 
|---|
| 247 | fn default() -> Self { | 
|---|
| 248 | Self::new() | 
|---|
| 249 | } | 
|---|
| 250 | } | 
|---|
| 251 |  | 
|---|
| 252 | impl FrameDecoder { | 
|---|
| 253 | /// This will create a new decoder without allocating anything yet. | 
|---|
| 254 | /// init()/reset() will allocate all needed buffers if it is the first time this decoder is used | 
|---|
| 255 | /// else they just reset these buffers with not further allocations | 
|---|
| 256 | pub fn new() -> FrameDecoder { | 
|---|
| 257 | FrameDecoder { | 
|---|
| 258 | state: None, | 
|---|
| 259 | dicts: BTreeMap::new(), | 
|---|
| 260 | } | 
|---|
| 261 | } | 
|---|
| 262 |  | 
|---|
| 263 | /// init() will allocate all needed buffers if it is the first time this decoder is used | 
|---|
| 264 | /// else they just reset these buffers with not further allocations | 
|---|
| 265 | /// | 
|---|
| 266 | /// Note that all bytes currently in the decodebuffer from any previous frame will be lost. Collect them with collect()/collect_to_writer() | 
|---|
| 267 | /// | 
|---|
| 268 | /// equivalent to reset() | 
|---|
| 269 | pub fn init(&mut self, source: impl Read) -> Result<(), FrameDecoderError> { | 
|---|
| 270 | self.reset(source) | 
|---|
| 271 | } | 
|---|
| 272 |  | 
|---|
| 273 | /// reset() will allocate all needed buffers if it is the first time this decoder is used | 
|---|
| 274 | /// else they just reset these buffers with not further allocations | 
|---|
| 275 | /// | 
|---|
| 276 | /// Note that all bytes currently in the decodebuffer from any previous frame will be lost. Collect them with collect()/collect_to_writer() | 
|---|
| 277 | /// | 
|---|
| 278 | /// equivalent to init() | 
|---|
| 279 | pub fn reset(&mut self, source: impl Read) -> Result<(), FrameDecoderError> { | 
|---|
| 280 | use FrameDecoderError as err; | 
|---|
| 281 | let state = match &mut self.state { | 
|---|
| 282 | Some(s) => { | 
|---|
| 283 | s.reset(source)?; | 
|---|
| 284 | s | 
|---|
| 285 | } | 
|---|
| 286 | None => { | 
|---|
| 287 | self.state = Some(FrameDecoderState::new(source)?); | 
|---|
| 288 | self.state.as_mut().unwrap() | 
|---|
| 289 | } | 
|---|
| 290 | }; | 
|---|
| 291 | if let Some(dict_id) = state.frame.header.dictionary_id() { | 
|---|
| 292 | let dict = self | 
|---|
| 293 | .dicts | 
|---|
| 294 | .get(&dict_id) | 
|---|
| 295 | .ok_or(err::DictNotProvided { dict_id })?; | 
|---|
| 296 | state.decoder_scratch.init_from_dict(dict); | 
|---|
| 297 | state.using_dict = Some(dict_id); | 
|---|
| 298 | } | 
|---|
| 299 | Ok(()) | 
|---|
| 300 | } | 
|---|
| 301 |  | 
|---|
| 302 | /// Add a dict to the FrameDecoder that can be used when needed. The FrameDecoder uses the appropriate one dynamically | 
|---|
| 303 | pub fn add_dict(&mut self, dict: Dictionary) -> Result<(), FrameDecoderError> { | 
|---|
| 304 | self.dicts.insert(dict.id, dict); | 
|---|
| 305 | Ok(()) | 
|---|
| 306 | } | 
|---|
| 307 |  | 
|---|
| 308 | pub fn force_dict(&mut self, dict_id: u32) -> Result<(), FrameDecoderError> { | 
|---|
| 309 | use FrameDecoderError as err; | 
|---|
| 310 | let Some(state) = self.state.as_mut() else { | 
|---|
| 311 | return Err(err::NotYetInitialized); | 
|---|
| 312 | }; | 
|---|
| 313 |  | 
|---|
| 314 | let dict = self | 
|---|
| 315 | .dicts | 
|---|
| 316 | .get(&dict_id) | 
|---|
| 317 | .ok_or(err::DictNotProvided { dict_id })?; | 
|---|
| 318 | state.decoder_scratch.init_from_dict(dict); | 
|---|
| 319 | state.using_dict = Some(dict_id); | 
|---|
| 320 |  | 
|---|
| 321 | Ok(()) | 
|---|
| 322 | } | 
|---|
| 323 |  | 
|---|
| 324 | /// Returns how many bytes the frame contains after decompression | 
|---|
| 325 | pub fn content_size(&self) -> u64 { | 
|---|
| 326 | match &self.state { | 
|---|
| 327 | None => 0, | 
|---|
| 328 | Some(s) => s.frame.header.frame_content_size(), | 
|---|
| 329 | } | 
|---|
| 330 | } | 
|---|
| 331 |  | 
|---|
| 332 | /// Returns the checksum that was read from the data. Only available after all bytes have been read. It is the last 4 bytes of a zstd-frame | 
|---|
| 333 | pub fn get_checksum_from_data(&self) -> Option<u32> { | 
|---|
| 334 | let state = match &self.state { | 
|---|
| 335 | None => return None, | 
|---|
| 336 | Some(s) => s, | 
|---|
| 337 | }; | 
|---|
| 338 |  | 
|---|
| 339 | state.check_sum | 
|---|
| 340 | } | 
|---|
| 341 |  | 
|---|
| 342 | /// Returns the checksum that was calculated while decoding. | 
|---|
| 343 | /// Only a sensible value after all decoded bytes have been collected/read from the FrameDecoder | 
|---|
| 344 | #[ cfg(feature = "hash")] | 
|---|
| 345 | pub fn get_calculated_checksum(&self) -> Option<u32> { | 
|---|
| 346 | use core::hash::Hasher; | 
|---|
| 347 |  | 
|---|
| 348 | let state = match &self.state { | 
|---|
| 349 | None => return None, | 
|---|
| 350 | Some(s) => s, | 
|---|
| 351 | }; | 
|---|
| 352 | let cksum_64bit = state.decoder_scratch.buffer.hash.finish(); | 
|---|
| 353 | //truncate to lower 32bit because reasons... | 
|---|
| 354 | Some(cksum_64bit as u32) | 
|---|
| 355 | } | 
|---|
| 356 |  | 
|---|
| 357 | /// Counter for how many bytes have been consumed while decoding the frame | 
|---|
| 358 | pub fn bytes_read_from_source(&self) -> u64 { | 
|---|
| 359 | let state = match &self.state { | 
|---|
| 360 | None => return 0, | 
|---|
| 361 | Some(s) => s, | 
|---|
| 362 | }; | 
|---|
| 363 | state.bytes_read_counter | 
|---|
| 364 | } | 
|---|
| 365 |  | 
|---|
| 366 | /// Whether the current frames last block has been decoded yet | 
|---|
| 367 | /// If this returns true you can call the drain* functions to get all content | 
|---|
| 368 | /// (the read() function will drain automatically if this returns true) | 
|---|
| 369 | pub fn is_finished(&self) -> bool { | 
|---|
| 370 | let state = match &self.state { | 
|---|
| 371 | None => return true, | 
|---|
| 372 | Some(s) => s, | 
|---|
| 373 | }; | 
|---|
| 374 | if state.frame.header.descriptor.content_checksum_flag() { | 
|---|
| 375 | state.frame_finished && state.check_sum.is_some() | 
|---|
| 376 | } else { | 
|---|
| 377 | state.frame_finished | 
|---|
| 378 | } | 
|---|
| 379 | } | 
|---|
| 380 |  | 
|---|
| 381 | /// Counter for how many blocks have already been decoded | 
|---|
| 382 | pub fn blocks_decoded(&self) -> usize { | 
|---|
| 383 | let state = match &self.state { | 
|---|
| 384 | None => return 0, | 
|---|
| 385 | Some(s) => s, | 
|---|
| 386 | }; | 
|---|
| 387 | state.block_counter | 
|---|
| 388 | } | 
|---|
| 389 |  | 
|---|
| 390 | /// Decodes blocks from a reader. It requires that the framedecoder has been initialized first. | 
|---|
| 391 | /// The Strategy influences how many blocks will be decoded before the function returns | 
|---|
| 392 | /// This is important if you want to manage memory consumption carefully. If you don't care | 
|---|
| 393 | /// about that you can just choose the strategy "All" and have all blocks of the frame decoded into the buffer | 
|---|
| 394 | pub fn decode_blocks( | 
|---|
| 395 | &mut self, | 
|---|
| 396 | mut source: impl Read, | 
|---|
| 397 | strat: BlockDecodingStrategy, | 
|---|
| 398 | ) -> Result<bool, FrameDecoderError> { | 
|---|
| 399 | use FrameDecoderError as err; | 
|---|
| 400 | let state = self.state.as_mut().ok_or(err::NotYetInitialized)?; | 
|---|
| 401 |  | 
|---|
| 402 | let mut block_dec = decoding::block_decoder::new(); | 
|---|
| 403 |  | 
|---|
| 404 | let buffer_size_before = state.decoder_scratch.buffer.len(); | 
|---|
| 405 | let block_counter_before = state.block_counter; | 
|---|
| 406 | loop { | 
|---|
| 407 | vprintln!( "################"); | 
|---|
| 408 | vprintln!( "Next Block: {} ", state.block_counter); | 
|---|
| 409 | vprintln!( "################"); | 
|---|
| 410 | let (block_header, block_header_size) = block_dec | 
|---|
| 411 | .read_block_header(&mut source) | 
|---|
| 412 | .map_err(err::FailedToReadBlockHeader)?; | 
|---|
| 413 | state.bytes_read_counter += u64::from(block_header_size); | 
|---|
| 414 |  | 
|---|
| 415 | vprintln!(); | 
|---|
| 416 | vprintln!( | 
|---|
| 417 | "Found {}  block with size: {} , which will be of size: {} ", | 
|---|
| 418 | block_header.block_type, | 
|---|
| 419 | block_header.content_size, | 
|---|
| 420 | block_header.decompressed_size | 
|---|
| 421 | ); | 
|---|
| 422 |  | 
|---|
| 423 | let bytes_read_in_block_body = block_dec | 
|---|
| 424 | .decode_block_content(&block_header, &mut state.decoder_scratch, &mut source) | 
|---|
| 425 | .map_err(err::FailedToReadBlockBody)?; | 
|---|
| 426 | state.bytes_read_counter += bytes_read_in_block_body; | 
|---|
| 427 |  | 
|---|
| 428 | state.block_counter += 1; | 
|---|
| 429 |  | 
|---|
| 430 | vprintln!( "Output: {} ", state.decoder_scratch.buffer.len()); | 
|---|
| 431 |  | 
|---|
| 432 | if block_header.last_block { | 
|---|
| 433 | state.frame_finished = true; | 
|---|
| 434 | if state.frame.header.descriptor.content_checksum_flag() { | 
|---|
| 435 | let mut chksum = [0u8; 4]; | 
|---|
| 436 | source | 
|---|
| 437 | .read_exact(&mut chksum) | 
|---|
| 438 | .map_err(err::FailedToReadChecksum)?; | 
|---|
| 439 | state.bytes_read_counter += 4; | 
|---|
| 440 | let chksum = u32::from_le_bytes(chksum); | 
|---|
| 441 | state.check_sum = Some(chksum); | 
|---|
| 442 | } | 
|---|
| 443 | break; | 
|---|
| 444 | } | 
|---|
| 445 |  | 
|---|
| 446 | match strat { | 
|---|
| 447 | BlockDecodingStrategy::All => { /* keep going */ } | 
|---|
| 448 | BlockDecodingStrategy::UptoBlocks(n) => { | 
|---|
| 449 | if state.block_counter - block_counter_before >= n { | 
|---|
| 450 | break; | 
|---|
| 451 | } | 
|---|
| 452 | } | 
|---|
| 453 | BlockDecodingStrategy::UptoBytes(n) => { | 
|---|
| 454 | if state.decoder_scratch.buffer.len() - buffer_size_before >= n { | 
|---|
| 455 | break; | 
|---|
| 456 | } | 
|---|
| 457 | } | 
|---|
| 458 | } | 
|---|
| 459 | } | 
|---|
| 460 |  | 
|---|
| 461 | Ok(state.frame_finished) | 
|---|
| 462 | } | 
|---|
| 463 |  | 
|---|
| 464 | /// Collect bytes and retain window_size bytes while decoding is still going on. | 
|---|
| 465 | /// After decoding of the frame (is_finished() == true) has finished it will collect all remaining bytes | 
|---|
| 466 | pub fn collect(&mut self) -> Option<Vec<u8>> { | 
|---|
| 467 | let finished = self.is_finished(); | 
|---|
| 468 | let state = self.state.as_mut()?; | 
|---|
| 469 | if finished { | 
|---|
| 470 | Some(state.decoder_scratch.buffer.drain()) | 
|---|
| 471 | } else { | 
|---|
| 472 | state.decoder_scratch.buffer.drain_to_window_size() | 
|---|
| 473 | } | 
|---|
| 474 | } | 
|---|
| 475 |  | 
|---|
| 476 | /// Collect bytes and retain window_size bytes while decoding is still going on. | 
|---|
| 477 | /// After decoding of the frame (is_finished() == true) has finished it will collect all remaining bytes | 
|---|
| 478 | pub fn collect_to_writer(&mut self, w: impl Write) -> Result<usize, Error> { | 
|---|
| 479 | let finished = self.is_finished(); | 
|---|
| 480 | let state = match &mut self.state { | 
|---|
| 481 | None => return Ok(0), | 
|---|
| 482 | Some(s) => s, | 
|---|
| 483 | }; | 
|---|
| 484 | if finished { | 
|---|
| 485 | state.decoder_scratch.buffer.drain_to_writer(w) | 
|---|
| 486 | } else { | 
|---|
| 487 | state.decoder_scratch.buffer.drain_to_window_size_writer(w) | 
|---|
| 488 | } | 
|---|
| 489 | } | 
|---|
| 490 |  | 
|---|
| 491 | /// How many bytes can currently be collected from the decodebuffer, while decoding is going on this will be lower than the actual decodbuffer size | 
|---|
| 492 | /// because window_size bytes need to be retained for decoding. | 
|---|
| 493 | /// After decoding of the frame (is_finished() == true) has finished it will report all remaining bytes | 
|---|
| 494 | pub fn can_collect(&self) -> usize { | 
|---|
| 495 | let finished = self.is_finished(); | 
|---|
| 496 | let state = match &self.state { | 
|---|
| 497 | None => return 0, | 
|---|
| 498 | Some(s) => s, | 
|---|
| 499 | }; | 
|---|
| 500 | if finished { | 
|---|
| 501 | state.decoder_scratch.buffer.can_drain() | 
|---|
| 502 | } else { | 
|---|
| 503 | state | 
|---|
| 504 | .decoder_scratch | 
|---|
| 505 | .buffer | 
|---|
| 506 | .can_drain_to_window_size() | 
|---|
| 507 | .unwrap_or(0) | 
|---|
| 508 | } | 
|---|
| 509 | } | 
|---|
| 510 |  | 
|---|
| 511 | /// Decodes as many blocks as possible from the source slice and reads from the decodebuffer into the target slice | 
|---|
| 512 | /// The source slice may contain only parts of a frame but must contain at least one full block to make progress | 
|---|
| 513 | /// | 
|---|
| 514 | /// By all means use decode_blocks if you have a io.Reader available. This is just for compatibility with other decompressors | 
|---|
| 515 | /// which try to serve an old-style c api | 
|---|
| 516 | /// | 
|---|
| 517 | /// Returns (read, written), if read == 0 then the source did not contain a full block and further calls with the same | 
|---|
| 518 | /// input will not make any progress! | 
|---|
| 519 | /// | 
|---|
| 520 | /// Note that no kind of block can be bigger than 128kb. | 
|---|
| 521 | /// So to be safe use at least 128*1024 (max block content size) + 3 (block_header size) + 18 (max frame_header size) bytes as your source buffer | 
|---|
| 522 | /// | 
|---|
| 523 | /// You may call this function with an empty source after all bytes have been decoded. This is equivalent to just call decoder.read(&mut target) | 
|---|
| 524 | pub fn decode_from_to( | 
|---|
| 525 | &mut self, | 
|---|
| 526 | source: &[u8], | 
|---|
| 527 | target: &mut [u8], | 
|---|
| 528 | ) -> Result<(usize, usize), FrameDecoderError> { | 
|---|
| 529 | use FrameDecoderError as err; | 
|---|
| 530 | let bytes_read_at_start = match &self.state { | 
|---|
| 531 | Some(s) => s.bytes_read_counter, | 
|---|
| 532 | None => 0, | 
|---|
| 533 | }; | 
|---|
| 534 |  | 
|---|
| 535 | if !self.is_finished() || self.state.is_none() { | 
|---|
| 536 | let mut mt_source = source; | 
|---|
| 537 |  | 
|---|
| 538 | if self.state.is_none() { | 
|---|
| 539 | self.init(&mut mt_source)?; | 
|---|
| 540 | } | 
|---|
| 541 |  | 
|---|
| 542 | //pseudo block to scope "state" so we can borrow self again after the block | 
|---|
| 543 | { | 
|---|
| 544 | let state = match &mut self.state { | 
|---|
| 545 | Some(s) => s, | 
|---|
| 546 | None => panic!( "Bug in library"), | 
|---|
| 547 | }; | 
|---|
| 548 | let mut block_dec = decoding::block_decoder::new(); | 
|---|
| 549 |  | 
|---|
| 550 | if state.frame.header.descriptor.content_checksum_flag() | 
|---|
| 551 | && state.frame_finished | 
|---|
| 552 | && state.check_sum.is_none() | 
|---|
| 553 | { | 
|---|
| 554 | //this block is needed if the checksum were the only 4 bytes that were not included in the last decode_from_to call for a frame | 
|---|
| 555 | if mt_source.len() >= 4 { | 
|---|
| 556 | let chksum = mt_source[..4].try_into().expect( "optimized away"); | 
|---|
| 557 | state.bytes_read_counter += 4; | 
|---|
| 558 | let chksum = u32::from_le_bytes(chksum); | 
|---|
| 559 | state.check_sum = Some(chksum); | 
|---|
| 560 | } | 
|---|
| 561 | return Ok((4, 0)); | 
|---|
| 562 | } | 
|---|
| 563 |  | 
|---|
| 564 | loop { | 
|---|
| 565 | //check if there are enough bytes for the next header | 
|---|
| 566 | if mt_source.len() < 3 { | 
|---|
| 567 | break; | 
|---|
| 568 | } | 
|---|
| 569 | let (block_header, block_header_size) = block_dec | 
|---|
| 570 | .read_block_header(&mut mt_source) | 
|---|
| 571 | .map_err(err::FailedToReadBlockHeader)?; | 
|---|
| 572 |  | 
|---|
| 573 | // check the needed size for the block before updating counters. | 
|---|
| 574 | // If not enough bytes are in the source, the header will have to be read again, so act like we never read it in the first place | 
|---|
| 575 | if mt_source.len() < block_header.content_size as usize { | 
|---|
| 576 | break; | 
|---|
| 577 | } | 
|---|
| 578 | state.bytes_read_counter += u64::from(block_header_size); | 
|---|
| 579 |  | 
|---|
| 580 | let bytes_read_in_block_body = block_dec | 
|---|
| 581 | .decode_block_content( | 
|---|
| 582 | &block_header, | 
|---|
| 583 | &mut state.decoder_scratch, | 
|---|
| 584 | &mut mt_source, | 
|---|
| 585 | ) | 
|---|
| 586 | .map_err(err::FailedToReadBlockBody)?; | 
|---|
| 587 | state.bytes_read_counter += bytes_read_in_block_body; | 
|---|
| 588 | state.block_counter += 1; | 
|---|
| 589 |  | 
|---|
| 590 | if block_header.last_block { | 
|---|
| 591 | state.frame_finished = true; | 
|---|
| 592 | if state.frame.header.descriptor.content_checksum_flag() { | 
|---|
| 593 | //if there are enough bytes handle this here. Else the block at the start of this function will handle it at the next call | 
|---|
| 594 | if mt_source.len() >= 4 { | 
|---|
| 595 | let chksum = mt_source[..4].try_into().expect( "optimized away"); | 
|---|
| 596 | state.bytes_read_counter += 4; | 
|---|
| 597 | let chksum = u32::from_le_bytes(chksum); | 
|---|
| 598 | state.check_sum = Some(chksum); | 
|---|
| 599 | } | 
|---|
| 600 | } | 
|---|
| 601 | break; | 
|---|
| 602 | } | 
|---|
| 603 | } | 
|---|
| 604 | } | 
|---|
| 605 | } | 
|---|
| 606 |  | 
|---|
| 607 | let result_len = self.read(target).map_err(err::FailedToDrainDecodebuffer)?; | 
|---|
| 608 | let bytes_read_at_end = match &mut self.state { | 
|---|
| 609 | Some(s) => s.bytes_read_counter, | 
|---|
| 610 | None => panic!( "Bug in library"), | 
|---|
| 611 | }; | 
|---|
| 612 | let read_len = bytes_read_at_end - bytes_read_at_start; | 
|---|
| 613 | Ok((read_len as usize, result_len)) | 
|---|
| 614 | } | 
|---|
| 615 |  | 
|---|
| 616 | /// Decode multiple frames into the output slice. | 
|---|
| 617 | /// | 
|---|
| 618 | /// `input` must contain an exact number of frames. | 
|---|
| 619 | /// | 
|---|
| 620 | /// `output` must be large enough to hold the decompressed data. If you don't know | 
|---|
| 621 | /// how large the output will be, use [`FrameDecoder::decode_blocks`] instead. | 
|---|
| 622 | /// | 
|---|
| 623 | /// This calls [`FrameDecoder::init`], and all bytes currently in the decoder will be lost. | 
|---|
| 624 | /// | 
|---|
| 625 | /// Returns the number of bytes written to `output`. | 
|---|
| 626 | pub fn decode_all( | 
|---|
| 627 | &mut self, | 
|---|
| 628 | mut input: &[u8], | 
|---|
| 629 | mut output: &mut [u8], | 
|---|
| 630 | ) -> Result<usize, FrameDecoderError> { | 
|---|
| 631 | let mut total_bytes_written = 0; | 
|---|
| 632 | while !input.is_empty() { | 
|---|
| 633 | match self.init(&mut input) { | 
|---|
| 634 | Ok(_) => {} | 
|---|
| 635 | Err(FrameDecoderError::ReadFrameHeaderError( | 
|---|
| 636 | frame::ReadFrameHeaderError::SkipFrame { length, .. }, | 
|---|
| 637 | )) => { | 
|---|
| 638 | input = input | 
|---|
| 639 | .get(length as usize..) | 
|---|
| 640 | .ok_or(FrameDecoderError::FailedToSkipFrame)?; | 
|---|
| 641 | continue; | 
|---|
| 642 | } | 
|---|
| 643 | Err(e) => return Err(e), | 
|---|
| 644 | }; | 
|---|
| 645 | loop { | 
|---|
| 646 | self.decode_blocks(&mut input, BlockDecodingStrategy::UptoBlocks(1))?; | 
|---|
| 647 | let bytes_written = self | 
|---|
| 648 | .read(output) | 
|---|
| 649 | .map_err(FrameDecoderError::FailedToDrainDecodebuffer)?; | 
|---|
| 650 | output = &mut output[bytes_written..]; | 
|---|
| 651 | total_bytes_written += bytes_written; | 
|---|
| 652 | if self.can_collect() != 0 { | 
|---|
| 653 | return Err(FrameDecoderError::TargetTooSmall); | 
|---|
| 654 | } | 
|---|
| 655 | if self.is_finished() { | 
|---|
| 656 | break; | 
|---|
| 657 | } | 
|---|
| 658 | } | 
|---|
| 659 | } | 
|---|
| 660 |  | 
|---|
| 661 | Ok(total_bytes_written) | 
|---|
| 662 | } | 
|---|
| 663 |  | 
|---|
| 664 | /// Decode multiple frames into the extra capacity of the output vector. | 
|---|
| 665 | /// | 
|---|
| 666 | /// `input` must contain an exact number of frames. | 
|---|
| 667 | /// | 
|---|
| 668 | /// `output` must have enough extra capacity to hold the decompressed data. | 
|---|
| 669 | /// This function will not reallocate or grow the vector. If you don't know | 
|---|
| 670 | /// how large the output will be, use [`FrameDecoder::decode_blocks`] instead. | 
|---|
| 671 | /// | 
|---|
| 672 | /// This calls [`FrameDecoder::init`], and all bytes currently in the decoder will be lost. | 
|---|
| 673 | /// | 
|---|
| 674 | /// The length of the output vector is updated to include the decompressed data. | 
|---|
| 675 | /// The length is not changed if an error occurs. | 
|---|
| 676 | pub fn decode_all_to_vec( | 
|---|
| 677 | &mut self, | 
|---|
| 678 | input: &[u8], | 
|---|
| 679 | output: &mut Vec<u8>, | 
|---|
| 680 | ) -> Result<(), FrameDecoderError> { | 
|---|
| 681 | let len = output.len(); | 
|---|
| 682 | let cap = output.capacity(); | 
|---|
| 683 | output.resize(cap, 0); | 
|---|
| 684 | match self.decode_all(input, &mut output[len..]) { | 
|---|
| 685 | Ok(bytes_written) => { | 
|---|
| 686 | let new_len = core::cmp::min(len + bytes_written, cap); // Sanitizes `bytes_written`. | 
|---|
| 687 | output.resize(new_len, 0); | 
|---|
| 688 | Ok(()) | 
|---|
| 689 | } | 
|---|
| 690 | Err(e) => { | 
|---|
| 691 | output.resize(len, 0); | 
|---|
| 692 | Err(e) | 
|---|
| 693 | } | 
|---|
| 694 | } | 
|---|
| 695 | } | 
|---|
| 696 | } | 
|---|
| 697 |  | 
|---|
| 698 | /// Read bytes from the decode_buffer that are no longer needed. While the frame is not yet finished | 
|---|
| 699 | /// this will retain window_size bytes, else it will drain it completely | 
|---|
| 700 | impl Read for FrameDecoder { | 
|---|
| 701 | fn read(&mut self, target: &mut [u8]) -> Result<usize, Error> { | 
|---|
| 702 | let state: &mut FrameDecoderState = match &mut self.state { | 
|---|
| 703 | None => return Ok(0), | 
|---|
| 704 | Some(s: &mut FrameDecoderState) => s, | 
|---|
| 705 | }; | 
|---|
| 706 | if state.frame_finished { | 
|---|
| 707 | state.decoder_scratch.buffer.read_all(target) | 
|---|
| 708 | } else { | 
|---|
| 709 | state.decoder_scratch.buffer.read(buf:target) | 
|---|
| 710 | } | 
|---|
| 711 | } | 
|---|
| 712 | } | 
|---|
| 713 |  | 
|---|