pallet_balances: Add try_state for checking Holds and Freezes

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,27 @@ 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))]

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"));
+	});
+}