cadmpeg_core/decode/
budget.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Resource ceilings for one decode session, and the local work counter.
3//!
4//! A session counts input bytes, decompressed bytes, retained output bytes,
5//! entities, and recursion depth. A refusal in any of them fuses the session:
6//! every later charge fails and `finish_session` returns the refusal.
7
8use std::cell::{Cell, RefCell};
9use std::rc::Rc;
10
11use crate::CodecError;
12
13use super::error::{ResourceDimension, ResourceFailure, ResourceLimit};
14use super::policy::{
15    DecodePolicy, DECOMPRESSED_TOTAL_BASE, DECOMPRESSED_TOTAL_PER_INPUT_BYTE, RETAINED_BASE,
16    RETAINED_PER_INPUT_BYTE,
17};
18
19#[derive(Debug)]
20pub(super) struct DecodeBudget {
21    /// Policy applied to this budget.
22    policy: DecodePolicy,
23    input_bytes: Cell<u64>,
24    decompressed: Cell<u64>,
25    retained: Cell<u64>,
26    entities: Cell<u64>,
27    recursion_depth: Cell<u64>,
28    fuse: Cell<Option<ResourceLimit>>,
29    pub(super) trace: RefCell<Option<Vec<super::transaction::Event>>>,
30}
31
32impl DecodeBudget {
33    pub(super) fn record(&self, event: super::transaction::Event) {
34        if let Some(trace) = self.trace.borrow_mut().as_mut() {
35            super::transaction::record(trace, event);
36        }
37    }
38
39    pub(super) fn state(&self) -> super::transaction::BudgetState {
40        super::transaction::BudgetState {
41            policy: self.policy,
42            input: self.input_bytes.get(),
43            decompressed: self.decompressed.get(),
44            retained: self.retained.get(),
45            entities: self.entities.get(),
46            depth: self.recursion_depth.get(),
47            fuse: self.fuse.get(),
48        }
49    }
50
51    pub(super) fn restore(&self, state: &super::transaction::BudgetState) {
52        self.input_bytes.set(state.input);
53        self.decompressed.set(state.decompressed);
54        self.retained.set(state.retained);
55        self.entities.set(state.entities);
56        self.recursion_depth.set(state.depth);
57        self.fuse.set(state.fuse);
58    }
59
60    pub(super) fn replay_charge(
61        &self,
62        dimension: ResourceDimension,
63        amount: u64,
64        operation: &'static str,
65    ) -> Result<(), CodecError> {
66        match dimension {
67            ResourceDimension::InputBytes => self.charge_input(amount, operation),
68            ResourceDimension::DecompressedBytes => self.charge_decompressed(amount, operation),
69            ResourceDimension::RetainedBytes => {
70                self.retain_bytes(amount, operation).map_err(Into::into)
71            }
72            ResourceDimension::Entities => self.charge_entities(amount, operation),
73            ResourceDimension::RecursionDepth => self
74                .charge(
75                    dimension,
76                    &self.recursion_depth,
77                    self.policy.limits.max_recursion_depth,
78                    amount,
79                    operation,
80                )
81                .map_err(Into::into),
82        }
83    }
84
85    pub(super) fn leave(&self) {
86        if let Some(depth) = self.recursion_depth.get().checked_sub(1) {
87            self.recursion_depth.set(depth);
88        }
89    }
90
91    /// Returns the policy applied to this budget.
92    pub(super) fn policy(&self) -> &DecodePolicy {
93        &self.policy
94    }
95
96    pub(super) fn new(policy: DecodePolicy, input_bytes: u64) -> Self {
97        Self {
98            policy,
99            input_bytes: Cell::new(input_bytes),
100            decompressed: Cell::new(0),
101            retained: Cell::new(0),
102            entities: Cell::new(0),
103            recursion_depth: Cell::new(0),
104            fuse: Cell::new(None),
105            trace: RefCell::new(None),
106        }
107    }
108
109    pub(super) fn fused(&self) -> Option<ResourceLimit> {
110        self.fuse.get()
111    }
112
113    /// The smaller of a policy ceiling and `base + per_input_byte * input`.
114    pub(super) fn proportional_allowance(
115        &self,
116        ceiling: u64,
117        base: u64,
118        per_input_byte: u64,
119    ) -> u64 {
120        let Some(proportional) = per_input_byte
121            .checked_mul(self.input_bytes.get())
122            .and_then(|amount| base.checked_add(amount))
123        else {
124            return ceiling;
125        };
126        ceiling.min(proportional)
127    }
128
129    pub(super) fn decompression_allowance(&self) -> u64 {
130        self.proportional_allowance(
131            self.policy.limits.max_decompressed_bytes_total,
132            DECOMPRESSED_TOTAL_BASE,
133            DECOMPRESSED_TOTAL_PER_INPUT_BYTE,
134        )
135    }
136
137    fn retained_allowance(&self) -> u64 {
138        self.proportional_allowance(
139            self.policy.limits.max_retained_bytes,
140            RETAINED_BASE,
141            RETAINED_PER_INPUT_BYTE,
142        )
143    }
144
145    pub(super) fn input_bytes(&self) -> u64 {
146        self.input_bytes.get()
147    }
148
149    pub(super) fn charge_input(
150        &self,
151        amount: u64,
152        operation: &'static str,
153    ) -> Result<(), CodecError> {
154        self.charge(
155            ResourceDimension::InputBytes,
156            &self.input_bytes,
157            self.policy.limits.max_input_bytes,
158            amount,
159            operation,
160        )
161        .map_err(Into::into)
162    }
163
164    pub(super) fn charge_decompressed(
165        &self,
166        amount: u64,
167        operation: &'static str,
168    ) -> Result<(), CodecError> {
169        self.charge(
170            ResourceDimension::DecompressedBytes,
171            &self.decompressed,
172            self.decompression_allowance(),
173            amount,
174            operation,
175        )
176        .map_err(Into::into)
177    }
178
179    pub(super) fn decompressed_used(&self) -> u64 {
180        self.decompressed.get()
181    }
182
183    fn charge(
184        &self,
185        dimension: ResourceDimension,
186        used: &Cell<u64>,
187        limit: u64,
188        amount: u64,
189        operation: &'static str,
190    ) -> Result<(), ResourceLimit> {
191        self.record(super::transaction::Event::Charge(
192            dimension, amount, operation,
193        ));
194        if let Some(resource) = self.fuse.get() {
195            return Err(resource);
196        }
197        let before = used.get();
198        #[cfg(feature = "test-support")]
199        let limit = super::test_support::charge_limit(
200            &self.policy.limits,
201            dimension,
202            operation,
203            before,
204            amount,
205        )
206        .map_or(limit, |probe| limit.min(probe));
207        if limit
208            .checked_sub(before)
209            .is_none_or(|remaining| amount > remaining)
210        {
211            return Err(self.refuse_limit(
212                dimension,
213                ResourceFailure::BudgetExceeded,
214                limit,
215                before,
216                amount,
217                operation,
218            ));
219        }
220        used.set(before + amount);
221        #[cfg(feature = "test-support")]
222        super::test_support::record_charge(dimension, operation, amount);
223        Ok(())
224    }
225
226    pub(super) fn refuse_limit(
227        &self,
228        dimension: ResourceDimension,
229        reason: ResourceFailure,
230        limit: u64,
231        used: u64,
232        additional: u64,
233        operation: &'static str,
234    ) -> ResourceLimit {
235        let resource = ResourceLimit {
236            dimension,
237            reason,
238            limit,
239            used,
240            additional,
241            operation,
242        };
243        self.fuse.set(Some(resource));
244        resource
245    }
246
247    pub(super) fn refuse(
248        &self,
249        dimension: ResourceDimension,
250        reason: ResourceFailure,
251        limit: u64,
252        used: u64,
253        additional: u64,
254        operation: &'static str,
255    ) -> CodecError {
256        self.refuse_limit(dimension, reason, limit, used, additional, operation)
257            .into()
258    }
259
260    /// Charges bytes that decoded output keeps for the rest of the session:
261    /// native record text, arena buffers, and retained source byte copies.
262    pub(super) fn retain_bytes(
263        &self,
264        bytes: u64,
265        operation: &'static str,
266    ) -> Result<(), ResourceLimit> {
267        self.charge(
268            ResourceDimension::RetainedBytes,
269            &self.retained,
270            self.retained_allowance(),
271            bytes,
272            operation,
273        )
274    }
275
276    /// Retained bytes the session can still charge before `retain_bytes` refuses.
277    pub(super) fn retained_bytes_available(&self) -> u64 {
278        let available = self.retained_available();
279        self.record(super::transaction::Event::Available(available));
280        available
281    }
282
283    pub(super) fn retained_available(&self) -> u64 {
284        if self.fuse.get().is_some() {
285            0
286        } else {
287            let Some(available) = self.retained_allowance().checked_sub(self.retained.get()) else {
288                return 0;
289            };
290            available
291        }
292    }
293
294    /// Reports an allocator refusal for `requested` bytes or slots.
295    pub(super) fn retained_allocation_failed(
296        &self,
297        requested: u64,
298        operation: &'static str,
299    ) -> CodecError {
300        self.refuse_limit(
301            ResourceDimension::RetainedBytes,
302            ResourceFailure::AllocationFailed,
303            self.retained_allowance(),
304            self.retained.get(),
305            requested,
306            operation,
307        )
308        .into()
309    }
310
311    pub(super) fn charge_entities(
312        &self,
313        count: u64,
314        operation: &'static str,
315    ) -> Result<(), CodecError> {
316        self.charge(
317            ResourceDimension::Entities,
318            &self.entities,
319            self.policy.limits.max_entities,
320            count,
321            operation,
322        )
323        .map_err(Into::into)
324    }
325
326    pub(super) fn enter_nested(
327        &self,
328        operation: &'static str,
329    ) -> Result<DepthGuard<'_>, ResourceLimit> {
330        self.charge(
331            ResourceDimension::RecursionDepth,
332            &self.recursion_depth,
333            self.policy.limits.max_recursion_depth,
334            1,
335            operation,
336        )?;
337        Ok(DepthGuard { budget: self })
338    }
339}
340
341/// A guard holding one unit of recursive nesting depth.
342#[derive(Debug)]
343pub struct DepthGuard<'a> {
344    budget: &'a DecodeBudget,
345}
346
347impl Drop for DepthGuard<'_> {
348    fn drop(&mut self) {
349        let Some(depth) = self.budget.recursion_depth.get().checked_sub(1) else {
350            return;
351        };
352        self.budget.record(super::transaction::Event::Leave);
353        self.budget.recursion_depth.set(depth);
354    }
355}
356
357/// A local work counter for one bounded search or iteration.
358///
359/// Exhaustion is sticky: once a charge is refused, every later charge is
360/// refused. A caller that runs out skips the step it bounds and reports the
361/// skip; exhaustion never refuses the decode session.
362#[derive(Debug)]
363pub struct WorkBudget {
364    limit: usize,
365    account: Rc<WorkAccount>,
366    recursion_depth: Cell<usize>,
367}
368
369/// Remaining units of one counter, and the counter that pays for its work.
370#[derive(Debug)]
371struct WorkAccount {
372    remaining: Cell<Option<usize>>,
373    parent: Option<Rc<WorkAccount>>,
374}
375
376/// RAII guard for one recursive frame of a [`WorkBudget`].
377#[derive(Debug)]
378pub struct WorkBudgetRecursionGuard<'budget> {
379    budget: &'budget WorkBudget,
380}
381
382impl Drop for WorkBudgetRecursionGuard<'_> {
383    fn drop(&mut self) {
384        if let Some(depth) = self.budget.recursion_depth.get().checked_sub(1) {
385            self.budget.recursion_depth.set(depth);
386        }
387    }
388}
389
390impl WorkBudget {
391    /// The recursion ceiling of [`WorkBudget::recursion_guard`].
392    pub const RECURSION_DEPTH: usize = 256;
393
394    /// Creates a work counter with `limit` units.
395    #[must_use]
396    pub fn new(limit: usize) -> Self {
397        Self::with_parent(limit, None)
398    }
399
400    fn with_parent(limit: usize, parent: Option<Rc<WorkAccount>>) -> Self {
401        Self {
402            limit,
403            account: Rc::new(WorkAccount {
404                remaining: Cell::new(Some(limit)),
405                parent,
406            }),
407            recursion_depth: Cell::new(0),
408        }
409    }
410
411    fn ancestors(&self) -> impl Iterator<Item = &WorkAccount> {
412        std::iter::successors(self.account.parent.as_deref(), |account| {
413            account.parent.as_deref()
414        })
415    }
416
417    /// Charges one work unit, returning false after exhaustion.
418    pub fn charge(&self) -> bool {
419        self.charge_by(1)
420    }
421
422    /// Charges several work units, with sticky exhaustion on refusal.
423    ///
424    /// A charging slice also charges every counter it was drawn from. A
425    /// charge that one of them cannot pay exhausts this counter only.
426    pub fn charge_by(&self, work: usize) -> bool {
427        let charged = self.debit(work);
428        #[cfg(feature = "test-support")]
429        if charged {
430            super::test_support::record_work(work);
431        }
432        charged
433    }
434
435    fn debit(&self, work: usize) -> bool {
436        let Some(remaining) = self.account.remaining.get() else {
437            return false;
438        };
439        let affordable = |account: &WorkAccount| {
440            account
441                .remaining
442                .get()
443                .is_some_and(|remaining| work <= remaining)
444        };
445        if work > remaining || !self.ancestors().all(affordable) {
446            self.exhaust();
447            return false;
448        }
449        self.account.remaining.set(Some(remaining - work));
450        for account in self.ancestors() {
451            account
452                .remaining
453                .set(account.remaining.get().map(|remaining| remaining - work));
454        }
455        true
456    }
457
458    /// Returns whether this counter has refused a charge.
459    pub fn exhausted(&self) -> bool {
460        self.account.remaining.get().is_none()
461    }
462
463    /// Marks this counter exhausted. The counters it was drawn from are not.
464    pub fn exhaust(&self) {
465        self.account.remaining.set(None);
466    }
467
468    /// Returns the work units a charge can still use: the smallest remainder
469    /// of this counter and every counter it was drawn from.
470    pub fn remaining(&self) -> usize {
471        std::iter::once(&*self.account)
472            .chain(self.ancestors())
473            .map(|account| account.remaining.get().unwrap_or(0))
474            .min()
475            .unwrap_or(0)
476    }
477
478    /// Returns work charged to this counter and the remainder forfeited by
479    /// exhaustion.
480    pub fn consumed(&self) -> usize {
481        self.limit - self.account.remaining.get().unwrap_or(0)
482    }
483
484    /// Enters one recursive frame, refusing past [`WorkBudget::RECURSION_DEPTH`].
485    /// The refusal exhausts the counter; it does not touch a decode session.
486    pub fn recursion_guard(&self) -> Result<WorkBudgetRecursionGuard<'_>, BudgetExhausted> {
487        let depth = self.recursion_depth.get();
488        if depth >= Self::RECURSION_DEPTH {
489            self.exhaust();
490            return Err(BudgetExhausted);
491        }
492        self.recursion_depth.set(depth + 1);
493        Ok(WorkBudgetRecursionGuard { budget: self })
494    }
495
496    /// Creates a child counter capped by this counter's remainder.
497    ///
498    /// The child's charges do not reach this counter until
499    /// [`WorkBudget::consume_child`] transfers them.
500    #[must_use]
501    pub fn child_slice(&self, limit: usize) -> WorkBudget {
502        WorkBudget::new(limit.min(self.remaining()))
503    }
504
505    /// Creates a child counter of at most `limit` units whose every charge is
506    /// also charged to this counter, so the work below a counter never
507    /// exceeds its limit.
508    ///
509    /// Exhausting the child, by its own limit or by an ancestor that cannot
510    /// pay, leaves this counter usable.
511    #[must_use]
512    pub fn charging_slice(&self, limit: usize) -> WorkBudget {
513        WorkBudget::with_parent(limit.min(self.remaining()), Some(Rc::clone(&self.account)))
514    }
515
516    /// Charges this counter for work consumed by a child counter.
517    ///
518    /// # Errors
519    ///
520    /// [`BudgetExhausted`] when the child's work is above this counter's
521    /// remainder. The counter marks itself exhausted in that case.
522    pub fn consume_child(&self, child: &WorkBudget) -> Result<(), BudgetExhausted> {
523        if self.debit(child.consumed()) {
524            Ok(())
525        } else {
526            Err(BudgetExhausted)
527        }
528    }
529}
530
531/// A work budget reached its limit.
532///
533/// The budget marks itself exhausted when it refuses, so this error names the
534/// event and carries no state of its own.
535#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
536#[error("work budget exhausted")]
537pub struct BudgetExhausted;
538
539/// The work a step over `items` owes a [`WorkBudget`].
540///
541/// A step that walks no item still runs, so it owes one unit. Charging zero
542/// would leave an unbounded number of empty steps free of the budget, and the
543/// walk that makes them would not be bounded by it.
544#[must_use]
545pub const fn work_units(items: usize) -> usize {
546    match items {
547        0 => 1,
548        items => items,
549    }
550}
551
552#[cfg(test)]
553mod tests {
554    use super::{work_units, DecodeBudget, WorkBudget, RETAINED_BASE, RETAINED_PER_INPUT_BYTE};
555    use crate::decode::{DecodePolicy, ResourceDimension, ResourceFailure};
556
557    fn descend(budget: &WorkBudget, depth: usize) -> usize {
558        let Ok(_guard) = budget.recursion_guard() else {
559            return depth;
560        };
561        descend(budget, depth + 1)
562    }
563
564    #[test]
565    fn recursion_guard_is_bounded_and_exhausts_the_counter() {
566        let budget = WorkBudget::new(10_000);
567        assert_eq!(descend(&budget, 0), WorkBudget::RECURSION_DEPTH);
568        assert!(budget.exhausted());
569    }
570
571    #[test]
572    fn a_step_over_no_item_owes_exactly_one_unit() {
573        assert_eq!(work_units(0), 1);
574        for items in [1, 2, 7, usize::MAX] {
575            assert_eq!(work_units(items), items);
576        }
577
578        let budget = WorkBudget::new(1);
579        assert!(budget.charge_by(work_units(0)));
580        assert_eq!(budget.consumed(), 1);
581        assert!(!budget.charge_by(work_units(0)));
582        assert!(budget.exhausted());
583    }
584
585    #[test]
586    fn child_consumption_transfers_to_the_parent_once() {
587        let parent = WorkBudget::new(3);
588        let child = parent.child_slice(2);
589        assert!(child.charge());
590        assert!(parent.consume_child(&child).is_ok());
591        assert_eq!(parent.remaining(), 2);
592        let greedy = WorkBudget::new(5);
593        assert!(greedy.charge_by(5));
594        assert!(parent.consume_child(&greedy).is_err());
595        assert!(parent.exhausted());
596    }
597
598    #[test]
599    fn charging_slices_charge_every_ancestor_and_exhaust_locally() {
600        let root = WorkBudget::new(10);
601        let child = root.charging_slice(8);
602        let grandchild = child.charging_slice(100);
603        assert_eq!(grandchild.remaining(), 8);
604        assert!(grandchild.charge_by(3));
605        assert_eq!((root.remaining(), child.remaining()), (7, 5));
606        assert_eq!(grandchild.consumed(), 3);
607
608        assert!(root.charge_by(5));
609        assert_eq!(grandchild.remaining(), 2);
610        assert!(!grandchild.charge_by(3));
611        assert!(grandchild.exhausted());
612        assert!(!child.exhausted() && !root.exhausted());
613        assert_eq!((root.remaining(), child.remaining()), (2, 2));
614
615        let sibling = child.charging_slice(1);
616        assert!(sibling.charge());
617        assert!(!sibling.charge());
618        assert!(!child.exhausted());
619        assert_eq!(root.remaining(), 1);
620
621        root.exhaust();
622        assert!(!child.charge());
623        assert!(child.exhausted());
624        assert_eq!(root.consumed(), 10);
625    }
626
627    #[test]
628    fn retained_allocation_refusal_fuses_with_current_usage() {
629        let session = DecodeBudget::new(DecodePolicy::default(), 1);
630        session
631            .retain_bytes(3, "copy retained")
632            .expect("retained charge fits the default session allowance");
633        let crate::CodecError::ResourceLimit(limit) =
634            session.retained_allocation_failed(5, "copy retained")
635        else {
636            panic!("retained allocation must produce a resource refusal");
637        };
638        assert_eq!(limit.dimension, ResourceDimension::RetainedBytes);
639        assert_eq!(limit.reason, ResourceFailure::AllocationFailed);
640        assert_eq!(limit.used, 3);
641        assert_eq!(limit.additional, 5);
642        assert_eq!(session.fused(), Some(limit));
643    }
644
645    #[test]
646    fn retained_bytes_available_reports_proportional_allowance() {
647        let policy = DecodePolicy::default();
648        let input = 3;
649        let expected = policy
650            .limits
651            .max_retained_bytes
652            .min(RETAINED_BASE + RETAINED_PER_INPUT_BYTE * input);
653        let session = DecodeBudget::new(policy, input);
654        assert_eq!(session.retained_bytes_available(), expected);
655    }
656
657    #[test]
658    fn retained_bytes_available_reports_policy_ceiling() {
659        let mut policy = DecodePolicy::default();
660        policy.limits.max_retained_bytes = RETAINED_BASE / 2;
661        let session = DecodeBudget::new(policy, 3);
662        assert_eq!(
663            session.retained_bytes_available(),
664            policy.limits.max_retained_bytes
665        );
666    }
667
668    #[test]
669    fn retained_bytes_available_reads_without_charging_or_fusing() {
670        let policy = DecodePolicy::default();
671        let input = 3;
672        let allowance = policy
673            .limits
674            .max_retained_bytes
675            .min(RETAINED_BASE + RETAINED_PER_INPUT_BYTE * input);
676        let session = DecodeBudget::new(policy, input);
677        let retained = 7;
678        assert!(session.retain_bytes(retained, "retain").is_ok());
679        for _ in 0..3 {
680            assert_eq!(session.retained_bytes_available(), allowance - retained);
681            assert_eq!(session.fused(), None);
682        }
683        assert!(session
684            .retain_bytes(allowance - retained, "retain remaining")
685            .is_ok());
686        assert_eq!(session.retained_bytes_available(), 0);
687        assert_eq!(session.fused(), None);
688    }
689
690    #[test]
691    fn retained_bytes_available_is_zero_after_any_dimension_fuses() {
692        let mut policy = DecodePolicy::default();
693        policy.limits.max_retained_bytes = RETAINED_BASE / 2;
694        let retained_session = DecodeBudget::new(policy, 3);
695        assert!(retained_session
696            .retain_bytes(policy.limits.max_retained_bytes + 1, "retain over limit")
697            .is_err());
698        assert!(retained_session.fused().is_some());
699        assert_eq!(retained_session.retained_bytes_available(), 0);
700
701        let entity_session = DecodeBudget::new(policy, 3);
702        assert!(entity_session
703            .charge_entities(policy.limits.max_entities + 1, "entities over limit")
704            .is_err());
705        assert!(entity_session.fused().is_some());
706        assert_eq!(entity_session.retained_bytes_available(), 0);
707    }
708
709    #[test]
710    fn zero_charge_refuses_when_usage_already_exceeds_limit() {
711        let mut policy = DecodePolicy::default();
712        policy.limits.max_entities = 1;
713        let session = DecodeBudget::new(policy, 1);
714        session.entities.set(2);
715        let error = session.charge_entities(0, "entities");
716        assert!(matches!(
717            error,
718            Err(crate::CodecError::ResourceLimit(limit))
719                if limit.dimension == ResourceDimension::Entities
720                    && limit.reason == ResourceFailure::BudgetExceeded
721                    && limit.used == 2
722                    && limit.additional == 0
723        ));
724    }
725}
cadmpeg_core/decode/budget.rs

cadmpeg_core/decode/
budget.rs
