cadmpeg_core/decode/
input.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Bounded input acquisition under the caller's session budget.
3
4use std::io::Read;
5
6use super::{u64_from_index, DecodeContext, ResourceDimension, ResourceFailure};
7use crate::CodecError;
8
9const INPUT_CHUNK: usize = 8192;
10
11impl DecodeContext<'_> {
12    /// Reads at most `length` bytes and admits each copied byte before growth.
13    /// Input storage is governed by the input-byte dimension.
14    pub fn read_input_prefix(
15        &self,
16        reader: &mut dyn Read,
17        length: usize,
18    ) -> Result<Vec<u8>, CodecError> {
19        let mut bytes = Vec::new();
20        self.extend_input_prefix(reader, &mut bytes, length)?;
21        Ok(bytes)
22    }
23
24    /// Extends admitted input storage up to a total bounded length.
25    pub fn extend_input_prefix(
26        &self,
27        reader: &mut dyn Read,
28        bytes: &mut Vec<u8>,
29        length: usize,
30    ) -> Result<(), CodecError> {
31        let mut chunk = [0_u8; INPUT_CHUNK];
32        while bytes.len() < length {
33            let count = (length - bytes.len()).min(chunk.len());
34            let read = reader.read(&mut chunk[..count]).map_err(CodecError::Io)?;
35            if read == 0 {
36                break;
37            }
38            self.budget
39                .charge_input(u64_from_index(read), "read input prefix")?;
40
41            bytes.try_reserve_exact(read).map_err(|_| {
42                self.budget.refuse(
43                    ResourceDimension::InputBytes,
44                    ResourceFailure::AllocationFailed,
45                    self.policy().limits.max_input_bytes,
46                    self.budget.input_bytes(),
47                    u64_from_index(read),
48                    "input prefix storage",
49                )
50            })?;
51            bytes.extend_from_slice(&chunk[..read]);
52        }
53        Ok(())
54    }
55
56    /// Acquires the remaining input without reading or charging the prefix twice.
57    pub fn complete_input(
58        &self,
59        reader: &mut dyn Read,
60        bytes: &mut Vec<u8>,
61    ) -> Result<(), CodecError> {
62        let max = self.policy().limits.max_input_bytes;
63        loop {
64            let remaining = max
65                .checked_sub(u64_from_index(bytes.len()))
66                .ok_or_else(|| self.refuse_input_limit(0, "complete input prefix"))?;
67            let count = if remaining >= u64_from_index(INPUT_CHUNK) {
68                INPUT_CHUNK
69            } else {
70                usize::try_from(remaining)
71                    .map_err(|_| self.refuse_input_limit(remaining, "address input chunk"))?
72            };
73            let length = bytes
74                .len()
75                .checked_add(count)
76                .ok_or_else(|| self.refuse_input_limit(max, "address complete input"))?;
77            self.extend_input_prefix(reader, bytes, length)?;
78            if count == 0 || bytes.len() < length {
79                break;
80            }
81        }
82
83        if reader.read(&mut [0_u8; 1])? != 0 {
84            return Err(self.refuse_input_limit(1, "complete input"));
85        }
86        Ok(())
87    }
88
89    /// Raises a typed input-byte refusal without changing its dimension.
90    pub fn refuse_input_limit(&self, additional: u64, operation: &'static str) -> CodecError {
91        self.budget.refuse(
92            ResourceDimension::InputBytes,
93            ResourceFailure::BudgetExceeded,
94            self.policy().limits.max_input_bytes,
95            self.budget.input_bytes(),
96            additional,
97            operation,
98        )
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ResourceDimension};
105    use crate::CodecError;
106    use std::io::Cursor;
107
108    #[test]
109    fn input_prefix_refusal_keeps_input_dimension() {
110        let arena = DecodeArena::new();
111        let mut policy = DecodePolicy::service();
112        policy.limits.max_input_bytes = 2;
113        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
114        assert!(matches!(ctx.read_input_prefix(&mut Cursor::new(b"abc"), 3),
115            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
116                && limit.used == 0 && limit.additional == 3));
117    }
118
119    #[test]
120    fn input_prefix_reads_share_one_input_budget() {
121        let arena = DecodeArena::new();
122        let mut policy = DecodePolicy::service();
123        policy.limits.max_input_bytes = 2;
124        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
125        let mut reader = Cursor::new(b"abc");
126        assert_eq!(
127            ctx.read_input_prefix(&mut reader, 2).expect("prefix"),
128            b"ab"
129        );
130        assert!(matches!(ctx.read_input_prefix(&mut reader, 1),
131            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
132                && limit.used == 2 && limit.additional == 1));
133    }
134
135    #[test]
136    fn complete_input_preserves_large_policy_and_prefix() {
137        let arena = DecodeArena::new();
138        let mut policy = DecodePolicy::service();
139        policy.limits.max_input_bytes = u64::MAX;
140        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
141        let source = vec![b'x'; 9000];
142        let mut reader = Cursor::new(&source);
143        let mut bytes = ctx.read_input_prefix(&mut reader, 7).expect("prefix");
144        ctx.complete_input(&mut reader, &mut bytes)
145            .expect("complete input");
146        assert_eq!(bytes, source);
147        assert_eq!(ctx.budget.input_bytes(), 9000);
148    }
149
150    #[test]
151    fn complete_input_refusal_keeps_input_dimension() {
152        let arena = DecodeArena::new();
153        let mut policy = DecodePolicy::service();
154        policy.limits.max_input_bytes = 2;
155        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
156        let mut bytes = Vec::new();
157        assert!(
158            matches!(ctx.complete_input(&mut Cursor::new(b"abc"), &mut bytes),
159            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
160                && limit.used == 2 && limit.additional == 1)
161        );
162        assert_eq!(bytes, b"ab");
163    }
164}
cadmpeg_core/decode/input.rs

cadmpeg_core/decode/
input.rs
