Allow SetBalance to handle error when trying to kill acount with re…

…ference counter. (#10826)

* bug found

* fix logic

* a little simpler

* add test

frame/balances/src/lib.rs

@@ -314,25 +314,32 @@ pub mod pallet {
 			let new_free = if wipeout { Zero::zero() } else { new_free };
 			let new_reserved = if wipeout { Zero::zero() } else { new_reserved };
 
-			let (free, reserved) = Self::mutate_account(&who, |account| {
-				if new_free > account.free {
-					mem::drop(PositiveImbalance::<T, I>::new(new_free - account.free));
-				} else if new_free < account.free {
-					mem::drop(NegativeImbalance::<T, I>::new(account.free - new_free));
-				}
-
-				if new_reserved > account.reserved {
-					mem::drop(PositiveImbalance::<T, I>::new(new_reserved - account.reserved));
-				} else if new_reserved < account.reserved {
-					mem::drop(NegativeImbalance::<T, I>::new(account.reserved - new_reserved));
-				}
+			// First we try to modify the account's balance to the forced balance.
+			let (old_free, old_reserved) = Self::mutate_account(&who, |account| {
+				let old_free = account.free;
+				let old_reserved = account.reserved;
 
 				account.free = new_free;
 				account.reserved = new_reserved;
 
-				(account.free, account.reserved)
+				(old_free, old_reserved)
 			})?;
-			Self::deposit_event(Event::BalanceSet { who, free, reserved });
+
+			// This will adjust the total issuance, which was not done by the `mutate_account`
+			// above.
+			if new_free > old_free {
+				mem::drop(PositiveImbalance::<T, I>::new(new_free - old_free));
+			} else if new_free < old_free {
+				mem::drop(NegativeImbalance::<T, I>::new(old_free - new_free));
+			}
+
+			if new_reserved > old_reserved {
+				mem::drop(PositiveImbalance::<T, I>::new(new_reserved - old_reserved));
+			} else if new_reserved < old_reserved {
+				mem::drop(NegativeImbalance::<T, I>::new(old_reserved - new_reserved));
+			}
+
+			Self::deposit_event(Event::BalanceSet { who, free: new_free, reserved: new_reserved });
 			Ok(().into())
 		}
 

frame/balances/src/tests.rs

@@ -1239,5 +1239,29 @@ macro_rules! decl_tests {
 				assert_eq!(Balances::free_balance(&3), 25);
 			});
 		}
+
+		#[test]
+		fn set_balance_handles_killing_account() {
+			<$ext_builder>::default().build().execute_with(|| {
+					let _ = Balances::deposit_creating(&1, 111);
+					assert_ok!(frame_system::Pallet::<Test>::inc_consumers(&1));
+					assert_noop!(
+						Balances::set_balance(Origin::root(), 1, 0, 0),
+						DispatchError::ConsumerRemaining,
+					);
+			});
+		}
+
+		#[test]
+		fn set_balance_handles_total_issuance() {
+			<$ext_builder>::default().build().execute_with(|| {
+					let old_total_issuance = Balances::total_issuance();
+					assert_ok!(Balances::set_balance(Origin::root(), 1337, 69, 42));
+					assert_eq!(Balances::total_issuance(), old_total_issuance + 69 + 42);
+					assert_eq!(Balances::total_balance(&1337), 69 + 42);
+					assert_eq!(Balances::free_balance(&1337), 69);
+					assert_eq!(Balances::reserved_balance(&1337), 42);
+			});
+		}
 	}
 }