Move dust collection hook to outside of account mutate (#8087)

* Move dust collection hook to outside of account mutate

* Fix dust cleanup in nested mutates.

* Fixes

* Fixes

* Apply suggestions from code review

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

* dust removal reentrancy test case integration (#8133)

* dust removal reentrancy test case integration

* Update frame/balances/src/tests_reentrancy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Update frame/balances/src/tests_reentrancy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Update frame/balances/src/tests_reentrancy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Update frame/balances/src/tests_reentrancy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Update frame/balances/src/tests_reentrancy.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* dust removal reentrancy test case integration | removed dependency on tests.rs

* dust removal reentrancy test case integration | formatt correction

* dust removal reentrancy test case integration | formatt correction

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: RK <r.raajey@gmail.com>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

substrate/frame/balances/src/lib.rs

@@ -151,6 +151,7 @@
 mod tests;
 mod tests_local;
 mod tests_composite;
+mod tests_reentrancy;
 mod benchmarking;
 pub mod weights;
 
@@ -618,6 +619,17 @@ impl Default for Releases {
 	}
 }
 
+pub struct DustCleaner<T: Config<I>, I: 'static = ()>(Option<(T::AccountId, NegativeImbalance<T, I>)>);
+
+impl<T: Config<I>, I: 'static> Drop for DustCleaner<T, I> {
+	fn drop(&mut self) {
+		if let Some((who, dust)) = self.0.take() {
+			Module::<T, I>::deposit_event(Event::DustLost(who, dust.peek()));
+			T::DustRemoval::on_unbalanced(dust);
+		}
+	}
+}
+
 impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// Get the free balance of an account.
 	pub fn free_balance(who: impl sp_std::borrow::Borrow<T::AccountId>) -> T::Balance {
@@ -646,25 +658,27 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		T::AccountStore::get(&who)
 	}
 
-	/// Places the `free` and `reserved` parts of `new` into `account`. Also does any steps needed
-	/// after mutating an account. This includes DustRemoval unbalancing, in the case than the `new`
-	/// account's total balance is non-zero but below ED.
+	/// Handles any steps needed after mutating an account.
 	///
