pallet-treasury: Ensure we respect max_amount for spend across batch calls (#13468)

* `pallet-treasury`: Ensure we respect `max_amount` for spend across batch calls

When calling `spend` the origin defines the `max_amount` of tokens it is allowed to spend. The
problem is that someone can send a `batch(spend, spend)` to circumvent this restriction as we don't
check across different calls that the `max_amount` is respected. This pull request fixes this
behavior by introducing a so-called dispatch context. This dispatch context is created once per
outer most `dispatch` call. For more information see the docs in this pr. The treasury then uses
this dispatch context to attach information about already spent funds per `max_amount` (we assume
that each origin has a different `max_amount` configured). So, a `batch(spend, spend)` is now
checked to stay inside the allowed spending bounds.

Fixes: https://github.com/paritytech/substrate/issues/13167

* Import `Box` for wasm

* FMT

substrate/Cargo.lock

@@ -2383,6 +2383,7 @@ version = "4.0.0-dev"
 dependencies = [
  "assert_matches",
  "bitflags",
+ "environmental",
  "frame-metadata",
  "frame-support-procedural",
  "frame-system",
@@ -6705,6 +6706,7 @@ dependencies = [
  "frame-system",
  "impl-trait-for-tuples",
  "pallet-balances",
+ "pallet-utility",
  "parity-scale-codec",
  "scale-info",
  "serde",

substrate/frame/support/Cargo.toml

@@ -38,6 +38,7 @@ smallvec = "1.8.0"
 log = { version = "0.4.17", default-features = false }
 sp-core-hashing-proc-macro = { version = "5.0.0", path = "../../primitives/core/hashing/proc-macro" }
 k256 = { version = "0.11.5", default-features = false, features = ["ecdsa"] }
+environmental = { version = "1.1.4", default-features = false }
 
 [dev-dependencies]
 serde_json = "1.0.85"
@@ -67,6 +68,7 @@ std = [
 	"sp-weights/std",
 	"frame-support-procedural/std",
 	"log/std",
+	"environmental/std",
 ]
 runtime-benchmarks = []
 try-runtime = []

substrate/frame/support/procedural/src/pallet/expand/call.rs

@@ -333,22 +333,24 @@ pub fn expand_call(def: &mut Def) -> proc_macro2::TokenStream {
 				self,
 				origin: Self::RuntimeOrigin
 			) -> #frame_support::dispatch::DispatchResultWithPostInfo {
-				match self {
-					#(
-						Self::#fn_name { #( #args_name_pattern, )* } => {
-							#frame_support::sp_tracing::enter_span!(
-								#frame_support::sp_tracing::trace_span!(stringify!(#fn_name))
-							);
-							#maybe_allow_attrs
-							<#pallet_ident<#type_use_gen>>::#fn_name(origin, #( #args_name, )* )
-								.map(Into::into).map_err(Into::into)
+				#frame_support::dispatch_context::run_in_context(|| {
+					match self {
+						#(
+							Self::#fn_name { #( #args_name_pattern, )* } => {
+								#frame_support::sp_tracing::enter_span!(
+									#frame_support::sp_tracing::trace_span!(stringify!(#fn_name))
+								);
+								#maybe_allow_attrs
+								<#pallet_ident<#type_use_gen>>::#fn_name(origin, #( #args_name, )* )
+									.map(Into::into).map_err(Into::into)
+							},
+						)*
+						Self::__Ignore(_, _) => {
+							let _ = origin; // Use origin for empty Call enum
+							unreachable!("__PhantomItem cannot be used.");
 						},
-					)*
-					Self::__Ignore(_, _) => {
-						let _ = origin; // Use origin for empty Call enum
-						unreachable!("__PhantomItem cannot be used.");
-					},
-				}
+					}
+				})
 			}
 		}
 

substrate/frame/support/src/dispatch_context.rs

