cadmpeg_core/decode/
zstd.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Window-admitted Zstandard streams.
3
4use super::{DecodeContext, View};
5use crate::CodecError;
6use zstd::zstd_safe::{self, DCtx, DParameter, InBuffer, OutBuffer};
7
8/// The largest window the linked decoder can open on this target. A larger
9/// frame is outside the reader, not over a session ceiling.
10const MAX_DECODER_WINDOW: u64 = 1_u64 << if usize::BITS == 32 { 30 } else { 31 };
11
12/// One Zstandard decoder over an admitted input.
13pub struct ZstdDecoder<'ctx, 'input> {
14    decoder: DCtx<'static>,
15    ctx: &'ctx DecodeContext<'input>,
16    input: InBuffer<'input>,
17    finished: bool,
18    window_bytes: u64,
19}
20
21impl<'input> DecodeContext<'input> {
22    /// Inspects all frame windows before allocating a decoder. The output
23    /// each frame's decoder buffers, its window or its declared content size
24    /// when that is smaller, must fit the per-expansion decompression
25    /// allowance, and the declared window must fit the decoder's. Either
26    /// excess is [`CodecError::NotImplemented`] for this stream and leaves the
27    /// session open, since nothing was allocated.
28    pub fn open_zstd<'ctx>(
29        &'ctx self,
30        source: View<'input>,
31    ) -> Result<ZstdDecoder<'ctx, 'input>, CodecError> {
32        let mut offset = 0;
33        let mut largest_window = 0_u64;
34        let allowance = self.per_expand_allowance();
35        while offset < source.window().len() {
36            let frame = source
37                .child(source.start() + offset, source.end())
38                .ok_or_else(|| {
39                    CodecError::Malformed("Zstandard frame offset escapes input".into())
40                })?;
41            let FrameWindow { declared, buffered } = frame_window(frame)?;
42            if buffered > allowance {
43                return Err(CodecError::NotImplemented(format!(
44                    "Zstandard window of {buffered} bytes exceeds the \
45                     {allowance}-byte per-expansion allowance"
46                )));
47            }
48            if declared > MAX_DECODER_WINDOW {
49                return Err(window_unsupported(declared));
50            }
51            largest_window = largest_window.max(declared);
52            let size = zstd_safe::find_frame_compressed_size(frame.window())
53                .map_err(|code| zstd_error(self, code, declared))?;
54            if size == 0 || size > frame.window().len() {
55                return Err(CodecError::Malformed("invalid Zstandard frame size".into()));
56            }
57            offset += size;
58        }
59        if offset == 0 {
60            return Err(CodecError::Malformed("empty Zstandard stream".into()));
61        }
62        let mut decoder = DCtx::try_create().ok_or_else(|| {
63            self.budget
64                .retained_allocation_failed(largest_window, "Zstandard context")
65        })?;
66        // Round the inspected maximum up to the next permitted log, with the
67        // decoder's format minimum of 1 KiB. Each frame's buffered output was
68        // checked against the exact allowance above before allocation.
69        let minimum_window = largest_window.max(1024);
70        let window_log = u64::BITS - (minimum_window - 1).leading_zeros();
71        decoder
72            .set_parameter(DParameter::WindowLogMax(window_log))
73            .map_err(|code| zstd_error(self, code, largest_window))?;
74        Ok(ZstdDecoder {
75            decoder,
76            ctx: self,
77            input: InBuffer {
78                src: source.window(),
79                pos: 0,
80            },
81            finished: false,
82            window_bytes: largest_window,
83        })
84    }
85}
86
87impl ZstdDecoder<'_, '_> {
88    /// Reads decoded output into caller-owned storage.
89    pub fn read_chunk(&mut self, bytes: &mut [u8]) -> Result<usize, CodecError> {
90        if self.finished || bytes.is_empty() {
91            return Ok(0);
92        }
93        loop {
94            let before = self.input.pos;
95            let mut chunk = [0_u8; 16 * 1024];
96            let capacity = bytes.len().min(chunk.len());
97            let mut output = OutBuffer::around(&mut chunk[..capacity]);
98            let remaining = self
99                .decoder
100                .decompress_stream(&mut output, &mut self.input)
101                .map_err(|code| zstd_error(self.ctx, code, self.window_bytes))?;
102            let produced = output.pos();
103
104            bytes[..produced].copy_from_slice(&chunk[..produced]);
105            if remaining == 0 && self.input.pos == self.input.src.len() {
106                self.finished = true;
107            }
108            if produced != 0 || self.finished {
109                return Ok(produced);
110            }
111            if before == self.input.pos {
112                return Err(CodecError::Malformed("truncated Zstandard stream".into()));
113            }
114        }
115    }
116}
117
118fn window_unsupported(window: u64) -> CodecError {
119    CodecError::NotImplemented(format!(
120        "Zstandard window of {window} bytes exceeds the {MAX_DECODER_WINDOW}-byte decoder window"
121    ))
122}
123
124fn zstd_error(ctx: &DecodeContext<'_>, code: usize, window: u64) -> CodecError {
125    // SAFETY: the code is the error result returned by the same zstd library.
126    let error = unsafe { zstd_safe::zstd_sys::ZSTD_getErrorCode(code) };
127    match error {
128        zstd_safe::zstd_sys::ZSTD_ErrorCode::ZSTD_error_memory_allocation => ctx
129            .budget
130            .retained_allocation_failed(window, "Zstandard decoder allocation"),
131        zstd_safe::zstd_sys::ZSTD_ErrorCode::ZSTD_error_frameParameter_windowTooLarge => {
132            window_unsupported(window)
133        }
134        _ => CodecError::malformed(format_args!(
135            "invalid Zstandard stream: {}",
136            zstd_safe::get_error_name(code)
137        )),
138    }
139}
140
141/// One frame's declared window and the output its decoder buffers: the
142/// window, or the declared content size when that is smaller.
143struct FrameWindow {
144    declared: u64,
145    buffered: u64,
146}
147
148fn frame_window(mut frame: View<'_>) -> Result<FrameWindow, CodecError> {
149    let magic = frame.req_u32_le()?;
150    if (0x184d_2a50..=0x184d_2a5f).contains(&magic) {
151        frame.req_u32_le()?;
152        return Ok(FrameWindow {
153            declared: 0,
154            buffered: 0,
155        });
156    }
157    if magic != 0xfd2f_b528 {
158        return Err(CodecError::Malformed("invalid Zstandard magic".into()));
159    }
160    let flags = frame.req_u8()?;
161    if flags & 8 != 0 {
162        return Err(CodecError::Malformed(
163            "reserved Zstandard frame flag".into(),
164        ));
165    }
166    let single_segment = flags & 32 != 0;
167    let window = if single_segment {
168        None
169    } else {
170        let descriptor = frame.req_u8()?;
171        let base = 1_u64 << (10 + u32::from(descriptor >> 3));
172        Some(base + (base / 8) * u64::from(descriptor & 7))
173    };
174    let dictionary_bytes = match flags & 3 {
175        0 => 0,
176        1 => 1,
177        2 => 2,
178        _ => 4,
179    };
180    frame.req_take(dictionary_bytes)?;
181    let content_size = match flags >> 6 {
182        0 if single_segment => Some(u64::from(frame.req_u8()?)),
183        0 => None,
184        1 => Some(u64::from(frame.req_u16_le()?) + 256),
185        2 => Some(u64::from(frame.req_u32_le()?)),
186        _ => Some(frame.req_u64_le()?),
187    };
188    let declared = window.or(content_size).unwrap_or(0);
189    Ok(FrameWindow {
190        declared,
191        buffered: content_size.map_or(declared, |size| size.min(declared)),
192    })
193}
194
195#[cfg(test)]
196mod tests {
197    use super::MAX_DECODER_WINDOW;
198    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ExpandSpec};
199    use crate::CodecError;
200
201    /// A frame header with a non-single-segment window of `window` bytes,
202    /// followed by one empty raw block.
203    fn frame_with_window(window: u64) -> Vec<u8> {
204        let exponent = u8::try_from(window.trailing_zeros() - 10).expect("window exponent");
205        vec![0x28, 0xb5, 0x2f, 0xfd, 0, exponent << 3, 1, 0, 0]
206    }
207
208    #[test]
209    fn zstd_window_past_the_per_expand_allowance_is_unsupported_and_leaves_the_session_open() {
210        let window = 32 << 20;
211        let bytes = frame_with_window(window);
212        let policy = DecodePolicy::service();
213        let arena = DecodeArena::new();
214        let (ctx, root) = DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
215        let allowance = ctx.per_expand_allowance();
216        assert!(allowance < window);
217        assert!(window < policy.limits.max_decompressed_bytes_per_expand);
218        let Err(CodecError::NotImplemented(message)) = ctx.open_zstd(root) else {
219            panic!("a window above the allowance is unsupported");
220        };
221        assert_eq!(
222            message,
223            format!(
224                "Zstandard window of {window} bytes exceeds the \
225                 {allowance}-byte per-expansion allowance"
226            )
227        );
228        assert_eq!(ctx.resource_refusal(), None);
229        let mut writer = ctx
230            .begin_expand(ExpandSpec::Exact(1))
231            .expect("later expansions are admitted");
232        writer.write(b"x").expect("write");
233        writer.finalize().expect("finalize");
234        assert!(ctx.finish_session().is_ok());
235    }
236
237    #[test]
238    fn zstd_window_is_admitted_by_a_smaller_declared_content_size() {
239        let window = 32 << 20;
240        let content = [7_u8; 256];
241        // A 32 MiB window descriptor, a two-byte content size of 256 bytes,
242        // and one raw last block holding the content.
243        let mut bytes = vec![
244            0x28,
245            0xb5,
246            0x2f,
247            0xfd,
248            0x40,
249            15 << 3,
250            0,
251            0,
252            0x01,
253            0x08,
254            0x00,
255        ];
256        bytes.extend_from_slice(&content);
257        let policy = DecodePolicy::service();
258        let arena = DecodeArena::new();
259        let (ctx, root) = DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
260        assert!(ctx.per_expand_allowance() < window);
261        let mut decoder = ctx.open_zstd(root).expect("the content size fits");
262        let mut output = Vec::new();
263        let mut chunk = [0; 64];
264        loop {
265            let read = decoder.read_chunk(&mut chunk).expect("read");
266            if read == 0 {
267                break;
268            }
269            output.extend_from_slice(&chunk[..read]);
270        }
271        assert_eq!(output, content);
272        assert_eq!(ctx.resource_refusal(), None);
273    }
274
275    #[test]
276    fn zstd_window_above_the_decoder_is_unsupported_and_leaves_the_session_open() {
277        let window = MAX_DECODER_WINDOW * 2;
278        let mut bytes = frame_with_window(window);
279        // Each input byte raises the per-expand allowance by 256 bytes.
280        bytes.resize(usize::try_from(window / 256).expect("padding length"), 0);
281        let mut policy = DecodePolicy::service();
282        policy.limits.max_decompressed_bytes_per_expand = window;
283        let arena = DecodeArena::new();
284        let (ctx, root) = DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
285        assert!(ctx.per_expand_allowance() >= window);
286        assert!(matches!(
287            ctx.open_zstd(root),
288            Err(CodecError::NotImplemented(message)) if message.contains("decoder window")
289        ));
290        assert_eq!(ctx.resource_refusal(), None);
291        assert!(ctx.finish_session().is_ok());
292    }
293
294    #[test]
295    fn zstd_checks_each_window_before_decoder_allocation() {
296        let mut bytes = zstd::bulk::compress(b"prefix", 0).expect("frame");
297        bytes.extend_from_slice(&[0x28, 0xb5, 0x2f, 0xfd, 0, 0xa0, 1, 0, 0]);
298        let arena = DecodeArena::new();
299        let (ctx, root) =
300            DecodeContext::from_root_bytes(&bytes, &arena, &DecodePolicy::service()).expect("root");
301        assert!(
302            matches!(ctx.open_zstd(root), Err(CodecError::NotImplemented(message)) if message.starts_with(&format!("Zstandard window of {} bytes", 1_u64 << 30)) && message.ends_with("per-expansion allowance"))
303        );
304        assert_eq!(ctx.resource_refusal(), None);
305    }
306
307    #[test]
308    fn zstd_decodes_a_frame() {
309        let bytes = zstd::bulk::compress(b"payload", 0).expect("frame");
310        let arena = DecodeArena::new();
311        let (ctx, root) =
312            DecodeContext::from_root_bytes(&bytes, &arena, &DecodePolicy::service()).expect("root");
313        let mut decoder = ctx.open_zstd(root).expect("decoder");
314        let mut output = [0; 7];
315        assert_eq!(decoder.read_chunk(&mut output).expect("read"), 7);
316        assert_eq!(&output, b"payload");
317    }
318}
cadmpeg_core/decode/zstd.rs

cadmpeg_core/decode/
zstd.rs
