Fix fungible unbalanced trait (paritytech#12569)

* Fix fungible unbalanced trait * Add simple decrease_balance test Signed-off-by: Oliver Tale-Yazdi <[email protected]> * Fix decrease_balance_at_most * Fix decrease_balance_at_most in fungibles * Rename free_balanceto balance_on_free * Use reducible_balance instead of balance_on_free Signed-off-by: Oliver Tale-Yazdi <[email protected]> Co-authored-by: Oliver Tale-Yazdi <[email protected]>
ZentaChain · Nov 2, 2022 · 5683513 · 5683513
1 parent 8fb4138
commit 5683513
Show file tree

Hide file tree

Showing 4 changed files with 175 additions and 11 deletions.
diff --git a/frame/balances/src/lib.rs b/frame/balances/src/lib.rs
@@ -1149,15 +1149,19 @@ impl<T: Config<I>, I: 'static> fungible::Transfer<T::AccountId> for Pallet<T, I>
 
 impl<T: Config<I>, I: 'static> fungible::Unbalanced<T::AccountId> for Pallet<T, I> {
 	fn set_balance(who: &T::AccountId, amount: Self::Balance) -> DispatchResult {
-		Self::mutate_account(who, |account| {
-			account.free = amount;
+		Self::mutate_account(who, |account| -> DispatchResult {
+			// fungibles::Unbalanced::decrease_balance didn't check account.reserved
+			// free = new_balance - reserved
+			account.free =
+				amount.checked_sub(&account.reserved).ok_or(ArithmeticError::Underflow)?;
 			Self::deposit_event(Event::BalanceSet {
 				who: who.clone(),
 				free: account.free,
 				reserved: account.reserved,
 			});
-		})?;
-		Ok(())
+
+			Ok(())
+		})?
 	}
 
 	fn set_total_issuance(amount: Self::Balance) {

diff --git a/frame/balances/src/tests.rs b/frame/balances/src/tests.rs
@@ -24,7 +24,7 @@ macro_rules! decl_tests {
 	($test:ty, $ext_builder:ty, $existential_deposit:expr) => {
 
 		use crate::*;
-		use sp_runtime::{ArithmeticError, FixedPointNumber, traits::{SignedExtension, BadOrigin}};
+		use sp_runtime::{ArithmeticError, TokenError, FixedPointNumber, traits::{SignedExtension, BadOrigin}};
 		use frame_support::{
 			assert_noop, assert_storage_noop, assert_ok, assert_err,
 			traits::{
@@ -1298,5 +1298,163 @@ macro_rules! decl_tests {
 					assert_eq!(Balances::reserved_balance(&1337), 42);
 			});
 		}
+
+		#[test]
+		fn fungible_unbalanced_trait_set_balance_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 0);
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 100);
+
+				assert_ok!(Balances::reserve(&1337, 60));
+				assert_eq!(Balances::free_balance(1337) , 40);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+
+				assert_noop!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 0), ArithmeticError::Underflow);
+
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 60));
+				assert_eq!(Balances::free_balance(1337) , 0);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_set_total_issuance_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_eq!(<Balances as fungible::Inspect<_>>::total_issuance(), 0);
+				<Balances as fungible::Unbalanced<_>>::set_total_issuance(100);
+				assert_eq!(<Balances as fungible::Inspect<_>>::total_issuance(), 100);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_decrease_balance_simple_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				// An Account that starts at 100
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				// and reserves 50
+				assert_ok!(Balances::reserve(&1337, 50));
+				// and is decreased by 20
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::decrease_balance(&1337, 20));
+				// should end up at 80.
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 80);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_decrease_balance_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 100);
+
+				assert_noop!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance(&1337, 101),
+					TokenError::NoFunds
+				);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance(&1337, 100),
+					Ok(100)
+				);
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 0);
+
+				// free: 40, reserved: 60
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_ok!(Balances::reserve(&1337, 60));
+				assert_eq!(Balances::free_balance(1337) , 40);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+				assert_noop!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance(&1337, 41),
+					TokenError::NoFunds
+				);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance(&1337, 40),
+					Ok(40)
+				);
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 60);
+				assert_eq!(Balances::free_balance(1337), 0);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_decrease_balance_at_most_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 100);
+
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance_at_most(&1337, 101),
+					100
+				);
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 0);
+
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance_at_most(&1337, 100),
+					100
+				);
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 0);
+
+				// free: 40, reserved: 60
+				assert_ok!(<Balances as fungible::Unbalanced<_>>::set_balance(&1337, 100));
+				assert_ok!(Balances::reserve(&1337, 60));
+				assert_eq!(Balances::free_balance(1337) , 40);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance_at_most(&1337, 0),
+					0
+				);
+				assert_eq!(Balances::free_balance(1337) , 40);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance_at_most(&1337, 10),
+					10
+				);
+				assert_eq!(Balances::free_balance(1337), 30);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::decrease_balance_at_most(&1337, 200),
+					30
+				);
+				assert_eq!(<Balances as fungible::Inspect<_>>::balance(&1337), 60);
+				assert_eq!(Balances::free_balance(1337), 0);
+				assert_eq!(Balances::reserved_balance(1337), 60);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_increase_balance_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_noop!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance(&1337, 0),
+					TokenError::BelowMinimum
+				);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance(&1337, 1),
+					Ok(1)
+				);
+				assert_noop!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance(&1337, u64::MAX),
+					ArithmeticError::Overflow
+				);
+			});
+		}
+
+		#[test]
+		fn fungible_unbalanced_trait_increase_balance_at_most_works() {
+			<$ext_builder>::default().build().execute_with(|| {
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance_at_most(&1337, 0),
+					0
+				);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance_at_most(&1337, 1),
+					1
+				);
+				assert_eq!(
+					<Balances as fungible::Unbalanced<_>>::increase_balance_at_most(&1337, u64::MAX),
+					u64::MAX - 1
+				);
+			});
+		}
 	}
 }
