Move dust collection hook to outside of account mutate (paritytech#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 (paritytech#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>

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.
+	///
+	/// This includes DustRemoval unbalancing, in the case than the `new` account's total balance
+	/// is non-zero but below ED.
 	///
-	/// Returns the final free balance, iff the account was previously of total balance zero, known
-	/// as its "endowment".
+	/// 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)?;
-
-				// 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(())
-			})
-		})?;
+					|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(
+							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)

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))
 						]
 					);
 				});

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))
 				]
 			);
 		});