Extrinsic to restore corrupt staking ledgers (#3706)

This PR adds a new extrinsic `Call::restore_ledger ` gated by
`StakingAdmin` origin that restores a corrupted staking ledger. This
extrinsic will be used to recover ledgers that were affected by the
issue discussed in
#3245.

The extrinsic will re-write the storage items associated with a stash
account provided as input parameter. The data used to reset the ledger
can be either i) fetched on-chain or ii) partially/totally set by the
input parameters of the call.

In order to use on-chain data to restore the staking locks, we need a
way to read the current lock in the balances pallet. This PR adds a
`InspectLockableCurrency` trait and implements it in the pallet
balances. An alternative would be to tightly couple staking with the
pallet balances but that's inelegant (an example of how it would look
like in [this
branch](https://github.com/paritytech/polkadot-sdk/tree/gpestana/ledger-badstate-clean_tightly)).

More details on the type of corruptions and corresponding fixes
https://hackmd.io/DLb5jEYWSmmvqXC9ae4yRg?view#/

We verified that the `Call::restore_ledger` does fix all current
corrupted ledgers in Polkadot and Kusama. You can verify it here
https://hackmd.io/v-XNrEoGRpe7APR-EZGhOA.

**Changes introduced**
- Adds `Call::restore_ledger ` extrinsic to recover a corrupted ledger;
- Adds trait `frame_support::traits::currency::InspectLockableCurrency`
to allow external pallets to read current locks given an account and
lock ID;
- Implements the `InspectLockableCurrency` in the pallet-balances.
- Adds staking locks try-runtime checks
(#3751)

**Todo**
- [x] benchmark `Call::restore_ledger`
- [x] throughout testing of all ledger recovering cases
- [x] consider adding the staking locks try-runtime checks to this PR
(#3751)
- [x] simulate restoring all ledgers
(https://hackmd.io/Dsa2tvhISNSs7zcqriTaxQ?view) in Polkadot and Kusama
using chopsticks -- https://hackmd.io/v-XNrEoGRpe7APR-EZGhOA

Related to #3245
Closes #3751

---------

Co-authored-by: command-bot <>

prdoc/pr_3706.prdoc

@@ -0,0 +1,20 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Extrinsic to restore corrupted staking ledgers
+
+doc:
+  - audience: Runtime User
+    description: |
+      This PR adds a new extrinsic `Call::restore_ledger ` gated by `StakingAdmin` origin that restores a corrupted staking ledger. This extrinsic will be used to recover ledgers that were affected by the issue discussed in https://github.com/paritytech/polkadot-sdk/issues/3245.
+      The extrinsic will re-write the storage items associated with a stash account provided as input parameter. The data used to reset the ledger can be either i) fetched on-chain or ii) partially/totally set by the input parameters of the call.
+
+      Changes introduced:
+        - Adds `Call::restore_ledger ` extrinsic to recover a corrupted ledger;
+        - Adds trait `frame_support::traits::currency::InspectLockableCurrency` to allow external pallets to read current locks given an account and lock ID;
+        - Implements the `InspectLockableCurrency` in the pallet-balances.
+        - Adds staking locks try-runtime checks (https://github.com/paritytech/polkadot-sdk/issues/3751)
+
+crates: 
+ - name: pallet-staking
+ - name: pallet-balances

substrate/frame/balances/src/impl_currency.rs

@@ -28,8 +28,8 @@ use frame_support::{
 		tokens::{fungible, BalanceStatus as Status, Fortitude::Polite, Precision::BestEffort},
 		Currency, DefensiveSaturating, ExistenceRequirement,
 		ExistenceRequirement::AllowDeath,
-		Get, Imbalance, LockIdentifier, LockableCurrency, NamedReservableCurrency,
-		ReservableCurrency, SignedImbalance, TryDrop, WithdrawReasons,
+		Get, Imbalance, InspectLockableCurrency, LockIdentifier, LockableCurrency,
+		NamedReservableCurrency, ReservableCurrency, SignedImbalance, TryDrop, WithdrawReasons,
 	},
 };
 use frame_system::pallet_prelude::BlockNumberFor;
@@ -918,3 +918,12 @@ where
 		Self::update_locks(who, &locks[..]);
 	}
 }
