cadmpeg_core/decode/
budget.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Unified resource accounting for one decode session.
3
4use std::cell::Cell;
5
6use crate::CodecError;
7
8use super::error::{
9    ErrorContext, LimitScope, ResourceDimension, ResourceFailure, ResourceLimit, SourceLocation,
10};
11use super::policy::{
12    DecodePolicy, DECOMPRESSED_TOTAL_BASE, DECOMPRESSED_TOTAL_PER_INPUT_BYTE, MATERIALIZED_BASE,
13    MATERIALIZED_PER_INPUT_BYTE, RETAINED_BASE, RETAINED_PER_INPUT_BYTE,
14};
15
16#[derive(Debug)]
17pub(crate) struct DecodeBudget {
18    policy: DecodePolicy,
19    input_bytes: u64,
20    decompressed: Cell<u64>,
21    materialized: Cell<u64>,
22    retained: Cell<u64>,
23    entities: Cell<u64>,
24    collection_items: Cell<u64>,
25    recursion_depth: Cell<u64>,
26    work: Cell<u64>,
27    fuse: Cell<Option<ResourceLimit>>,
28}
29
30impl DecodeBudget {
31    pub(crate) fn new(policy: DecodePolicy, input_bytes: u64) -> Self {
32        Self {
33            policy,
34            input_bytes,
35            decompressed: Cell::new(0),
36            materialized: Cell::new(0),
37            retained: Cell::new(0),
38            entities: Cell::new(0),
39            collection_items: Cell::new(0),
40            recursion_depth: Cell::new(0),
41            work: Cell::new(0),
42            fuse: Cell::new(None),
43        }
44    }
45
46    pub(crate) fn fused(&self) -> Option<ResourceLimit> {
47        self.fuse.get()
48    }
49
50    pub(crate) fn decompression_allowance(&self) -> u64 {
51        let proportional = DECOMPRESSED_TOTAL_BASE
52            .saturating_add(DECOMPRESSED_TOTAL_PER_INPUT_BYTE.saturating_mul(self.input_bytes));
53        self.policy
54            .limits
55            .max_decompressed_bytes_total
56            .min(proportional)
57    }
58
59    pub(crate) fn input_bytes(&self) -> u64 {
60        self.input_bytes
61    }
62
63    pub(crate) fn charge_decompressed(
64        &self,
65        amount: u64,
66        operation: &'static str,
67        location: Option<SourceLocation>,
68    ) -> Result<(), CodecError> {
69        self.charge(
70            ResourceDimension::DecompressedBytes,
71            &self.decompressed,
72            self.decompression_allowance(),
73            amount,
74            operation,
75            location,
76        )
77    }
78
79    pub(crate) fn decompressed_used(&self) -> u64 {
80        self.decompressed.get()
81    }
82
83    fn materialized_allowance(&self) -> u64 {
84        self.policy.limits.max_materialized_bytes.min(
85            MATERIALIZED_BASE
86                .saturating_add(MATERIALIZED_PER_INPUT_BYTE.saturating_mul(self.input_bytes)),
87        )
88    }
89
90    fn retained_allowance(&self) -> u64 {
91        self.policy.limits.max_retained_bytes.min(
92            RETAINED_BASE.saturating_add(RETAINED_PER_INPUT_BYTE.saturating_mul(self.input_bytes)),
93        )
94    }
95
96    fn charge(
97        &self,
98        dimension: ResourceDimension,
99        used: &Cell<u64>,
100        limit: u64,
101        amount: u64,
102        operation: &'static str,
103        location: Option<SourceLocation>,
104    ) -> Result<(), CodecError> {
105        if let Some(resource) = self.fuse.get() {
106            return Err(CodecError::ResourceLimit(resource));
107        }
108        let before = used.get();
109        if amount > limit.saturating_sub(before) {
110            return Err(self.refuse(
111                dimension,
112                ResourceFailure::BudgetExceeded,
113                LimitScope::Global,
114                limit,
115                before,
116                amount,
117                operation,
118                location,
119            ));
120        }
121        used.set(before + amount);
122        Ok(())
123    }
124
125    #[allow(clippy::too_many_arguments)]
126    pub(crate) fn refuse(
127        &self,
128        dimension: ResourceDimension,
129        reason: ResourceFailure,
130        scope: LimitScope,
131        limit: u64,
132        used: u64,
133        additional: u64,
134        operation: &'static str,
135        location: Option<SourceLocation>,
136    ) -> CodecError {
137        let resource = ResourceLimit {
138            dimension,
139            reason,
140            scope,
141            limit,
142            used,
143            additional,
144            context: ErrorContext {
145                operation,
146                location,
147            },
148        };
149        self.fuse.set(Some(resource));
150        CodecError::ResourceLimit(resource)
151    }
152
153    pub(crate) fn reserve_scoped(
154        &self,
155        bytes: u64,
156        operation: &'static str,
157        location: Option<SourceLocation>,
158    ) -> Result<ScopedReservation<'_>, CodecError> {
159        self.charge(
160            ResourceDimension::MaterializedBytes,
161            &self.materialized,
162            self.materialized_allowance(),
163            bytes,
164            operation,
165            location,
166        )?;
167        Ok(ScopedReservation {
168            budget: self,
169            bytes,
170            operation,
171            location,
172        })
173    }
174
175    pub(crate) fn charge_retained(
176        &self,
177        bytes: u64,
178        operation: &'static str,
179        location: Option<SourceLocation>,
180    ) -> Result<(), CodecError> {
181        self.charge(
182            ResourceDimension::RetainedBytes,
183            &self.retained,
184            self.retained_allowance(),
185            bytes,
186            operation,
187            location,
188        )
189    }
190
191    pub(crate) fn charge_entities(
192        &self,
193        count: u64,
194        operation: &'static str,
195    ) -> Result<(), CodecError> {
196        self.charge(
197            ResourceDimension::Entities,
198            &self.entities,
199            self.policy.limits.max_entities,
200            count,
201            operation,
202            None,
203        )
204    }
205
206    pub(crate) fn charge_collection_items(
207        &self,
208        count: u64,
209        operation: &'static str,
210    ) -> Result<(), CodecError> {
211        self.charge(
212            ResourceDimension::CollectionItems,
213            &self.collection_items,
214            self.policy.limits.max_collection_items,
215            count,
216            operation,
217            None,
218        )
219    }
220
221    pub(crate) fn enter_nested(
222        &self,
223        operation: &'static str,
224        location: Option<SourceLocation>,
225    ) -> Result<DepthGuard<'_>, CodecError> {
226        self.charge(
227            ResourceDimension::RecursionDepth,
228            &self.recursion_depth,
229            self.policy.limits.max_recursion_depth,
230            1,
231            operation,
232            location,
233        )?;
234        Ok(DepthGuard { budget: self })
235    }
236
237    pub(crate) fn charge_work(
238        &self,
239        units: u64,
240        operation: &'static str,
241    ) -> Result<(), CodecError> {
242        self.charge(
243            ResourceDimension::WorkUnits,
244            &self.work,
245            self.policy.limits.max_work_units,
246            units,
247            operation,
248            None,
249        )
250    }
251}
252
253/// A temporary materialization charged until it is dropped or committed.
254#[derive(Debug)]
255pub struct ScopedReservation<'a> {
256    budget: &'a DecodeBudget,
257    bytes: u64,
258    operation: &'static str,
259    location: Option<SourceLocation>,
260}
261
262impl ScopedReservation<'_> {
263    /// Increases the live temporary reservation.
264    pub fn grow(&mut self, bytes: u64) -> Result<(), CodecError> {
265        self.budget.charge(
266            ResourceDimension::MaterializedBytes,
267            &self.budget.materialized,
268            self.budget.materialized_allowance(),
269            bytes,
270            self.operation,
271            self.location,
272        )?;
273        self.bytes = self.bytes.saturating_add(bytes);
274        Ok(())
275    }
276
277    /// Converts the temporary reservation into session-retained bytes.
278    pub fn commit(self) -> Result<(), CodecError> {
279        self.budget
280            .charge_retained(self.bytes, self.operation, self.location)
281    }
282
283    /// Returns the currently reserved byte count.
284    pub fn bytes(&self) -> u64 {
285        self.bytes
286    }
287}
288
289impl Drop for ScopedReservation<'_> {
290    fn drop(&mut self) {
291        self.budget
292            .materialized
293            .set(self.budget.materialized.get().saturating_sub(self.bytes));
294    }
295}
296
297/// A guard holding one unit of recursive nesting depth.
298#[derive(Debug)]
299pub struct DepthGuard<'a> {
300    budget: &'a DecodeBudget,
301}
302
303impl Drop for DepthGuard<'_> {
304    fn drop(&mut self) {
305        self.budget
306            .recursion_depth
307            .set(self.budget.recursion_depth.get().saturating_sub(1));
308    }
309}
310
311/// A sticky, context-free local work counter.
312#[derive(Debug)]
313pub struct WorkBudget<'a> {
314    limit: usize,
315    remaining: Cell<usize>,
316    exhausted: Cell<bool>,
317    session: Option<&'a DecodeBudget>,
318}
319
320impl WorkBudget<'static> {
321    /// Creates an independent local work slice.
322    pub const fn new(limit: usize) -> Self {
323        Self {
324            limit,
325            remaining: Cell::new(limit),
326            exhausted: Cell::new(false),
327            session: None,
328        }
329    }
330}
331
332impl<'a> WorkBudget<'a> {
333    pub(crate) fn for_session(limit: u64, session: &'a DecodeBudget) -> Self {
334        let limit = if limit > usize::MAX as u64 {
335            usize::MAX
336        } else {
337            limit as usize
338        };
339        Self {
340            limit,
341            remaining: Cell::new(limit),
342            exhausted: Cell::new(false),
343            session: Some(session),
344        }
345    }
346
347    /// Charges one work unit, returning false after exhaustion.
348    pub fn charge(&self) -> bool {
349        self.charge_by(1)
350    }
351
352    /// Charges several work units, with sticky exhaustion on refusal.
353    pub fn charge_by(&self, work: usize) -> bool {
354        if self.exhausted.get() {
355            return false;
356        }
357        let remaining = self.remaining.get();
358        if work > remaining {
359            self.exhausted.set(true);
360            false
361        } else {
362            if let Some(session) = self.session {
363                if session.charge_work(work as u64, "work_budget").is_err() {
364                    self.exhausted.set(true);
365                    return false;
366                }
367            }
368            self.remaining.set(remaining - work);
369            true
370        }
371    }
372
373    /// Returns whether this slice has refused a charge.
374    pub fn exhausted(&self) -> bool {
375        self.exhausted.get()
376    }
377
378    /// Marks this slice exhausted without consuming additional session work.
379    pub fn exhaust(&self) {
380        self.exhausted.set(true);
381    }
382
383    /// Returns unconsumed work units.
384    pub fn remaining(&self) -> usize {
385        self.remaining.get()
386    }
387
388    /// Returns work units consumed by successful charges.
389    pub fn consumed(&self) -> usize {
390        self.limit.saturating_sub(self.remaining.get())
391    }
392
393    /// Creates an independent child slice capped by this budget's remainder.
394    pub fn child_slice(&self, limit: usize) -> WorkBudget<'static> {
395        WorkBudget::new(limit.min(self.remaining()))
396    }
397
398    /// Creates a child slice capped by this budget's remainder and attached to its session.
399    pub fn session_child_slice(&self, limit: usize) -> WorkBudget<'_> {
400        WorkBudget {
401            limit: limit.min(self.remaining()),
402            remaining: Cell::new(limit.min(self.remaining())),
403            exhausted: Cell::new(false),
404            session: self.session,
405        }
406    }
407
408    /// Charges this budget for work consumed by a child slice.
409    pub fn consume_child(&self, child: &WorkBudget<'_>) -> bool {
410        self.charge_by(child.consumed())
411    }
412
413    /// Classifies this local budget's refusal as a resource limit.
414    pub fn refuse(&self, operation: &'static str) -> CodecError {
415        if let Some(resource) = self.session.and_then(DecodeBudget::fused) {
416            return CodecError::ResourceLimit(resource);
417        }
418        local_limit_error(
419            "work_units",
420            self.limit as u64,
421            self.consumed().saturating_add(1) as u64,
422            operation,
423            None,
424        )
425    }
426}
427
428/// Builds a correctly classified refusal for a codec-local ceiling.
429pub fn refuse_local_limit(
430    what: &'static str,
431    limit: u64,
432    requested: u64,
433    location: Option<SourceLocation>,
434) -> CodecError {
435    local_limit_error(what, limit, requested, what, location)
436}
437
438/// Allocates `count` copies of `value` with `try_reserve_exact` and no panic on OOM.
439///
440/// Use [`DecodeContext::alloc_filled`] when a decode session can also charge
441/// collection items. This helper is for call sites that only need the
442/// allocation bound.
443pub fn alloc_filled<T: Clone>(
444    count: usize,
445    value: T,
446    operation: &'static str,
447) -> Result<Vec<T>, CodecError> {
448    let mut out = Vec::new();
449    out.try_reserve_exact(count)
450        .map_err(|_| refuse_local_limit(operation, count as u64, count as u64, None))?;
451    out.resize(count, value);
452    Ok(out)
453}
454
455fn local_limit_error(
456    what: &'static str,
457    limit: u64,
458    requested: u64,
459    operation: &'static str,
460    location: Option<SourceLocation>,
461) -> CodecError {
462    CodecError::ResourceLimit(ResourceLimit {
463        dimension: ResourceDimension::Codec(what),
464        reason: ResourceFailure::BudgetExceeded,
465        scope: LimitScope::Global,
466        limit,
467        used: requested.min(limit),
468        additional: requested.saturating_sub(limit),
469        context: ErrorContext {
470            operation,
471            location,
472        },
473    })
474}
cadmpeg_core/decode/budget.rs

cadmpeg_core/decode/
budget.rs
