subspace_runtime_primitives/
extension.rs

1#[cfg(feature = "runtime-benchmarks")]
2pub mod benchmarking;
3
4use crate::utility::{MaybeNestedCall, nested_call_iter};
5use crate::weights::balance_transfer_check_extension::WeightInfo as SubstrateWeightInfo;
6use core::marker::PhantomData;
7use frame_support::RuntimeDebugNoBound;
8use frame_support::pallet_prelude::Weight;
9use frame_system::Config;
10use frame_system::pallet_prelude::{OriginFor, RuntimeCallFor};
11use pallet_balances::Call as BalancesCall;
12use parity_scale_codec::{Decode, Encode};
13use scale_info::TypeInfo;
14use scale_info::prelude::fmt;
15use sp_runtime::DispatchResult;
16use sp_runtime::traits::{
17    AsSystemOriginSigner, DispatchInfoOf, DispatchOriginOf, Dispatchable, PostDispatchInfoOf,
18    TransactionExtension, ValidateResult,
19};
20use sp_runtime::transaction_validity::{
21    InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction,
22};
23
24/// Maximum number of calls we benchmarked for.
25const MAXIMUM_NUMBER_OF_CALLS: u32 = 5_000;
26
27/// Weights for the balance transfer check extension.
28pub trait WeightInfo {
29    fn balance_transfer_check_multiple(c: u32) -> Weight;
30    fn balance_transfer_check_utility(c: u32) -> Weight;
31    fn balance_transfer_check_multisig(c: u32) -> Weight;
32}
33
34/// Trait to convert Runtime call to possible Balance call.
35pub trait MaybeBalancesCall<Runtime>
36where
37    Runtime: pallet_balances::Config,
38{
39    fn maybe_balance_call(&self) -> Option<&BalancesCall<Runtime>>;
40}
41
42/// Trait to check if the Balance transfers are enabled.
43pub trait BalanceTransferChecks {
44    fn is_balance_transferable() -> bool;
45}
46
47/// Disable balance transfers, if configured in the runtime.
48#[derive(Debug, Encode, Decode, Clone, Eq, PartialEq, TypeInfo)]
49pub struct BalanceTransferCheckExtension<Runtime>(PhantomData<Runtime>);
50
51impl<Runtime> Default for BalanceTransferCheckExtension<Runtime>
52where
53    Runtime: BalanceTransferChecks + pallet_balances::Config,
54    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
55{
56    fn default() -> Self {
57        Self(PhantomData)
58    }
59}
60
61impl<Runtime> BalanceTransferCheckExtension<Runtime>
62where
63    Runtime: BalanceTransferChecks + pallet_balances::Config,
64    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
65{
66    fn do_validate_signed(
67        call: &RuntimeCallFor<Runtime>,
68    ) -> Result<(ValidTransaction, u32), TransactionValidityError> {
69        if Runtime::is_balance_transferable() {
70            return Ok((ValidTransaction::default(), 0));
71        }
72
73        // Disable normal balance transfers.
74        let (contains_balance_call, calls) = Self::contains_balance_transfer(call);
75        if contains_balance_call {
76            Err(InvalidTransaction::Call.into())
77        } else {
78            Ok((ValidTransaction::default(), calls))
79        }
80    }
81
82    fn contains_balance_transfer(call: &RuntimeCallFor<Runtime>) -> (bool, u32) {
83        let mut calls = 0;
84        for call in nested_call_iter::<Runtime>(call) {
85            calls += 1;
86            // Any other calls might contain nested calls, so we can only return early if we find a
87            // balance transfer call.
88            if let Some(balance_call) = call.maybe_balance_call()
89                && matches!(
90                    balance_call,
91                    BalancesCall::transfer_allow_death { .. }
92                        | BalancesCall::transfer_keep_alive { .. }
93                        | BalancesCall::transfer_all { .. }
94                )
95            {
96                return (true, calls);
97            }
98        }
99
100        (false, calls)
101    }
102
103    fn get_weights(n: u32) -> Weight {
104        SubstrateWeightInfo::<Runtime>::balance_transfer_check_multisig(n)
105            .max(SubstrateWeightInfo::<Runtime>::balance_transfer_check_multiple(n))
106            .max(SubstrateWeightInfo::<Runtime>::balance_transfer_check_utility(n))
107    }
108}
109
110/// Data passed from prepare to post_dispatch.
111#[derive(RuntimeDebugNoBound)]
112pub enum Pre {
113    Refund(Weight),
114}
115
116/// Data passed from validate to prepare.
117#[derive(RuntimeDebugNoBound)]
118pub enum Val {
119    FullRefund,
120    PartialRefund(Option<u32>),
121}
122
123impl<Runtime> TransactionExtension<RuntimeCallFor<Runtime>>
124    for BalanceTransferCheckExtension<Runtime>
125where
126    Runtime: Config
127        + pallet_balances::Config
128        + scale_info::TypeInfo
129        + fmt::Debug
130        + Send
131        + Sync
132        + BalanceTransferChecks,
133    <RuntimeCallFor<Runtime> as Dispatchable>::RuntimeOrigin:
134        AsSystemOriginSigner<<Runtime as Config>::AccountId> + Clone,
135    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
136{
137    const IDENTIFIER: &'static str = "BalanceTransferCheckExtension";
138    type Implicit = ();
139    type Val = Val;
140    type Pre = Pre;
141
142    fn weight(&self, _call: &RuntimeCallFor<Runtime>) -> Weight {
143        Self::get_weights(MAXIMUM_NUMBER_OF_CALLS)
144    }
145
146    fn validate(
147        &self,
148        origin: OriginFor<Runtime>,
149        call: &RuntimeCallFor<Runtime>,
150        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
151        _len: usize,
152        _self_implicit: Self::Implicit,
153        _inherited_implication: &impl Encode,
154        _source: TransactionSource,
155    ) -> ValidateResult<Self::Val, RuntimeCallFor<Runtime>> {
156        let (validity, val) = if origin.as_system_origin_signer().is_some() {
157            let (valid, maybe_calls) =
158                Self::do_validate_signed(call).map(|(valid, calls)| (valid, Some(calls)))?;
159            (valid, Val::PartialRefund(maybe_calls))
160        } else {
161            (ValidTransaction::default(), Val::FullRefund)
162        };
163
164        Ok((validity, val, origin))
165    }
166
167    fn prepare(
168        self,
169        val: Self::Val,
170        _origin: &DispatchOriginOf<RuntimeCallFor<Runtime>>,
171        _call: &RuntimeCallFor<Runtime>,
172        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
173        _len: usize,
174    ) -> Result<Self::Pre, TransactionValidityError> {
175        let total_weight = Self::get_weights(MAXIMUM_NUMBER_OF_CALLS);
176        match val {
177            // not a signed transaction, so return full refund.
178            Val::FullRefund => Ok(Pre::Refund(total_weight)),
179
180            // signed transaction with a minimum of one read weight,
181            // so refund any extra call weight
182            Val::PartialRefund(maybe_calls) => {
183                let actual_weights = Self::get_weights(maybe_calls.unwrap_or(0));
184                Ok(Pre::Refund(total_weight.saturating_sub(actual_weights)))
185            }
186        }
187    }
188
189    fn post_dispatch_details(
190        pre: Self::Pre,
191        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
192        _post_info: &PostDispatchInfoOf<RuntimeCallFor<Runtime>>,
193        _len: usize,
194        _result: &DispatchResult,
195    ) -> Result<Weight, TransactionValidityError> {
196        let Pre::Refund(weight) = pre;
197        Ok(weight)
198    }
199}
subspace_runtime_primitives/extension.rs

subspace_runtime_primitives/
extension.rs