+
+impl<T: Config<I>, I: 'static> InspectLockableCurrency<T::AccountId> for Pallet<T, I> {
+	fn balance_locked(id: LockIdentifier, who: &T::AccountId) -> Self::Balance {
+		Self::locks(who)
+			.into_iter()
+			.filter(|l| l.id == id)
+			.fold(Zero::zero(), |acc, l| acc + l.amount)
+	}
+}

substrate/frame/balances/src/tests/currency_tests.rs

@@ -24,8 +24,8 @@ use frame_support::{
 		BalanceStatus::{Free, Reserved},
 		Currency,
 		ExistenceRequirement::{self, AllowDeath, KeepAlive},
-		Hooks, LockIdentifier, LockableCurrency, NamedReservableCurrency, ReservableCurrency,
-		WithdrawReasons,
+		Hooks, InspectLockableCurrency, LockIdentifier, LockableCurrency, NamedReservableCurrency,
+		ReservableCurrency, WithdrawReasons,
 	},
 	StorageNoopGuard,
 };
@@ -88,6 +88,24 @@ fn basic_locking_should_work() {
 		});
 }
 
+#[test]
+fn inspect_lock_should_work() {
+	ExtBuilder::default()
+		.existential_deposit(1)
+		.monied(true)
+		.build_and_execute_with(|| {
+			Balances::set_lock(ID_1, &1, 10, WithdrawReasons::all());
+			Balances::set_lock(ID_2, &1, 10, WithdrawReasons::all());
+			Balances::set_lock(ID_1, &2, 20, WithdrawReasons::all());
+
+			assert_eq!(<Balances as InspectLockableCurrency<_>>::balance_locked(ID_1, &1), 10);
+			assert_eq!(<Balances as InspectLockableCurrency<_>>::balance_locked(ID_2, &1), 10);
+			assert_eq!(<Balances as InspectLockableCurrency<_>>::balance_locked(ID_1, &2), 20);
+			assert_eq!(<Balances as InspectLockableCurrency<_>>::balance_locked(ID_2, &2), 0);
+			assert_eq!(<Balances as InspectLockableCurrency<_>>::balance_locked(ID_1, &3), 0);
+		})
+}
+
 #[test]
 fn account_should_be_reaped() {
 	ExtBuilder::default()

substrate/frame/staking/Cargo.toml

@@ -40,10 +40,10 @@ frame-benchmarking = { path = "../benchmarking", default-features = false, optio
 rand_chacha = { version = "0.2", default-features = false, optional = true }
 
 [dev-dependencies]
+pallet-balances = { path = "../balances" }
 sp-tracing = { path = "../../primitives/tracing" }
 sp-core = { path = "../../primitives/core" }
 sp-npos-elections = { path = "../../primitives/npos-elections" }
-pallet-balances = { path = "../balances" }
 pallet-timestamp = { path = "../timestamp" }
 pallet-staking-reward-curve = { path = "reward-curve" }
 pallet-bags-list = { path = "../bags-list" }

substrate/frame/staking/src/benchmarking.rs

@@ -953,6 +953,15 @@ benchmarks! {
 		assert_eq!(MinCommission::<T>::get(), Perbill::from_percent(100));
 	}
 
+	restore_ledger {
+		let (stash, controller) = create_stash_controller::<T>(0, 100, RewardDestination::Staked)?;
+		// corrupt ledger.
+		Ledger::<T>::remove(controller);
+	}: _(RawOrigin::Root, stash.clone(), None, None, None)
+	verify {
+		assert_eq!(Staking::<T>::inspect_bond_state(&stash), Ok(LedgerIntegrityState::Ok));
+	}
+
 	impl_benchmark_test_suite!(
 		Staking,
 		crate::mock::ExtBuilder::default().has_stakers(true),

substrate/frame/staking/src/lib.rs

@@ -491,6 +491,20 @@ pub struct StakingLedger<T: Config> {
 	controller: Option<T::AccountId>,
 }
 
+/// State of a ledger with regards with its data and metadata integrity.
+#[derive(PartialEq, Debug)]
+enum LedgerIntegrityState {
+	/// Ledger, bond and corresponding staking lock is OK.
+	Ok,
+	/// Ledger and/or bond is corrupted. This means that the bond has a ledger with a different
+	/// stash than the bonded stash.
+	Corrupted,
+	/// Ledger was corrupted and it has been killed.
+	CorruptedKilled,
+	/// Ledger and bond are OK, however the ledger's stash lock is out of sync.
+	LockCorrupted,
+}
+
 impl<T: Config> StakingLedger<T> {
 	/// Remove entries from `unlocking` that are sufficiently old and reduce the
 	/// total by the sum of their balances.

substrate/frame/staking/src/mock.rs

@@ -25,8 +25,8 @@ use frame_election_provider_support::{
 use frame_support::{
 	assert_ok, derive_impl, ord_parameter_types, parameter_types,
 	traits::{
-		ConstU64, Currency, EitherOfDiverse, FindAuthor, Get, Hooks, Imbalance, OnUnbalanced,
-		OneSessionHandler,
+		ConstU64, Currency, EitherOfDiverse, FindAuthor, Get, Hooks, Imbalance, LockableCurrency,
+		OnUnbalanced, OneSessionHandler, WithdrawReasons,
 	},
 	weights::constants::RocksDbWeight,
 };
@@ -786,55 +786,86 @@ pub(crate) fn bond_controller_stash(controller: AccountId, stash: AccountId) ->
 	Ok(())
 }
 
+// simulates `set_controller` without corrupted ledger checks for testing purposes.
+pub(crate) fn set_controller_no_checks(stash: &AccountId) {
+	let controller = Bonded::<Test>::get(stash).expect("testing stash should be bonded");
+	let ledger = Ledger::<Test>::get(&controller).expect("testing ledger should exist");
+
+	Ledger::<Test>::remove(&controller);
+	Ledger::<Test>::insert(stash, ledger);
+	Bonded::<Test>::insert(stash, stash);
+}
+
+// simulates `bond_extra` without corrupted ledger checks for testing purposes.
+pub(crate) fn bond_extra_no_checks(stash: &AccountId, amount: Balance) {
+	let controller = Bonded::<Test>::get(stash).expect("bond must exist to bond_extra");
+	let mut ledger = Ledger::<Test>::get(&controller).expect("ledger must exist to bond_extra");
+
+	let new_total = ledger.total + amount;
+	Balances::set_lock(crate::STAKING_ID, stash, new_total, WithdrawReasons::all());
+	ledger.total = new_total;
+	ledger.active = new_total;
+	Ledger::<Test>::insert(controller, ledger);
+}
+
 pub(crate) fn setup_double_bonded_ledgers() {
-	assert_ok!(Staking::bond(RuntimeOrigin::signed(1), 10, RewardDestination::Staked));
-	assert_ok!(Staking::bond(RuntimeOrigin::signed(2), 20, RewardDestination::Staked));
-	assert_ok!(Staking::bond(RuntimeOrigin::signed(3), 20, RewardDestination::Staked));
+	let init_ledgers = Ledger::<Test>::iter().count();
+
+	let _ = Balances::make_free_balance_be(&333, 2000);
+	let _ = Balances::make_free_balance_be(&444, 2000);
+	let _ = Balances::make_free_balance_be(&555, 2000);
+	let _ = Balances::make_free_balance_be(&777, 2000);
+
+	assert_ok!(Staking::bond(RuntimeOrigin::signed(333), 10, RewardDestination::Staked));
+	assert_ok!(Staking::bond(RuntimeOrigin::signed(444), 20, RewardDestination::Staked));
+	assert_ok!(Staking::bond(RuntimeOrigin::signed(555), 20, RewardDestination::Staked));
 	// not relevant to the test case, but ensures try-runtime checks pass.
-	[1, 2, 3]
+	[333, 444, 555]
 		.iter()
 		.for_each(|s| Payee::<Test>::insert(s, RewardDestination::Staked));
 
 	// we want to test the case where a controller can also be a stash of another ledger.
 	// for that, we change the controller/stash bonding so that:
-	// * 2 becomes controller of 1.
-	// * 3 becomes controller of 2.
-	// * 4 becomes controller of 3.
-	let ledger_1 = Ledger::<Test>::get(1).unwrap();
-	let ledger_2 = Ledger::<Test>::get(2).unwrap();
-	let ledger_3 = Ledger::<Test>::get(3).unwrap();
-
-	// 4 becomes controller of 3.
-	Bonded::<Test>::mutate(3, |controller| *controller = Some(4));
-	Ledger::<Test>::insert(4, ledger_3);
-
-	// 3 becomes controller of 2.
-	Bonded::<Test>::mutate(2, |controller| *controller = Some(3));
-	Ledger::<Test>::insert(3, ledger_2);
-
-	// 2 becomes controller of 1
-	Bonded::<Test>::mutate(1, |controller| *controller = Some(2));
-	Ledger::<Test>::insert(2, ledger_1);
-	// 1 is not controller anymore.
-	Ledger::<Test>::remove(1);
+	// * 444 becomes controller of 333.
+	// * 555 becomes controller of 444.
+	// * 777 becomes controller of 555.
+	let ledger_333 = Ledger::<Test>::get(333).unwrap();
+	let ledger_444 = Ledger::<Test>::get(444).unwrap();
+	let ledger_555 = Ledger::<Test>::get(555).unwrap();
+
+	// 777 becomes controller of 555.
+	Bonded::<Test>::mutate(555, |controller| *controller = Some(777));
+	Ledger::<Test>::insert(777, ledger_555);
+
+	// 555 becomes controller of 444.
+	Bonded::<Test>::mutate(444, |controller| *controller = Some(555));
+	Ledger::<Test>::insert(555, ledger_444);
+
+	// 444 becomes controller of 333.
+	Bonded::<Test>::mutate(333, |controller| *controller = Some(444));
+	Ledger::<Test>::insert(444, ledger_333);
+
+	// 333 is not controller anymore.
+	Ledger::<Test>::remove(333);
 
 	// checks. now we have:
-	// * 3 ledgers
-	assert_eq!(Ledger::<Test>::iter().count(), 3);
-	// * stash 1 has controller 2.
-	assert_eq!(Bonded::<Test>::get(1), Some(2));
-	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(1)), Some(2));
-	assert_eq!(Ledger::<Test>::get(2).unwrap().stash, 1);
-
-	// * stash 2 has controller 3.
-	assert_eq!(Bonded::<Test>::get(2), Some(3));
-	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(2)), Some(3));
-	assert_eq!(Ledger::<Test>::get(3).unwrap().stash, 2);
-
-	// * stash 3 has controller 4.
-	assert_eq!(Bonded::<Test>::get(3), Some(4));
-	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(3)), Some(4));
-	assert_eq!(Ledger::<Test>::get(4).unwrap().stash, 3);
+	// * +3 ledgers
+	assert_eq!(Ledger::<Test>::iter().count(), 3 + init_ledgers);
+
+	// * stash 333 has controller 444.
+	assert_eq!(Bonded::<Test>::get(333), Some(444));
+	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(333)), Some(444));
+	assert_eq!(Ledger::<Test>::get(444).unwrap().stash, 333);
+
+	// * stash 444 has controller 555.
+	assert_eq!(Bonded::<Test>::get(444), Some(555));
+	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(444)), Some(555));
+	assert_eq!(Ledger::<Test>::get(555).unwrap().stash, 444);
+
+	// * stash 555 has controller 777.
+	assert_eq!(Bonded::<Test>::get(555), Some(777));
+	assert_eq!(StakingLedger::<Test>::paired_account(StakingAccount::Stash(555)), Some(777));
+	assert_eq!(Ledger::<Test>::get(777).unwrap().stash, 555);
 }
 
 #[macro_export]

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

