fix: unwanted delegator state removal when tx fails (#266)

* fix: unwanted delegator state removal when tx fails

* feat: add v6 migration

* feat: add v6 pre check, rm expect

* fix: dont panic

Author: wischli

pallets/parachain-staking/src/lib.rs

@@ -172,6 +172,7 @@ mod types;
 use frame_support::pallet;
 
 pub use crate::{default_weights::WeightInfo, pallet::*};
+use types::ReplacedDelegator;
 
 #[pallet]
 pub mod pallet {
@@ -1391,9 +1392,17 @@ pub mod pallet {
 		) -> DispatchResultWithPostInfo {
 			let acc = ensure_signed(origin)?;
 			let collator = T::Lookup::lookup(collator)?;
+
+			// check balance
+			ensure!(
+				pallet_balances::Pallet::<T>::free_balance(acc.clone()) >= amount.into(),
+				pallet_balances::Error::<T>::InsufficientBalance
+			);
+
 			// first delegation
 			ensure!(<DelegatorState<T>>::get(&acc).is_none(), Error::<T>::AlreadyDelegating);
 			ensure!(amount >= T::MinDelegatorStake::get(), Error::<T>::NomStakeBelowMin);
+
 			// cannot be a collator candidate and delegator with same AccountId
 			ensure!(!Self::is_active_candidate(&acc).is_some(), Error::<T>::CandidateExists);
 			ensure!(
@@ -1431,12 +1440,15 @@ pub mod pallet {
 				.map_err(|_| Error::<T>::MaxCollatorsPerDelegatorExceeded)?;
 
 			let old_total = state.total;
-			// update state and potentially kick a delegator with less staked amount
-			state = if num_delegations_pre_insertion == T::MaxDelegatorsPerCollator::get() {
+
+			// update state and potentially prepare kicking a delegator with less staked
+			// amount
+			let (state, maybe_kicked_delegator) = if num_delegations_pre_insertion == T::MaxDelegatorsPerCollator::get()
+			{
 				Self::do_update_delegator(delegation, state)?
 			} else {
 				state.total = state.total.saturating_add(amount);
-				state
+				(state, None)
 			};
 			let new_total = state.total;
 
@@ -1446,11 +1458,16 @@ pub mod pallet {
 				Self::update_top_candidates(collator.clone(), old_total, new_total);
 			}
 
+			// *** No Fail beyond this point ***
+
 			// update states
 			<CandidatePool<T>>::insert(&collator, state);
 			<DelegatorState<T>>::insert(&acc, delegator_state);
 			<LastDelegation<T>>::insert(&acc, delegation_counter);
 
+			// update or clear storage of potentially kicked delegator
+			Self::update_kicked_delegator_storage(maybe_kicked_delegator);
+
 			// update candidates for next round
 			let (num_collators, num_delegators, _, _) = Self::update_total_stake();
 
@@ -1520,6 +1537,14 @@ pub mod pallet {
 			let acc = ensure_signed(origin)?;
 			let collator = T::Lookup::lookup(collator)?;
 			let mut delegator = <DelegatorState<T>>::get(&acc).ok_or(Error::<T>::NotYetDelegating)?;
+
+			// check balance
+			ensure!(
+				pallet_balances::Pallet::<T>::free_balance(acc.clone())
+					>= delegator.total.saturating_add(amount).into(),
+				pallet_balances::Error::<T>::InsufficientBalance
+			);
+
 			// delegation after first
 			ensure!(amount >= T::MinDelegation::get(), Error::<T>::DelegationBelowMin);
 			ensure!(
@@ -1535,9 +1560,10 @@ pub mod pallet {
 			let num_delegations_pre_insertion: u32 = state.delegators.len().saturated_into();
 			ensure!(!state.is_leaving(), Error::<T>::CannotDelegateIfLeaving);
 
-			// attempt to insert delegator and check for uniqueness
-			// NOTE: excess is handled below because we support replacing a delegator with
-			// fewer stake
+			// attempt to insert delegation, check for uniqueness and update total delegated
+			// amount
+			// NOTE: excess is handled below because we support replacing a delegator
+			// with fewer stake
 			ensure!(
 				delegator
 					.add_delegation(Stake {
@@ -1561,12 +1587,14 @@ pub mod pallet {
 
 			let old_total = state.total;
 
-			// update state and potentially kick a delegator with less staked amount
-			state = if num_delegations_pre_insertion == T::MaxDelegatorsPerCollator::get() {
+			// update state and potentially prepare kicking a delegator with less staked
+			// amount
+			let (state, maybe_kicked_delegator) = if num_delegations_pre_insertion == T::MaxDelegatorsPerCollator::get()
+			{
 				Self::do_update_delegator(delegation, state)?
 			} else {
 				state.total = state.total.saturating_add(amount);
-				state
+				(state, None)
 			};
 			let new_total = state.total;
 
@@ -1576,11 +1604,16 @@ pub mod pallet {
 				Self::update_top_candidates(collator.clone(), old_total, new_total);
 			}
 
+			// *** No Fail beyond this point ***
+
 			// Update states
 			<CandidatePool<T>>::insert(&collator, state);
 			<DelegatorState<T>>::insert(&acc, delegator);
 			<LastDelegation<T>>::insert(&acc, delegation_counter);
 
+			// update or clear storage of potentially kicked delegator
+			Self::update_kicked_delegator_storage(maybe_kicked_delegator);
+
 			// update candidates for next round
 			let (num_collators, num_delegators, _, _) = Self::update_total_stake();
 
@@ -2012,19 +2045,49 @@ pub mod pallet {
 			Ok(())
 		}
 
-		fn kick_delegator(delegation: &StakeOf<T>, collator: &T::AccountId) -> DispatchResult {
+		/// Check for remaining delegations of the delegator which has been
+		/// removed from the given collator.
+		///
+		/// Returns the removed delegator's address and
+		/// * Either the updated delegator state if other delegations are still
+		///   remaining
+		/// * Or `None`, signalling the delegator state should be cleared once
+		///   the transaction cannot fail anymore.
+		fn prep_kick_delegator(
+			delegation: &StakeOf<T>,
+			collator: &T::AccountId,
+		) -> Result<ReplacedDelegator<T>, DispatchError> {
 			let mut state = <DelegatorState<T>>::get(&delegation.owner).ok_or(Error::<T>::DelegatorNotFound)?;
 			state.rm_delegation(collator);
+
 			// we don't unlock immediately
 			Self::prep_unstake(&delegation.owner, delegation.amount, true)?;
 
-			// clear storage if no delegations are remaining
+			// return state if not empty for later removal after all checks have passed
 			if state.delegations.is_empty() {
-				<DelegatorState<T>>::remove(&delegation.owner);
+				Ok(ReplacedDelegator {
+					who: delegation.owner.clone(),
+					state: None,
+				})
 			} else {
-				<DelegatorState<T>>::insert(&delegation.owner, state);
+				Ok(ReplacedDelegator {
+					who: delegation.owner.clone(),
+					state: Some(state),
+				})
+			}
+		}
+
+		/// Either clear the storage of a kicked delegator or update its
+		/// delegation state if it still contains other delegations.
+		fn update_kicked_delegator_storage(delegator: Option<ReplacedDelegator<T>>) {
+			match delegator {
+				Some(ReplacedDelegator {
+					who,
+					state: Some(state),
+				}) => DelegatorState::<T>::insert(who, state),
+				Some(ReplacedDelegator { who, .. }) => DelegatorState::<T>::remove(who),
+				_ => (),
 			}
-			Ok(())
 		}
 
 		/// Select the top `MaxSelectedCandidates` many collators in terms of
@@ -2125,7 +2188,9 @@ pub mod pallet {
 		/// amount is at most the minimum staked value of the original delegator
 		/// set, an error is returned.
 		///
-		/// Returns the old delegation that is updated, if any.
+		/// Returns a tuple which contains the updated candidate state as well
+		/// as the potentially replaced delegation which will be used later when
+		/// updating the storage of the replaced delegator.
 		///
 		/// Emits `DelegationReplaced` if the stake exceeds one of the current
 		/// delegations.
@@ -2135,10 +2200,17 @@ pub mod pallet {
 		/// bounded by `MaxDelegatorsPerCollator`.
 		/// - Reads/Writes: 0
 		/// # </weight>
+		#[allow(clippy::type_complexity)]
 		fn do_update_delegator(
 			stake: Stake<T::AccountId, BalanceOf<T>>,
 			mut state: Candidate<T::AccountId, BalanceOf<T>, T::MaxDelegatorsPerCollator>,
-		) -> Result<CandidateOf<T, T::MaxDelegatorsPerCollator>, DispatchError> {
+		) -> Result<
+			(
+				CandidateOf<T, T::MaxDelegatorsPerCollator>,
+				Option<ReplacedDelegator<T>>,
+			),
+			DispatchError,
+		> {
 			// attempt to replace the last element of the set
 			let stake_to_remove = state
 				.delegators
@@ -2159,7 +2231,7 @@ pub mod pallet {
 				state.total = state.total.saturating_sub(stake_to_remove.amount);
 
 				// update storage of kicked delegator
-				Self::kick_delegator(&stake_to_remove, &state.id)?;
+				let kicked_delegator = Self::prep_kick_delegator(&stake_to_remove, &state.id)?;
 
 				Self::deposit_event(Event::DelegationReplaced(
 					stake.owner,
@@ -2169,9 +2241,11 @@ pub mod pallet {
 					state.id.clone(),
 					state.total,
 				));
-			}
 
-			Ok(state)
+				Ok((state, Some(kicked_delegator)))
+			} else {
+				Ok((state, None))
+			}
 		}
 
 		/// Either set or increase the BalanceLock of target account to

pallets/parachain-staking/src/migrations.rs

@@ -30,6 +30,7 @@ mod v2;
 mod v3;
 mod v4;
 mod v5;
+mod v6;
 
 /// A trait that allows version migrators to access the underlying pallet's
 /// context, e.g., its Config trait.
@@ -55,6 +56,7 @@ pub enum StakingStorageVersion {
 	V3_0_0, // Update InflationConfig
 	V4,     // Sort TopCandidates and parachain-stakings by amount
 	V5,     // Remove SelectedCandidates, Count Candidates
+	V6,     // Fix delegator replacement bug
 }
 
 #[cfg(feature = "try-runtime")]
@@ -87,7 +89,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V2_0_0 => v3::pre_migrate::<T>(),
 			Self::V3_0_0 => v4::pre_migrate::<T>(),
 			Self::V4 => v5::pre_migrate::<T>(),
-			Self::V5 => Ok(()),
+			Self::V5 => v6::pre_migrate::<T>(),
+			Self::V6 => Ok(()),
 		}
 	}
 
@@ -98,7 +101,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V2_0_0 => v3::migrate::<T>(),
 			Self::V3_0_0 => v4::migrate::<T>(),
 			Self::V4 => v5::migrate::<T>(),
-			Self::V5 => Weight::zero(),
+			Self::V5 => v6::migrate::<T>(),
+			Self::V6 => Weight::zero(),
 		}
 	}
 
@@ -111,7 +115,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V2_0_0 => v3::post_migrate::<T>(),
 			Self::V3_0_0 => v4::post_migrate::<T>(),
 			Self::V4 => v5::post_migrate::<T>(),
-			Self::V5 => Ok(()),
+			Self::V5 => v6::post_migrate::<T>(),
+			Self::V6 => Ok(()),
 		}
 	}
 }
@@ -133,7 +138,8 @@ impl<T: Config> StakingStorageMigrator<T> {
 			// Migration happens naturally, no need to point to the latest version
 			StakingStorageVersion::V3_0_0 => Some(StakingStorageVersion::V4),
 			StakingStorageVersion::V4 => Some(StakingStorageVersion::V5),
-			StakingStorageVersion::V5 => None,
+			StakingStorageVersion::V5 => Some(StakingStorageVersion::V6),
+			StakingStorageVersion::V6 => None,
 		}
 	}
 

pallets/parachain-staking/src/migrations/v6.rs

@@ -0,0 +1,95 @@
+// KILT Blockchain – https://botlabs.org
+// Copyright (C) 2019-2021 BOTLabs GmbH
+
+// The KILT Blockchain is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// The KILT Blockchain is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+// If you feel like getting in touch with us, you can do so at info@botlabs.org
+
+use frame_support::{dispatch::Weight, traits::Get};
+
+use crate::{
+	migrations::StakingStorageVersion,
+	types::{BalanceOf, Delegator},
+	CandidatePool, Config, DelegatorState, StorageVersion,
+};
+
+#[cfg(feature = "try-runtime")]
+pub(crate) fn pre_migrate<T: Config>() -> Result<(), &'static str> {
+	use crate::types::Stake;
+
+	assert_eq!(StorageVersion::<T>::get(), StakingStorageVersion::V5);
+
+	// ensure each delegator has at most one delegation (ideally should have exactly
+	// one)
+	let mut corrupt_delegation: u32 = 0;
+	for candidate in CandidatePool::<T>::iter_values() {
+		for delegation in candidate.delegators.into_iter() {
+			if let Some(state) = DelegatorState::<T>::get(delegation.owner.clone()) {
+				assert_eq!(state.delegations.len(), 1);
+				if let Some(Stake::<T::AccountId, BalanceOf<T>> { amount, owner }) =
+					state.delegations.into_bounded_vec().into_inner().get(0)
+				{
+					assert_eq!(amount, &state.total);
+					assert_eq!(owner, &candidate.id);
+				}
+			} else {
+				corrupt_delegation = corrupt_delegation.saturating_add(1);
+			}
+		}
+	}
+
+	log::info!("Found {} corrupt delegations", corrupt_delegation);
+	Ok(())
+}
+
+pub(crate) fn migrate<T: Config>() -> Weight {
+	log::info!("Migrating staking to StakingStorageVersion::V6");
+
+	let mut reads = 0u64;
+	let mut writes = 1u64;
+
+	// iter candidate pool and check whether any delegator has a cleared state
+	for candidate in CandidatePool::<T>::iter_values() {
+		reads = reads.saturating_add(1u64);
+		for delegation in candidate.delegators.into_iter() {
+			// we do not have to mutate existing entries since MaxCollatorsPerDelegator = 1
+			if !DelegatorState::<T>::contains_key(delegation.owner.clone()) {
+				if let Ok(delegator) = Delegator::try_new(candidate.id.clone(), delegation.amount) {
+					DelegatorState::<T>::insert(delegation.owner, delegator);
+					writes = writes.saturating_add(1u64);
+				}
+			}
+			reads = reads.saturating_add(1u64);
+		}
+	}
+
+	// update storage version
+	StorageVersion::<T>::put(StakingStorageVersion::V6);
+	log::info!("Completed staking migration to StakingStorageVersion::V6");
+
+	T::DbWeight::get().reads_writes(reads, writes)
+}
+
+#[cfg(feature = "try-runtime")]
+pub(crate) fn post_migrate<T: Config>() -> Result<(), &'static str> {
+	assert_eq!(StorageVersion::<T>::get(), StakingStorageVersion::V6);
+
+	for candidate in CandidatePool::<T>::iter_values() {
+		for delegation in candidate.delegators.into_iter() {
+			assert!(DelegatorState::<T>::contains_key(delegation.owner));
+		}
+	}
+
+	Ok(())
+}