[staking] Currency Migration and Overstake fix (#7763)

## Summary
The existing fungible migration code has an issue when handling
partially unbonding accounts, leaving them in an inconsistent state.
These changes fix it by properly withdrawing overstake from unlock
chunks.

This PR also removes the `withdraw_overstake` extrinsic from
pallet-staking, as this scenario could only occur before the fungible
migration. With fungibles, over-staking is no longer possible.

## TODO
- [ ] Backport to stable2503.

---------

Co-authored-by: cmd[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: 2 people

polkadot/runtime/westend/src/tests.rs

@@ -109,6 +109,7 @@ fn check_treasury_pallet_id() {
 #[cfg(all(test, feature = "try-runtime"))]
 mod remote_tests {
 	use super::*;
+	use frame_support::traits::{TryState, TryStateSelect::All};
 	use frame_try_runtime::{runtime_decl_for_try_runtime::TryRuntime, UpgradeCheckSelect};
 	use remote_externalities::{
 		Builder, Mode, OfflineConfig, OnlineConfig, SnapshotConfig, Transport,
@@ -242,6 +243,10 @@ mod remote_tests {
 				unexpected_errors
 			);
 		});
+
+		ext.execute_with(|| {
+			AllPalletsWithSystem::try_state(System::block_number(), All).unwrap();
+		});
 	}
 
 	#[tokio::test]
@@ -313,6 +318,10 @@ mod remote_tests {
 				force_withdraw_acc
 			);
 		});
+
+		ext.execute_with(|| {
+			AllPalletsWithSystem::try_state(System::block_number(), All).unwrap();
+		});
 	}
 }
 

prdoc/pr_7763.prdoc

@@ -0,0 +1,12 @@
+title: '[staking] Currency Migration and Overstake fix'
+doc:
+- audience: Runtime Dev
+  description: |-
+    - Fixes Currency to fungible migration with force-unstake of partial unbonding accounts.
+    - Removes the extrinsic `withdraw_overstake` which is not useful post fungibe migration of pallet-staking.
+
+crates:
+- name: westend-runtime
+  bump: major
+- name: pallet-staking
+  bump: major

substrate/frame/staking/src/lib.rs

@@ -632,6 +632,52 @@ impl<T: Config> StakingLedger<T> {
 		}
 	}
 
+	/// Sets ledger total to the `new_total`.
+	///
+	/// Removes entries from `unlocking` upto `amount` starting from the oldest first.
+	fn update_total_stake(mut self, new_total: BalanceOf<T>) -> Self {
+		let old_total = self.total;
+		self.total = new_total;
+		debug_assert!(
+			new_total <= old_total,
+			"new_total {:?} must be <= old_total {:?}",
+			new_total,
+			old_total
+		);
+
+		let to_withdraw = old_total.defensive_saturating_sub(new_total);
+		// accumulator to keep track of how much is withdrawn.
+		// First we take out from active.
+		let mut withdrawn = BalanceOf::<T>::zero();
+
+		// first we try to remove stake from active
+		if self.active >= to_withdraw {
+			self.active -= to_withdraw;
+			return self
+		} else {
+			withdrawn += self.active;
+			self.active = BalanceOf::<T>::zero();
+		}
+
+		// start removing from the oldest chunk.
+		while let Some(last) = self.unlocking.last_mut() {
+			if withdrawn.defensive_saturating_add(last.value) <= to_withdraw {
+				withdrawn += last.value;
+				self.unlocking.pop();
+			} else {
+				let diff = to_withdraw.defensive_saturating_sub(withdrawn);
+				withdrawn += diff;
+				last.value -= diff;
+			}
+
+			if withdrawn >= to_withdraw {
+				break
+			}
+		}
+
+		self
+	}
+
 	/// Re-bond funds that were scheduled for unlocking.
 	///
 	/// Returns the updated ledger, and the amount actually rebonded.

substrate/frame/staking/src/pallet/impls.rs