@@ -27,8 +27,8 @@ use frame_support::{
 	dispatch::WithPostDispatchInfo,
 	pallet_prelude::*,
 	traits::{
-		Currency, Defensive, DefensiveSaturating, EstimateNextNewSession, Get, Imbalance, Len,
-		OnUnbalanced, TryCollect, UnixTime,
+		Currency, Defensive, DefensiveSaturating, EstimateNextNewSession, Get, Imbalance,
+		InspectLockableCurrency, Len, OnUnbalanced, TryCollect, UnixTime,
 	},
 	weights::Weight,
 };
@@ -50,8 +50,8 @@ use sp_std::prelude::*;
 use crate::{
 	election_size_tracker::StaticTracker, log, slashing, weights::WeightInfo, ActiveEraInfo,
 	BalanceOf, EraInfo, EraPayout, Exposure, ExposureOf, Forcing, IndividualExposure,
-	MaxNominationsOf, MaxWinnersOf, Nominations, NominationsQuota, PositiveImbalanceOf,
-	RewardDestination, SessionInterface, StakingLedger, ValidatorPrefs,
+	LedgerIntegrityState, MaxNominationsOf, MaxWinnersOf, Nominations, NominationsQuota,
+	PositiveImbalanceOf, RewardDestination, SessionInterface, StakingLedger, ValidatorPrefs,
 };
 
 use super::pallet::*;
