From ad0011a1bd562d857c92eb133268b2b933f18f4e Mon Sep 17 00:00:00 2001
From: Wei Tang <hi@that.world>
Date: Sat, 15 Jun 2019 10:26:22 +0200
Subject: [PATCH] Maintain TotalIssurance integrity in case of sudo set_balance
 (#2871)

* Maintain TotalInsurance integrity in case of sudo set_balance

* Fix set_balance docs

* Update impl_version
---
 substrate/node/runtime/src/lib.rs  |  2 +-
 substrate/srml/balances/src/lib.rs | 34 +++++++++++++++++++++++-------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/substrate/node/runtime/src/lib.rs b/substrate/node/runtime/src/lib.rs
index 0c5a2da317..3719130eda 100644
--- a/substrate/node/runtime/src/lib.rs
+++ b/substrate/node/runtime/src/lib.rs
@@ -59,7 +59,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	impl_name: create_runtime_str!("substrate-node"),
 	authoring_version: 10,
 	spec_version: 94,
-	impl_version: 94,
+	impl_version: 95,
 	apis: RUNTIME_API_VERSIONS,
 };
 
diff --git a/substrate/srml/balances/src/lib.rs b/substrate/srml/balances/src/lib.rs
index 3c3762833b..72ec997206 100644
--- a/substrate/srml/balances/src/lib.rs
+++ b/substrate/srml/balances/src/lib.rs
@@ -142,11 +142,15 @@
 //! ## Genesis config
 //!
 //! The Balances module depends on the [`GenesisConfig`](./struct.GenesisConfig.html).
+//!
+//! ## Assumptions
+//!
+//! * Total issued balanced of all accounts should be less than `Trait::Balance::max_value()`.
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use rstd::prelude::*;
-use rstd::{cmp, result};
+use rstd::{cmp, result, mem};
 use parity_codec::{Codec, Encode, Decode};
 use srml_support::{StorageValue, StorageMap, Parameter, decl_event, decl_storage, decl_module};
 use srml_support::traits::{
@@ -373,10 +377,10 @@ decl_module! {
 
 		/// Set the balances of a given account.
 		///
-		/// This will alter `FreeBalance` and `ReservedBalance` in storage.
+		/// This will alter `FreeBalance` and `ReservedBalance` in storage. it will
+		/// also decrease the total issuance of the system (`TotalIssuance`).
 		/// If the new free or reserved balance is below the existential deposit,
-		/// it will also decrease the total issuance of the system (`TotalIssuance`)
-		/// and reset the account nonce (`system::AccountNonce`).
+		/// it will reset the account nonce (`system::AccountNonce`).
 		///
 		/// The dispatch origin for this call is `root`.
 		///
@@ -386,12 +390,26 @@ decl_module! {
 		/// # </weight>
 		fn set_balance(
 			who: <T::Lookup as StaticLookup>::Source,
-			#[compact] free: T::Balance,
-			#[compact] reserved: T::Balance
+			#[compact] new_free: T::Balance,
+			#[compact] new_reserved: T::Balance
 		) {
 			let who = T::Lookup::lookup(who)?;
-			Self::set_free_balance(&who, free);
-			Self::set_reserved_balance(&who, reserved);
+
+			let current_free = <FreeBalance<T, I>>::get(&who);
+			if new_free > current_free {
+				mem::drop(PositiveImbalance::<T, I>::new(new_free - current_free));
+			} else if new_free < current_free {
+				mem::drop(NegativeImbalance::<T, I>::new(current_free - new_free));
+			}
+			Self::set_free_balance(&who, new_free);
+
+			let current_reserved = <ReservedBalance<T, I>>::get(&who);
+			if new_reserved > current_reserved {
+				mem::drop(PositiveImbalance::<T, I>::new(new_reserved - current_reserved));
+			} else if new_reserved < current_reserved {
+				mem::drop(NegativeImbalance::<T, I>::new(current_reserved - new_reserved));
+			}
+			Self::set_reserved_balance(&who, new_reserved);
 		}
 	}
 }