-	/// Returns the final free balance, iff the account was previously of total balance zero, known
-	/// as its "endowment".
+	/// This includes DustRemoval unbalancing, in the case than the `new` account's total balance
+	/// is non-zero but below ED.
+	///
+	/// Returns two values:
+	/// - `Some` containing the the `new` account, iff the account has sufficient balance.
+	/// - `Some` containing the dust to be dropped, iff some dust should be dropped.
 	fn post_mutation(
-		who: &T::AccountId,
+		_who: &T::AccountId,
 		new: AccountData<T::Balance>,
-	) -> Option<AccountData<T::Balance>> {
+	) -> (Option<AccountData<T::Balance>>, Option<NegativeImbalance<T, I>>) {
 		let total = new.total();
 		if total < T::ExistentialDeposit::get() {
-			if !total.is_zero() {
-				T::DustRemoval::on_unbalanced(NegativeImbalance::new(total));
-				Self::deposit_event(Event::DustLost(who.clone(), total));
+			if total.is_zero() {
+				(None, None)
+			} else {
+				(None, Some(NegativeImbalance::new(total)))
 			}
-			None
 		} else {
-			Some(new)
+			(Some(new), None)
 		}
 	}
 
@@ -696,19 +710,46 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		who: &T::AccountId,
 		f: impl FnOnce(&mut AccountData<T::Balance>, bool) -> Result<R, E>
 	) -> Result<R, E> {
-		T::AccountStore::try_mutate_exists(who, |maybe_account| {
+		Self::try_mutate_account_with_dust(who, f)
+			.map(|(result, dust_cleaner)| {
+				drop(dust_cleaner);
+				result
+			})
+	}
+
+	/// Mutate an account to some new value, or delete it entirely with `None`. Will enforce
+	/// `ExistentialDeposit` law, annulling the account as needed. This will do nothing if the
+	/// result of `f` is an `Err`.
+	///
+	/// It returns both the result from the closure, and an optional `DustCleaner` instance which
+	/// should be dropped once it is known that all nested mutates that could affect storage items
+	/// what the dust handler touches have completed.
+	///
+	/// NOTE: Doesn't do any preparatory work for creating a new account, so should only be used
+	/// when it is known that the account already exists.
+	///
+	/// NOTE: LOW-LEVEL: This will not attempt to maintain total issuance. It is expected that
+	/// the caller will do this.
+	fn try_mutate_account_with_dust<R, E: From<StoredMapError>>(
+		who: &T::AccountId,
+		f: impl FnOnce(&mut AccountData<T::Balance>, bool) -> Result<R, E>
+	) -> Result<(R, DustCleaner<T, I>), E> {
+		let result = T::AccountStore::try_mutate_exists(who, |maybe_account| {
 			let is_new = maybe_account.is_none();
 			let mut account = maybe_account.take().unwrap_or_default();
 			f(&mut account, is_new).map(move |result| {
 				let maybe_endowed = if is_new { Some(account.free) } else { None };
-				*maybe_account = Self::post_mutation(who, account);
-				(maybe_endowed, result)
+				let maybe_account_maybe_dust = Self::post_mutation(who, account);
+				*maybe_account = maybe_account_maybe_dust.0;
+				(maybe_endowed, maybe_account_maybe_dust.1, result)
 			})
-		}).map(|(maybe_endowed, result)| {
+		});
+		result.map(|(maybe_endowed, maybe_dust, result)| {
 			if let Some(endowed) = maybe_endowed {
 				Self::deposit_event(Event::Endowed(who.clone(), endowed));
 			}
-			result
+			let dust_cleaner = DustCleaner(maybe_dust.map(|dust| (who.clone(), dust)));
+			(result, dust_cleaner)
 		})
 	}
 
@@ -772,7 +813,7 @@ mod imbalances {
 	/// funds have been created without any equal and opposite accounting.
 	#[must_use]
 	#[derive(RuntimeDebug, PartialEq, Eq)]
-	pub struct PositiveImbalance<T: Config<I>, I: 'static>(T::Balance);
+	pub struct PositiveImbalance<T: Config<I>, I: 'static = ()>(T::Balance);
 
 	impl<T: Config<I>, I: 'static> PositiveImbalance<T, I> {
 		/// Create a new positive imbalance from a balance.
@@ -785,7 +826,7 @@ mod imbalances {
 	/// funds have been destroyed without any equal and opposite accounting.
 	#[must_use]
 	#[derive(RuntimeDebug, PartialEq, Eq)]
-	pub struct NegativeImbalance<T: Config<I>, I: 'static>(T::Balance);
+	pub struct NegativeImbalance<T: Config<I>, I: 'static = ()>(T::Balance);
 
 	impl<T: Config<I>, I: 'static> NegativeImbalance<T, I> {
 		/// Create a new negative imbalance from a balance.
@@ -1001,34 +1042,40 @@ impl<T: Config<I>, I: 'static> Currency<T::AccountId> for Pallet<T, I> where
 	) -> DispatchResult {
 		if value.is_zero() || transactor == dest { return Ok(()) }
 
-		Self::try_mutate_account(dest, |to_account, _| -> DispatchResult {
-			Self::try_mutate_account(transactor, |from_account, _| -> DispatchResult {
-				from_account.free = from_account.free.checked_sub(&value)
-					.ok_or(Error::<T, I>::InsufficientBalance)?;
-
-				// NOTE: total stake being stored in the same type means that this could never overflow
-				// but better to be safe than sorry.
-				to_account.free = to_account.free.checked_add(&value).ok_or(Error::<T, I>::Overflow)?;
-
-				let ed = T::ExistentialDeposit::get();
-				ensure!(to_account.total() >= ed, Error::<T, I>::ExistentialDeposit);
-
-				Self::ensure_can_withdraw(
+		Self::try_mutate_account_with_dust(
+			dest,
+			|to_account, _| -> Result<DustCleaner<T, I>, DispatchError> {
+				Self::try_mutate_account_with_dust(
 					transactor,
-					value,
-					WithdrawReasons::TRANSFER,
-					from_account.free,
-				).map_err(|_| Error::<T, I>::LiquidityRestrictions)?;
+					|from_account, _| -> DispatchResult {
+						from_account.free = from_account.free.checked_sub(&value)
+							.ok_or(Error::<T, I>::InsufficientBalance)?;
 
-				// TODO: This is over-conservative. There may now be other providers, and this pallet
-				//   may not even be a provider.
-				let allow_death = existence_requirement == ExistenceRequirement::AllowDeath;
-				let allow_death = allow_death && !system::Pallet::<T>::is_provider_required(transactor);
-				ensure!(allow_death || from_account.free >= ed, Error::<T, I>::KeepAlive);
+						// NOTE: total stake being stored in the same type means that this could never overflow
+						// but better to be safe than sorry.
+						to_account.free = to_account.free.checked_add(&value).ok_or(Error::<T, I>::Overflow)?;
 
-				Ok(())
-			})
-		})?;
+						let ed = T::ExistentialDeposit::get();
+						ensure!(to_account.total() >= ed, Error::<T, I>::ExistentialDeposit);
+
+						Self::ensure_can_withdraw(
+							transactor,
+							value,
+							WithdrawReasons::TRANSFER,
+							from_account.free,
+						).map_err(|_| Error::<T, I>::LiquidityRestrictions)?;
+
+						// TODO: This is over-conservative. There may now be other providers, and this pallet
+						//   may not even be a provider.
+						let allow_death = existence_requirement == ExistenceRequirement::AllowDeath;
+						let allow_death = allow_death && !system::Pallet::<T>::is_provider_required(transactor);
+						ensure!(allow_death || from_account.free >= ed, Error::<T, I>::KeepAlive);
+
+						Ok(())
+					}
+				).map(|(_, maybe_dust_cleaner)| maybe_dust_cleaner)
+			}
+		)?;
 
 		// Emit transfer event.
 		Self::deposit_event(Event::Transfer(transactor.clone(), dest.clone(), value));
@@ -1322,18 +1369,28 @@ impl<T: Config<I>, I: 'static> ReservableCurrency<T::AccountId> for Pallet<T, I>
 			};
 		}
 
-		let actual = Self::try_mutate_account(beneficiary, |to_account, is_new|-> Result<Self::Balance, DispatchError> {
-			ensure!(!is_new, Error::<T, I>::DeadAccount);
-			Self::try_mutate_account(slashed, |from_account, _| -> Result<Self::Balance, DispatchError> {
-				let actual = cmp::min(from_account.reserved, value);
-				match status {
-					Status::Free => to_account.free = to_account.free.checked_add(&actual).ok_or(Error::<T, I>::Overflow)?,
-					Status::Reserved => to_account.reserved = to_account.reserved.checked_add(&actual).ok_or(Error::<T, I>::Overflow)?,
-				}
-				from_account.reserved -= actual;
-				Ok(actual)
-			})
-		})?;
+		let ((actual, _maybe_one_dust), _maybe_other_dust) = Self::try_mutate_account_with_dust(
+			beneficiary,
+			|to_account, is_new| -> Result<(Self::Balance, DustCleaner<T, I>), DispatchError> {
+				ensure!(!is_new, Error::<T, I>::DeadAccount);
+				Self::try_mutate_account_with_dust(
+					slashed,
+					|from_account, _| -> Result<Self::Balance, DispatchError> {
+						let actual = cmp::min(from_account.reserved, value);
+						match status {
+							Status::Free => to_account.free = to_account.free
+								.checked_add(&actual)
+								.ok_or(Error::<T, I>::Overflow)?,
+							Status::Reserved => to_account.reserved = to_account.reserved
+								.checked_add(&actual)
+								.ok_or(Error::<T, I>::Overflow)?,
+						}
+						from_account.reserved -= actual;
+						Ok(actual)
+					}
+				)
+			}
+		)?;
 
 		Self::deposit_event(Event::ReserveRepatriated(slashed.clone(), beneficiary.clone(), actual, status));
 		Ok(value - actual)

substrate/frame/balances/src/tests.rs

@@ -732,8 +732,8 @@ macro_rules! decl_tests {
 					assert_eq!(
 						events(),
 						[
+							Event::frame_system(system::Event::KilledAccount(1)),
 							Event::pallet_balances(crate::Event::DustLost(1, 99)),
-							Event::frame_system(system::Event::KilledAccount(1))
 						]
 					);
 				});

substrate/frame/balances/src/tests_local.rs

@@ -184,8 +184,8 @@ fn emit_events_with_no_existential_deposit_suicide_with_dust() {
 			assert_eq!(
 				events(),
 				[
+					Event::frame_system(system::Event::KilledAccount(1)),
 					Event::pallet_balances(crate::Event::DustLost(1, 1)),
-					Event::frame_system(system::Event::KilledAccount(1))
 				]
 			);
 		});

