feat(tokens): OnDeposit, OnTransfer, OnSlash hooks (#815)

* feat(tokens): OnDeposit, OnTransfer, OnSlash hooks

* test(tokens): mutation hooks (OnDeposit, OnSlash, OnTransfer)

Author: daniel-savu

asset-registry/src/mock/para.rs

@@ -93,6 +93,9 @@ impl orml_tokens::Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = ();
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type ReserveIdentifier = [u8; 8];
 	type MaxReserves = ();
 	type MaxLocks = ConstU32<50>;

currencies/src/mock.rs

@@ -81,6 +81,9 @@ impl orml_tokens::Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = orml_tokens::TransferDust<Runtime, DustAccount>;
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type MaxLocks = ConstU32<100_000>;
 	type MaxReserves = ConstU32<100_000>;
 	type ReserveIdentifier = ReserveIdentifier;

payments/src/mock.rs

@@ -99,6 +99,9 @@ impl orml_tokens::Config for Test {
 	type Event = Event;
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = ();
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type WeightInfo = ();
 	type MaxLocks = MaxLocks;
 	type DustRemovalWhitelist = MockDustRemovalWhitelist;

tokens/src/lib.rs

@@ -66,7 +66,7 @@ use sp_std::{cmp, convert::Infallible, marker, prelude::*, vec::Vec};
 
 use orml_traits::{
 	arithmetic::{self, Signed},
-	currency::TransferAll,
+	currency::{OnDeposit, OnSlash, OnTransfer, TransferAll},
 	BalanceStatus, GetByKey, Happened, LockIdentifier, MultiCurrency, MultiCurrencyExtended, MultiLockableCurrency,
 	MultiReservableCurrency, NamedMultiReservableCurrency, OnDust,
 };
@@ -173,6 +173,8 @@ pub use module::*;
 
 #[frame_support::pallet]
 pub mod module {
+	use orml_traits::currency::{OnDeposit, OnSlash, OnTransfer};
+
 	use super::*;
 
 	#[pallet::config]
@@ -216,6 +218,15 @@ pub mod module {
 		/// Handler to burn or transfer account's dust
 		type OnDust: OnDust<Self::AccountId, Self::CurrencyId, Self::Balance>;
 
+		/// Hook to run before slashing an account.
+		type OnSlash: OnSlash<Self::AccountId, Self::CurrencyId, Self::Balance>;
+
+		/// Hook to run before depositing into an account.
+		type OnDeposit: OnDeposit<Self::AccountId, Self::CurrencyId, Self::Balance>;
+
+		/// Hook to run before transferring from an account to another.
+		type OnTransfer: OnTransfer<Self::AccountId, Self::CurrencyId, Self::Balance>;
+
 		/// Handler for when an account was created
 		type OnNewTokenAccount: Happened<(Self::AccountId, Self::CurrencyId)>;
 
@@ -894,6 +905,7 @@ impl<T: Config> Pallet<T> {
 			return Ok(());
 		}
 
+		T::OnTransfer::on_transfer(currency_id, from, to, amount)?;
 		Self::try_mutate_account(to, currency_id, |to_account, _existed| -> DispatchResult {
 			Self::try_mutate_account(from, currency_id, |from_account, _existed| -> DispatchResult {
 				from_account.free = from_account
@@ -1019,6 +1031,7 @@ impl<T: Config> Pallet<T> {
 			return Ok(());
 		}
 
+		T::OnDeposit::on_deposit(currency_id, who, amount)?;
 		Self::try_mutate_account(who, currency_id, |account, existed| -> DispatchResult {
 			if require_existed {
 				ensure!(existed, Error::<T>::DeadAccount);
@@ -1114,6 +1127,7 @@ impl<T: Config> MultiCurrency<T::AccountId> for Pallet<T> {
 			return amount;
 		}
 
+		T::OnSlash::on_slash(currency_id, who, amount);
 		let account = Self::accounts(who, currency_id);
 		let free_slashed_amount = account.free.min(amount);
 		// Cannot underflow because free_slashed_amount can never be greater than amount
@@ -1280,6 +1294,7 @@ impl<T: Config> MultiReservableCurrency<T::AccountId> for Pallet<T> {
 			return value;
 		}
 
+		T::OnSlash::on_slash(currency_id, who, value);
 		let reserved_balance = Self::reserved_balance(currency_id, who);
 		let actual = reserved_balance.min(value);
 		Self::mutate_account(who, currency_id, |account, _| {

tokens/src/mock.rs

@@ -268,6 +268,55 @@ impl Happened<(AccountId, CurrencyId)> for TrackKilledAccounts {
 	}
 }
 
+thread_local! {
+	pub static ON_SLASH_CALLS: RefCell<u32> = RefCell::new(0);
+	pub static ON_DEPOSIT_CALLS: RefCell<u32> = RefCell::new(0);
+	pub static ON_TRANSFER_CALLS: RefCell<u32> = RefCell::new(0);
+}
+
+pub struct OnSlashHook<T>(marker::PhantomData<T>);
+impl<T: Config> OnSlash<T::AccountId, CurrencyId, Balance> for OnSlashHook<T> {
+	fn on_slash(_currency_id: CurrencyId, _account_id: &T::AccountId, _amount: Balance) {
+		ON_SLASH_CALLS.with(|cell| *cell.borrow_mut() += 1);
+	}
+}
+impl<T: Config> OnSlashHook<T> {
+	pub fn calls() -> u32 {
+		ON_SLASH_CALLS.with(|accounts| accounts.borrow().clone())
+	}
+}
+
+pub struct OnDepositHook<T>(marker::PhantomData<T>);
+impl<T: Config> OnDeposit<T::AccountId, CurrencyId, Balance> for OnDepositHook<T> {
+	fn on_deposit(_currency_id: CurrencyId, _account_id: &T::AccountId, _amount: Balance) -> DispatchResult {
+		ON_DEPOSIT_CALLS.with(|cell| *cell.borrow_mut() += 1);
+		Ok(())
+	}
+}
+impl<T: Config> OnDepositHook<T> {
+	pub fn calls() -> u32 {
+		ON_DEPOSIT_CALLS.with(|accounts| accounts.borrow().clone())
+	}
+}
+
+pub struct OnTransferHook<T>(marker::PhantomData<T>);
+impl<T: Config> OnTransfer<T::AccountId, CurrencyId, Balance> for OnTransferHook<T> {
+	fn on_transfer(
+		_currency_id: CurrencyId,
+		_from: &T::AccountId,
+		_to: &T::AccountId,
+		_amount: Balance,
+	) -> DispatchResult {
+		ON_TRANSFER_CALLS.with(|cell| *cell.borrow_mut() += 1);
+		Ok(())
+	}
+}
+impl<T: Config> OnTransferHook<T> {
+	pub fn calls() -> u32 {
+		ON_TRANSFER_CALLS.with(|accounts| accounts.borrow().clone())
+	}
+}
+
 parameter_types! {
 	pub DustReceiver: AccountId = PalletId(*b"orml/dst").into_account_truncating();
 }
@@ -280,6 +329,9 @@ impl Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = TransferDust<Runtime, DustReceiver>;
+	type OnSlash = OnSlashHook<Runtime>;
+	type OnDeposit = OnDepositHook<Runtime>;
+	type OnTransfer = OnTransferHook<Runtime>;
 	type OnNewTokenAccount = TrackCreatedAccounts;
 	type OnKilledTokenAccount = TrackKilledAccounts;
 	type MaxLocks = ConstU32<2>;

tokens/src/tests.rs

@@ -1167,3 +1167,61 @@ fn lifecycle_callbacks_are_activated() {
 		assert_eq!(TrackKilledAccounts::accounts(), vec![(ALICE, BTC)]);
 	})
 }
+
+// *************************************************
+// tests for mutation hooks (OnDeposit, OnTransfer)
+// (tests for the OnSlash hook can be found in `./tests_multicurrency.rs`)
+// *************************************************
+
+#[test]
+fn deposit_hook_works() {
+	ExtBuilder::default().build().execute_with(|| {
+		let initial_hook_calls = OnDepositHook::<Runtime>::calls();
+		assert_ok!(Tokens::do_deposit(DOT, &CHARLIE, 0, false, true),);
+		assert_eq!(OnDepositHook::<Runtime>::calls(), initial_hook_calls);
+
+		assert_ok!(Tokens::do_deposit(DOT, &CHARLIE, 100, false, true),);
+		assert_eq!(OnDepositHook::<Runtime>::calls(), initial_hook_calls + 1);
+
+		// The hook must be called even if the actual deposit ends up failing
+		assert_noop!(
+			Tokens::do_deposit(DOT, &BOB, 1, false, true),
+			Error::<Runtime>::ExistentialDeposit
+		);
+		assert_eq!(OnDepositHook::<Runtime>::calls(), initial_hook_calls + 2);
+	});
+}
+
+#[test]
+fn transfer_hook_works() {
+	ExtBuilder::default()
+		.balances(vec![(ALICE, DOT, 100)])
+		.build()
+		.execute_with(|| {
+			let initial_hook_calls = OnTransferHook::<Runtime>::calls();
+			assert_ok!(Tokens::do_transfer(
+				DOT,
+				&ALICE,
+				&CHARLIE,
+				0,
+				ExistenceRequirement::AllowDeath
+			),);
+			assert_eq!(OnTransferHook::<Runtime>::calls(), initial_hook_calls);
+
+			assert_ok!(Tokens::do_transfer(
+				DOT,
+				&ALICE,
+				&CHARLIE,
+				10,
+				ExistenceRequirement::AllowDeath
+			));
+			assert_eq!(OnTransferHook::<Runtime>::calls(), initial_hook_calls + 1);
+
+			// The hook must be called even if the actual transfer ends up failing
+			assert_noop!(
+				Tokens::do_transfer(DOT, &ALICE, &BOB, 1, ExistenceRequirement::AllowDeath),
+				Error::<Runtime>::ExistentialDeposit
+			);
+			assert_eq!(OnTransferHook::<Runtime>::calls(), initial_hook_calls + 2);
+		});
+}

tokens/src/tests_multicurrency.rs

@@ -731,3 +731,72 @@ fn named_multi_reservable_repatriate_all_reserved_named_works() {
 			}));
 		});
 }
+
+#[test]
+fn slash_hook_works() {
+	ExtBuilder::default()
+		.balances(vec![(ALICE, DOT, 100)])
+		.build()
+		.execute_with(|| {
+			let initial_hook_calls = OnSlashHook::<Runtime>::calls();
+
+			// slashing zero tokens is a no-op
+			assert_eq!(Tokens::slash(DOT, &ALICE, 0), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_hook_calls);
+
+			assert_eq!(Tokens::slash(DOT, &ALICE, 50), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_hook_calls + 1);
+
+			// `slash` calls the hook even if no amount was slashed
+			assert_eq!(Tokens::slash(DOT, &ALICE, 100), 50);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_hook_calls + 2);
+		});
+}
+
+#[test]
+fn slash_hook_works_for_reserved() {
+	ExtBuilder::default()
+		.balances(vec![(ALICE, DOT, 100)])
+		.build()
+		.execute_with(|| {
+			let initial_slash_hook_calls = OnSlashHook::<Runtime>::calls();
+
+			assert_ok!(Tokens::reserve(DOT, &ALICE, 50));
+			// slashing zero tokens is a no-op
+			assert_eq!(Tokens::slash_reserved(DOT, &ALICE, 0), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls);
+
+			assert_eq!(Tokens::slash_reserved(DOT, &ALICE, 50), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls + 1);
+
+			// `slash_reserved` calls the hook even if no amount was slashed
+			assert_eq!(Tokens::slash_reserved(DOT, &ALICE, 50), 50);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls + 2);
+		});
+}
+
+#[test]
+fn slash_hook_works_for_reserved_named() {
+	ExtBuilder::default()
+		.balances(vec![(ALICE, DOT, 100)])
+		.build()
+		.execute_with(|| {
+			let initial_slash_hook_calls = OnSlashHook::<Runtime>::calls();
+
+			assert_ok!(Tokens::reserve_named(&RID_1, DOT, &ALICE, 10));
+			// slashing zero tokens is a no-op
+			assert_eq!(Tokens::slash_reserved_named(&RID_1, DOT, &ALICE, 0), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls);
+
+			assert_eq!(Tokens::slash_reserved_named(&RID_1, DOT, &ALICE, 10), 0);
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls + 1);
+
+			// `slash_reserved_named` calls `slash_reserved` under-the-hood with a
+			// value to slash based on the account's balance. Because the account's
+			// balance is currently zero, `slash_reserved` will be a no-op and
+			// the OnSlash hook will not be called.
+			assert_eq!(Tokens::slash_reserved_named(&RID_1, DOT, &ALICE, 50), 50);
+			// Same value as previously because of the no-op
+			assert_eq!(OnSlashHook::<Runtime>::calls(), initial_slash_hook_calls + 1);
+		});
+}

traits/src/currency.rs

@@ -657,3 +657,34 @@ impl<AccountId> TransferAll<AccountId> for Tuple {
 		Ok(())
 	}
 }
+
+/// Hook to run before slashing an account.
+pub trait OnSlash<AccountId, CurrencyId, Balance> {
+	fn on_slash(currency_id: CurrencyId, who: &AccountId, amount: Balance);
+}
+
+impl<AccountId, CurrencyId, Balance> OnSlash<AccountId, CurrencyId, Balance> for () {
+	fn on_slash(_: CurrencyId, _: &AccountId, _: Balance) {}
+}
+
+/// Hook to run before depositing into an account.
+pub trait OnDeposit<AccountId, CurrencyId, Balance> {
+	fn on_deposit(currency_id: CurrencyId, who: &AccountId, amount: Balance) -> DispatchResult;
+}
+
+impl<AccountId, CurrencyId, Balance> OnDeposit<AccountId, CurrencyId, Balance> for () {
+	fn on_deposit(_: CurrencyId, _: &AccountId, _: Balance) -> DispatchResult {
+		Ok(())
+	}
+}
+
+/// Hook to run before transferring from an account to another.
+pub trait OnTransfer<AccountId, CurrencyId, Balance> {
+	fn on_transfer(currency_id: CurrencyId, from: &AccountId, to: &AccountId, amount: Balance) -> DispatchResult;
+}
+
+impl<AccountId, CurrencyId, Balance> OnTransfer<AccountId, CurrencyId, Balance> for () {
+	fn on_transfer(_: CurrencyId, _: &AccountId, _: &AccountId, _: Balance) -> DispatchResult {
+		Ok(())
+	}
+}

xtokens/src/mock/para.rs

@@ -84,6 +84,9 @@ impl orml_tokens::Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = ();
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type MaxLocks = ConstU32<50>;
 	type MaxReserves = ConstU32<50>;
 	type ReserveIdentifier = [u8; 8];

xtokens/src/mock/para_relative_view.rs

@@ -87,6 +87,9 @@ impl orml_tokens::Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = ();
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type MaxLocks = ConstU32<50>;
 	type MaxReserves = ConstU32<50>;
 	type ReserveIdentifier = [u8; 8];

xtokens/src/mock/para_teleport.rs

@@ -85,6 +85,9 @@ impl orml_tokens::Config for Runtime {
 	type WeightInfo = ();
 	type ExistentialDeposits = ExistentialDeposits;
 	type OnDust = ();
+	type OnSlash = ();
+	type OnDeposit = ();
+	type OnTransfer = ();
 	type MaxLocks = ConstU32<50>;
 	type MaxReserves = ConstU32<50>;
 	type ReserveIdentifier = [u8; 8];