cadmpeg_ir/decode/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Decode state, decompression limits, and session lifecycle.
3
4use std::cell::Cell;
5
6use crate::codec::{CodecError, DecodeResult, ReadSeek};
7use crate::report::StrictConsequence;
8
9use super::arena::DecodeArena;
10use super::error::{
11    ErrorContext, LimitScope, ResourceDimension, ResourceFailure, ResourceLimit, SourceLocation,
12};
13use super::policy::{
14    DecodeMode, DecodePolicy, DECOMPRESSED_PER_EXPAND_BASE, DECOMPRESSED_PER_EXPAND_PER_INPUT_BYTE,
15    DECOMPRESSED_TOTAL_BASE, DECOMPRESSED_TOTAL_PER_INPUT_BYTE,
16};
17use super::space::{ByteRange, SpaceId};
18use super::view::View;
19
20/// Initial per-expand reservation clamp: an attacker's declared size cannot
21/// force a large up-front reservation before any output is produced.
22const RESERVE_CLAMP: u64 = 8 * 1024 * 1024;
23
24/// Shared monotonic decode state.
25#[derive(Debug)]
26pub struct DecodeContext<'a> {
27    arena: &'a DecodeArena,
28    policy: DecodePolicy,
29    container_only: bool,
30    input_bytes: u64,
31    decompressed_bytes: Cell<u64>,
32    next_space: Cell<u32>,
33    fuse: Cell<Option<ResourceLimit>>,
34}
35
36impl<'a> DecodeContext<'a> {
37    /// Reads the root input under `max_input_bytes`, copies it into the arena,
38    /// registers the root space, establishes input-proportional allowances,
39    /// and returns the context and root view.
40    pub fn read_root(
41        reader: &mut dyn ReadSeek,
42        arena: &'a DecodeArena,
43        policy: &DecodePolicy,
44    ) -> Result<(Self, View<'a>), CodecError> {
45        let max = policy.limits.max_input_bytes;
46        let cap = max.saturating_add(1);
47        let mut buffer: Vec<u8> = Vec::new();
48        let mut chunk = [0u8; 8192];
49        loop {
50            let remaining = cap.saturating_sub(buffer.len() as u64);
51            if remaining == 0 {
52                break;
53            }
54            let want = usize::try_from(remaining.min(chunk.len() as u64)).unwrap_or(chunk.len());
55            let read = reader.read(&mut chunk[..want]).map_err(CodecError::Io)?;
56            if read == 0 {
57                break;
58            }
59            buffer
60                .try_reserve(read)
61                .map_err(|_| root_error(ResourceFailure::AllocationFailed, max, read as u64))?;
62            buffer.extend_from_slice(&chunk[..read]);
63        }
64        if buffer.len() as u64 > max {
65            return Err(root_error(
66                ResourceFailure::BudgetExceeded,
67                max,
68                buffer.len() as u64,
69            ));
70        }
71        let bytes = arena.alloc(buffer.into_boxed_slice());
72        Self::from_bytes(bytes, arena, policy)
73    }
74
75    /// Builds a context over caller-owned root bytes, for fuzz targets and
76    /// tests. The arena still backs any expansions produced during decode.
77    pub fn from_root_bytes(
78        bytes: &'a [u8],
79        arena: &'a DecodeArena,
80        policy: &DecodePolicy,
81    ) -> Result<(Self, View<'a>), CodecError> {
82        Self::from_bytes(bytes, arena, policy)
83    }
84
85    fn from_bytes(
86        bytes: &'a [u8],
87        arena: &'a DecodeArena,
88        policy: &DecodePolicy,
89    ) -> Result<(Self, View<'a>), CodecError> {
90        let length = bytes.len() as u64;
91        if length > policy.limits.max_input_bytes {
92            return Err(root_error(
93                ResourceFailure::BudgetExceeded,
94                policy.limits.max_input_bytes,
95                length,
96            ));
97        }
98        let ctx = DecodeContext {
99            arena,
100            policy: *policy,
101            container_only: false,
102            input_bytes: length,
103            decompressed_bytes: Cell::new(0),
104            next_space: Cell::new(1),
105            fuse: Cell::new(None),
106        };
107        Ok((ctx, View::over_space(bytes, SpaceId::ROOT)))
108    }
109
110    /// Returns the decode policy in force.
111    pub fn policy(&self) -> &DecodePolicy {
112        &self.policy
113    }
114
115    /// Returns whether the caller requested container-only decoding.
116    pub fn container_only(&self) -> bool {
117        self.container_only
118    }
119
120    /// Records the caller's container-only request before decoding begins.
121    pub fn set_container_only(&mut self, value: bool) {
122        self.container_only = value;
123    }
124
125    fn decompression_allowance(&self) -> u64 {
126        let proportional = DECOMPRESSED_TOTAL_BASE
127            .saturating_add(DECOMPRESSED_TOTAL_PER_INPUT_BYTE.saturating_mul(self.input_bytes));
128        self.policy
129            .limits
130            .max_decompressed_bytes_total
131            .min(proportional)
132    }
133
134    fn per_expand_allowance(&self) -> u64 {
135        let proportional = DECOMPRESSED_PER_EXPAND_BASE.saturating_add(
136            DECOMPRESSED_PER_EXPAND_PER_INPUT_BYTE.saturating_mul(self.input_bytes),
137        );
138        self.policy
139            .limits
140            .max_decompressed_bytes_per_expand
141            .min(proportional)
142    }
143
144    fn allocate_space(&self) -> SpaceId {
145        let index = self.next_space.get();
146        self.next_space.set(index.saturating_add(1));
147        SpaceId::from_index(index)
148    }
149
150    fn charge_decompressed(
151        &self,
152        scope: LimitScope,
153        amount: u64,
154        operation: &'static str,
155        location: Option<SourceLocation>,
156    ) -> Result<(), CodecError> {
157        if let Some(limit) = self.fuse.get() {
158            return Err(CodecError::ResourceLimit(limit));
159        }
160        let allowance = self.decompression_allowance();
161        let used = self.decompressed_bytes.get();
162        if amount > allowance.saturating_sub(used) {
163            return Err(self.fuse(
164                ResourceFailure::BudgetExceeded,
165                scope,
166                amount,
167                operation,
168                location,
169            ));
170        }
171        self.decompressed_bytes.set(used.saturating_add(amount));
172        Ok(())
173    }
174
175    /// Records a permanent fuse and returns the resource error to propagate.
176    fn fuse(
177        &self,
178        reason: ResourceFailure,
179        scope: LimitScope,
180        amount: u64,
181        operation: &'static str,
182        location: Option<SourceLocation>,
183    ) -> CodecError {
184        let limit = match scope {
185            LimitScope::Global => self.decompression_allowance(),
186            LimitScope::PerExpand => self.per_expand_allowance(),
187        };
188        let used = self.decompressed_bytes.get();
189        let resource = ResourceLimit {
190            dimension: ResourceDimension::DecompressedBytes,
191            reason,
192            scope,
193            limit,
194            used,
195            additional: amount,
196            context: ErrorContext {
197                operation,
198                location,
199            },
200        };
201        self.fuse.set(Some(resource));
202        CodecError::ResourceLimit(resource)
203    }
204
205    // --- decompression ------------------------------------------------------
206
207    /// Begins an expansion whose output is charged incrementally and becomes
208    /// available only after successful finalization.
209    pub fn begin_expand(
210        &self,
211        source: View<'_>,
212        spec: ExpandSpec,
213    ) -> Result<ExpandWriter<'_, 'a>, CodecError> {
214        if let Some(limit) = self.fuse.get() {
215            return Err(CodecError::ResourceLimit(limit));
216        }
217        let mut buffer: Vec<u8> = Vec::new();
218        let reserve = match spec {
219            ExpandSpec::Exact(size) => size.min(RESERVE_CLAMP),
220            ExpandSpec::Unknown => 0,
221        };
222        if reserve > 0 {
223            let reserve = usize::try_from(reserve).unwrap_or(usize::MAX);
224            buffer.try_reserve(reserve).map_err(|_| {
225                self.fuse(
226                    ResourceFailure::AllocationFailed,
227                    LimitScope::PerExpand,
228                    reserve as u64,
229                    "begin_expand",
230                    Some(source.location()),
231                )
232            })?;
233        }
234        Ok(ExpandWriter {
235            ctx: self,
236            spec,
237            location: source.location(),
238            buffer,
239            written: 0,
240        })
241    }
242
243    /// Copies several input extents into one derived view.
244    pub fn concat_views(&self, inputs: &[View<'_>]) -> Result<View<'a>, CodecError> {
245        if let Some(limit) = self.fuse.get() {
246            return Err(CodecError::ResourceLimit(limit));
247        }
248        let total = inputs.iter().try_fold(0usize, |total, view| {
249            total.checked_add(view.window().len()).ok_or_else(|| {
250                CodecError::Io(std::io::Error::other("concatenated view is too large"))
251            })
252        })?;
253        let mut buffer = Vec::new();
254        buffer.try_reserve_exact(total).map_err(|_| {
255            CodecError::Io(std::io::Error::other("concatenated view allocation failed"))
256        })?;
257        for view in inputs {
258            buffer.extend_from_slice(view.window());
259        }
260        let bytes = self.arena.alloc(buffer.into_boxed_slice());
261        let space = self.allocate_space();
262        Ok(View::over_space(bytes, space))
263    }
264
265    /// Registers a stored (uncompressed) child range as a space that borrows
266    /// the parent bytes without copying.
267    ///
268    /// `range` is expressed in the parent view's own space coordinates and must
269    /// lie within the parent window; a range that escapes the parent is refused
270    /// here, at the request site, exactly as [`View::child`] refuses. No bytes
271    /// are copied. It is the archive-entry counterpart of [`DecodeContext::begin_expand`] —
272    /// stored ZIP entries take this path, compressed ones take the expander.
273    /// Registration still refuses on a fused context so a stored entry cannot be
274    /// admitted after a refusal.
275    pub fn register_slice<'v>(
276        &self,
277        parent: View<'v>,
278        range: ByteRange,
279    ) -> Result<View<'v>, CodecError> {
280        if let Some(limit) = self.fuse.get() {
281            return Err(CodecError::ResourceLimit(limit));
282        }
283        let start = usize::try_from(range.start).ok();
284        let end = usize::try_from(range.end).ok();
285        let child = start
286            .zip(end)
287            .and_then(|(start, end)| parent.child(start, end))
288            .ok_or_else(|| {
289                CodecError::Malformed(format!(
290                    "stored slice [{}, {}) escapes parent space {}",
291                    range.start,
292                    range.end,
293                    parent.space().index()
294                ))
295            })?;
296        let space = self.allocate_space();
297        Ok(View::over_space(child.window(), space))
298    }
299
300    // --- lifecycle ----------------------------------------------------------
301
302    /// Closes an inspection, returning a fused resource error even when codec
303    /// code swallowed the charge that caused it.
304    pub(crate) fn finish_inspection<T>(
305        self,
306        result: Result<T, CodecError>,
307    ) -> Result<T, CodecError> {
308        if let Some(limit) = self.fuse.get() {
309            return Err(CodecError::ResourceLimit(limit));
310        }
311        result
312    }
313
314    /// Closes a decode, returning a fused resource error even when codec code
315    /// swallowed the charge that caused it.
316    pub fn finish(
317        self,
318        result: Result<DecodeResult, CodecError>,
319    ) -> Result<DecodeResult, CodecError> {
320        if let Some(limit) = self.fuse.get() {
321            return Err(CodecError::ResourceLimit(limit));
322        }
323        let result = result?;
324        if self.policy.mode == DecodeMode::Strict && !result.report.container_only {
325            if let Some(loss) = result
326                .report
327                .losses
328                .iter()
329                .find(|loss| loss.code.strict_consequence() == StrictConsequence::Reject)
330            {
331                return Err(CodecError::Malformed(format!(
332                    "strict mode rejects {}: {}",
333                    loss.code, loss.message
334                )));
335            }
336        }
337        Ok(result)
338    }
339}
340
341/// Builds the root-input resource error before a context exists.
342fn root_error(reason: ResourceFailure, limit: u64, used: u64) -> CodecError {
343    CodecError::ResourceLimit(ResourceLimit {
344        dimension: ResourceDimension::InputBytes,
345        reason,
346        scope: LimitScope::Global,
347        limit,
348        used,
349        additional: used.saturating_sub(limit),
350        context: ErrorContext {
351            operation: "read_root",
352            location: None,
353        },
354    })
355}
356
357/// How much output an expansion is expected to produce.
358#[derive(Debug, Clone, Copy, PartialEq, Eq)]
359pub enum ExpandSpec {
360    /// A declared exact size, enforced per-write and at finalize.
361    Exact(u64),
362    /// No trustworthy declared size: the decompression limits apply.
363    Unknown,
364}
365
366/// Writes decompressed output under incremental charging.
367#[derive(Debug)]
368pub struct ExpandWriter<'ctx, 'a> {
369    ctx: &'ctx DecodeContext<'a>,
370    spec: ExpandSpec,
371    location: SourceLocation,
372    buffer: Vec<u8>,
373    written: u64,
374}
375
376impl<'a> ExpandWriter<'_, 'a> {
377    /// Appends decompressed output, charging before it is retained.
378    pub fn write(&mut self, data: &[u8]) -> Result<(), CodecError> {
379        let len = data.len() as u64;
380        let new_written = self.written.saturating_add(len);
381        match self.spec {
382            ExpandSpec::Exact(size) if new_written > size => {
383                return Err(CodecError::Malformed(format!(
384                    "expansion exceeded declared exact size {size}"
385                )))
386            }
387            _ => {}
388        }
389        let per_expand = self.ctx.per_expand_allowance();
390        if new_written > per_expand {
391            return Err(self.ctx.fuse(
392                ResourceFailure::BudgetExceeded,
393                LimitScope::PerExpand,
394                len,
395                "expand_write",
396                Some(self.location),
397            ));
398        }
399        self.ctx.charge_decompressed(
400            LimitScope::Global,
401            len,
402            "expand_write",
403            Some(self.location),
404        )?;
405        self.buffer.try_reserve(data.len()).map_err(|_| {
406            self.ctx.fuse(
407                ResourceFailure::AllocationFailed,
408                LimitScope::PerExpand,
409                len,
410                "expand_write",
411                Some(self.location),
412            )
413        })?;
414        self.buffer.extend_from_slice(data);
415        self.written = new_written;
416        Ok(())
417    }
418
419    /// Finalizes the expansion, stores it in the arena, and registers its space.
420    pub fn finalize(self) -> Result<View<'a>, CodecError> {
421        if let ExpandSpec::Exact(size) = self.spec {
422            if self.written != size {
423                return Err(CodecError::Malformed(format!(
424                    "expansion produced {} of declared exact {size} bytes",
425                    self.written
426                )));
427            }
428        }
429        let bytes = self.ctx.arena.alloc(self.buffer.into_boxed_slice());
430        let space = self.ctx.allocate_space();
431        Ok(View::over_space(bytes, space))
432    }
433
434    /// Returns how many bytes have been written so far.
435    pub fn written(&self) -> u64 {
436        self.written
437    }
438}
cadmpeg_ir/decode/context.rs

cadmpeg_ir/decode/
context.rs
