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            self.charge_work(u64_from_index(count), "read input prefix")?;
35            let read = reader.read(&mut chunk[..count]).map_err(CodecError::Io)?;
36            if read == 0 {
37                break;
38            }
39            self.budget
40                .charge_input(u64_from_index(read), "read input prefix")?;
41            self.charge_work(u64_from_index(read), "copy input prefix")?;
42            bytes.try_reserve_exact(read).map_err(|_| {
43                self.budget.refuse(
44                    ResourceDimension::InputBytes,
45                    ResourceFailure::AllocationFailed,
46                    self.policy().limits.max_input_bytes,
47                    self.budget.input_bytes(),
48                    u64_from_index(read),
49                    "input prefix storage",
50                )
51            })?;
52            bytes.extend_from_slice(&chunk[..read]);
53        }
54        Ok(())
55    }
56
57    /// Acquires the remaining input without reading or charging the prefix twice.
58    pub fn complete_input(
59        &self,
60        reader: &mut dyn Read,
61        bytes: &mut Vec<u8>,
62    ) -> Result<(), CodecError> {
63        let max = self.policy().limits.max_input_bytes;
64        loop {
65            let remaining = max
66                .checked_sub(u64_from_index(bytes.len()))
67                .ok_or_else(|| self.refuse_input_limit(0, "complete input prefix"))?;
68            let count = if remaining >= u64_from_index(INPUT_CHUNK) {
69                INPUT_CHUNK
70            } else {
71                usize::try_from(remaining).map_err(|_| {
72                    self.refuse_codec_limit(
73                        "address input chunk",
74                        u64_from_index(INPUT_CHUNK),
75                        remaining,
76                    )
77                })?
78            };
79            let length = bytes.len().checked_add(count).ok_or_else(|| {
80                self.refuse_codec_limit("address complete input", u64_from_index(usize::MAX), max)
81            })?;
82            self.extend_input_prefix(reader, bytes, length)?;
83            if count == 0 || bytes.len() < length {
84                break;
85            }
86        }
87        self.charge_work(1, "input end probe")?;
88        if reader.read(&mut [0_u8; 1])? != 0 {
89            return Err(self.refuse_input_limit(1, "complete input"));
90        }
91        Ok(())
92    }
93
94    /// Raises a typed input-byte refusal without changing its dimension.
95    pub fn refuse_input_limit(&self, additional: u64, operation: &'static str) -> CodecError {
96        self.budget.refuse(
97            ResourceDimension::InputBytes,
98            ResourceFailure::BudgetExceeded,
99            self.policy().limits.max_input_bytes,
100            self.budget.input_bytes(),
101            additional,
102            operation,
103        )
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ResourceDimension};
110    use crate::CodecError;
111    use std::io::Cursor;
112
113    #[test]
114    fn raw_input_storage_does_not_consume_decoded_collection_slots() {
115        let arena = DecodeArena::new();
116        let mut policy = DecodePolicy::service();
117        // One arena-registry entry; payload bytes are not decoded items.
118        policy.limits.max_collection_items = 1;
119        let source = vec![b'x'; 9000];
120        let (ctx, root) =
121            DecodeContext::read_root(&mut Cursor::new(&source), &arena, &policy, false)
122                .expect("physical input is governed by input bytes");
123        assert_eq!(root.window(), source);
124        ctx.finish_session().expect("no collection refusal");
125    }
126
127    #[test]
128    fn input_prefix_refusal_keeps_input_dimension() {
129        let arena = DecodeArena::new();
130        let mut policy = DecodePolicy::service();
131        policy.limits.max_input_bytes = 2;
132        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
133        assert!(matches!(ctx.read_input_prefix(&mut Cursor::new(b"abc"), 3),
134            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
135                && limit.used == 0 && limit.additional == 3));
136    }
137
138    #[test]
139    fn input_prefix_reads_share_one_input_budget() {
140        let arena = DecodeArena::new();
141        let mut policy = DecodePolicy::service();
142        policy.limits.max_input_bytes = 2;
143        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
144        let mut reader = Cursor::new(b"abc");
145        assert_eq!(
146            ctx.read_input_prefix(&mut reader, 2).expect("prefix"),
147            b"ab"
148        );
149        assert!(matches!(ctx.read_input_prefix(&mut reader, 1),
150            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
151                && limit.used == 2 && limit.additional == 1));
152    }
153
154    #[test]
155    fn complete_input_preserves_large_policy_and_prefix() {
156        let arena = DecodeArena::new();
157        let mut policy = DecodePolicy::service();
158        policy.limits.max_input_bytes = u64::MAX;
159        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
160        let source = vec![b'x'; 9000];
161        let mut reader = Cursor::new(&source);
162        let mut bytes = ctx.read_input_prefix(&mut reader, 7).expect("prefix");
163        ctx.complete_input(&mut reader, &mut bytes)
164            .expect("complete input");
165        assert_eq!(bytes, source);
166        assert_eq!(ctx.budget.input_bytes(), 9000);
167    }
168
169    #[test]
170    fn complete_input_refusal_keeps_input_dimension() {
171        let arena = DecodeArena::new();
172        let mut policy = DecodePolicy::service();
173        policy.limits.max_input_bytes = 2;
174        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root");
175        let mut bytes = Vec::new();
176        assert!(
177            matches!(ctx.complete_input(&mut Cursor::new(b"abc"), &mut bytes),
178            Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::InputBytes
179                && limit.used == 2 && limit.additional == 1)
180        );
181        assert_eq!(bytes, b"ab");
182    }
183}
cadmpeg_core/decode/input.rs

cadmpeg_core/decode/
input.rs