@@ -84,6 +84,38 @@ impl<T: Config> Pallet<T> {
 		StakingLedger::<T>::paired_account(Stash(stash.clone()))
 	}
 
+	/// Inspects and returns the corruption state of a ledger and bond, if any.
+	///
+	/// Note: all operations in this method access directly the `Bonded` and `Ledger` storage maps
+	/// instead of using the [`StakingLedger`] API since the bond and/or ledger may be corrupted.
+	pub(crate) fn inspect_bond_state(
+		stash: &T::AccountId,
+	) -> Result<LedgerIntegrityState, Error<T>> {
+		let lock = T::Currency::balance_locked(crate::STAKING_ID, &stash);
+
+		let controller = <Bonded<T>>::get(stash).ok_or_else(|| {
+			if lock == Zero::zero() {
+				Error::<T>::NotStash
+			} else {
+				Error::<T>::BadState
+			}
+		})?;
+
+		match Ledger::<T>::get(controller) {
+			Some(ledger) =>
+				if ledger.stash != *stash {
+					Ok(LedgerIntegrityState::Corrupted)
+				} else {
+					if lock != ledger.total {
+						Ok(LedgerIntegrityState::LockCorrupted)
+					} else {
+						Ok(LedgerIntegrityState::Ok)
+					}
+				},
+			None => Ok(LedgerIntegrityState::CorruptedKilled),
+		}
+	}
+
 	/// The total balance that can be slashed from a stash account as of right now.
 	pub fn slashable_balance_of(stash: &T::AccountId) -> BalanceOf<T> {
 		// Weight note: consider making the stake accessible through stash.
@@ -1837,12 +1869,12 @@ impl<T: Config> Pallet<T> {
 			"VoterList contains non-staker"
 		);
 
+		Self::check_ledgers()?;
 		Self::check_bonded_consistency()?;
 		Self::check_payees()?;
 		Self::check_nominators()?;
 		Self::check_exposures()?;
 		Self::check_paged_exposures()?;
-		Self::check_ledgers()?;
 		Self::check_count()
 	}
 