@@ -0,0 +1,232 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Provides functions to interact with the dispatch context.
+//!
+//! A Dispatch context is created by calling [`run_in_context`] and then the given closure will be
+//! executed in this dispatch context. Everyting run in this `closure` will have access to the same
+//! dispatch context. This also applies to nested calls of [`run_in_context`]. The dispatch context
+//! can be used to store and retrieve information locally in this context. The dispatch context can
+//! be accessed by using [`with_context`]. This function will execute the given closure and give it
+//! access to the value stored in the dispatch context.
+//!
+//! # FRAME integration
+//!
+//! The FRAME macros implement [`UnfilteredDispatchable`](crate::traits::UnfilteredDispatchable) for
+//! each pallet `Call` enum. Part of this implementation is the call to [`run_in_context`], so that
+//! each call to
+//! [`UnfilteredDispatchable::dispatch_bypass_filter`](crate::traits::UnfilteredDispatchable::dispatch_bypass_filter)
+//! or [`Dispatchable::dispatch`](crate::dispatch::Dispatchable::dispatch) will run in a dispatch
+//! context.
+//!
+//! # Example
+//!
+//! ```
+//! use frame_support::dispatch_context::{with_context, run_in_context};
+//!
+//! // Not executed in a dispatch context, so it should return `None`.
+//! assert!(with_context::<(), _>(|_| println!("Hello")).is_none());
+//!
+//! // Run it in a dispatch context and `with_context` returns `Some(_)`.
+//! run_in_context(|| {
+//!     assert!(with_context::<(), _>(|_| println!("Hello")).is_some());
+//! });
+//!
+//! #[derive(Default)]
+//! struct CustomContext(i32);
+//!
+//! run_in_context(|| {
+//!     with_context::<CustomContext, _>(|v| {
+//!         // Intitialize the value to the default value.
+//!         assert_eq!(0, v.or_default().0);
+//!         v.or_default().0 = 10;
+//!     });
+//!
+//!     with_context::<CustomContext, _>(|v| {
+//!         // We are still in the same context and can still access the set value.
+//!         assert_eq!(10, v.or_default().0);
+//!     });
+//!
+//!     run_in_context(|| {
+//!         with_context::<CustomContext, _>(|v| {
+//!             // A nested call of `run_in_context` stays in the same dispatch context
+//!             assert_eq!(10, v.or_default().0);
+//!         })
+//!     })
+//! });
+//!
+//! run_in_context(|| {
+//!     with_context::<CustomContext, _>(|v| {
+//!         // We left the other context and created a new one, so we should be back
+//!         // to our default value.
+//!         assert_eq!(0, v.or_default().0);
+//!     });
+//! });
+//! ```
+//!
+//! In your pallet you will only have to use [`with_context`], because as described above
+//! [`run_in_context`] will be handled by FRAME for you.
+
+use sp_std::{
+	any::{Any, TypeId},
+	boxed::Box,
+	collections::btree_map::{BTreeMap, Entry},
+};
+
+environmental::environmental!(DISPATCH_CONTEXT: BTreeMap<TypeId, Box<dyn Any>>);
+
+/// Abstraction over some optional value `T` that is stored in the dispatch context.
+pub struct Value<'a, T> {
+	value: Option<&'a mut T>,
+	new_value: Option<T>,
+}
+
+impl<T> Value<'_, T> {
+	/// Get the value as reference.
+	pub fn get(&self) -> Option<&T> {
+		self.new_value.as_ref().or_else(|| self.value.as_ref().map(|v| *v as &T))
+	}
+
+	/// Get the value as mutable reference.
+	pub fn get_mut(&mut self) -> Option<&mut T> {
+		self.new_value.as_mut().or_else(|| self.value.as_mut().map(|v| *v as &mut T))
+	}
+
+	/// Set to the given value.
+	///
+	/// [`Self::get`] and [`Self::get_mut`] will return `new_value` afterwards.
+	pub fn set(&mut self, new_value: T) {
+		self.value = None;
+		self.new_value = Some(new_value);
+	}
+
+	/// Returns a mutable reference to the value.
+	///
+	/// If the internal value isn't initialized, this will set it to [`Default::default()`] before
+	/// returning the mutable reference.
+	pub fn or_default(&mut self) -> &mut T
+	where
+		T: Default,
+	{
+		if let Some(v) = &mut self.value {
+			return v
+		}
+
+		self.new_value.get_or_insert_with(|| Default::default())
+	}
+
+	/// Clear the internal value.
+	///
+	/// [`Self::get`] and [`Self::get_mut`] will return `None` afterwards.
+	pub fn clear(&mut self) {
+		self.new_value = None;
+		self.value = None;
+	}
+}
+
+/// Runs the given `callback` in the dispatch context and gives access to some user defined value.
+///
+/// Passes the a mutable reference of [`Value`] to the callback. The value will be of type `T` and
+/// is identified using the [`TypeId`] of `T`. This means that `T` should be some unique type to
+/// make the value unique. If no value is set yet [`Value::get()`] and [`Value::get_mut()`] will
+/// return `None`. It is totally valid to have some `T` that is shared between different callers to
+/// have access to the same value.
+///
+/// Returns `None` if the current context is not a dispatch context. To create a context it is
+/// required to call [`run_in_context`] with the closure to execute in this context. So, for example
+/// in tests it could be that there isn't any dispatch context or when calling a dispatchable like a
+/// normal Rust function from some FRAME hook.
+pub fn with_context<T: 'static, R>(callback: impl FnOnce(&mut Value<T>) -> R) -> Option<R> {
+	DISPATCH_CONTEXT::with(|c| match c.entry(TypeId::of::<T>()) {
+		Entry::Occupied(mut o) => {
+			let value = o.get_mut().downcast_mut::<T>();
+
+			if value.is_none() {
+				log::error!(
+					"Failed to downcast value for type {} in dispatch context!",
+					sp_std::any::type_name::<T>(),
+				);
+			}
+
+			let mut value = Value { value, new_value: None };
+			let res = callback(&mut value);
+
+			if value.value.is_none() && value.new_value.is_none() {
+				o.remove();
+			} else if let Some(new_value) = value.new_value {
+				o.insert(Box::new(new_value) as Box<_>);
+			}
+
+			res
+		},
+		Entry::Vacant(v) => {
+			let mut value = Value { value: None, new_value: None };
+
+			let res = callback(&mut value);
+
+			if let Some(new_value) = value.new_value {
+				v.insert(Box::new(new_value) as Box<_>);
+			}
+
+			res
+		},
+	})
+}
+
+/// Run the given closure `run` in a dispatch context.
+///
+/// Nested calls to this function will execute `run` in the same dispatch context as the initial
+/// call to this function. In other words, all nested calls of this function will be done in the
+/// same dispatch context.
+pub fn run_in_context<R>(run: impl FnOnce() -> R) -> R {
+	DISPATCH_CONTEXT::using_once(&mut Default::default(), run)
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn dispatch_context_works() {
+		// No context, so we don't execute
+		assert!(with_context::<(), _>(|_| ()).is_none());
+
+		let ret = run_in_context(|| with_context::<(), _>(|_| 1).unwrap());
+		assert_eq!(1, ret);
+
+		#[derive(Default)]
+		struct Context(i32);
+
+		let res = run_in_context(|| {
+			with_context::<Context, _>(|v| {
+				assert_eq!(0, v.or_default().0);
+
+				v.or_default().0 = 100;
+			});
+
+			run_in_context(|| {
+				run_in_context(|| {
+					run_in_context(|| with_context::<Context, _>(|v| v.or_default().0).unwrap())
+				})
+			})
+		});
+
+		// Ensure that the initial value set in the context is also accessible after nesting the
+		// `run_in_context` calls.
+		assert_eq!(100, res);
+	}
+}

