Benchmarks for balance transfer check extension

crates/pallet-domains/src/extensions.rs

@@ -3,6 +3,7 @@
 use crate::pallet::Call as DomainsCall;
 use crate::weights::WeightInfo;
 use crate::{Config, FraudProofFor, OpaqueBundleOf, Origin, Pallet as Domains, SingletonReceiptOf};
+use frame_support::ensure;
 use frame_support::pallet_prelude::{PhantomData, TypeInfo};
 use frame_support::weights::Weight;
 use frame_system::pallet_prelude::RuntimeCallFor;
@@ -28,6 +29,12 @@ where
     fn maybe_domains_call(&self) -> Option<&DomainsCall<Runtime>>;
 }
 
+/// Trait to check if the Domains are enabled on Consensus.
+pub trait DomainsCheck {
+    /// Check if the domains are enabled on Runtime.
+    fn is_domains_enabled() -> bool;
+}
+
 /// Extensions for pallet-domains unsigned extrinsics.
 #[derive(Encode, Decode, Clone, Eq, PartialEq, TypeInfo)]
 pub struct DomainsExtension<Runtime>(PhantomData<Runtime>);
@@ -168,7 +175,7 @@ where
 
 impl<Runtime> TransactionExtension<RuntimeCallFor<Runtime>> for DomainsExtension<Runtime>
 where
-    Runtime: Config + scale_info::TypeInfo + fmt::Debug + Send + Sync,
+    Runtime: Config + scale_info::TypeInfo + fmt::Debug + Send + Sync + DomainsCheck,
     <RuntimeCallFor<Runtime> as Dispatchable>::RuntimeOrigin:
         AsSystemOriginSigner<<Runtime as frame_system::Config>::AccountId> + From<Origin> + Clone,
     RuntimeCallFor<Runtime>: MaybeDomainsCall<Runtime>,
@@ -180,19 +187,24 @@ where
 
     fn weight(&self, call: &Runtime::RuntimeCall) -> Weight {
         // This extension only apply to the following 3 calls thus only return weight for them.
-        match call.maybe_domains_call() {
+        let maybe_weight = match call.maybe_domains_call() {
             Some(DomainsCall::submit_bundle { .. }) => {
-                <Runtime as Config>::WeightInfo::validate_submit_bundle()
+                Some(<Runtime as Config>::WeightInfo::validate_submit_bundle())
             }
-            Some(DomainsCall::submit_fraud_proof { fraud_proof }) => {
+            Some(DomainsCall::submit_fraud_proof { fraud_proof }) => Some(
                 <Runtime as Config>::WeightInfo::fraud_proof_pre_check()
-                    .saturating_add(fraud_proof_verification_weights::<_, _, _, _>(fraud_proof))
-            }
+                    .saturating_add(fraud_proof_verification_weights::<_, _, _, _>(fraud_proof)),
+            ),
             Some(DomainsCall::submit_receipt { .. }) => {
-                <Runtime as Config>::WeightInfo::validate_singleton_receipt()
+                Some(<Runtime as Config>::WeightInfo::validate_singleton_receipt())
             }
-            _ => Weight::zero(),
-        }
+            _ => None,
+        };
+
+        // There is one additional runtime read to check if the domains are enabled
+        maybe_weight
+            .and_then(|weight| weight.checked_add(&Runtime::DbWeight::get().reads(1)))
+            .unwrap_or(Weight::zero())
     }
 
     fn validate(
@@ -215,6 +227,9 @@ where
             None => return Ok((ValidTransaction::default(), (), origin)),
         };
 