@@ -1356,18 +1356,20 @@ impl<T: Config> Pallet<T> {
 		} else {
 			// if we are here, it means we cannot hold all user stake. We will do a force withdraw
 			// from ledger, but that's okay since anyways user do not have funds for it.
-			let force_withdraw = staked.saturating_sub(max_hold);
-
-			// we ignore if active is 0. It implies the locked amount is not actively staked. The
-			// account can still get away from potential slash but we can't do much better here.
-			StakingLedger {
-				total: max_hold,
-				active: ledger.active.saturating_sub(force_withdraw),
-				// we are not changing the stash, so we can keep the stash.
-				..ledger
-			}
-			.update()?;
-			force_withdraw
+
+			let old_total = ledger.total;
+			// update ledger with total stake as max_hold.
+			let updated_ledger = ledger.update_total_stake(max_hold);
+
+			// new total stake in ledger.
+			let new_total = updated_ledger.total;
+			debug_assert_eq!(new_total, max_hold);
+
+			// update ledger in storage.
+			updated_ledger.update()?;
+
+			// return the diff
+			old_total.defensive_saturating_sub(new_total)
 		};
 
 		// Get rid of the extra consumer we used to have with OldCurrency.
@@ -1377,20 +1379,31 @@ impl<T: Config> Pallet<T> {
 		Ok(())
 	}
 
+	// These are system accounts and don’t normally hold funds, so migration isn’t strictly
+	// necessary. However, this is a good opportunity to clean up the extra consumer/providers that
+	// were previously used.
 	fn do_migrate_virtual_staker(stash: &T::AccountId) -> DispatchResult {
-		// Funds for virtual stakers not managed/held by this pallet. We only need to clear
-		// the extra consumer we used to have with OldCurrency.
-		frame_system::Pallet::<T>::dec_consumers(&stash);
+		let consumer_count = frame_system::Pallet::<T>::consumers(stash);
+		// fail early if no consumers.
+		ensure!(consumer_count > 0, Error::<T>::AlreadyMigrated);
+
+		// provider/consumer ref count has been a mess (inconsistent), and some of these accounts
+		// accumulated upto 2 consumers. But if it's more than 2, we simply fail to not allow
+		// this migration to be called multiple times.
+		ensure!(consumer_count <= 2, Error::<T>::BadState);
+
+		// get rid of the consumers
+		for _ in 0..consumer_count {
+			frame_system::Pallet::<T>::dec_consumers(&stash);
+		}
 
-		// The delegation system that manages the virtual staker needed to increment provider
-		// previously because of the consumer needed by this pallet. In reality, this stash
-		// is just a key for managing the ledger and the account does not need to hold any
-		// balance or exist. We decrement this provider.
+		// get the current count of providers
 		let actual_providers = frame_system::Pallet::<T>::providers(stash);
 
 		let expected_providers =
-			// provider is expected to be 1 but someone can always transfer some free funds to
-			// these accounts, increasing the provider.
+			// We expect these accounts to have only one provider, and hold no balance. However, if
+			// someone mischievously sends some funds to these accounts, they may have an additional
+			// provider, which we can safely ignore.
 			if asset::free_to_stake::<T>(&stash) >= asset::existential_deposit::<T>() {
 				2
 			} else {
@@ -1403,10 +1416,9 @@ impl<T: Config> Pallet<T> {
 		// if actual provider is less than expected, it is already migrated.
 		ensure!(actual_providers == expected_providers, Error::<T>::AlreadyMigrated);
 
-		// dec provider
 		let _ = frame_system::Pallet::<T>::dec_providers(&stash)?;
 
-		return Ok(())
+		Ok(())
 	}
 }
 

substrate/frame/staking/src/pallet/mod.rs

@@ -2607,37 +2607,5 @@ pub mod pallet {
 
 			Ok(Pays::No.into())
 		}
-
-		/// Adjusts the staking ledger by withdrawing any excess staked amount.
-		///
-		/// This function corrects cases where a user's recorded stake in the ledger
-		/// exceeds their actual staked funds. This situation can arise due to cases such as
-		/// external slashing by another pallet, leading to an inconsistency between the ledger
-		/// and the actual stake.
-		#[pallet::call_index(32)]
-		#[pallet::weight(T::DbWeight::get().reads_writes(2, 1))]
-		pub fn withdraw_overstake(origin: OriginFor<T>, stash: T::AccountId) -> DispatchResult {
-			let _ = ensure_signed(origin)?;
-
-			let ledger = Self::ledger(Stash(stash.clone()))?;
-			let actual_stake = asset::staked::<T>(&stash);
-			let force_withdraw_amount = ledger.total.defensive_saturating_sub(actual_stake);
-
-			// ensure there is something to force unstake.
-			ensure!(!force_withdraw_amount.is_zero(), Error::<T>::BoundNotMet);
-
-			// we ignore if active is 0. It implies the locked amount is not actively staked. The
-			// account can still get away from potential slash, but we can't do much better here.
-			StakingLedger {
-				total: actual_stake,
-				active: ledger.active.saturating_sub(force_withdraw_amount),
-				..ledger
-			}
-			.update()?;
-
-			Self::deposit_event(Event::<T>::Withdrawn { stash, amount: force_withdraw_amount });
-
-			Ok(())
-		}
 	}
 }