substrate/frame/support/src/lib.rs

@@ -78,6 +78,7 @@ pub mod inherent;
 #[macro_use]
 pub mod error;
 pub mod crypto;
+pub mod dispatch_context;
 pub mod instances;
 pub mod migrations;
 pub mod traits;

substrate/frame/support/test/tests/pallet.rs

@@ -16,9 +16,12 @@
 // limitations under the License.
 
 use frame_support::{
+	assert_ok,
 	dispatch::{
-		DispatchClass, DispatchInfo, GetDispatchInfo, Parameter, Pays, UnfilteredDispatchable,
+		DispatchClass, DispatchInfo, Dispatchable, GetDispatchInfo, Parameter, Pays,
+		UnfilteredDispatchable,
 	},
+	dispatch_context::with_context,
 	pallet_prelude::{StorageInfoTrait, ValueQuery},
 	storage::unhashed,
 	traits::{
@@ -102,6 +105,7 @@ pub mod pallet {
 	use super::*;
 	use frame_support::pallet_prelude::*;
 	use frame_system::pallet_prelude::*;
+	use sp_runtime::DispatchResult;
 
 	type BalanceOf<T> = <T as Config>::Balance;
 
@@ -227,6 +231,12 @@ pub mod pallet {
 		pub fn foo_no_post_info(_origin: OriginFor<T>) -> DispatchResult {
 			Ok(())
 		}
+
+		#[pallet::call_index(3)]
+		#[pallet::weight(1)]
+		pub fn check_for_dispatch_context(_origin: OriginFor<T>) -> DispatchResult {
+			with_context::<(), _>(|_| ()).ok_or_else(|| DispatchError::Unavailable)
+		}
 	}
 
 	#[pallet::error]
@@ -713,7 +723,7 @@ fn call_expand() {
 	assert_eq!(call_foo.get_call_name(), "foo");
 	assert_eq!(
 		pallet::Call::<Runtime>::get_call_names(),
-		&["foo", "foo_storage_layer", "foo_no_post_info"],
+		&["foo", "foo_storage_layer", "foo_no_post_info", "check_for_dispatch_context"],
 	);
 }
 
@@ -1933,3 +1943,21 @@ fn test_storage_alias() {
 		);
 	})
 }