+        // ensure domains are enabled
+        ensure!(Runtime::is_domains_enabled(), InvalidTransaction::Call);
+
         let validity = match domains_call {
             DomainsCall::submit_bundle { opaque_bundle } => {
                 Self::do_validate_submit_bundle(opaque_bundle, source)?

crates/subspace-malicious-operator/src/malicious_bundle_producer.rs

@@ -28,7 +28,8 @@ use sp_transaction_pool::runtime_api::TaggedTransactionQueue;
 use std::error::Error;
 use std::sync::Arc;
 use subspace_core_primitives::pot::PotOutput;
-use subspace_runtime::{DisablePallets, Runtime, RuntimeCall, SignedExtra, UncheckedExtrinsic};
+use subspace_runtime::{Runtime, RuntimeCall, SignedExtra, UncheckedExtrinsic};
+use subspace_runtime_primitives::extension::BalanceTransferCheckExtension;
 use subspace_runtime_primitives::opaque::Block as CBlock;
 use subspace_runtime_primitives::{AccountId, Balance, BlockHashFor, HeaderFor, Nonce};
 
@@ -415,7 +416,7 @@ fn get_singed_extra(best_number: u64, immortal: bool, nonce: Nonce) -> SignedExt
         frame_system::CheckNonce::<Runtime>::from(nonce.into()),
         frame_system::CheckWeight::<Runtime>::new(),
         pallet_transaction_payment::ChargeTransactionPayment::<Runtime>::from(0u128),
-        DisablePallets,
+        BalanceTransferCheckExtension::<Runtime>::default(),
         pallet_subspace::extensions::SubspaceExtension::<Runtime>::new(),
         pallet_domains::extensions::DomainsExtension::<Runtime>::new(),
         pallet_messenger::extensions::MessengerExtension::<Runtime>::new(),

crates/subspace-runtime-primitives/Cargo.toml

@@ -17,8 +17,10 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
 parity-scale-codec = { workspace = true, features = ["derive"] }
+frame-benchmarking = { workspace = true, optional = true }
 frame-support.workspace = true
 frame-system.workspace = true
+pallet-balances.workspace = true
 pallet-multisig.workspace = true
 pallet-transaction-payment.workspace = true
 pallet-utility.workspace = true
@@ -35,6 +37,7 @@ std = [
     "parity-scale-codec/std",
     "frame-support/std",
     "frame-system/std",
+    "pallet-balances/std",
     "pallet-multisig/std",
     "pallet-transaction-payment/std",
     "pallet-utility/std",
@@ -47,3 +50,10 @@ std = [
 testing = [
     "sp-io"
 ]
+
+runtime-benchmarks = [
+    "frame-benchmarking",
+    "frame-benchmarking/runtime-benchmarks",
+    "frame-support/runtime-benchmarks",
+    "frame-system/runtime-benchmarks",
+]

crates/subspace-runtime-primitives/src/extension.rs

@@ -0,0 +1,187 @@
+#[cfg(feature = "runtime-benchmarks")]
+pub mod benchmarking;
+pub mod weights;
+
+use crate::extension::weights::WeightInfo as SubstrateWeightInfo;
+use crate::utility::{nested_call_iter, MaybeNestedCall};
+use core::marker::PhantomData;
+use frame_support::pallet_prelude::Weight;
+use frame_support::RuntimeDebugNoBound;
+use frame_system::pallet_prelude::{OriginFor, RuntimeCallFor};
+use frame_system::Config;
+use pallet_balances::Call as BalancesCall;
+use parity_scale_codec::{Decode, Encode};
+use scale_info::prelude::fmt;
+use scale_info::TypeInfo;
+use sp_runtime::traits::{
+    AsSystemOriginSigner, DispatchInfoOf, DispatchOriginOf, Dispatchable, PostDispatchInfoOf,
+    TransactionExtension, ValidateResult,
+};
+use sp_runtime::transaction_validity::{
+    InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction,
+};
+use sp_runtime::DispatchResult;
+
+/// Maximum number of calls we benchmarked for.
+const MAXIMUM_NUMBER_OF_CALLS: u32 = 1000;
+
+/// Weights for the balance transfer check extension.
+pub trait WeightInfo {
+    fn balance_transfer_check_mixed(c: u32) -> Weight;
+    fn balance_transfer_check_utility(c: u32) -> Weight;
+    fn balance_transfer_check_multisig(c: u32) -> Weight;
+}
+
+/// Trait to convert Runtime call to possible Balance call.
+pub trait MaybeBalancesCall<Runtime>
+where
+    Runtime: pallet_balances::Config,
+{
+    fn maybe_balance_call(&self) -> Option<&BalancesCall<Runtime>>;
+}
+
+/// Trait to check if the Balance transfers are enabled.
+pub trait BalanceTransferChecks {
+    fn is_balance_transferable() -> bool;
+}
+
+/// Disable balance transfers, if configured in the runtime.
+#[derive(Debug, Encode, Decode, Clone, Eq, PartialEq, TypeInfo)]
+pub struct BalanceTransferCheckExtension<Runtime>(PhantomData<Runtime>);
+
+impl<Runtime> Default for BalanceTransferCheckExtension<Runtime>
+where
+    Runtime: BalanceTransferChecks + pallet_balances::Config,
+    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
+{
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<Runtime> BalanceTransferCheckExtension<Runtime>
+where
+    Runtime: BalanceTransferChecks + pallet_balances::Config,
+    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
+{
+    fn do_validate_signed(
+        call: &RuntimeCallFor<Runtime>,
+    ) -> Result<(ValidTransaction, u32), TransactionValidityError> {
+        if Runtime::is_balance_transferable() {
+            return Ok((ValidTransaction::default(), 0));
+        }
+
+        // Disable normal balance transfers.
+        let (contains_balance_call, calls) = Self::contains_balance_transfer(call);
+        if contains_balance_call {
+            Err(InvalidTransaction::Call.into())
+        } else {
+            Ok((ValidTransaction::default(), calls))
+        }
+    }
+
+    fn contains_balance_transfer(call: &RuntimeCallFor<Runtime>) -> (bool, u32) {
+        let mut calls = 0;
+        for call in nested_call_iter::<Runtime>(call) {
+            calls += 1;
+            // Any other calls might contain nested calls, so we can only return early if we find a
+            // balance transfer call.
+            if let Some(balance_call) = call.maybe_balance_call()
+                && matches!(
+                    balance_call,
+                    BalancesCall::transfer_allow_death { .. }
+                        | BalancesCall::transfer_keep_alive { .. }
+                        | BalancesCall::transfer_all { .. }
+                )
+            {
+                return (true, calls);
+            }
+        }
+
+        (false, calls)
+    }
+
+    fn get_weights(n: u32) -> Weight {
+        SubstrateWeightInfo::<Runtime>::balance_transfer_check_multisig(n)
+            .max(SubstrateWeightInfo::<Runtime>::balance_transfer_check_mixed(n))
+            .max(SubstrateWeightInfo::<Runtime>::balance_transfer_check_utility(n))
+    }
+}
+
+/// Data passed from prepare to post_dispatch.
+#[derive(RuntimeDebugNoBound)]
+pub enum Pre {
+    Refund(Weight),
+}
+
+impl<Runtime> TransactionExtension<RuntimeCallFor<Runtime>>
+    for BalanceTransferCheckExtension<Runtime>
+where
+    Runtime: Config
+        + pallet_balances::Config
+        + scale_info::TypeInfo
+        + fmt::Debug
+        + Send
+        + Sync
+        + BalanceTransferChecks,
+    <RuntimeCallFor<Runtime> as Dispatchable>::RuntimeOrigin:
+        AsSystemOriginSigner<<Runtime as Config>::AccountId> + Clone,
+    RuntimeCallFor<Runtime>: MaybeBalancesCall<Runtime> + MaybeNestedCall<Runtime>,
+{
+    const IDENTIFIER: &'static str = "BalanceTransferCheckExtension";
+    type Implicit = ();
+    type Val = Option<u32>;
+    type Pre = Pre;
+
+    fn weight(&self, _call: &RuntimeCallFor<Runtime>) -> Weight {
+        Self::get_weights(MAXIMUM_NUMBER_OF_CALLS)
+    }
+
+    fn validate(
+        &self,
+        origin: OriginFor<Runtime>,
+        call: &RuntimeCallFor<Runtime>,
+        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _len: usize,
+        _self_implicit: Self::Implicit,
+        _inherited_implication: &impl Encode,
+        _source: TransactionSource,
+    ) -> ValidateResult<Self::Val, RuntimeCallFor<Runtime>> {
+        let (validity, maybe_calls) = if origin.as_system_origin_signer().is_some() {
+            Self::do_validate_signed(call).map(|(valid, calls)| (valid, Some(calls)))?
+        } else {
+            (ValidTransaction::default(), None)
+        };
+
+        Ok((validity, maybe_calls, origin))
+    }
+
+    fn prepare(
+        self,
+        val: Self::Val,
+        _origin: &DispatchOriginOf<RuntimeCallFor<Runtime>>,
+        _call: &RuntimeCallFor<Runtime>,
+        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _len: usize,
+    ) -> Result<Self::Pre, TransactionValidityError> {
+        let assigned_weight = Self::get_weights(MAXIMUM_NUMBER_OF_CALLS);
+        match val {
+            None => Ok(Pre::Refund(assigned_weight)),
+            Some(calls) => {
+                let actual_weights = Self::get_weights(calls);
+                Ok(Pre::Refund(assigned_weight.saturating_sub(actual_weights)))
+            }
+        }
+    }
+
+    fn post_dispatch_details(
+        pre: Self::Pre,
+        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _post_info: &PostDispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _len: usize,
+        _result: &DispatchResult,
+    ) -> Result<Weight, TransactionValidityError> {
+        let Pre::Refund(weight) = pre;
+        Ok(weight)
+    }
+}