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

cadmpeg_core/decode/
budget.rs
