feat: balance check as signed extension

.codecov.yml

@@ -12,3 +12,4 @@ coverage:
     patch: 
       default:
         threshold: 5%
+        informational: true

pallets/transaction-multi-payment/src/lib.rs

@@ -54,6 +54,11 @@ use primitives::asset::AssetPair;
 use primitives::traits::AMM;
 use primitives::{Amount, AssetId, Balance, CORE_ASSET_ID};
 
+use codec::{Decode, Encode};
+use frame_support::sp_runtime::traits::SignedExtension;
+use frame_support::sp_runtime::transaction_validity::{TransactionValidity, ValidTransaction};
+use frame_support::traits::IsSubType;
+
 type NegativeImbalanceOf<C, T> = <C as Currency<<T as frame_system::Config>::AccountId>>::NegativeImbalance;
 // Re-export pallet items so that they can be accessed from the crate namespace.
 pub use pallet::*;
@@ -391,6 +396,13 @@ impl<T: Config> Pallet<T> {
 		let capped_weight: Weight = weight.min(T::BlockWeights::get().max_block);
 		<T as Config>::WeightToFee::calc(&capped_weight)
 	}
+
+	fn check_balance(account: &T::AccountId, currency: AssetId) -> Result<(), Error<T>> {
+		if T::MultiCurrency::free_balance(currency, account) == Balance::zero() {
+			return Err(Error::<T>::ZeroBalance.into());
+		};
+		Ok(())
+	}
 }
 
 use crate::traits::PaymentSwapResult;
@@ -512,3 +524,51 @@ where
 		Ok(())
 	}
 }
+
+/// Signed extension that checks for the `set_currency` call and in that case, it checks the account balance
+#[derive(Encode, Decode, Clone, Eq, PartialEq)]
+pub struct CurrencyBalanceCheck<T: Config + Send + Sync>(PhantomData<T>);
+
+impl<T: Config + Send + Sync> sp_std::fmt::Debug for CurrencyBalanceCheck<T> {
+	fn fmt(&self, f: &mut sp_std::fmt::Formatter) -> sp_std::fmt::Result {
+		write!(f, "CurrencyBalanceCheck")
+	}
+}
+
+impl<T: Config + Send + Sync> SignedExtension for CurrencyBalanceCheck<T>
+where
+	<T as frame_system::Config>::Call: IsSubType<Call<T>>,
+{
+	const IDENTIFIER: &'static str = "CurrencyBalanceCheck";
+	type AccountId = T::AccountId;
+	type Call = <T as frame_system::Config>::Call;
+	type AdditionalSigned = ();
+	type Pre = ();
+
+	fn additional_signed(&self) -> sp_std::result::Result<(), TransactionValidityError> {
+		Ok(())
+	}
+
+	fn validate(
+		&self,
+		who: &Self::AccountId,
+		call: &Self::Call,
+		_info: &DispatchInfoOf<Self::Call>,
+		_len: usize,
+	) -> TransactionValidity {
+		match call.is_sub_type() {
+			Some(Call::set_currency(currency)) => match Pallet::<T>::check_balance(who, *currency) {
+				Ok(_) => Ok(ValidTransaction::default()),
+				Err(error) => InvalidTransaction::Custom(error.as_u8()).into(),
+			},
+			_ => Ok(Default::default()),
+		}
+	}
+}
+
+impl<T: Config + Send + Sync> CurrencyBalanceCheck<T> {
+	#[cfg_attr(feature = "cargo-clippy", allow(clippy::new_without_default))]
+	pub fn new() -> Self {
+		Self(sp_std::marker::PhantomData)
+	}
+}

pallets/transaction-multi-payment/src/tests.rs

@@ -20,10 +20,13 @@ use frame_support::{assert_noop, assert_ok};
 use pallet_transaction_payment::ChargeTransactionPayment;
 use sp_runtime::traits::SignedExtension;
 
+use crate::CurrencyBalanceCheck;
+use frame_support::sp_runtime::transaction_validity::{InvalidTransaction, ValidTransaction};
 use frame_support::weights::DispatchInfo;
 use orml_traits::MultiCurrency;
 use pallet_balances::Call as BalancesCall;
 use primitives::Price;
+use sp_std::marker::PhantomData;
 
 const CALL: &<Test as frame_system::Config>::Call = &Call::Balances(BalancesCall::transfer(2, 69));
 
@@ -432,3 +435,43 @@ fn fee_payment_non_native_insufficient_balance_with_no_pool() {
 			assert_eq!(Tokens::free_balance(SUPPORTED_CURRENCY_WITH_BALANCE, &CHARLIE), 457);
 		});
 }
+
+#[test]
+fn check_balance_extension_works() {
+	const CHARLIE: AccountId = 5;
+
+	ExtBuilder::default()
+		.account_tokens(CHARLIE, SUPPORTED_CURRENCY_WITH_BALANCE, 1000)
+		.build()
+		.execute_with(|| {
+			let call = <crate::Call<Test>>::set_currency(SUPPORTED_CURRENCY_WITH_BALANCE).into();
+			let info = DispatchInfo::default();
+
+			assert_eq!(
+				CurrencyBalanceCheck::<Test>(PhantomData).validate(&CHARLIE, &call, &info, 150),
+				Ok(ValidTransaction::default())
+			);
+
+			let call = <crate::Call<Test>>::add_currency(SUPPORTED_CURRENCY_WITH_BALANCE, Price::from(1)).into();
+
+			assert_eq!(
+				CurrencyBalanceCheck::<Test>(PhantomData).validate(&CHARLIE, &call, &info, 150),
+				Ok(ValidTransaction::default())
+			);
+		});
+}
+
+#[test]
+fn check_balance_extension_fails() {
+	const NOT_CHARLIE: AccountId = 6;
+
+	ExtBuilder::default().build().execute_with(|| {
+		let call = <crate::Call<Test>>::set_currency(SUPPORTED_CURRENCY_WITH_BALANCE).into();
+		let info = DispatchInfo::default();
+
+		assert_eq!(
+			CurrencyBalanceCheck::<Test>(PhantomData).validate(&NOT_CHARLIE, &call, &info, 150),
+			InvalidTransaction::Custom(Error::<Test>::ZeroBalance.as_u8()).into()
+		);
+	});
+}

runtime/src/lib.rs

@@ -884,6 +884,7 @@ pub type SignedExtra = (
 	frame_system::CheckNonce<Runtime>,
 	frame_system::CheckWeight<Runtime>,
 	pallet_transaction_payment::ChargeTransactionPayment<Runtime>,
+	pallet_transaction_multi_payment::CurrencyBalanceCheck<Runtime>,
 );
 /// Unchecked extrinsic type as expected by this runtime.
 pub type UncheckedExtrinsic = generic::UncheckedExtrinsic<Address, Call, Signature, SignedExtra>;