+
+#[test]
+fn test_dispatch_context() {
+	TestExternalities::default().execute_with(|| {
+		// By default there is no context
+		assert!(with_context::<(), _>(|_| ()).is_none());
+
+		// When not using `dispatch`, there should be no dispatch context
+		assert_eq!(
+			DispatchError::Unavailable,
+			Example::check_for_dispatch_context(RuntimeOrigin::root()).unwrap_err(),
+		);
+
+		// When using `dispatch`, there should be a dispatch context
+		assert_ok!(RuntimeCall::from(pallet::Call::<Runtime>::check_for_dispatch_context {})
+			.dispatch(RuntimeOrigin::root()));
+	});
+}

substrate/frame/treasury/Cargo.toml

@@ -30,6 +30,7 @@ sp-std = { version = "5.0.0", default-features = false, path = "../../primitives
 [dev-dependencies]
 sp-core = { version = "7.0.0", path = "../../primitives/core" }
 sp-io = { version = "7.0.0", path = "../../primitives/io" }
+pallet-utility = { version = "4.0.0-dev", path = "../utility" }
 
 [features]
 default = ["std"]

substrate/frame/treasury/src/lib.rs

@@ -67,10 +67,10 @@ use codec::{Decode, Encode, MaxEncodedLen};
 use scale_info::TypeInfo;
 
 use sp_runtime::{
-	traits::{AccountIdConversion, Saturating, StaticLookup, Zero},
+	traits::{AccountIdConversion, CheckedAdd, Saturating, StaticLookup, Zero},
 	Permill, RuntimeDebug,
 };
-use sp_std::prelude::*;
+use sp_std::{collections::btree_map::BTreeMap, prelude::*};
 
 use frame_support::{
 	print,
@@ -136,7 +136,7 @@ pub struct Proposal<AccountId, Balance> {
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
-	use frame_support::pallet_prelude::*;
+	use frame_support::{dispatch_context::with_context, pallet_prelude::*};
 	use frame_system::pallet_prelude::*;
 
 	#[pallet::pallet]
@@ -339,6 +339,11 @@ pub mod pallet {
 		}
 	}
 
+	#[derive(Default)]
+	struct SpendContext<Balance> {
+		spend_in_context: BTreeMap<Balance, Balance>,
+	}
+
 	#[pallet::call]
 	impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		/// Put forward a suggestion for spending. A deposit proportional to the value
@@ -433,9 +438,29 @@ pub mod pallet {
 			beneficiary: AccountIdLookupOf<T>,
 		) -> DispatchResult {
 			let max_amount = T::SpendOrigin::ensure_origin(origin)?;
-			let beneficiary = T::Lookup::lookup(beneficiary)?;
-
 			ensure!(amount <= max_amount, Error::<T, I>::InsufficientPermission);
+
+			with_context::<SpendContext<BalanceOf<T, I>>, _>(|v| {
+				let context = v.or_default();
+
+				// We group based on `max_amount`, to dinstinguish between different kind of
+				// origins. (assumes that all origins have different `max_amount`)
+				//
+				// Worst case is that we reject some "valid" request.
+				let spend = context.spend_in_context.entry(max_amount).or_default();
+
+				// Ensure that we don't overflow nor use more than `max_amount`
+				if spend.checked_add(&amount).map(|s| s > max_amount).unwrap_or(true) {
+					Err(Error::<T, I>::InsufficientPermission)
+				} else {
+					*spend = spend.saturating_add(amount);
+
+					Ok(())
+				}
+			})
+			.unwrap_or(Ok(()))?;
+
+			let beneficiary = T::Lookup::lookup(beneficiary)?;
 			let proposal_index = Self::proposal_count();
 			Approvals::<T, I>::try_append(proposal_index)
 				.map_err(|_| Error::<T, I>::TooManyApprovals)?;

substrate/frame/treasury/src/tests.rs

@@ -22,11 +22,11 @@
 use sp_core::H256;
 use sp_runtime::{
 	testing::Header,
-	traits::{BadOrigin, BlakeTwo256, IdentityLookup},
+	traits::{BadOrigin, BlakeTwo256, Dispatchable, IdentityLookup},
 };
 
 use frame_support::{
-	assert_noop, assert_ok,
+	assert_err_ignore_postinfo, assert_noop, assert_ok,
 	pallet_prelude::GenesisBuild,
 	parameter_types,
 	traits::{ConstU32, ConstU64, OnInitialize},
@@ -38,6 +38,8 @@ use crate as treasury;
 
 type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
 type Block = frame_system::mocking::MockBlock<Test>;
+type UtilityCall = pallet_utility::Call<Test>;
+type TreasuryCall = crate::Call<Test>;
 
 frame_support::construct_runtime!(
 	pub enum Test where
@@ -48,6 +50,7 @@ frame_support::construct_runtime!(
 		System: frame_system::{Pallet, Call, Config, Storage, Event<T>},
 		Balances: pallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>},
 		Treasury: treasury::{Pallet, Call, Storage, Config, Event<T>},
+		Utility: pallet_utility,
 	}
 );
 
@@ -88,6 +91,14 @@ impl pallet_balances::Config for Test {
 	type AccountStore = System;
 	type WeightInfo = ();
 }
+
+impl pallet_utility::Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type RuntimeCall = RuntimeCall;
+	type PalletsOrigin = OriginCaller;
+	type WeightInfo = ();
+}
+
 parameter_types! {
 	pub const ProposalBond: Permill = Permill::from_percent(5);
 	pub const Burn: Permill = Permill::from_percent(50);
@@ -470,3 +481,28 @@ fn remove_already_removed_approval_fails() {
 		);
 	});
 }
+
+#[test]
+fn spending_in_batch_respects_max_total() {
+	new_test_ext().execute_with(|| {
+		// Respect the `max_total` for the given origin.
+		assert_ok!(RuntimeCall::from(UtilityCall::batch_all {
+			calls: vec![
+				RuntimeCall::from(TreasuryCall::spend { amount: 2, beneficiary: 100 }),
+				RuntimeCall::from(TreasuryCall::spend { amount: 2, beneficiary: 101 })
+			]
+		})
+		.dispatch(RuntimeOrigin::signed(10)));
+
+		assert_err_ignore_postinfo!(
+			RuntimeCall::from(UtilityCall::batch_all {
+				calls: vec![
+					RuntimeCall::from(TreasuryCall::spend { amount: 2, beneficiary: 100 }),
+					RuntimeCall::from(TreasuryCall::spend { amount: 4, beneficiary: 101 })
+				]
+			})
+			.dispatch(RuntimeOrigin::signed(10)),
+			Error::<Test, _>::InsufficientPermission
+		);
+	})
+}