diff --git a/frame/support/src/traits/tokens/fungible/balanced.rs b/frame/support/src/traits/tokens/fungible/balanced.rs
@@ -164,7 +164,7 @@ pub trait Unbalanced<AccountId>: Inspect<AccountId> {
 		amount: Self::Balance,
 	) -> Result<Self::Balance, DispatchError> {
 		let old_balance = Self::balance(who);
-		let (mut new_balance, mut amount) = if old_balance < amount {
+		let (mut new_balance, mut amount) = if Self::reducible_balance(who, false) < amount {
 			return Err(TokenError::NoFunds.into())
 		} else {
 			(old_balance - amount, amount)
@@ -186,8 +186,9 @@ pub trait Unbalanced<AccountId>: Inspect<AccountId> {
 	/// Return the imbalance by which the account was reduced.
 	fn decrease_balance_at_most(who: &AccountId, amount: Self::Balance) -> Self::Balance {
 		let old_balance = Self::balance(who);
-		let (mut new_balance, mut amount) = if old_balance < amount {
-			(Zero::zero(), old_balance)
+		let old_free_balance = Self::reducible_balance(who, false);
+		let (mut new_balance, mut amount) = if old_free_balance < amount {
+			(old_balance.saturating_sub(old_free_balance), old_free_balance)
 		} else {
 			(old_balance - amount, amount)
 		};

diff --git a/frame/support/src/traits/tokens/fungibles/balanced.rs b/frame/support/src/traits/tokens/fungibles/balanced.rs
@@ -185,7 +185,7 @@ pub trait Unbalanced<AccountId>: Inspect<AccountId> {
 		amount: Self::Balance,
 	) -> Result<Self::Balance, DispatchError> {
 		let old_balance = Self::balance(asset, who);
-		let (mut new_balance, mut amount) = if old_balance < amount {
+		let (mut new_balance, mut amount) = if Self::reducible_balance(asset, who, false) < amount {
 			return Err(TokenError::NoFunds.into())
 		} else {
 			(old_balance - amount, amount)
@@ -211,8 +211,9 @@ pub trait Unbalanced<AccountId>: Inspect<AccountId> {
 		amount: Self::Balance,
 	) -> Self::Balance {
 		let old_balance = Self::balance(asset, who);
-		let (mut new_balance, mut amount) = if old_balance < amount {
-			(Zero::zero(), old_balance)
+		let old_free_balance = Self::reducible_balance(asset, who, false);
+		let (mut new_balance, mut amount) = if old_free_balance < amount {
+			(old_balance.saturating_sub(old_free_balance), old_free_balance)
 		} else {
 			(old_balance - amount, amount)
 		};