staging_xcm/
double_encoded.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Polkadot.
3
4// Polkadot is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Polkadot is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
16
17use crate::{DECODE_ALL_ERR_MSG, MAX_XCM_DECODE_DEPTH, RECURSION_LIMIT};
18use alloc::vec::Vec;
19use codec::{Decode, DecodeLimit, DecodeWithMemTracking, Encode};
20use sp_runtime::Saturating;
21
22pub(crate) const DECODE_MAX_DEPTH_MSG: &str =
23	"Depth limit exceeded while decoding DoubleEncoded object";
24pub(crate) const DECODE_RECURSION_LIMIT_MSG: &str =
25	"Recursion limit exceeded while decoding DoubleEncoded object";
26
27environmental::environmental!(nesting_count: u32);
28
29fn descend_ref_and_check_depth(
30	depth: &mut u32,
31	depth_limit: u32,
32	err_msg: &'static str,
33) -> Result<(), codec::Error> {
34	depth.saturating_inc();
35	if *depth > depth_limit {
36		return Err(err_msg.into());
37	}
38	Ok(())
39}
40
41/// `Input` implementation used for recursively decoding nested `DoubleEncoded` structures.
42///
43/// One instance of this input corresponds to one `DoubleEncoded` structure being decoded.
44/// For nested `DoubleEncoded` structures, the decoding logic will create an equal number of
45/// `NestedInput`s chained through the `downstream_input` field. For example:
46/// ```ignore
47/// NestedInput {
48/// 	downstream_input: NestedInput {
49/// 		downstream_input: NestedInput { downstream_input: ... }
50/// 	}
51/// }
52/// ```
53///
54/// Has the following behaviors:
55/// - propagates the memory allocation notifications to the downstream input. This way the
56///   downstream input will have a full picture of the entire heap memory used by the top-level
57///   decoded object, including the double encoded structures nested within it
58/// - doesn't propagate the depth related notifications to the downstream input: acts as if the
59///   double encoded structure doesn't add depth to the top-level decoded object
60/// - keeps track of and limits the decoding depth of the `DoubleEncoded` structure currently
61///   decoded.
62struct NestedInput<'a> {
63	downstream_input: &'a mut dyn codec::Input,
64	encoded: &'a [u8],
65	depth: u32,
66}
67
68impl<'a> codec::Input for NestedInput<'a> {
69	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
70		self.encoded.remaining_len()
71	}
72
73	fn read(&mut self, into: &mut [u8]) -> Result<(), codec::Error> {
74		self.encoded.read(into)
75	}
76
77	fn read_byte(&mut self) -> Result<u8, codec::Error> {
78		self.encoded.read_byte()
79	}
80
81	fn descend_ref(&mut self) -> Result<(), codec::Error> {
82		descend_ref_and_check_depth(&mut self.depth, MAX_XCM_DECODE_DEPTH, DECODE_MAX_DEPTH_MSG)
83	}
84
85	fn ascend_ref(&mut self) {
86		self.depth.saturating_dec();
87	}
88
89	fn on_before_alloc_mem(&mut self, size: usize) -> Result<(), codec::Error> {
90		self.downstream_input.on_before_alloc_mem(size)
91	}
92}
93
94/// Wrapper around the encoded and decoded versions of a value.
95/// Caches the decoded value once computed.
96#[derive(Encode, DecodeWithMemTracking, scale_info::TypeInfo)]
97#[codec(encode_bound())]
98#[codec(decode_with_mem_tracking_bound(T: Decode))]
99#[scale_info(bounds(), skip_type_params(T))]
100#[scale_info(replace_segment("staging_xcm", "xcm"))]
101#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
102pub struct DoubleEncoded<T> {
103	encoded: Vec<u8>,
104	#[codec(skip)]
105	decoded: Option<T>,
106}
107
108impl<T> Decode for DoubleEncoded<T>
109where
110	T: Decode,
111{
112	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
113		let mut obj = Self { encoded: Vec::<u8>::decode(input)?, decoded: None };
114
115		// If it's a local call, we also decode the inner double encoded object,
116		// in order to make sure that its heap memory is accounted for.
117		nesting_count::using_once(&mut 0, || {
118			nesting_count::with(|count| {
119				descend_ref_and_check_depth(
120					count,
121					RECURSION_LIMIT as u32,
122					DECODE_RECURSION_LIMIT_MSG,
123				)
124			})
125			.unwrap_or(Err("Could not access nesting_count env variable".into()))?;
126
127			let mut nested_input =
128				NestedInput { downstream_input: input, encoded: &obj.encoded[..], depth: 0 };
129			let decoded = T::decode(&mut nested_input)?;
130			// If we didn't manage to consume any byte, this is a remote call, and it can't
131			// be decoded locally.
132			if nested_input.encoded.len() == obj.encoded.len() {
133				let _ = nesting_count::with(|count| {
134					count.saturating_dec();
135				});
136
137				return Ok(obj);
138			}
139			obj.decoded = Some(decoded);
140
141			// We need to also make sure that we consumed all the input data, but we can't use
142			// `decode_all()`, because it only accepts a byte slice as input.
143			if !nested_input.encoded.is_empty() {
144				return Err(DECODE_ALL_ERR_MSG.into());
145			}
146
147			let _ = nesting_count::with(|count| {
148				count.saturating_dec();
149			});
150
151			Ok(obj)
152		})
153	}
154}
155
156impl<T> Clone for DoubleEncoded<T> {
157	fn clone(&self) -> Self {
158		Self { encoded: self.encoded.clone(), decoded: None }
159	}
160}
161
162impl<T> PartialEq for DoubleEncoded<T> {
163	fn eq(&self, other: &Self) -> bool {
164		self.encoded.eq(&other.encoded)
165	}
166}
167impl<T> Eq for DoubleEncoded<T> {}
168
169impl<T> core::fmt::Debug for DoubleEncoded<T> {
170	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
171		array_bytes::bytes2hex("0x", &self.encoded).fmt(f)
172	}
173}
174
175impl<T> From<Vec<u8>> for DoubleEncoded<T> {
176	fn from(encoded: Vec<u8>) -> Self {
177		Self { encoded, decoded: None }
178	}
179}
180
181impl<T> DoubleEncoded<T> {
182	pub fn encoded(&self) -> &[u8] {
183		&self.encoded
184	}
185
186	/// Converts a `DoubleEncoded<T>` into a `DoubleEncoded<S>`, dropping the decoded value.
187	pub fn transmute_encoded<S>(self) -> DoubleEncoded<S> {
188		DoubleEncoded { encoded: self.encoded, decoded: None }
189	}
190}
191
192impl<T: Decode> DoubleEncoded<T> {
193	/// Decode the inner encoded value and store it.
194	/// Returns a reference to the value in case of success and `Err(())` in case the decoding
195	/// fails.
196	pub fn ensure_decoded(&mut self) -> Result<&T, ()> {
197		if self.decoded.is_none() {
198			self.decoded =
199				T::decode_all_with_depth_limit(MAX_XCM_DECODE_DEPTH, &mut &self.encoded[..]).ok();
200		}
201		self.decoded.as_ref().ok_or(())
202	}
203
204	/// Provides an API similar to `TryInto` that allows fallible conversion to the inner value
205	/// type. `TryInto` implementation would collide with std blanket implementation based on
206	/// `TryFrom`.
207	pub fn try_into(mut self) -> Result<T, ()> {
208		self.ensure_decoded()?;
209		self.decoded.ok_or(())
210	}
211}
212
213#[cfg(test)]
214mod tests {
215	use super::*;
216
217	#[test]
218	fn ensure_decoded_works() {
219		let val: u64 = 42;
220		let mut encoded: DoubleEncoded<_> = Encode::encode(&val).into();
221		assert_eq!(encoded.ensure_decoded(), Ok(&val));
222	}
223
224	#[test]
225	fn try_into_works() {
226		let val: u64 = 42;
227		let encoded: DoubleEncoded<_> = Encode::encode(&val).into();
228		assert_eq!(encoded.try_into(), Ok(val));
229	}
230}
staging_xcm/double_encoded.rs

staging_xcm/
double_encoded.rs