@@ -1851,6 +1883,7 @@ impl<T: Config> Pallet<T> {
 	/// * A bonded (stash, controller) pair should have only one associated ledger. I.e. if the
 	///   ledger is bonded by stash, the controller account must not bond a different ledger.
 	/// * A bonded (stash, controller) pair must have an associated ledger.
+	///
 	/// NOTE: these checks result in warnings only. Once
 	/// <https://github.com/paritytech/polkadot-sdk/issues/3245> is resolved, turn warns into check
 	/// failures.
@@ -1945,19 +1978,18 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Invariants:
-	/// * `ledger.controller` is not stored in the storage (but populated at retrieval).
 	/// * Stake consistency: ledger.total == ledger.active + sum(ledger.unlocking).
-	/// * The controller keying the ledger and the ledger stash matches the state of the `Bonded`
-	/// storage.
+	/// * The ledger's controller and stash matches the associated `Bonded` tuple.
+	/// * Staking locked funds for every bonded stash should be the same as its ledger's total.
+	/// * Staking ledger and bond are not corrupted.
 	fn check_ledgers() -> Result<(), TryRuntimeError> {
 		Bonded::<T>::iter()
 			.map(|(stash, ctrl)| {
-				// `ledger.controller` is never stored in raw storage.
-				let raw = Ledger::<T>::get(stash).unwrap_or_else(|| {
-					Ledger::<T>::get(ctrl.clone())
-						.expect("try_check: bonded stash/ctrl does not have an associated ledger")
-				});
-				ensure!(raw.controller.is_none(), "raw storage controller should be None");
+				// ensure locks consistency.
+				ensure!(
+					Self::inspect_bond_state(&stash) == Ok(LedgerIntegrityState::Ok),
+					"bond, ledger and/or staking lock inconsistent for a bonded stash."
+				);
 
 				// ensure ledger consistency.
 				Self::ensure_ledger_consistent(ctrl)