parity_scale_codec/
depth_limit.rs

1// Copyright 2017, 2018 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
15use crate::{Decode, Error, Input};
16
17/// The error message returned when depth limit is reached.
18const DECODE_MAX_DEPTH_MSG: &str = "Maximum recursion depth reached when decoding";
19
20/// Extension trait to [`Decode`] for decoding with a maximum recursion depth.
21pub trait DecodeLimit: Sized {
22	/// Decode `Self` with the given maximum recursion depth and advance `input` by the number of
23	/// bytes consumed.
24	///
25	/// If `limit` is hit, an error is returned.
26	fn decode_with_depth_limit<I: Input>(limit: u32, input: &mut I) -> Result<Self, Error>;
27
28	/// Decode `Self` and consume all of the given input data.
29	///
30	/// If not all data is consumed or `limit` is hit, an error is returned.
31	fn decode_all_with_depth_limit(limit: u32, input: &mut &[u8]) -> Result<Self, Error>;
32}
33
34struct DepthTrackingInput<'a, I> {
35	input: &'a mut I,
36	depth: u32,
37	max_depth: u32,
38}
39
40impl<'a, I: Input> Input for DepthTrackingInput<'a, I> {
41	fn remaining_len(&mut self) -> Result<Option<usize>, Error> {
42		self.input.remaining_len()
43	}
44
45	fn read(&mut self, into: &mut [u8]) -> Result<(), Error> {
46		self.input.read(into)
47	}
48
49	fn read_byte(&mut self) -> Result<u8, Error> {
50		self.input.read_byte()
51	}
52
53	fn descend_ref(&mut self) -> Result<(), Error> {
54		self.input.descend_ref()?;
55		self.depth += 1;
56		if self.depth > self.max_depth {
57			Err(DECODE_MAX_DEPTH_MSG.into())
58		} else {
59			Ok(())
60		}
61	}
62
63	fn ascend_ref(&mut self) {
64		self.input.ascend_ref();
65		self.depth -= 1;
66	}
67}
68
69impl<T: Decode> DecodeLimit for T {
70	fn decode_all_with_depth_limit(limit: u32, input: &mut &[u8]) -> Result<Self, Error> {
71		let t = <Self as DecodeLimit>::decode_with_depth_limit(limit, input)?;
72
73		if input.is_empty() {
74			Ok(t)
75		} else {
76			Err(crate::decode_all::DECODE_ALL_ERR_MSG.into())
77		}
78	}
79
80	fn decode_with_depth_limit<I: Input>(limit: u32, input: &mut I) -> Result<Self, Error> {
81		let mut input = DepthTrackingInput { input, depth: 0, max_depth: limit };
82		T::decode(&mut input)
83	}
84}
85
86#[cfg(test)]
87mod tests {
88	use super::*;
89	use crate::Encode;
90
91	#[test]
92	fn decode_limit_works() {
93		type NestedVec = Vec<Vec<Vec<Vec<u8>>>>;
94		let nested: NestedVec = vec![vec![vec![vec![1]]]];
95		let encoded = nested.encode();
96
97		let decoded = NestedVec::decode_with_depth_limit(3, &mut encoded.as_slice()).unwrap();
98		assert_eq!(decoded, nested);
99		assert!(NestedVec::decode_with_depth_limit(2, &mut encoded.as_slice()).is_err());
100	}
101
102	#[test]
103	fn decode_limit_advances_input() {
104		type NestedVec = Vec<Vec<Vec<Vec<u8>>>>;
105		let nested: NestedVec = vec![vec![vec![vec![1]]]];
106		let encoded = nested.encode();
107		let encoded_slice = &mut encoded.as_slice();
108
109		let decoded = Vec::<u8>::decode_with_depth_limit(1, encoded_slice).unwrap();
110		assert_eq!(decoded, vec![4]);
111		assert!(NestedVec::decode_with_depth_limit(3, encoded_slice).is_err());
112	}
113
114	#[test]
115	fn decode_all_with_limit_advances_input() {
116		type NestedVec = Vec<Vec<Vec<Vec<u8>>>>;
117		let nested: NestedVec = vec![vec![vec![vec![1]]]];
118		let mut encoded = NestedVec::encode(&nested);
119
120		let decoded = NestedVec::decode_all_with_depth_limit(3, &mut encoded.as_slice()).unwrap();
121		assert_eq!(decoded, nested);
122
123		encoded.extend(&[1, 2, 3, 4, 5, 6]);
124		assert_eq!(
125			NestedVec::decode_all_with_depth_limit(3, &mut encoded.as_slice())
126				.unwrap_err()
127				.to_string(),
128			"Input buffer has still data left after decoding!",
129		);
130	}
131}
parity_scale_codec/depth_limit.rs

parity_scale_codec/
depth_limit.rs
