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/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
+		);
+	})
+}