Fix account call for freeze/lock inconsistency (#1276)

* Fix account call for freeze/lock inconsistency

* Finish up fix

* change

* Comments

* Review comment

pallets/dapp-staking-v3/src/lib.rs

@@ -386,6 +386,8 @@ pub mod pallet {
         NoExpiredEntries,
         /// Force call is not allowed in production.
         ForceNotAllowed,
+        /// Account doesn't have the freeze inconsistency
+        AccountNotInconsistent, // TODO: can be removed after call `fix_account` is removed
     }
 
     /// General information about dApp staking protocol state.
@@ -931,28 +933,7 @@ pub mod pallet {
             Self::ensure_pallet_enabled()?;
             let account = ensure_signed(origin)?;
 
-            let mut ledger = Ledger::<T>::get(&account);
-
-            let current_block = frame_system::Pallet::<T>::block_number();
-            let amount = ledger.claim_unlocked(current_block.saturated_into());
-            ensure!(amount > Zero::zero(), Error::<T>::NoUnlockedChunksToClaim);
-
-            // In case it's full unlock, account is exiting dApp staking, ensure all storage is cleaned up.
-            let removed_entries = if ledger.is_empty() {
-                let _ = StakerInfo::<T>::clear_prefix(&account, ledger.contract_stake_count, None);
-                ledger.contract_stake_count
-            } else {
-                0
-            };
-
-            Self::update_ledger(&account, ledger)?;
-            CurrentEraInfo::<T>::mutate(|era_info| {
-                era_info.unlocking_removed(amount);
-            });
-
-            Self::deposit_event(Event::<T>::ClaimedUnlocked { account, amount });
-
-            Ok(Some(T::WeightInfo::claim_unlocked(removed_entries)).into())
+            Self::internal_claim_unlocked(account)
         }
 
         #[pallet::call_index(10)]
@@ -1603,6 +1584,50 @@ pub mod pallet {
 
             Ok(())
         }
+
+        /// A call used to fix accounts with inconsistent state, where frozen balance is actually higher than what's available.
+        ///
+        /// The approach is as simple as possible:
+        /// 1. Caller provides an account to fix.
+        /// 2. If account is eligible for the fix, all unlocking chunks are modified to be claimable immediately.
+        /// 3. The `claim_unlocked` call is executed using the provided account as the origin.
+        /// 4. All states are updated accordingly, and the account is no longer in an inconsistent state.
+        ///
+        /// The benchmarked weight of the `claim_unlocked` call is used as a base, and additional overestimated weight is added.
+        /// Call doesn't touch any storage items that aren't already touched by the `claim_unlocked` call, hence the simplified approach.
+        #[pallet::call_index(100)]
+        #[pallet::weight(T::DbWeight::get().reads_writes(4, 1))]
+        pub fn fix_account(
+            origin: OriginFor<T>,
+            account: T::AccountId,
+        ) -> DispatchResultWithPostInfo {
+            Self::ensure_pallet_enabled()?;
+            ensure_signed(origin)?;
+
+            let mut ledger = Ledger::<T>::get(&account);
+            let locked_amount_ledger = ledger.total_locked_amount();
+            let total_balance = T::Currency::total_balance(&account);
+
+            if locked_amount_ledger > total_balance {
+                // 1. Modify all unlocking chunks so they can be unlocked immediately.
+                let current_block: BlockNumber =
+                    frame_system::Pallet::<T>::block_number().saturated_into();
+                ledger
+                    .unlocking
+                    .iter_mut()
+                    .for_each(|chunk| chunk.unlock_block = current_block);
+                Ledger::<T>::insert(&account, ledger);
+
+                // 2. Execute the unlock call, clearing all of the unlocking chunks.
+                Self::internal_claim_unlocked(account)?;
+
+                // 3. In case of success, ensure no fee is paid.
+                Ok(Pays::No.into())
+            } else {
+                // The above logic is designed for a specific scenario and cannot be used otherwise.
+                Err(Error::<T>::AccountNotInconsistent.into())
+            }
+        }
     }
 
     impl<T: Config> Pallet<T> {
@@ -2125,5 +2150,30 @@ pub mod pallet {
             // It could end up being less than this weight, but this won't occur often enough to be important.
             T::WeightInfo::on_idle_cleanup()
         }
+
+        fn internal_claim_unlocked(account: T::AccountId) -> DispatchResultWithPostInfo {
+            let mut ledger = Ledger::<T>::get(&account);
+
+            let current_block = frame_system::Pallet::<T>::block_number();
+            let amount = ledger.claim_unlocked(current_block.saturated_into());
+            ensure!(amount > Zero::zero(), Error::<T>::NoUnlockedChunksToClaim);
+
+            // In case it's full unlock, account is exiting dApp staking, ensure all storage is cleaned up.
+            let removed_entries = if ledger.is_empty() {
+                let _ = StakerInfo::<T>::clear_prefix(&account, ledger.contract_stake_count, None);
+                ledger.contract_stake_count
+            } else {
+                0
+            };
+
+            Self::update_ledger(&account, ledger)?;
+            CurrentEraInfo::<T>::mutate(|era_info| {
+                era_info.unlocking_removed(amount);
+            });
+
+            Self::deposit_event(Event::<T>::ClaimedUnlocked { account, amount });
+
+            Ok(Some(T::WeightInfo::claim_unlocked(removed_entries)).into())
+        }
     }
 }

