Skip to content

Commit

Permalink
pallet-treasury: Ensure we respect max_amount for spend across ba…
Browse files Browse the repository at this point in the history
…tch calls (paritytech#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: paritytech#13167

* Import `Box` for wasm

* FMT
  • Loading branch information
bkchr authored and nathanwhit committed Jul 19, 2023
1 parent 6508e13 commit e7cc6b7
Show file tree
Hide file tree
Showing 9 changed files with 353 additions and 24 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions frame/support/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -67,6 +68,7 @@ std = [
"sp-weights/std",
"frame-support-procedural/std",
"log/std",
"environmental/std",
]
runtime-benchmarks = []
try-runtime = []
Expand Down
32 changes: 17 additions & 15 deletions frame/support/procedural/src/pallet/expand/call.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
},
}
}
})
}
}

Expand Down
232 changes: 232 additions & 0 deletions frame/support/src/dispatch_context.rs
Original file line number Diff line number Diff line change
@@ -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);
}
}
1 change: 1 addition & 0 deletions frame/support/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
32 changes: 30 additions & 2 deletions frame/support/test/tests/pallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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"],
);
}

Expand Down Expand Up @@ -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()));
});
}
1 change: 1 addition & 0 deletions frame/treasury/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
Loading

0 comments on commit e7cc6b7

Please sign in to comment.