From f56208c93fbd58f2e412d103dad6d90cd65e5606 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Fri, 17 May 2024 13:35:15 +0200
Subject: [PATCH] pallet_balances: Add `try_state` for checking `Holds` and
 `Freezes` (#4490)

Co-authored-by: command-bot <>
---
 substrate/frame/balances/src/lib.rs           | 25 ++++++++++++++-
 .../frame/balances/src/tests/general_tests.rs | 32 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/substrate/frame/balances/src/lib.rs b/substrate/frame/balances/src/lib.rs
index 8d904d3d21b8a..56eb81b49e2db 100644
--- a/substrate/frame/balances/src/lib.rs
+++ b/substrate/frame/balances/src/lib.rs
@@ -542,8 +542,8 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
-		#[cfg(not(feature = "insecure_zero_ed"))]
 		fn integrity_test() {
+			#[cfg(not(feature = "insecure_zero_ed"))]
 			assert!(
 				!<T as Config<I>>::ExistentialDeposit::get().is_zero(),
 				"The existential deposit must be greater than zero!"
@@ -555,6 +555,29 @@ pub mod pallet {
 				T::MaxFreezes::get(), <T::RuntimeFreezeReason as VariantCount>::VARIANT_COUNT,
 			);
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
+			Holds::<T, I>::iter_keys().try_for_each(|k| {
+				if Holds::<T, I>::decode_len(k).unwrap_or(0) >
+					T::RuntimeHoldReason::VARIANT_COUNT as usize
+				{
+					Err("Found `Hold` with too many elements")
+				} else {
+					Ok(())
+				}
+			})?;
+
+			Freezes::<T, I>::iter_keys().try_for_each(|k| {
+				if Freezes::<T, I>::decode_len(k).unwrap_or(0) > T::MaxFreezes::get() as usize {
+					Err("Found `Freeze` with too many elements")
+				} else {
+					Ok(())
+				}
+			})?;
+
+			Ok(())
+		}
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
diff --git a/substrate/frame/balances/src/tests/general_tests.rs b/substrate/frame/balances/src/tests/general_tests.rs
index 0f3e015d0a892..a855fae5616af 100644
--- a/substrate/frame/balances/src/tests/general_tests.rs
+++ b/substrate/frame/balances/src/tests/general_tests.rs
@@ -109,3 +109,35 @@ fn regression_historic_acc_does_not_evaporate_reserve() {
 		});
 	});
 }
+
+#[cfg(feature = "try-runtime")]
+#[test]
+fn try_state_works() {
+	use crate::{Config, Freezes, Holds};
+	use frame_support::{
+		storage,
+		traits::{Get, Hooks, VariantCount},
+	};
+
+	ExtBuilder::default().build_and_execute_with(|| {
+		storage::unhashed::put(
+			&Holds::<Test>::hashed_key_for(1),
+			&vec![0u8; <Test as Config>::RuntimeHoldReason::VARIANT_COUNT as usize + 1],
+		);
+
+		assert!(format!("{:?}", Balances::try_state(0).unwrap_err())
+			.contains("Found `Hold` with too many elements"));
+	});
+
+	ExtBuilder::default().build_and_execute_with(|| {
+		let max_freezes: u32 = <Test as Config>::MaxFreezes::get();
+
+		storage::unhashed::put(
+			&Freezes::<Test>::hashed_key_for(1),
+			&vec![0u8; max_freezes as usize + 1],
+		);
+
+		assert!(format!("{:?}", Balances::try_state(0).unwrap_err())
+			.contains("Found `Freeze` with too many elements"));
+	});
+}