pallets/dapp-staking-v3/src/test/tests.rs

@@ -3302,3 +3302,75 @@ fn lock_correctly_considers_unlocking_amount() {
         );
     })
 }
+
+#[test]
+fn fix_account_scenarios_work() {
+    ExtBuilder::build().execute_with(|| {
+        // 1. Lock some amount correctly, unstake it, try to fix it, and ensure the call fails
+        let (account_1, lock_1) = (1, 100);
+        assert_lock(account_1, lock_1);
+        assert_noop!(
+            DappStaking::fix_account(RuntimeOrigin::signed(11), account_1),
+            Error::<Test>::AccountNotInconsistent
+        );
+
+        assert_unlock(account_1, lock_1);
+        assert_noop!(
+            DappStaking::fix_account(RuntimeOrigin::signed(11), account_1),
+            Error::<Test>::AccountNotInconsistent
+        );
+
+        // 2. Reproduce the issue where the account has more frozen than balance
+        let (account_2, unlock_2) = (2, 13);
+        let lock_2 = Balances::total_balance(&account_2);
+        assert_lock(account_2, lock_2);
+        assert_unlock(account_2, unlock_2);
+
+        // With the fix implemented, the scenario needs to be reproduced by hand.
+        // Account calls `lock`, specifying the amount that is undergoing the unlocking process.
+        // It can be either more or less, it doesn't matter for the test or the issue.
+
+        // But first, a sanity check.
+        assert_noop!(
+            DappStaking::lock(RuntimeOrigin::signed(account_2), unlock_2),
+            Error::<Test>::ZeroAmount,
+        );
+
+        // Now reproduce the incorrect lock/freeze operation.
+        let mut ledger = Ledger::<Test>::get(&account_2);
+        ledger.add_lock_amount(unlock_2);
+        assert_ok!(DappStaking::update_ledger(&account_2, ledger));
+        use crate::CurrentEraInfo;
+        CurrentEraInfo::<Test>::mutate(|era_info| {
+            era_info.add_locked(unlock_2);
+        });
+        assert!(
+            Balances::free_balance(&account_2)
+                < Ledger::<Test>::get(&account_2).total_locked_amount(),
+            "Sanity check."
+        );
+
+        // Now fix the account
+        assert_ok!(DappStaking::fix_account(
+            RuntimeOrigin::signed(11),
+            account_2
+        ));
+        System::assert_last_event(RuntimeEvent::DappStaking(Event::ClaimedUnlocked {
+            account: account_2,
+            amount: unlock_2,
+        }));
+
+        // Post-fix checks
+        assert_eq!(
+            Balances::free_balance(&account_2),
+            Ledger::<Test>::get(&account_2).total_locked_amount(),
+            "After the fix, balances should be equal."
+        );
+
+        // Cannot fix the same account again.
+        assert_noop!(
+            DappStaking::fix_account(RuntimeOrigin::signed(11), account_2),
+            Error::<Test>::AccountNotInconsistent
+        );
+    })
+}