cadmpeg_core/decode/
zstd.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Window-admitted Zstandard streams with scoped decoder workspace.
3
4use super::{
5    u64_from_index, DecodeContext, ResourceDimension, ResourceFailure, ScopedReservation, View,
6};
7use crate::CodecError;
8use zstd::zstd_safe::{self, DCtx, DParameter, InBuffer, OutBuffer};
9
10const BLOCK_BYTES: u64 = 128 * 1024;
11const FIXED_STATE_BYTES: u64 = 1024 * 1024;
12const COPY_PADDING_BYTES: u64 = 64;
13const MAX_DECODER_WINDOW: u64 = 1_u64 << if usize::BITS == 32 { 30 } else { 31 };
14
15/// One decoder and its live workspace reservation. The decoder is dropped
16/// before the reservation that admits its allocation.
17pub struct ZstdDecoder<'ctx, 'input> {
18    decoder: DCtx<'static>,
19    _workspace: ScopedReservation<'ctx>,
20    ctx: &'ctx DecodeContext<'input>,
21    input: InBuffer<'input>,
22    finished: bool,
23    workspace_bytes: u64,
24    window_bytes: u64,
25}
26
27impl<'input> DecodeContext<'input> {
28    /// Inspects all frame windows before allocating a decoder. The largest
29    /// frame window must fit the materialized-byte and per-expansion ceilings.
30    pub fn open_zstd<'ctx>(
31        &'ctx self,
32        source: View<'input>,
33    ) -> Result<ZstdDecoder<'ctx, 'input>, CodecError> {
34        self.charge_work(
35            u64_from_index(source.window().len()),
36            "inspect Zstandard frames",
37        )?;
38        self.charge_work(
39            u64_from_index(source.window().len()),
40            "Zstandard compressed input",
41        )?;
42        let mut offset = 0;
43        let mut largest_window = 0_u64;
44        let ceiling = self
45            .policy()
46            .limits
47            .max_materialized_bytes
48            .min(self.policy().limits.max_decompressed_bytes_per_expand)
49            .min(MAX_DECODER_WINDOW);
50        while offset < source.window().len() {
51            self.charge_work(1, "inspect Zstandard frame header")?;
52            let frame = source
53                .child(source.start() + offset, source.end())
54                .ok_or_else(|| {
55                    CodecError::Malformed("Zstandard frame offset escapes input".into())
56                })?;
57            let window = frame_window(frame)?;
58            if window > ceiling {
59                return Err(self.budget.refuse(
60                    ResourceDimension::MaterializedBytes,
61                    ResourceFailure::BudgetExceeded,
62                    ceiling,
63                    0,
64                    window,
65                    "Zstandard window",
66                ));
67            }
68            largest_window = largest_window.max(window);
69            let size = zstd_safe::find_frame_compressed_size(frame.window())
70                .map_err(|code| zstd_error(self, code, 0, window))?;
71            if size == 0 || size > frame.window().len() {
72                return Err(CodecError::Malformed("invalid Zstandard frame size".into()));
73            }
74            offset += size;
75        }
76        if offset == 0 {
77            return Err(CodecError::Malformed("empty Zstandard stream".into()));
78        }
79        // ZSTD_estimateDStreamSize in zstd 1.5.7 bounds streaming storage by
80        // DCtx state + window + three 128 KiB blocks + 2*WILDCOPY_OVERLENGTH.
81        // The fixed state, including entropy tables and the literal scratch
82        // buffer, is below 1 MiB. Dictionaries and legacy frames are not opened.
83        let workspace_bytes = largest_window
84            .checked_add(3 * BLOCK_BYTES + FIXED_STATE_BYTES + COPY_PADDING_BYTES)
85            .ok_or_else(|| {
86                CodecError::from(
87                    self.budget
88                        .scoped_size_overflow_limit("Zstandard workspace"),
89                )
90            })?;
91        let workspace = self.reserve_scoped(workspace_bytes, "Zstandard workspace")?;
92        let mut decoder = DCtx::try_create().ok_or_else(|| {
93            self.budget
94                .scoped_allocation_failed(workspace_bytes, "Zstandard context")
95        })?;
96        // Round the inspected maximum up to the next permitted log, with the
97        // decoder's format minimum of 1 KiB. Each actual frame was checked
98        // against the exact ceiling above before allocation.
99        let minimum_window = largest_window.max(1024);
100        let window_log = u64::BITS - (minimum_window - 1).leading_zeros();
101        decoder
102            .set_parameter(DParameter::WindowLogMax(window_log))
103            .map_err(|code| zstd_error(self, code, workspace_bytes, largest_window))?;
104        Ok(ZstdDecoder {
105            decoder,
106            _workspace: workspace,
107            ctx: self,
108            input: InBuffer {
109                src: source.window(),
110                pos: 0,
111            },
112            finished: false,
113            workspace_bytes,
114            window_bytes: largest_window,
115        })
116    }
117}
118
119impl ZstdDecoder<'_, '_> {
120    /// Reads admitted output into caller-owned storage and charges step/copy work.
121    pub fn read_chunk(&mut self, bytes: &mut [u8]) -> Result<usize, CodecError> {
122        if self.finished || bytes.is_empty() {
123            return Ok(0);
124        }
125        loop {
126            self.ctx
127                .charge_work(u64_from_index(bytes.len()), "Zstandard decode step")?;
128            let before = self.input.pos;
129            let mut chunk = [0_u8; 16 * 1024];
130            let capacity = bytes.len().min(chunk.len());
131            let mut output = OutBuffer::around(&mut chunk[..capacity]);
132            let remaining = self
133                .decoder
134                .decompress_stream(&mut output, &mut self.input)
135                .map_err(|code| {
136                    zstd_error(self.ctx, code, self.workspace_bytes, self.window_bytes)
137                })?;
138            let produced = output.pos();
139            self.ctx
140                .charge_work(u64_from_index(produced), "Zstandard output copy")?;
141            bytes[..produced].copy_from_slice(&chunk[..produced]);
142            if remaining == 0 && self.input.pos == self.input.src.len() {
143                self.finished = true;
144            }
145            if produced != 0 || self.finished {
146                return Ok(produced);
147            }
148            if before == self.input.pos {
149                return Err(CodecError::Malformed("truncated Zstandard stream".into()));
150            }
151        }
152    }
153}
154
155fn zstd_error(ctx: &DecodeContext<'_>, code: usize, bytes: u64, window: u64) -> CodecError {
156    // SAFETY: the code is the error result returned by the same zstd library.
157    let error = unsafe { zstd_safe::zstd_sys::ZSTD_getErrorCode(code) };
158    match error {
159        zstd_safe::zstd_sys::ZSTD_ErrorCode::ZSTD_error_memory_allocation => ctx
160            .budget
161            .scoped_allocation_failed(bytes, "Zstandard decoder allocation"),
162        zstd_safe::zstd_sys::ZSTD_ErrorCode::ZSTD_error_frameParameter_windowTooLarge => {
163            ctx.budget.refuse(
164                ResourceDimension::MaterializedBytes,
165                ResourceFailure::BudgetExceeded,
166                ctx.policy()
167                    .limits
168                    .max_materialized_bytes
169                    .min(ctx.policy().limits.max_decompressed_bytes_per_expand)
170                    .min(MAX_DECODER_WINDOW),
171                0,
172                window,
173                "Zstandard decoder window",
174            )
175        }
176        _ => CodecError::malformed(format_args!(
177            "invalid Zstandard stream: {}",
178            zstd_safe::get_error_name(code)
179        )),
180    }
181}
182
183fn frame_window(mut frame: View<'_>) -> Result<u64, CodecError> {
184    let magic = frame.req_u32_le()?;
185    if (0x184d_2a50..=0x184d_2a5f).contains(&magic) {
186        frame.req_u32_le()?;
187        return Ok(0);
188    }
189    if magic != 0xfd2f_b528 {
190        return Err(CodecError::Malformed("invalid Zstandard magic".into()));
191    }
192    let flags = frame.req_u8()?;
193    if flags & 8 != 0 {
194        return Err(CodecError::Malformed(
195            "reserved Zstandard frame flag".into(),
196        ));
197    }
198    let single_segment = flags & 32 != 0;
199    let window = if single_segment {
200        0
201    } else {
202        let descriptor = frame.req_u8()?;
203        let base = 1_u64 << (10 + u32::from(descriptor >> 3));
204        base + (base / 8) * u64::from(descriptor & 7)
205    };
206    let dictionary_bytes = match flags & 3 {
207        0 => 0,
208        1 => 1,
209        2 => 2,
210        _ => 4,
211    };
212    frame.req_take(dictionary_bytes)?;
213    let content_size = match flags >> 6 {
214        0 if single_segment => u64::from(frame.req_u8()?),
215        0 => 0,
216        1 => u64::from(frame.req_u16_le()?) + 256,
217        2 => u64::from(frame.req_u32_le()?),
218        _ => frame.req_u64_le()?,
219    };
220    Ok(if single_segment { content_size } else { window })
221}
222
223#[cfg(test)]
224mod tests {
225    use super::{BLOCK_BYTES, COPY_PADDING_BYTES, FIXED_STATE_BYTES};
226    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ResourceDimension};
227    use crate::CodecError;
228
229    #[test]
230    fn zstd_checks_each_window_before_decoder_allocation() {
231        let mut bytes = zstd::bulk::compress(b"prefix", 0).expect("frame");
232        bytes.extend_from_slice(&[0x28, 0xb5, 0x2f, 0xfd, 0, 0xa0, 1, 0, 0]);
233        let arena = DecodeArena::new();
234        let (ctx, root) =
235            DecodeContext::from_root_bytes(&bytes, &arena, &DecodePolicy::service()).expect("root");
236        assert!(
237            matches!(ctx.open_zstd(root), Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::MaterializedBytes && limit.operation == "Zstandard window" && limit.additional == 1 << 30)
238        );
239    }
240
241    #[test]
242    fn zstd_workspace_stays_reserved_until_decoder_drop() {
243        let bytes = zstd::bulk::compress(b"payload", 0).expect("frame");
244        let arena = DecodeArena::new();
245        let mut policy = DecodePolicy::service();
246        policy.limits.max_materialized_bytes =
247            7 + 3 * BLOCK_BYTES + FIXED_STATE_BYTES + COPY_PADDING_BYTES;
248        let (ctx, root) = DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
249        let mut decoder = ctx.open_zstd(root).expect("decoder");
250        assert_eq!(
251            decoder.workspace_bytes,
252            policy.limits.max_materialized_bytes
253        );
254        let mut output = [0; 7];
255        assert_eq!(decoder.read_chunk(&mut output).expect("read"), 7);
256        assert_eq!(&output, b"payload");
257        drop(decoder);
258        ctx.reserve_scoped(policy.limits.max_materialized_bytes, "reuse workspace")
259            .expect("released workspace");
260    }
261
262    #[test]
263    fn zstd_workspace_and_work_limits_remain_typed() {
264        let bytes = zstd::bulk::compress(b"payload", 0).expect("frame");
265        for dimension in [
266            ResourceDimension::MaterializedBytes,
267            ResourceDimension::WorkUnits,
268        ] {
269            let arena = DecodeArena::new();
270            let mut policy = DecodePolicy::service();
271            if dimension == ResourceDimension::MaterializedBytes {
272                policy.limits.max_materialized_bytes = 0;
273            } else {
274                policy.limits.max_work_units = 0;
275            }
276            let (ctx, root) =
277                DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
278            assert!(
279                matches!(ctx.open_zstd(root), Err(CodecError::ResourceLimit(limit)) if limit.dimension == dimension)
280            );
281        }
282    }
283}
cadmpeg_core/decode/zstd.rs

cadmpeg_core/decode/
zstd.rs
