Replace Overflow-Prone Operations in Staking with Saturating Arithmetic (#2115)

* Improve a few doc string

* Replace overflow-prone operation with saturating.

* Remove whitespace.

* Update wasm; Bump spec;

* Bump impl again.

* Fix review comments.

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: 44,
-	impl_version: 44,
+	impl_version: 45,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/staking/src/phragmen.rs

@@ -18,69 +18,69 @@
 
 use rstd::prelude::*;
 use primitives::Perquintill;
-use primitives::traits::{Zero, As, Bounded, CheckedMul, CheckedSub};
+use primitives::traits::{Zero, As, Bounded, CheckedMul, Saturating};
 use parity_codec::{HasCompact, Encode, Decode};
 use crate::{Exposure, BalanceOf, Trait, ValidatorPrefs, IndividualExposure};
 
 
-// Configure the behavior of the Phragmen election.
-// Might be deprecated.
+/// Configure the behavior of the Phragmen election.
+/// Might be deprecated.
 pub struct ElectionConfig<Balance: HasCompact> {
-	// Perform equalise?.
+	/// Perform equalise?.
 	pub equalise: bool,
-	// Number of equalise iterations.
+	/// Number of equalise iterations.
 	pub iterations: usize,
-	// Tolerance of max change per equalise iteration.
+	/// Tolerance of max change per equalise iteration.
 	pub tolerance: Balance,
 }
 
-// Wrapper around validation candidates some metadata.
+/// Wrapper around validation candidates some metadata.
 #[derive(Clone, Encode, Decode, Default)]
 #[cfg_attr(feature = "std", derive(Debug))]
 pub struct Candidate<AccountId, Balance: HasCompact> {
-	// The validator's account
+	/// The validator's account
 	pub who: AccountId,
-	// Exposure struct, holding info about the value that the validator has in stake.
+	/// Exposure struct, holding info about the value that the validator has in stake.
 	pub exposure: Exposure<AccountId, Balance>,
-	// Intermediary value used to sort candidates.
+	/// Intermediary value used to sort candidates.
 	pub score: Perquintill,
-	// Accumulator of the stake of this candidate based on received votes.
+	/// Accumulator of the stake of this candidate based on received votes.
 	approval_stake: Balance,
-	// Flag for being elected.
+	/// Flag for being elected.
 	elected: bool,
-	// This is most often equal to `Exposure.total` but not always. Needed for [`equalise`]
+	/// This is most often equal to `Exposure.total` but not always. Needed for [`equalise`]
 	backing_stake: Balance
 }
 
-// Wrapper around the nomination info of a single nominator for a group of validators.
+/// Wrapper around the nomination info of a single nominator for a group of validators.
 #[derive(Clone, Encode, Decode, Default)]
 #[cfg_attr(feature = "std", derive(Debug))]
 pub struct Nominator<AccountId, Balance: HasCompact> {
-	// The nominator's account.
+	/// The nominator's account.
 	who: AccountId,
-	// List of validators proposed by this nominator.
+	/// List of validators proposed by this nominator.
 	edges: Vec<Edge<AccountId, Balance>>,
-	// the stake amount proposed by the nominator as a part of the vote.
+	/// the stake amount proposed by the nominator as a part of the vote.
 	budget: Balance,
-	// Incremented each time a nominee that this nominator voted for has been elected.
+	/// Incremented each time a nominee that this nominator voted for has been elected.
 	load: Perquintill,
 }
 
-// Wrapper around a nominator vote and the load of that vote.
+/// Wrapper around a nominator vote and the load of that vote.
 #[derive(Clone, Encode, Decode, Default)]
 #[cfg_attr(feature = "std", derive(Debug))]
 pub struct Edge<AccountId, Balance: HasCompact> {
-	// Account being voted for
+	/// Account being voted for
 	who: AccountId,
-	// Load of this vote.
+	/// Load of this vote.
 	load: Perquintill,
-	// Final backing stake of this vote.
+	/// Final backing stake of this vote.
 	backing_stake: Balance,
-	// Index of the candidate stored in the 'candidates' vector
+	/// Index of the candidate stored in the 'candidates' vector
 	candidate_index: usize,
-	// Index of the candidate stored in the 'elected_candidates' vector. Used only with equalise.
+	/// Index of the candidate stored in the 'elected_candidates' vector. Used only with equalise.
 	elected_idx: usize,
-	// Indicates if this edge is a vote for an elected candidate. Used only with equalise.
+	/// Indicates if this edge is a vote for an elected candidate. Used only with equalise.
 	elected: bool,
 }
 
@@ -88,7 +88,8 @@ pub struct Edge<AccountId, Balance: HasCompact> {
 ///
 /// Reference implementation: https://github.com/w3f/consensus
 ///
-/// Returns a vector of elected candidates
+/// Returns an Option of elected candidates, if election is performed.
+/// Returns None if not enough candidates exist.
 pub fn elect<T: Trait + 'static, FR, FN, FV, FS>(
 		get_rounds: FR,
 		get_validators: FV,
@@ -138,7 +139,7 @@ pub fn elect<T: Trait + 'static, FR, FN, FV, FS>(
 		let mut edges: Vec<Edge<T::AccountId, BalanceOf<T>>> = Vec::with_capacity(nominees.len());
 		for n in &nominees {
 			if let Some(idx) = candidates.iter_mut().position(|i| i.who == *n) {
-				candidates[idx].approval_stake += nominator_stake;
+				candidates[idx].approval_stake = candidates[idx].approval_stake.saturating_add(nominator_stake);
 				edges.push(Edge { who: n.clone(), candidate_index: idx, ..Default::default() });
 			}
 		}
@@ -173,8 +174,10 @@ pub fn elect<T: Trait + 'static, FR, FN, FV, FS>(
 				for e in &n.edges {
 					let c = &mut candidates[e.candidate_index];
 					if !c.elected {
-						let temp = n.budget.as_() * *n.load / c.approval_stake.as_();
-						c.score = Perquintill::from_quintillionths(*c.score + temp);
+						// Note: This seems to never overflow, ok to be safe though
+						let temp = n.budget.as_().saturating_mul(*n.load) / c.approval_stake.as_();
+						// Note: This seems to never overflow, ok to be safe though
+						c.score = Perquintill::from_quintillionths((*c.score).saturating_add(temp));
 					}
 				}
 			}
@@ -206,12 +209,12 @@ pub fn elect<T: Trait + 'static, FR, FN, FV, FS>(
 				// if the target of this vote is among the winners, otherwise let go.
 				if let Some(c) = elected_candidates.iter_mut().find(|c| c.who == e.who) {
 					e.elected = true;
-					// NOTE: for now, always divide last to avoid collapse to zero.
-					e.backing_stake = <BalanceOf<T>>::sa((n.budget.as_() * *e.load) / *n.load);
-					c.backing_stake += e.backing_stake;
+					// NOTE: always divide last to avoid collapse to zero.
+					e.backing_stake = <BalanceOf<T>>::sa(n.budget.as_().saturating_mul(*e.load) / *n.load);
+					c.backing_stake = c.backing_stake.saturating_add(e.backing_stake);
 					if c.who != n.who {
 						// Only update the exposure if this vote is from some other account.
-						c.exposure.total += e.backing_stake;
+						c.exposure.total = c.exposure.total.saturating_add(e.backing_stake);
 						c.exposure.others.push(
 							IndividualExposure { who: n.who.clone(), value: e.backing_stake }
 						);
@@ -256,7 +259,7 @@ pub fn elect<T: Trait + 'static, FR, FN, FV, FS>(
 				let nominator = n.who.clone();
 				for e in &mut n.edges {
 					if let Some(c) = elected_candidates.iter_mut().find(|c| c.who == e.who && c.who != nominator) {
-						c.exposure.total += n.budget;
+						c.exposure.total = c.exposure.total.saturating_add(n.budget);
 						c.exposure.others.push(
 							IndividualExposure { who: n.who.clone(), value: n.budget }
 						);
@@ -284,7 +287,7 @@ pub fn equalise<T: Trait + 'static>(
 	if elected_edges.len() == 0 { return <BalanceOf<T>>::zero(); }
 	let stake_used = elected_edges
 		.iter()
-		.fold(<BalanceOf<T>>::zero(), |s, e| s + e.backing_stake);
+		.fold(<BalanceOf<T>>::zero(), |s, e| s.saturating_add(e.backing_stake));
 	let backed_stakes = elected_edges
 		.iter()
 		.map(|e| elected_candidates[e.elected_idx].backing_stake)
@@ -304,9 +307,9 @@ pub fn equalise<T: Trait + 'static>(
 		let min_stake = *backed_stakes
 			.iter()
 			.min()
-			.expect("vector with positive length will have a max; qed");
-		difference = max_stake - min_stake;
-		difference += nominator.budget - stake_used;
+			.expect("vector with positive length will have a min; qed");
+		difference = max_stake.saturating_sub(min_stake);
+		difference = difference.saturating_add(nominator.budget.saturating_sub(stake_used));
 		if difference < tolerance {
 			return difference;
 		}
@@ -330,25 +333,29 @@ pub fn equalise<T: Trait + 'static>(
 	let budget = nominator.budget;
 	elected_edges.iter_mut().enumerate().for_each(|(idx, e)| {
 		let stake = elected_candidates[e.elected_idx].backing_stake;
-
 		let stake_mul = stake.checked_mul(&<BalanceOf<T>>::sa(idx as u64)).unwrap_or(<BalanceOf<T>>::max_value());
-		let stake_sub = stake_mul.checked_sub(&cumulative_stake).unwrap_or_default();
+		let stake_sub = stake_mul.saturating_sub(cumulative_stake);
 		if stake_sub > budget {
-			last_index = idx.clone().checked_sub(1).unwrap_or(0);
+			last_index = idx.checked_sub(1).unwrap_or(0);
 			return
 		}
-		cumulative_stake += stake;
+		cumulative_stake = cumulative_stake.saturating_add(stake);
 	});
 
 	let last_stake = elected_candidates[elected_edges[last_index].elected_idx].backing_stake;
 	let split_ways = last_index + 1;
-	let excess = nominator.budget + cumulative_stake - last_stake * <BalanceOf<T>>::sa(split_ways as u64);
+	let excess = nominator.budget
+		.saturating_add(cumulative_stake)
+		.saturating_sub(
+			last_stake.checked_mul(&<BalanceOf<T>>::sa(split_ways as u64))
+				.unwrap_or(<BalanceOf<T>>::max_value())
+		);
 	let nominator_address = nominator.who.clone();
 	elected_edges.iter_mut().take(split_ways).for_each(|e| {
 		let c = &mut elected_candidates[e.elected_idx];
-		e.backing_stake = excess / <BalanceOf<T>>::sa(split_ways as u64) + last_stake - c.backing_stake;
-		c.exposure.total += e.backing_stake;
-		c.backing_stake += e.backing_stake;
+		e.backing_stake = (excess / <BalanceOf<T>>::sa(split_ways as u64)).saturating_add(last_stake) - c.backing_stake;
+		c.exposure.total = c.exposure.total.saturating_add(e.backing_stake);
+		c.backing_stake = c.backing_stake.saturating_add(e.backing_stake);
 		if let Some(i_expo) = c.exposure.others.iter_mut().find(|i| i.who == nominator_address) {
 			i_expo.value = e.backing_stake;
 		}

substrate/srml/staking/src/tests.rs

@@ -728,7 +728,7 @@ fn nominators_also_get_slashed() {
 		assert_ok!(Staking::nominate(Origin::signed(2), vec![20, 10]));
 
 		// new era, pay rewards,
-		System::set_block_number(2);
+		System::set_block_number(1);
 		Session::check_rotate_session(System::block_number());
 
 		// Nominator stash didn't collect any.
@@ -1885,3 +1885,105 @@ fn phragmen_chooses_correct_validators() {
 		assert_eq!(Session::validators().len(), 1);
 	})
 }
+
+#[test]
+fn phragmen_should_not_overflow_validators() {
+	with_externalities(&mut ExtBuilder::default()
+	.nominate(false)
+	.build()
+	, || {
+		let bond_validator = |a, b| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
+		};
+		let bond_nominator = |a, b, v| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::nominate(Origin::signed(a), v));
+		};
+
+		for i in 1..=8 {
+			let _ = Balances::make_free_balance_be(&i, u64::max_value());
+		}
+
+		let _ = Staking::chill(Origin::signed(10));
+		let _ = Staking::chill(Origin::signed(20));
+
+		bond_validator(2, u64::max_value());
+		bond_validator(4, u64::max_value());
+
+		bond_nominator(6, u64::max_value()/2, vec![1, 3]);
+		bond_nominator(8, u64::max_value()/2, vec![1, 3]);
+
+		System::set_block_number(2);
+		Session::check_rotate_session(System::block_number());
+
+		assert_eq!(Session::validators(), vec![4, 2]);
+	})
+}
+
+#[test]
+fn phragmen_should_not_overflow_nominators() {
+	with_externalities(&mut ExtBuilder::default()
+	.nominate(false)
+	.build()
+	, || {
+		let bond_validator = |a, b| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
+		};
+		let bond_nominator = |a, b, v| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::nominate(Origin::signed(a), v));
+		};
+
+		let _ = Staking::chill(Origin::signed(10));
+		let _ = Staking::chill(Origin::signed(20));
+
+		for i in 1..=8 {
+			let _ = Balances::make_free_balance_be(&i, u64::max_value());
+		}
+
+		bond_validator(2, u64::max_value()/2);
+		bond_validator(4, u64::max_value()/2);
+
+		bond_nominator(6, u64::max_value(), vec![1, 3]);
+		bond_nominator(8, u64::max_value(), vec![1, 3]);
+
+		System::set_block_number(2);
+		Session::check_rotate_session(System::block_number());
+
+		assert_eq!(Session::validators(), vec![4, 2]);
+	})
+}
+
+#[test]
+fn phragmen_should_not_overflow_ultimate() {
+	with_externalities(&mut ExtBuilder::default()
+	.nominate(false)
+	.build()
+	, || {
+		let bond_validator = |a, b| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::validate(Origin::signed(a), ValidatorPrefs::default()));
+		};
+		let bond_nominator = |a, b, v| {
+			assert_ok!(Staking::bond(Origin::signed(a-1), a, b, RewardDestination::Controller));
+			assert_ok!(Staking::nominate(Origin::signed(a), v));
+		};
+
+		for i in 1..=8 {
+			let _ = Balances::make_free_balance_be(&i, u64::max_value());
+		}
+
+		bond_validator(2, u64::max_value());
+		bond_validator(4, u64::max_value());
+
+		bond_nominator(6, u64::max_value(), vec![1, 3]);
+		bond_nominator(8, u64::max_value(), vec![1, 3]);
+
+		System::set_block_number(2);
+		Session::check_rotate_session(System::block_number());
+
+		assert_eq!(Session::validators(), vec![4, 2]);
+	})
+}

substrate/srml/support/src/traits.rs

@@ -310,7 +310,7 @@ pub trait Currency<AccountId> {
 		value: Self::Balance,
 	) -> result::Result<(), &'static str>;
 
-		/// Deducts up to `value` from the combined balance of `who`, preferring to deduct from the
+	/// Deducts up to `value` from the combined balance of `who`, preferring to deduct from the
 	/// free balance. This function cannot fail.
 	///
 	/// The resulting imbalance is the first item of the tuple returned.