substrate/frame/staking/src/tests.rs

@@ -49,7 +49,7 @@ use sp_runtime::{
 };
 use sp_staking::{
 	offence::{OffenceDetails, OnOffenceHandler},
-	SessionIndex, Stake, StakingInterface,
+	SessionIndex, StakingInterface,
 };
 use substrate_test_utils::assert_eq_uvec;
 
@@ -5099,72 +5099,6 @@ fn restricted_accounts_can_only_withdraw() {
 	})
 }
 
-#[test]
-fn permissionless_withdraw_overstake() {
-	ExtBuilder::default().build_and_execute(|| {
-		// Given Alice, Bob and Charlie with some stake.
-		let alice = 301;
-		let bob = 302;
-		let charlie = 303;
-		let _ = Balances::make_free_balance_be(&alice, 500);
-		let _ = Balances::make_free_balance_be(&bob, 500);
-		let _ = Balances::make_free_balance_be(&charlie, 500);
-		assert_ok!(Staking::bond(RuntimeOrigin::signed(alice), 100, RewardDestination::Staked));
-		assert_ok!(Staking::bond(RuntimeOrigin::signed(bob), 100, RewardDestination::Staked));
-		assert_ok!(Staking::bond(RuntimeOrigin::signed(charlie), 100, RewardDestination::Staked));
-
-		// WHEN: charlie is partially unbonding.
-		assert_ok!(Staking::unbond(RuntimeOrigin::signed(charlie), 90));
-		let charlie_ledger = StakingLedger::<Test>::get(StakingAccount::Stash(charlie)).unwrap();
-
-		// AND: alice and charlie ledger having higher value than actual stake.
-		Ledger::<Test>::insert(alice, StakingLedger::<Test>::new(alice, 200));
-		Ledger::<Test>::insert(
-			charlie,
-			StakingLedger { stash: charlie, total: 200, active: 200 - 90, ..charlie_ledger },
-		);
-
-		// THEN overstake can be permissionlessly withdrawn.
-		System::reset_events();
-
-		// Alice stake is corrected.
-		assert_eq!(
-			<Staking as StakingInterface>::stake(&alice).unwrap(),
-			Stake { total: 200, active: 200 }
-		);
-		assert_ok!(Staking::withdraw_overstake(RuntimeOrigin::signed(1), alice));
-		assert_eq!(
-			<Staking as StakingInterface>::stake(&alice).unwrap(),
-			Stake { total: 100, active: 100 }
-		);
-
-		// Charlie who is partially withdrawing also gets their stake corrected.
-		assert_eq!(
-			<Staking as StakingInterface>::stake(&charlie).unwrap(),
-			Stake { total: 200, active: 110 }
-		);
-		assert_ok!(Staking::withdraw_overstake(RuntimeOrigin::signed(1), charlie));
-		assert_eq!(
-			<Staking as StakingInterface>::stake(&charlie).unwrap(),
-			Stake { total: 200 - 100, active: 110 - 100 }
-		);
-
-		assert_eq!(
-			staking_events_since_last_call(),
-			vec![
-				Event::Withdrawn { stash: alice, amount: 200 - 100 },
-				Event::Withdrawn { stash: charlie, amount: 200 - 100 }
-			]
-		);
-
-		// but Bob ledger is fine and that cannot be withdrawn.
-		assert_noop!(
-			Staking::withdraw_overstake(RuntimeOrigin::signed(1), bob),
-			Error::<Test>::BoundNotMet
-		);
-	});
-}
-
 mod election_data_provider {
 	use super::*;
 	use frame_election_provider_support::ElectionDataProvider;
@@ -9071,6 +9005,92 @@ mod hold_migration {
 		});
 	}
 
