From e7cc6b7fbf6d1bba914a19670910270e66088da9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastian=20K=C3=B6cher?= Date: Mon, 27 Feb 2023 18:49:16 +0100 Subject: [PATCH] `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 --- Cargo.lock | 2 + frame/support/Cargo.toml | 2 + .../procedural/src/pallet/expand/call.rs | 32 +-- frame/support/src/dispatch_context.rs | 232 ++++++++++++++++++ frame/support/src/lib.rs | 1 + frame/support/test/tests/pallet.rs | 32 ++- frame/treasury/Cargo.toml | 1 + frame/treasury/src/lib.rs | 35 ++- frame/treasury/src/tests.rs | 40 ++- 9 files changed, 353 insertions(+), 24 deletions(-) create mode 100644 frame/support/src/dispatch_context.rs diff --git a/Cargo.lock b/Cargo.lock index 92db97a36c943..df6843364331a 100644 --- a/Cargo.lock +++ b/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", diff --git a/frame/support/Cargo.toml b/frame/support/Cargo.toml index 4f62ae42ef78f..008d4b3f26015 100644 --- a/frame/support/Cargo.toml +++ b/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 = [] diff --git a/frame/support/procedural/src/pallet/expand/call.rs b/frame/support/procedural/src/pallet/expand/call.rs index 72367dc3957d9..3db454eb6211b 100644 --- a/frame/support/procedural/src/pallet/expand/call.rs +++ b/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."); - }, - } + } + }) } } diff --git a/frame/support/src/dispatch_context.rs b/frame/support/src/dispatch_context.rs new file mode 100644 index 0000000000000..31278ea9f8194 --- /dev/null +++ b/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::(|v| { +//! // Intitialize the value to the default value. +//! assert_eq!(0, v.or_default().0); +//! v.or_default().0 = 10; +//! }); +//! +//! with_context::(|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::(|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::(|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>); + +/// 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, +} + +impl 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(callback: impl FnOnce(&mut Value) -> R) -> Option { + DISPATCH_CONTEXT::with(|c| match c.entry(TypeId::of::()) { + Entry::Occupied(mut o) => { + let value = o.get_mut().downcast_mut::(); + + if value.is_none() { + log::error!( + "Failed to downcast value for type {} in dispatch context!", + sp_std::any::type_name::(), + ); + } + + 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(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::(|v| { + assert_eq!(0, v.or_default().0); + + v.or_default().0 = 100; + }); + + run_in_context(|| { + run_in_context(|| { + run_in_context(|| with_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); + } +} diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs index c11c902d64dd2..4e9bc32b0043b 100644 --- a/frame/support/src/lib.rs +++ b/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; diff --git a/frame/support/test/tests/pallet.rs b/frame/support/test/tests/pallet.rs index 7dd9aa0ddaf66..8ce85fa506348 100644 --- a/frame/support/test/tests/pallet.rs +++ b/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 = ::Balance; @@ -227,6 +231,12 @@ pub mod pallet { pub fn foo_no_post_info(_origin: OriginFor) -> DispatchResult { Ok(()) } + + #[pallet::call_index(3)] + #[pallet::weight(1)] + pub fn check_for_dispatch_context(_origin: OriginFor) -> 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::::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::::check_for_dispatch_context {}) + .dispatch(RuntimeOrigin::root())); + }); +} diff --git a/frame/treasury/Cargo.toml b/frame/treasury/Cargo.toml index 1fd1d7b3d4620..23fcb7944bfb7 100644 --- a/frame/treasury/Cargo.toml +++ b/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"] diff --git a/frame/treasury/src/lib.rs b/frame/treasury/src/lib.rs index 2e22063eb3ea3..04e989373c8cd 100644 --- a/frame/treasury/src/lib.rs +++ b/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 { #[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 { + spend_in_context: BTreeMap, + } + #[pallet::call] impl, I: 'static> Pallet { /// Put forward a suggestion for spending. A deposit proportional to the value @@ -433,9 +438,29 @@ pub mod pallet { beneficiary: AccountIdLookupOf, ) -> DispatchResult { let max_amount = T::SpendOrigin::ensure_origin(origin)?; - let beneficiary = T::Lookup::lookup(beneficiary)?; - ensure!(amount <= max_amount, Error::::InsufficientPermission); + + with_context::>, _>(|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::::InsufficientPermission) + } else { + *spend = spend.saturating_add(amount); + + Ok(()) + } + }) + .unwrap_or(Ok(()))?; + + let beneficiary = T::Lookup::lookup(beneficiary)?; let proposal_index = Self::proposal_count(); Approvals::::try_append(proposal_index) .map_err(|_| Error::::TooManyApprovals)?; diff --git a/frame/treasury/src/tests.rs b/frame/treasury/src/tests.rs index c9305dfd15037..24d2d01f92f8a 100644 --- a/frame/treasury/src/tests.rs +++ b/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; type Block = frame_system::mocking::MockBlock; +type UtilityCall = pallet_utility::Call; +type TreasuryCall = crate::Call; frame_support::construct_runtime!( pub enum Test where @@ -48,6 +50,7 @@ frame_support::construct_runtime!( System: frame_system::{Pallet, Call, Config, Storage, Event}, Balances: pallet_balances::{Pallet, Call, Storage, Config, Event}, Treasury: treasury::{Pallet, Call, Storage, Config, Event}, + 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::::InsufficientPermission + ); + }) +}