zstd_complete/decoding/
multi_frame_decoder.rs

1//! Bounded high-level decoding of concatenated Zstandard frames.
2
3use core::{fmt, num::NonZeroUsize};
4use std::io::{self, Read};
5
6use super::{
7    errors::{FrameDecoderError, ReadFrameHeaderError},
8    BlockDecodingStrategy, Dictionary, FrameDecoder, DEFAULT_MAX_WINDOW_SIZE,
9};
10
11/// Policy for RFC 8878 skippable frames encountered between data frames.
12#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
13pub enum SkippableFramePolicy {
14    /// Consume the skippable payload without exposing it to the decoded stream.
15    #[default]
16    Skip,
17    /// Return an error as soon as a skippable-frame header is encountered.
18    Reject,
19}
20
21/// Resource and format policy for [`MultiFrameDecoder`].
22#[derive(Clone, Debug)]
23pub struct MultiFrameDecoderOptions {
24    skippable_frames: SkippableFramePolicy,
25    max_frames: Option<NonZeroUsize>,
26    max_skippable_frame_size: usize,
27    max_window_size: u64,
28    max_decoded_bytes: Option<NonZeroUsize>,
29}
30
31impl MultiFrameDecoderOptions {
32    /// Creates a bounded default policy.
33    pub const fn new() -> Self {
34        Self {
35            skippable_frames: SkippableFramePolicy::Skip,
36            max_frames: None,
37            max_skippable_frame_size: 16 * 1024 * 1024,
38            max_window_size: DEFAULT_MAX_WINDOW_SIZE,
39            max_decoded_bytes: None,
40        }
41    }
42
43    pub const fn skippable_frame_policy(&self) -> SkippableFramePolicy {
44        self.skippable_frames
45    }
46
47    pub const fn max_frames(&self) -> Option<NonZeroUsize> {
48        self.max_frames
49    }
50
51    pub const fn max_skippable_frame_size(&self) -> usize {
52        self.max_skippable_frame_size
53    }
54
55    pub const fn max_window_size(&self) -> u64 {
56        self.max_window_size
57    }
58
59    pub const fn max_decoded_bytes(&self) -> Option<NonZeroUsize> {
60        self.max_decoded_bytes
61    }
62
63    pub const fn with_skippable_frame_policy(mut self, policy: SkippableFramePolicy) -> Self {
64        self.skippable_frames = policy;
65        self
66    }
67
68    pub const fn with_max_frames(mut self, limit: Option<NonZeroUsize>) -> Self {
69        self.max_frames = limit;
70        self
71    }
72
73    pub const fn with_max_skippable_frame_size(mut self, bytes: usize) -> Self {
74        self.max_skippable_frame_size = bytes;
75        self
76    }
77
78    pub const fn with_max_window_size(mut self, bytes: u64) -> Self {
79        self.max_window_size = bytes;
80        self
81    }
82
83    pub const fn with_max_decoded_bytes(mut self, limit: Option<NonZeroUsize>) -> Self {
84        self.max_decoded_bytes = limit;
85        self
86    }
87}
88
89impl Default for MultiFrameDecoderOptions {
90    fn default() -> Self {
91        Self::new()
92    }
93}
94
95/// A typed failure retained inside the [`io::Error`] returned by [`Read`].
96#[derive(Debug)]
97#[non_exhaustive]
98pub enum MultiFrameDecoderError {
99    Frame(FrameDecoderError),
100    Io(io::Error),
101    SkippableFrameRejected { magic_number: u32, length: u32 },
102    SkippableFrameTooLarge { length: u32, limit: usize },
103    FrameLimitExceeded { limit: usize },
104    DecodedSizeLimitExceeded { limit: usize },
105}
106
107impl fmt::Display for MultiFrameDecoderError {
108    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
109        match self {
110            Self::Frame(error) => error.fmt(formatter),
111            Self::Io(error) => error.fmt(formatter),
112            Self::SkippableFrameRejected {
113                magic_number,
114                length,
115            } => write!(
116                formatter,
117                "skippable frame 0x{magic_number:X} with {length} payload bytes was rejected"
118            ),
119            Self::SkippableFrameTooLarge { length, limit } => write!(
120                formatter,
121                "skippable frame has {length} payload bytes, exceeding the {limit}-byte limit"
122            ),
123            Self::FrameLimitExceeded { limit } => {
124                write!(formatter, "archive exceeds the {limit}-frame limit")
125            }
126            Self::DecodedSizeLimitExceeded { limit } => write!(
127                formatter,
128                "decoded archive exceeds the {limit}-byte output limit"
129            ),
130        }
131    }
132}
133
134impl std::error::Error for MultiFrameDecoderError {
135    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
136        match self {
137            Self::Frame(error) => Some(error),
138            Self::Io(error) => Some(error),
139            _ => None,
140        }
141    }
142}
143
144/// A `Read` decoder for concatenated data and skippable Zstandard frames.
145pub struct MultiFrameDecoder<R: Read> {
146    source: R,
147    decoder: FrameDecoder,
148    options: MultiFrameDecoderOptions,
149    state: DecoderState,
150    frame_count: usize,
151    decoded_bytes: usize,
152}
153
154#[derive(Clone, Copy, Debug, Eq, PartialEq)]
155enum DecoderState {
156    NeedFrame,
157    Decoding,
158    Finished,
159}
160
161impl<R: Read> MultiFrameDecoder<R> {
162    pub fn new(source: R) -> Self {
163        Self::with_options(source, MultiFrameDecoderOptions::new())
164    }
165
166    pub fn with_options(source: R, options: MultiFrameDecoderOptions) -> Self {
167        Self {
168            source,
169            decoder: FrameDecoder::new(),
170            options,
171            state: DecoderState::NeedFrame,
172            frame_count: 0,
173            decoded_bytes: 0,
174        }
175    }
176
177    /// Adds a formatted dictionary selected automatically by frame dictionary ID.
178    pub fn add_dictionary(&mut self, dictionary: Dictionary) -> Result<(), FrameDecoderError> {
179        self.decoder.add_dict(dictionary)
180    }
181
182    pub fn get_ref(&self) -> &R {
183        &self.source
184    }
185
186    pub fn get_mut(&mut self) -> &mut R {
187        &mut self.source
188    }
189
190    pub fn into_inner(self) -> R {
191        self.source
192    }
193
194    pub const fn frame_count(&self) -> usize {
195        self.frame_count
196    }
197
198    pub const fn decoded_bytes(&self) -> usize {
199        self.decoded_bytes
200    }
201
202    fn start_next_frame(&mut self) -> Result<bool, MultiFrameDecoderError> {
203        loop {
204            let Some(first) = read_one(&mut self.source)? else {
205                self.state = DecoderState::Finished;
206                return Ok(false);
207            };
208            if let Some(limit) = self.options.max_frames {
209                if self.frame_count >= limit.get() {
210                    return Err(MultiFrameDecoderError::FrameLimitExceeded { limit: limit.get() });
211                }
212            }
213            let prefix = io::Cursor::new([first]);
214            match self.decoder.reset_with_window_limit(
215                prefix.chain(&mut self.source),
216                self.options.max_window_size,
217            ) {
218                Ok(()) => {
219                    self.frame_count += 1;
220                    self.state = DecoderState::Decoding;
221                    return Ok(true);
222                }
223                Err(FrameDecoderError::ReadFrameHeaderError(ReadFrameHeaderError::SkipFrame {
224                    magic_number,
225                    length,
226                })) => {
227                    if self.options.skippable_frames == SkippableFramePolicy::Reject {
228                        return Err(MultiFrameDecoderError::SkippableFrameRejected {
229                            magic_number,
230                            length,
231                        });
232                    }
233                    if length as usize > self.options.max_skippable_frame_size {
234                        return Err(MultiFrameDecoderError::SkippableFrameTooLarge {
235                            length,
236                            limit: self.options.max_skippable_frame_size,
237                        });
238                    }
239                    discard_exact(&mut self.source, length as usize)?;
240                }
241                Err(error) => return Err(MultiFrameDecoderError::Frame(error)),
242            }
243        }
244    }
245
246    fn read_decoded(&mut self, target: &mut [u8]) -> Result<usize, MultiFrameDecoderError> {
247        let remaining = self.options.max_decoded_bytes.map_or(usize::MAX, |limit| {
248            limit.get().saturating_sub(self.decoded_bytes)
249        });
250
251        while self.decoder.can_collect() == 0 && !self.decoder.is_finished() {
252            self.decoder
253                .decode_blocks(&mut self.source, BlockDecodingStrategy::UptoBytes(1))
254                .map_err(MultiFrameDecoderError::Frame)?;
255        }
256
257        if remaining == 0 && self.decoder.can_collect() != 0 {
258            return Err(MultiFrameDecoderError::DecodedSizeLimitExceeded {
259                limit: self
260                    .options
261                    .max_decoded_bytes
262                    .expect("zero remaining requires a configured limit")
263                    .get(),
264            });
265        }
266
267        let requested = target.len().min(remaining);
268        let read = self
269            .decoder
270            .read(&mut target[..requested])
271            .map_err(MultiFrameDecoderError::Io)?;
272        self.decoded_bytes += read;
273        Ok(read)
274    }
275}
276
277impl<R: Read> Read for MultiFrameDecoder<R> {
278    fn read(&mut self, target: &mut [u8]) -> io::Result<usize> {
279        if target.is_empty() {
280            return Ok(0);
281        }
282
283        loop {
284            match self.state {
285                DecoderState::Finished => return Ok(0),
286                DecoderState::NeedFrame => {
287                    self.start_next_frame().map_err(io::Error::other)?;
288                }
289                DecoderState::Decoding => {
290                    let read = self.read_decoded(target).map_err(io::Error::other)?;
291                    if read != 0 {
292                        return Ok(read);
293                    }
294                    if self.decoder.is_finished() && self.decoder.can_collect() == 0 {
295                        self.state = DecoderState::NeedFrame;
296                    }
297                }
298            }
299        }
300    }
301}
302
303fn read_one(source: &mut impl Read) -> Result<Option<u8>, MultiFrameDecoderError> {
304    let mut byte = [0_u8; 1];
305    loop {
306        match source.read(&mut byte) {
307            Ok(0) => return Ok(None),
308            Ok(_) => return Ok(Some(byte[0])),
309            Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
310            Err(error) => return Err(MultiFrameDecoderError::Io(error)),
311        }
312    }
313}
314
315fn discard_exact(
316    source: &mut impl Read,
317    mut remaining: usize,
318) -> Result<(), MultiFrameDecoderError> {
319    let mut buffer = [0_u8; 8192];
320    while remaining != 0 {
321        let requested = remaining.min(buffer.len());
322        let read = source
323            .read(&mut buffer[..requested])
324            .map_err(MultiFrameDecoderError::Io)?;
325        if read == 0 {
326            return Err(MultiFrameDecoderError::Io(io::Error::new(
327                io::ErrorKind::UnexpectedEof,
328                "truncated skippable frame",
329            )));
330        }
331        remaining -= read;
332    }
333    Ok(())
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use crate::encoding::{encode_all, CompressionLevel, EncoderOptions};
340    use std::vec::Vec;
341
342    #[test]
343    fn reads_concatenated_frames_and_skippable_payloads() {
344        let first = b"first frame".repeat(100);
345        let second = b"second frame".repeat(100);
346        let mut archive = encode_all(
347            first.as_slice(),
348            EncoderOptions::new(CompressionLevel::FASTEST),
349        )
350        .unwrap();
351        archive.extend_from_slice(&0x184D2A50_u32.to_le_bytes());
352        archive.extend_from_slice(&4_u32.to_le_bytes());
353        archive.extend_from_slice(b"skip");
354        archive.extend_from_slice(
355            &encode_all(
356                second.as_slice(),
357                EncoderOptions::new(CompressionLevel::DEFAULT),
358            )
359            .unwrap(),
360        );
361
362        let mut decoder = MultiFrameDecoder::new(archive.as_slice());
363        let mut decoded = Vec::new();
364        decoder.read_to_end(&mut decoded).unwrap();
365        assert_eq!(decoded, [first, second].concat());
366        assert_eq!(decoder.frame_count(), 2);
367    }
368
369    #[test]
370    fn accepts_an_empty_archive() {
371        let mut decoded = Vec::new();
372        MultiFrameDecoder::new(&[][..])
373            .read_to_end(&mut decoded)
374            .unwrap();
375        assert!(decoded.is_empty());
376    }
377
378    #[test]
379    fn decoded_size_limit_is_enforced() {
380        let input = b"bounded decoded output".repeat(100);
381        let archive = encode_all(input.as_slice(), EncoderOptions::default()).unwrap();
382        let options = MultiFrameDecoderOptions::new()
383            .with_max_decoded_bytes(NonZeroUsize::new(input.len() - 1));
384        let mut decoder = MultiFrameDecoder::with_options(archive.as_slice(), options);
385        let error = io::copy(&mut decoder, &mut io::sink()).unwrap_err();
386        assert!(error
387            .get_ref()
388            .and_then(|error| error.downcast_ref::<MultiFrameDecoderError>())
389            .is_some_and(|error| matches!(
390                error,
391                MultiFrameDecoderError::DecodedSizeLimitExceeded { .. }
392            )));
393    }
394}
zstd_complete/decoding/multi_frame_decoder.rs

zstd_complete/decoding/
multi_frame_decoder.rs