substrate/frame/balances/src/tests_reentrancy.rs

@@ -0,0 +1,310 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2018-2020 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.
+
+//! Test setup for potential reentracy and lost updates of nested mutations.
+
+#![cfg(test)]
+
+use sp_runtime::{
+	traits::IdentityLookup,
+	testing::Header,
+};
+use sp_core::H256;
+use sp_io;
+use frame_support::parameter_types;
+use frame_support::traits::StorageMapShim;
+use frame_support::weights::{IdentityFee};
+use crate::{
+	self as pallet_balances,
+	Module, Config,
+};
+use pallet_transaction_payment::CurrencyAdapter;
+
+use crate::*;
+use frame_support::{
+	assert_ok,
+	traits::{
+		Currency, ReservableCurrency,
+	}
+};
+use frame_system::RawOrigin;
+
+type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
+type Block = frame_system::mocking::MockBlock<Test>;
+
+fn last_event() -> Event {
+	system::Module::<Test>::events().pop().expect("Event expected").event
+}
+
+frame_support::construct_runtime!(
+	pub enum Test where
+		Block = Block,
+		NodeBlock = Block,
+		UncheckedExtrinsic = UncheckedExtrinsic,
+	{
+		System: frame_system::{Module, Call, Config, Storage, Event<T>},
+		Balances: pallet_balances::{Module, Call, Storage, Config<T>, Event<T>},
+	}
+);
+
+parameter_types! {
+	pub const BlockHashCount: u64 = 250;
+	pub BlockWeights: frame_system::limits::BlockWeights =
+		frame_system::limits::BlockWeights::simple_max(1024);
+	pub static ExistentialDeposit: u64 = 0;
+}
+impl frame_system::Config for Test {
+	type BaseCallFilter = ();
+	type BlockWeights = BlockWeights;
+	type BlockLength = ();
+	type DbWeight = ();
+	type Origin = Origin;
+	type Index = u64;
+	type BlockNumber = u64;
+	type Call = Call;
+	type Hash = H256;
+	type Hashing = ::sp_runtime::traits::BlakeTwo256;
+	type AccountId = u64;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type Header = Header;
+	type Event = Event;
+	type BlockHashCount = BlockHashCount;
+	type Version = ();
+	type PalletInfo = PalletInfo;
+	type AccountData = ();
+	type OnNewAccount = ();
+	type OnKilledAccount = ();
+	type SystemWeightInfo = ();
+	type SS58Prefix = ();
+}
+parameter_types! {
+	pub const TransactionByteFee: u64 = 1;
+}
+impl pallet_transaction_payment::Config for Test {
+	type OnChargeTransaction = CurrencyAdapter<Module<Test>, ()>;
+	type TransactionByteFee = TransactionByteFee;
+	type WeightToFee = IdentityFee<u64>;
+	type FeeMultiplierUpdate = ();
+}
+
+pub struct OnDustRemoval;
+impl OnUnbalanced<NegativeImbalance<Test>> for OnDustRemoval {
+	fn on_nonzero_unbalanced(amount: NegativeImbalance<Test>) {
+		let _ = Balances::resolve_into_existing(&1, amount);
+	}
+}
+parameter_types! {
+	pub const MaxLocks: u32 = 50;
+}
+impl Config for Test {
+	type Balance = u64;
+	type DustRemoval = OnDustRemoval;
+	type Event = Event;
+	type ExistentialDeposit = ExistentialDeposit;
+	type AccountStore = StorageMapShim<
+		super::Account<Test>,
+		system::Provider<Test>,
+		u64,
+		super::AccountData<u64>,
+	>;
+	type MaxLocks = MaxLocks;
+	type WeightInfo = ();
+}
+
+pub struct ExtBuilder {
+	existential_deposit: u64,
+}
+impl Default for ExtBuilder {
+	fn default() -> Self {
+		Self {
+			existential_deposit: 1,
+		}
+	}
+}
+impl ExtBuilder {
+
+	pub fn existential_deposit(mut self, existential_deposit: u64) -> Self {
+		self.existential_deposit = existential_deposit;
+		self
+	}
+
+	pub fn set_associated_consts(&self) {
+		EXISTENTIAL_DEPOSIT.with(|v| *v.borrow_mut() = self.existential_deposit);
+	}
+
+	pub fn build(self) -> sp_io::TestExternalities {
+		self.set_associated_consts();
+		let mut t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
+		pallet_balances::GenesisConfig::<Test> {
+			balances: vec![],
+		}.assimilate_storage(&mut t).unwrap();
+		let mut ext = sp_io::TestExternalities::new(t);
+		ext.execute_with(|| System::set_block_number(1));
+		ext
+	}
+}
+
+#[test]
+fn transfer_dust_removal_tst1_should_work() {
+	ExtBuilder::default()
+		.existential_deposit(100)
+		.build()
+		.execute_with(|| {
+			// Verification of reentrancy in dust removal
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 1, 1000, 0));
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 2, 500, 0));
+
+			// In this transaction, account 2 free balance
+			// drops below existential balance
+			// and dust balance is removed from account 2
+			assert_ok!(Balances::transfer(RawOrigin::Signed(2).into(), 3, 450));
+
+			// As expected dust balance is removed.
+			assert_eq!(Balances::free_balance(&2), 0);
+
+			// As expected beneficiary account 3
+			// received the transfered fund.
+			assert_eq!(Balances::free_balance(&3), 450);
+
+			// Dust balance is deposited to account 1
+			// during the process of dust removal.
+			assert_eq!(Balances::free_balance(&1), 1050);
+
+			// Verify the events
+			// Number of events expected is 8
+			assert_eq!(System::events().len(), 11);
+
+			assert!(
+				System::events().iter().any(
+					|er|
+					er.event == Event::pallet_balances(
+						crate::Event::Transfer(2, 3, 450),
+					),
+				),
+			);
+
+			assert!(
+				System::events().iter().any(
+					|er|
+					er.event == Event::pallet_balances(
+						crate::Event::DustLost(2, 50)
+					),
+				),
+			);
+		}
+	);
+}
+
+#[test]
+fn transfer_dust_removal_tst2_should_work() {
+	ExtBuilder::default()
+		.existential_deposit(100)
+		.build()
+		.execute_with(|| {
+			// Verification of reentrancy in dust removal
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 1, 1000, 0));
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 2, 500, 0));
+
+			// In this transaction, account 2 free balance
+			// drops below existential balance
+			// and dust balance is removed from account 2
+			assert_ok!(Balances::transfer(RawOrigin::Signed(2).into(), 1, 450));
+
+			// As expected dust balance is removed.
+			assert_eq!(Balances::free_balance(&2), 0);
+
+			// Dust balance is deposited to account 1
+			// during the process of dust removal.
+			assert_eq!(Balances::free_balance(&1), 1500);
+
+			// Verify the events
+			// Number of events expected is 8
+			assert_eq!(System::events().len(), 9);
+
+			assert!(
+				System::events().iter().any(
+					|er|
+					er.event == Event::pallet_balances(
+						crate::Event::Transfer(2, 1, 450),
+					),
+				),
+			);
+
+			assert!(
+				System::events().iter().any(
+					|er|
+					er.event == Event::pallet_balances(
+						crate::Event::DustLost(2, 50),
+					),
+				),
+			);
+		}
+	);
+}
+
+#[test]
+fn repatriating_reserved_balance_dust_removal_should_work() {
+	ExtBuilder::default()
+		.existential_deposit(100)
+		.build()
+		.execute_with(|| {
+			// Verification of reentrancy in dust removal
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 1, 1000, 0));
+			assert_ok!(Balances::set_balance(RawOrigin::Root.into(), 2, 500, 0));
+
+			// Reserve a value on account 2,
+			// Such that free balance is lower than
+			// Exestintial deposit.
+			assert_ok!(Balances::reserve(&2, 450));
+
+			// Transfer of reserved fund from slashed account 2 to
+			// beneficiary account 1
+			assert_ok!(Balances::repatriate_reserved(&2, &1, 450, Status::Free), 0);
+
+			// Since free balance of account 2 is lower than
+			// existential deposit, dust amount is
+			// removed from the account 2
+			assert_eq!(Balances::reserved_balance(2), 0);
+			assert_eq!(Balances::free_balance(2), 0);
+
+			// account 1 is credited with reserved amount
+			// together with dust balance during dust
+			// removal.
+			assert_eq!(Balances::reserved_balance(1), 0);
+			assert_eq!(Balances::free_balance(1), 1500);
+
+			// Verify the events
+			// Number of events expected is 10
+			assert_eq!(System::events().len(), 10);
+
+			assert!(
+				System::events().iter().any(
+					|er|
+					er.event == Event::pallet_balances(
+						crate::Event::ReserveRepatriated(2, 1, 450, Status::Free),
+					),
+				),
+			);
+
+			assert_eq!(
+				last_event(),
+				Event::pallet_balances(crate::Event::DustLost(2, 50)),
+			);
+
+		}
+	);
+}