+	#[test]
+	fn overstaked_and_partially_unbonding() {
+		ExtBuilder::default().has_stakers(true).build_and_execute(|| {
+			// GIVEN alice who is a nominator with T::OldCurrency.
+			let alice = 300;
+			// 1000 + ED
+			let _ = Balances::make_free_balance_be(&alice, 1001);
+			let stake = 600;
+			let reserved_by_another_pallet = 400;
+			assert_ok!(Staking::bond(
+				RuntimeOrigin::signed(alice),
+				stake,
+				RewardDestination::Staked
+			));
+
+			// AND Alice is partially unbonding.
+			assert_ok!(Staking::unbond(RuntimeOrigin::signed(alice), 300));
+
+			// AND Alice has some funds reserved with another pallet.
+			assert_ok!(Balances::reserve(&alice, reserved_by_another_pallet));
+
+			// convert stake to T::OldCurrency.
+			testing_utils::migrate_to_old_currency::<Test>(alice);
+			assert_eq!(asset::staked::<Test>(&alice), 0);
+			assert_eq!(Balances::balance_locked(STAKING_ID, &alice), stake);
+
+			// ledger has correct amount staked.
+			assert_eq!(
+				<Staking as StakingInterface>::stake(&alice),
+				Ok(Stake { total: stake, active: stake - 300 })
+			);
+
+			// Alice becomes overstaked by withdrawing some staked balance.
+			assert_ok!(Balances::transfer_allow_death(
+				RuntimeOrigin::signed(alice),
+				10,
+				reserved_by_another_pallet
+			));
+
+			let expected_force_withdraw = reserved_by_another_pallet;
+
+			// ledger mutation would fail in this case before migration because of failing hold.
+			assert_noop!(
+				Staking::unbond(RuntimeOrigin::signed(alice), 100),
+				Error::<Test>::NotEnoughFunds
+			);
+
+			// clear events
+			System::reset_events();
+
+			// WHEN alice currency is migrated.
+			assert_ok!(Staking::migrate_currency(RuntimeOrigin::signed(1), alice));
+
+			// THEN
+			let expected_hold = stake - expected_force_withdraw;
+			// ensure no lock
+			assert_eq!(Balances::balance_locked(STAKING_ID, &alice), 0);
+			// ensure stake and hold are same.
+			assert_eq!(
+				<Staking as StakingInterface>::stake(&alice),
+				// expected stake is 0 since force withdrawn (400) is taken out completely of
+				// active stake.
+				Ok(Stake { total: expected_hold, active: 0 })
+			);
+
+			assert_eq!(asset::staked::<Test>(&alice), expected_hold);
+			// ensure events are emitted.
+			assert_eq!(
+				staking_events_since_last_call(),
+				vec![Event::CurrencyMigrated {
+					stash: alice,
+					force_withdraw: expected_force_withdraw
+				}]
+			);
+
+			// ensure cannot migrate again.
+			assert_noop!(
+				Staking::migrate_currency(RuntimeOrigin::signed(1), alice),
+				Error::<Test>::AlreadyMigrated
+			);
+
+			// unbond works after migration.
+			assert_ok!(Staking::unbond(RuntimeOrigin::signed(alice), 100));
+		});
+	}
+
 	#[test]
 	fn virtual_staker_consumer_provider_dec() {
 		// Ensure virtual stakers consumer and provider count is decremented.