fix(tokens): PostDeposit hook (#839)

* fix(tokens): PostDeposit hook

* test(tokens): posthook correctness

* fix(tokens): use `assert` instead of `ensure` in mock posthooks

Author: daniel-savu

tokens/src/lib.rs

@@ -1054,19 +1054,19 @@ impl<T: Config> Pallet<T> {
 				TotalIssuance::<T>::mutate(currency_id, |v| *v = new_total_issuance);
 			}
 			account.free = account.free.defensive_saturating_add(amount);
-
-			<T::CurrencyHooks as MutationHooks<T::AccountId, T::CurrencyId, T::Balance>>::PostDeposit::on_deposit(
-				currency_id,
-				who,
-				amount,
-			)?;
-			Self::deposit_event(Event::Deposited {
-				currency_id,
-				who: who.clone(),
-				amount,
-			});
 			Ok(())
-		})
+		})?;
+		<T::CurrencyHooks as MutationHooks<T::AccountId, T::CurrencyId, T::Balance>>::PostDeposit::on_deposit(
+			currency_id,
+			who,
+			amount,
+		)?;
+		Self::deposit_event(Event::Deposited {
+			currency_id,
+			who: who.clone(),
+			amount,
+		});
+		Ok(())
 	}
 }
 

tokens/src/mock.rs

@@ -327,8 +327,14 @@ impl<T: Config> PreDeposit<T> {
 
 pub struct PostDeposit<T>(marker::PhantomData<T>);
 impl<T: Config> OnDeposit<T::AccountId, T::CurrencyId, T::Balance> for PostDeposit<T> {
-	fn on_deposit(_currency_id: T::CurrencyId, _account_id: &T::AccountId, _amount: T::Balance) -> DispatchResult {
+	fn on_deposit(currency_id: T::CurrencyId, account_id: &T::AccountId, amount: T::Balance) -> DispatchResult {
 		ON_DEPOSIT_POSTHOOK_CALLS.with(|cell| *cell.borrow_mut() += 1);
+		let account_balance: AccountData<T::Balance> =
+			tokens::Pallet::<T>::accounts::<T::AccountId, T::CurrencyId>(account_id.clone(), currency_id);
+		assert!(
+			account_balance.free.ge(&amount),
+			"Posthook must run after the account balance is updated."
+		);
 		Ok(())
 	}
 }
@@ -359,12 +365,18 @@ impl<T: Config> PreTransfer<T> {
 pub struct PostTransfer<T>(marker::PhantomData<T>);
 impl<T: Config> OnTransfer<T::AccountId, T::CurrencyId, T::Balance> for PostTransfer<T> {
 	fn on_transfer(
-		_currency_id: T::CurrencyId,
+		currency_id: T::CurrencyId,
 		_from: &T::AccountId,
-		_to: &T::AccountId,
-		_amount: T::Balance,
+		to: &T::AccountId,
+		amount: T::Balance,
 	) -> DispatchResult {
 		ON_TRANSFER_POSTHOOK_CALLS.with(|cell| *cell.borrow_mut() += 1);
+		let account_balance: AccountData<T::Balance> =
+			tokens::Pallet::<T>::accounts::<T::AccountId, T::CurrencyId>(to.clone(), currency_id);
+		assert!(
+			account_balance.free.ge(&amount),
+			"Posthook must run after the account balance is updated."
+		);
 		Ok(())
 	}
 }

tokens/src/tests.rs

@@ -1200,6 +1200,17 @@ fn deposit_hooks_work() {
 	});
 }
 
+#[test]
+fn post_deposit_can_use_new_balance() {
+	ExtBuilder::default().build().execute_with(|| {
+		let initial_balance = Tokens::free_balance(DOT, &CHARLIE);
+		// The following will fail unless Charlie's new balance can be used by the hook,
+		// because `initial_balance + 100` is higher than Charlie's initial balance.
+		// If this fails, the posthook is called too soon.
+		assert_ok!(Tokens::do_deposit(DOT, &CHARLIE, initial_balance + 100, false, true),);
+	});
+}
+
 #[test]
 fn transfer_hooks_work() {
 	ExtBuilder::default()
@@ -1238,3 +1249,23 @@ fn transfer_hooks_work() {
 			assert_eq!(PostTransfer::<Runtime>::calls(), initial_posthook_calls + 1);
 		});
 }
+
+#[test]
+fn post_transfer_can_use_new_balance() {
+	ExtBuilder::default()
+		.balances(vec![(ALICE, DOT, 100)])
+		.build()
+		.execute_with(|| {
+			let initial_balance = Tokens::free_balance(DOT, &CHARLIE);
+			// The following will fail unless Charlie's new balance can be used by the hook,
+			// because `initial_balance + 100` is higher than Charlie's initial balance.
+			// If this fails, the posthook is called too soon.
+			assert_ok!(Tokens::do_transfer(
+				DOT,
+				&ALICE,
+				&CHARLIE,
+				initial_balance + 100,
+				ExistenceRequirement::AllowDeath
+			));
+		});
+}