parity_scale_codec/
counted_input.rs

1// Copyright 2017-2024 Parity Technologies
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15/// A wrapper for `Input` which tracks the number fo bytes that are successfully read.
16///
17/// If inner `Input` fails to read, the counter is not incremented.
18///
19/// It can count until `u64::MAX - 1` accurately then saturate.
20pub struct CountedInput<'a, I: crate::Input> {
21	input: &'a mut I,
22	counter: u64,
23}
24
25impl<'a, I: crate::Input> CountedInput<'a, I> {
26	/// Create a new `CountedInput` with the given input.
27	pub fn new(input: &'a mut I) -> Self {
28		Self { input, counter: 0 }
29	}
30
31	/// Get the number of bytes successfully read.
32	pub fn count(&self) -> u64 {
33		self.counter
34	}
35}
36
37impl<I: crate::Input> crate::Input for CountedInput<'_, I> {
38	fn remaining_len(&mut self) -> Result<Option<usize>, crate::Error> {
39		self.input.remaining_len()
40	}
41
42	fn read(&mut self, into: &mut [u8]) -> Result<(), crate::Error> {
43		self.input.read(into).inspect(|_r| {
44			self.counter = self.counter.saturating_add(into.len().try_into().unwrap_or(u64::MAX));
45		})
46	}
47
48	fn read_byte(&mut self) -> Result<u8, crate::Error> {
49		self.input.read_byte().inspect(|_r| {
50			self.counter = self.counter.saturating_add(1);
51		})
52	}
53
54	fn ascend_ref(&mut self) {
55		self.input.ascend_ref()
56	}
57
58	fn descend_ref(&mut self) -> Result<(), crate::Error> {
59		self.input.descend_ref()
60	}
61}
62
63#[cfg(test)]
64mod test {
65	use super::*;
66	use crate::Input;
67
68	#[test]
69	fn test_counted_input_input_impl() {
70		let mut input = &[1u8, 2, 3, 4, 5][..];
71		let mut counted_input = CountedInput::new(&mut input);
72
73		assert_eq!(counted_input.remaining_len().unwrap(), Some(5));
74		assert_eq!(counted_input.count(), 0);
75
76		counted_input.read_byte().unwrap();
77
78		assert_eq!(counted_input.remaining_len().unwrap(), Some(4));
79		assert_eq!(counted_input.count(), 1);
80
81		counted_input.read(&mut [0u8; 2][..]).unwrap();
82
83		assert_eq!(counted_input.remaining_len().unwrap(), Some(2));
84		assert_eq!(counted_input.count(), 3);
85
86		counted_input.ascend_ref();
87		counted_input.descend_ref().unwrap();
88
89		counted_input.read(&mut [0u8; 2][..]).unwrap();
90
91		assert_eq!(counted_input.remaining_len().unwrap(), Some(0));
92		assert_eq!(counted_input.count(), 5);
93
94		assert_eq!(counted_input.read_byte(), Err("Not enough data to fill buffer".into()));
95
96		assert_eq!(counted_input.remaining_len().unwrap(), Some(0));
97		assert_eq!(counted_input.count(), 5);
98
99		assert_eq!(
100			counted_input.read(&mut [0u8; 2][..]),
101			Err("Not enough data to fill buffer".into())
102		);
103
104		assert_eq!(counted_input.remaining_len().unwrap(), Some(0));
105		assert_eq!(counted_input.count(), 5);
106	}
107
108	#[test]
109	fn test_counted_input_max_count_read_byte() {
110		let max_exact_count = u64::MAX - 1;
111
112		let mut input = &[0u8; 1000][..];
113		let mut counted_input = CountedInput::new(&mut input);
114
115		counted_input.counter = max_exact_count;
116
117		assert_eq!(counted_input.count(), max_exact_count);
118
119		counted_input.read_byte().unwrap();
120
121		assert_eq!(counted_input.count(), u64::MAX);
122
123		counted_input.read_byte().unwrap();
124
125		assert_eq!(counted_input.count(), u64::MAX);
126	}
127
128	#[test]
129	fn test_counted_input_max_count_read() {
130		let max_exact_count = u64::MAX - 1;
131
132		let mut input = &[0u8; 1000][..];
133		let mut counted_input = CountedInput::new(&mut input);
134
135		counted_input.counter = max_exact_count;
136
137		assert_eq!(counted_input.count(), max_exact_count);
138
139		counted_input.read(&mut [0u8; 2][..]).unwrap();
140
141		assert_eq!(counted_input.count(), u64::MAX);
142
143		counted_input.read(&mut [0u8; 2][..]).unwrap();
144
145		assert_eq!(counted_input.count(), u64::MAX);
146	}
147}
parity_scale_codec/counted_input.rs

parity_scale_codec/
counted_input.rs
