Track nested memory used by xcm::DoubleEncoded

bridges/snowbridge/primitives/outbound-queue/src/v2/converter/convert.rs

@@ -284,9 +284,8 @@ where
 		let transact_call = match_expression!(self.peek(), Ok(Transact { call, .. }), call);
 		if let Some(transact_call) = transact_call {
 			let _ = self.next();
-			let transact =
-				ContractCall::decode_all(&mut transact_call.clone().into_encoded().as_slice())
-					.map_err(|_| TransactDecodeFailed)?;
+			let transact = ContractCall::decode_all(&mut transact_call.encoded())
+				.map_err(|_| TransactDecodeFailed)?;
 			match transact {
 				ContractCall::V1 { target, calldata, gas, value } => commands
 					.push(Command::CallContract { target: target.into(), calldata, gas, value }),

polkadot/xcm/src/double_encoded.rs

@@ -16,7 +16,11 @@
 
 use crate::MAX_XCM_DECODE_DEPTH;
 use alloc::vec::Vec;
-use codec::{Decode, DecodeLimit, DecodeWithMemTracking, Encode};
+use codec::{Decode, DecodeLimit, DecodeWithMemTracking, Encode, Input};
+
+use sp_runtime::nested_mem;
+
+const DECODE_ALL_ERR_MSG: &str = "Input buffer has still data left after decoding!";
 
 /// Wrapper around the encoded and decoded versions of a value.
 /// Caches the decoded value once computed.
@@ -29,7 +33,8 @@ use codec::{Decode, DecodeLimit, DecodeWithMemTracking, Encode};
 pub struct DoubleEncoded<T> {
 	encoded: Vec<u8>,
 	#[codec(skip)]
-	decoded: Option<T>,
+	#[cfg_attr(feature = "json-schema", schemars(skip))]
+	decoded: Option<(T, Option<nested_mem::DeallocationReminder>)>,
 }
 
 impl<T> Clone for DoubleEncoded<T> {
@@ -58,61 +63,61 @@ impl<T> From<Vec<u8>> for DoubleEncoded<T> {
 }
 
 impl<T> DoubleEncoded<T> {
-	pub fn into<S>(self) -> DoubleEncoded<S> {
-		DoubleEncoded::from(self)
+	pub fn encoded(&self) -> &[u8] {
+		&self.encoded
 	}
 
-	pub fn from<S>(e: DoubleEncoded<S>) -> Self {
-		Self { encoded: e.encoded, decoded: None }
-	}
-
-	/// Provides an API similar to `AsRef` that provides access to the inner value.
-	/// `AsRef` implementation would expect an `&Option<T>` return type.
-	pub fn as_ref(&self) -> Option<&T> {
-		self.decoded.as_ref()
-	}
-
-	/// Access the encoded data.
-	pub fn into_encoded(self) -> Vec<u8> {
-		self.encoded
+	/// Converts a `DoubleEncoded<T>` into a `DoubleEncoded<S>`, dropping the decoded value.
+	pub fn transmute_encoded<S>(self) -> DoubleEncoded<S> {
+		DoubleEncoded { encoded: self.encoded, decoded: None }
 	}
 }
 
 impl<T: Decode> DoubleEncoded<T> {
+	fn try_decode(&self) -> Result<(T, Option<nested_mem::DeallocationReminder>), codec::Error> {
+		nested_mem::decode_with_limiter(&mut &self.encoded[..], |mem_tracking_input| {
+			let decoded = T::decode_with_depth_limit(MAX_XCM_DECODE_DEPTH, mem_tracking_input)?;
+			if mem_tracking_input.remaining_len() != Ok(Some(0)) {
+				return Err(DECODE_ALL_ERR_MSG.into());
+			}
+			Ok(decoded)
+		})
+	}
+
 	/// Decode the inner encoded value and store it.
 	/// Returns a reference to the value in case of success and `Err(())` in case the decoding
 	/// fails.
-	pub fn ensure_decoded(&mut self) -> Result<&T, ()> {
+	pub fn ensure_decoded(&mut self) -> Result<&T, codec::Error> {
 		if self.decoded.is_none() {
-			self.decoded =
-				T::decode_all_with_depth_limit(MAX_XCM_DECODE_DEPTH, &mut &self.encoded[..]).ok();
+			self.decoded = Some(self.try_decode()?);
 		}
-		self.decoded.as_ref().ok_or(())
-	}
-
-	/// Move the decoded value out or (if not present) decode `encoded`.
-	pub fn take_decoded(&mut self) -> Result<T, ()> {
-		self.decoded
-			.take()
-			.or_else(|| {
-				T::decode_all_with_depth_limit(MAX_XCM_DECODE_DEPTH, &mut &self.encoded[..]).ok()
-			})
-			.ok_or(())
+		Ok(self
+			.decoded
+			.as_ref()
+			.map(|(decoded, _deallocation_reminder)| decoded)
+			.expect("The value has just been decoded"))
 	}
 
-	/// Provides an API similar to `TryInto` that allows fallible conversion to the inner value
-	/// type. `TryInto` implementation would collide with std blanket implementation based on
-	/// `TryFrom`.
-	pub fn try_into(mut self) -> Result<T, ()> {
+	/// Do something with the decoded value, consuming `self`.
+	pub fn try_using_decoded<F, R>(mut self, f: F) -> Result<R, codec::Error>
+	where
+		F: FnOnce(T) -> R,
+	{
 		self.ensure_decoded()?;
-		self.decoded.ok_or(())
+		let (decoded, _deallocation_reminder) =
+			self.decoded.expect("The value has just been decoded");
+		Ok(f(decoded))
 	}
 }
 
 #[cfg(test)]
 mod tests {
 	use super::*;
 
+	use sp_runtime::generic::DEFAULT_CALL_SIZE_LIMIT;
+
+	const DECODE_OOM_MSG: &str = "Heap memory limit exceeded while decoding";
+
 	#[test]
 	fn ensure_decoded_works() {
 		let val: u64 = 42;
@@ -121,16 +126,38 @@ mod tests {
 	}
 
 	#[test]
-	fn take_decoded_works() {
-		let val: u64 = 42;
-		let mut encoded: DoubleEncoded<_> = Encode::encode(&val).into();
-		assert_eq!(encoded.take_decoded(), Ok(val));
-	}
-
-	#[test]
-	fn try_into_works() {
-		let val: u64 = 42;
-		let encoded: DoubleEncoded<_> = Encode::encode(&val).into();
-		assert_eq!(encoded.try_into(), Ok(val));
+	fn try_using_decoded_works() {
+		let val_1 = vec![1; DEFAULT_CALL_SIZE_LIMIT - 1000];
+		let encoded_val_1: DoubleEncoded<Vec<u8>> = Encode::encode(&val_1).into();
+
+		assert_eq!(nested_mem::get_current_limit(), None);
+		nested_mem::using_limiter_once(|| {
+			assert_eq!(nested_mem::get_current_limit(), Some(DEFAULT_CALL_SIZE_LIMIT));
+			encoded_val_1
+				.try_using_decoded(|decoded_val| {
+					assert_eq!(nested_mem::get_current_limit(), Some(1000));
+					assert_eq!(decoded_val, val_1);
+
+					let val_2 = vec![2; 999];
+					let encoded_val_2: DoubleEncoded<Vec<u8>> = Encode::encode(&val_2).into();
+					let res = encoded_val_2.try_using_decoded(|decoded_val| {
+						assert_eq!(decoded_val, val_2);
+						assert_eq!(nested_mem::get_current_limit(), Some(1));
+					});
+					assert_eq!(res, Ok(()));
+					assert_eq!(nested_mem::get_current_limit(), Some(1000));
+
+					let val_2 = vec![2; 1000];
+					let encoded_val_2: DoubleEncoded<Vec<u8>> = Encode::encode(&val_2).into();
+					let res = encoded_val_2.try_using_decoded(|decoded_val| {
+						assert_eq!(decoded_val, val_2);
+					});
+					assert_eq!(res, Err(DECODE_OOM_MSG.into()));
+					assert_eq!(nested_mem::get_current_limit(), Some(1000));
+				})
+				.unwrap();
+			assert_eq!(nested_mem::get_current_limit(), Some(DEFAULT_CALL_SIZE_LIMIT));
+		});
+		assert_eq!(nested_mem::get_current_limit(), None);
 	}
 }

polkadot/xcm/src/v3/mod.rs

@@ -1161,7 +1161,7 @@ impl<Call> Instruction<Call> {
 			HrmpChannelClosing { initiator, sender, recipient } =>
 				HrmpChannelClosing { initiator, sender, recipient },
 			Transact { origin_kind, require_weight_at_most, call } =>
-				Transact { origin_kind, require_weight_at_most, call: call.into() },
+				Transact { origin_kind, require_weight_at_most, call: call.transmute_encoded() },
 			ReportError(response_info) => ReportError(response_info),
 			DepositAsset { assets, beneficiary } => DepositAsset { assets, beneficiary },
 			DepositReserveAsset { assets, dest, xcm } => DepositReserveAsset { assets, dest, xcm },

polkadot/xcm/src/v4/mod.rs

@@ -1122,7 +1122,7 @@ impl<Call> Instruction<Call> {
 			HrmpChannelClosing { initiator, sender, recipient } =>
 				HrmpChannelClosing { initiator, sender, recipient },
 			Transact { origin_kind, require_weight_at_most, call } =>
-				Transact { origin_kind, require_weight_at_most, call: call.into() },
+				Transact { origin_kind, require_weight_at_most, call: call.transmute_encoded() },
 			ReportError(response_info) => ReportError(response_info),
 			DepositAsset { assets, beneficiary } => DepositAsset { assets, beneficiary },
 			DepositReserveAsset { assets, dest, xcm } => DepositReserveAsset { assets, dest, xcm },
@@ -1306,7 +1306,7 @@ impl<Call: Decode + GetDispatchInfo> TryFrom<NewInstruction<Call>> for Instructi
 			Transact { origin_kind, mut call, fallback_max_weight } => {
 				// We first try to decode the call, if we can't, we use the fallback weight,
 				// if there's no fallback, we just return `Weight::MAX`.
-				let require_weight_at_most = match call.take_decoded() {
+				let require_weight_at_most = match call.ensure_decoded() {
 					Ok(decoded) => decoded.get_dispatch_info().call_weight,
 					Err(error) => {
 						let fallback_weight = fallback_max_weight.unwrap_or(Weight::MAX);
@@ -1319,7 +1319,7 @@ impl<Call: Decode + GetDispatchInfo> TryFrom<NewInstruction<Call>> for Instructi
 						fallback_weight
 					},
 				};
-				Self::Transact { origin_kind, require_weight_at_most, call: call.into() }
+				Self::Transact { origin_kind, require_weight_at_most, call }
 			},
 			ReportError(response_info) => Self::ReportError(QueryResponseInfo {
 				query_id: response_info.query_id,

polkadot/xcm/src/v5/mod.rs

@@ -1190,7 +1190,7 @@ impl<Call> Instruction<Call> {
 			HrmpChannelClosing { initiator, sender, recipient } =>
 				HrmpChannelClosing { initiator, sender, recipient },
 			Transact { origin_kind, call, fallback_max_weight } =>
-				Transact { origin_kind, call: call.into(), fallback_max_weight },
+				Transact { origin_kind, call: call.transmute_encoded(), fallback_max_weight },
 			ReportError(response_info) => ReportError(response_info),
 			DepositAsset { assets, beneficiary } => DepositAsset { assets, beneficiary },
 			DepositReserveAsset { assets, dest, xcm } => DepositReserveAsset { assets, dest, xcm },