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(
126        clippy::too_many_arguments,
127        reason = "Decode/encode helper keeps one parameter per independent arena, table, or control flag rather than a catch-all context struct."
128    )]
129    pub(crate) fn refuse(
130        &self,
131        dimension: ResourceDimension,
132        reason: ResourceFailure,
133        scope: LimitScope,
134        limit: u64,
135        used: u64,
136        additional: u64,
137        operation: &'static str,
138        location: Option<SourceLocation>,
139    ) -> CodecError {
140        let resource = ResourceLimit {
141            dimension,
142            reason,
143            scope,
144            limit,
145            used,
146            additional,
147            context: ErrorContext {
148                operation,
149                location,
150            },
151        };
152        self.fuse.set(Some(resource));
153        CodecError::ResourceLimit(resource)
154    }
155
156    pub(crate) fn reserve_scoped(
157        &self,
158        bytes: u64,
159        operation: &'static str,
160        location: Option<SourceLocation>,
161    ) -> Result<ScopedReservation<'_>, CodecError> {
162        self.charge(
163            ResourceDimension::MaterializedBytes,
164            &self.materialized,
165            self.materialized_allowance(),
166            bytes,
167            operation,
168            location,
169        )?;
170        Ok(ScopedReservation {
171            budget: self,
172            bytes,
173            operation,
174            location,
175        })
176    }
177
178    pub(crate) fn charge_retained(
179        &self,
180        bytes: u64,
181        operation: &'static str,
182        location: Option<SourceLocation>,
183    ) -> Result<(), CodecError> {
184        self.charge(
185            ResourceDimension::RetainedBytes,
186            &self.retained,
187            self.retained_allowance(),
188            bytes,
189            operation,
190            location,
191        )
192    }
193
194    pub(crate) fn charge_entities(
195        &self,
196        count: u64,
197        operation: &'static str,
198    ) -> Result<(), CodecError> {
199        self.charge(
200            ResourceDimension::Entities,
201            &self.entities,
202            self.policy.limits.max_entities,
203            count,
204            operation,
205            None,
206        )
207    }
208
209    pub(crate) fn charge_collection_items(
210        &self,
211        count: u64,
212        operation: &'static str,
213    ) -> Result<(), CodecError> {
214        self.charge(
215            ResourceDimension::CollectionItems,
216            &self.collection_items,
217            self.policy.limits.max_collection_items,
218            count,
219            operation,
220            None,
221        )
222    }
223
224    pub(crate) fn enter_nested(
225        &self,
226        operation: &'static str,
227        location: Option<SourceLocation>,
228    ) -> Result<DepthGuard<'_>, CodecError> {
229        self.charge(
230            ResourceDimension::RecursionDepth,
231            &self.recursion_depth,
232            self.policy.limits.max_recursion_depth,
233            1,
234            operation,
235            location,
236        )?;
237        Ok(DepthGuard { budget: self })
238    }
239
240    pub(crate) fn charge_work(
241        &self,
242        units: u64,
243        operation: &'static str,
244    ) -> Result<(), CodecError> {
245        self.charge(
246            ResourceDimension::WorkUnits,
247            &self.work,
248            self.policy.limits.max_work_units,
249            units,
250            operation,
251            None,
252        )
253    }
254}
255
256/// A temporary materialization charged until it is dropped or committed.
257#[derive(Debug)]
258pub struct ScopedReservation<'a> {
259    budget: &'a DecodeBudget,
260    bytes: u64,
261    operation: &'static str,
262    location: Option<SourceLocation>,
263}
264
265impl ScopedReservation<'_> {
266    /// Increases the live temporary reservation.
267    pub fn grow(&mut self, bytes: u64) -> Result<(), CodecError> {
268        self.budget.charge(
269            ResourceDimension::MaterializedBytes,
270            &self.budget.materialized,
271            self.budget.materialized_allowance(),
272            bytes,
273            self.operation,
274            self.location,
275        )?;
276        self.bytes = self.bytes.saturating_add(bytes);
277        Ok(())
278    }
279
280    /// Converts the temporary reservation into session-retained bytes.
281    pub fn commit(self) -> Result<(), CodecError> {
282        self.budget
283            .charge_retained(self.bytes, self.operation, self.location)
284    }
285
286    /// Returns the currently reserved byte count.
287    pub fn bytes(&self) -> u64 {
288        self.bytes
289    }
290}
291
292impl Drop for ScopedReservation<'_> {
293    fn drop(&mut self) {
294        self.budget
295            .materialized
296            .set(self.budget.materialized.get().saturating_sub(self.bytes));
297    }
298}
299
300/// A guard holding one unit of recursive nesting depth.
301#[derive(Debug)]
302pub struct DepthGuard<'a> {
303    budget: &'a DecodeBudget,
304}
305
306impl Drop for DepthGuard<'_> {
307    fn drop(&mut self) {
308        self.budget
309            .recursion_depth
310            .set(self.budget.recursion_depth.get().saturating_sub(1));
311    }
312}
313
314/// A sticky, context-free local work counter.
315#[derive(Debug)]
316pub struct WorkBudget<'a> {
317    limit: usize,
318    remaining: Cell<usize>,
319    exhausted: Cell<bool>,
320    session: Option<&'a DecodeBudget>,
321}
322
323impl WorkBudget<'static> {
324    /// Creates an independent local work slice.
325    pub const fn new(limit: usize) -> Self {
326        Self {
327            limit,
328            remaining: Cell::new(limit),
329            exhausted: Cell::new(false),
330            session: None,
331        }
332    }
333}
334
335impl<'a> WorkBudget<'a> {
336    pub(crate) fn for_session(limit: u64, session: &'a DecodeBudget) -> Self {
337        let limit = if limit > usize::MAX as u64 {
338            usize::MAX
339        } else {
340            limit as usize
341        };
342        Self {
343            limit,
344            remaining: Cell::new(limit),
345            exhausted: Cell::new(false),
346            session: Some(session),
347        }
348    }
349
350    /// Charges one work unit, returning false after exhaustion.
351    pub fn charge(&self) -> bool {
352        self.charge_by(1)
353    }
354
355    /// Charges several work units, with sticky exhaustion on refusal.
356    pub fn charge_by(&self, work: usize) -> bool {
357        if self.exhausted.get() {
358            return false;
359        }
360        let remaining = self.remaining.get();
361        if work > remaining {
362            self.exhausted.set(true);
363            false
364        } else {
365            if let Some(session) = self.session {
366                if session.charge_work(work as u64, "work_budget").is_err() {
367                    self.exhausted.set(true);
368                    return false;
369                }
370            }
371            self.remaining.set(remaining - work);
372            true
373        }
374    }
375
376    /// Returns whether this slice has refused a charge.
377    pub fn exhausted(&self) -> bool {
378        self.exhausted.get()
379    }
380
381    /// Marks this slice exhausted without consuming additional session work.
382    pub fn exhaust(&self) {
383        self.exhausted.set(true);
384    }
385
386    /// Returns unconsumed work units.
387    pub fn remaining(&self) -> usize {
388        self.remaining.get()
389    }
390
391    /// Returns work units consumed by successful charges.
392    pub fn consumed(&self) -> usize {
393        self.limit.saturating_sub(self.remaining.get())
394    }
395
396    /// Creates an independent child slice capped by this budget's remainder.
397    pub fn child_slice(&self, limit: usize) -> WorkBudget<'static> {
398        WorkBudget::new(limit.min(self.remaining()))
399    }
400
401    /// Creates a child slice capped by this budget's remainder and attached to its session.
402    pub fn session_child_slice(&self, limit: usize) -> WorkBudget<'_> {
403        WorkBudget {
404            limit: limit.min(self.remaining()),
405            remaining: Cell::new(limit.min(self.remaining())),
406            exhausted: Cell::new(false),
407            session: self.session,
408        }
409    }
410
411    /// Charges this budget for work consumed by a child slice.
412    pub fn consume_child(&self, child: &WorkBudget<'_>) -> bool {
413        self.charge_by(child.consumed())
414    }
415
416    /// Classifies this local budget's refusal as a resource limit.
417    pub fn refuse(&self, operation: &'static str) -> CodecError {
418        if let Some(resource) = self.session.and_then(DecodeBudget::fused) {
419            return CodecError::ResourceLimit(resource);
420        }
421        local_limit_error(
422            "work_units",
423            self.limit as u64,
424            self.consumed().saturating_add(1) as u64,
425            operation,
426            None,
427        )
428    }
429}
430
431/// Builds a correctly classified refusal for a codec-local ceiling.
432pub fn refuse_local_limit(
433    what: &'static str,
434    limit: u64,
435    requested: u64,
436    location: Option<SourceLocation>,
437) -> CodecError {
438    local_limit_error(what, limit, requested, what, location)
439}
440
441fn local_limit_error(
442    what: &'static str,
443    limit: u64,
444    requested: u64,
445    operation: &'static str,
446    location: Option<SourceLocation>,
447) -> CodecError {
448    CodecError::ResourceLimit(ResourceLimit {
449        dimension: ResourceDimension::Codec(what),
450        reason: ResourceFailure::BudgetExceeded,
451        scope: LimitScope::Global,
452        limit,
453        used: requested.min(limit),
454        additional: requested.saturating_sub(limit),
455        context: ErrorContext {
456            operation,
457            location,
458        },
459    })
460}
cadmpeg_core/decode/budget.rs

cadmpeg_core/decode/
budget.rs
