fix: add account check (#836)

bifrost-finance · Nov 25, 2022 · 3f0b32a · 3f0b32a
1 parent 40d0c54
commit 3f0b32a
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/pallets/slp/src/agents/parachain_staking_agent/agent.rs b/pallets/slp/src/agents/parachain_staking_agent/agent.rs
@@ -31,8 +31,9 @@ use crate::{
 	pallet::Error,
 	primitives::{Ledger, OneToManyDelegatorStatus, QueryId, BNC},
 	traits::StakingAgent,
-	AccountIdOf, BalanceOf, Config, CurrencyDelays, DelegatorLedgers, LedgerUpdateEntry,
-	MinimumsAndMaximums, Pallet, TimeUnit, Validators,
+	AccountIdOf, BalanceOf, Config, CurrencyDelays, DelegatorLedgers,
+	DelegatorsMultilocation2Index, LedgerUpdateEntry, MinimumsAndMaximums, Pallet, TimeUnit,
+	Validators,
 };
 use node_primitives::{CurrencyId, TokenSymbol, VtokenMintingOperator};
 use orml_traits::MultiCurrency;
@@ -910,8 +911,20 @@ impl<T: Config>
 		amount: BalanceOf<T>,
 		currency_id: CurrencyId,
 	) -> Result<(), Error<T>> {
+		// Ensure amount is greater than zero.
+		ensure!(!amount.is_zero(), Error::<T>::AmountZero);
+
+		// Check if from is one of our delegators. If not, return error.
+		DelegatorsMultilocation2Index::<T>::get(currency_id, from)
+			.ok_or(Error::<T>::DelegatorNotExist)?;
+
+		// Make sure the receiving account is the Exit_account from vtoken-minting module.
 		let from_account = Pallet::<T>::multilocation_to_account(from)?;
 		let to_account = Pallet::<T>::multilocation_to_account(to)?;
+
+		let (_, exit_account) = T::VtokenMinting::get_entrance_and_exit_accounts();
+		ensure!(to_account == exit_account, Error::<T>::InvalidAccount);
+
 		T::MultiCurrency::transfer(currency_id, &from_account, &to_account, amount)
 			.map_err(|_| Error::<T>::Unexpected)?;
 
@@ -927,8 +940,17 @@ impl<T: Config>
 		amount: BalanceOf<T>,
 		currency_id: CurrencyId,
 	) -> Result<(), Error<T>> {
+		// Make sure receiving account is one of the KSM delegators.
+		ensure!(
+			DelegatorsMultilocation2Index::<T>::contains_key(currency_id, to),
+			Error::<T>::DelegatorNotExist
+		);
+
+		// Make sure from account is the entrance account of vtoken-minting module.
 		let from_account = Pallet::<T>::multilocation_to_account(from)?;
 		let to_account = Pallet::<T>::multilocation_to_account(to)?;
+		let (entrance_account, _) = T::VtokenMinting::get_entrance_and_exit_accounts();
+		ensure!(from_account == entrance_account, Error::<T>::InvalidAccount);
 		T::MultiCurrency::transfer(currency_id, &from_account, &to_account, amount)
 			.map_err(|_| Error::<T>::Unexpected)?;
 

diff --git a/pallets/slp/src/tests/parachain_staking_tests.rs b/pallets/slp/src/tests/parachain_staking_tests.rs
@@ -690,6 +690,9 @@ fn parachain_staking_transfer_back_works() {
 			interior: X1(Junction::AccountId32 { network: Any, id: exit_account_id_32 }),
 		};
 
+		DelegatorsIndex2Multilocation::<Runtime>::insert(BNC, 0, subaccount_0_location.clone());
+		DelegatorsMultilocation2Index::<Runtime>::insert(BNC, subaccount_0_location.clone(), 0);
+
 		assert_ok!(Currencies::update_balance(Origin::root(), CHARLIE, BNC, 1000_000_000_000_000,));
 
 		assert_ok!(Slp::transfer_back(
@@ -712,6 +715,7 @@ fn parachain_staking_transfer_to_works() {
 	ExtBuilder::default().build().execute_with(|| {
 		// environment setup
 		parachain_staking_setup();
+
 		let entrance_account_id_32: [u8; 32] =
 			hex_literal::hex!["6d6f646c62662f76746b696e0000000000000000000000000000000000000000"]
 				.into();
@@ -721,6 +725,9 @@ fn parachain_staking_transfer_to_works() {
 			interior: X1(Junction::AccountId32 { network: Any, id: entrance_account_id_32 }),
 		};
 
+		DelegatorsIndex2Multilocation::<Runtime>::insert(BNC, 0, subaccount_0_location.clone());
+		DelegatorsMultilocation2Index::<Runtime>::insert(BNC, subaccount_0_location.clone(), 0);
+
 		assert_ok!(Currencies::update_balance(
 			Origin::root(),
 			entrance_account_id_32.into(),