Try State Hook for Beefy (#3246)

Part of: https://github.com/paritytech/polkadot-sdk/issues/239

Polkadot address: 12GyGD3QhT4i2JJpNzvMf96sxxBLWymz4RdGCxRH5Rj5agKW

prdoc/pr_3246.prdoc

@@ -0,0 +1,11 @@
+title: Try State Hook for Beefy.
+
+doc:
+  - audience: Runtime User
+    description: |
+      Invariants for storage items in the beefy pallet. Enforces the following Invariants:
+      1. `Authorities` should not exceed the `MaxAuthorities` capacity.
+      2. `NextAuthorities` should not exceed the `MaxAuthorities` capacity.
+      3. `ValidatorSetId` must be present in `SetIdSession`.
+crates:
+- name: pallet-beefy

substrate/frame/beefy/src/lib.rs

@@ -280,6 +280,14 @@ pub mod pallet {
 		}
 	}
 
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
+			Self::do_try_state()
+		}
+	}
+
 	#[pallet::validate_unsigned]
 	impl<T: Config> ValidateUnsigned for Pallet<T> {
 		type Call = Call<T>;
@@ -294,6 +302,60 @@ pub mod pallet {
 	}
 }
 
+#[cfg(any(feature = "try-runtime", test))]
+impl<T: Config> Pallet<T> {
+	/// Ensure the correctness of the state of this pallet.
+	///
+	/// This should be valid before or after each state transition of this pallet.
+	pub fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
+		Self::try_state_authorities()?;
+		Self::try_state_validators()?;
+
+		Ok(())
+	}
+
+	/// # Invariants
+	///
+	/// * `Authorities` should not exceed the `MaxAuthorities` capacity.
+	/// * `NextAuthorities` should not exceed the `MaxAuthorities` capacity.
+	fn try_state_authorities() -> Result<(), sp_runtime::TryRuntimeError> {
+		if let Some(authorities_len) = <Authorities<T>>::decode_len() {
+			ensure!(
+				authorities_len as u32 <= T::MaxAuthorities::get(),
+				"Authorities number exceeds what the pallet config allows."
+			);
+		} else {
+			return Err(sp_runtime::TryRuntimeError::Other(
+				"Failed to decode length of authorities",
+			));
+		}
+
+		if let Some(next_authorities_len) = <NextAuthorities<T>>::decode_len() {
+			ensure!(
+				next_authorities_len as u32 <= T::MaxAuthorities::get(),
+				"Next authorities number exceeds what the pallet config allows."
+			);
+		} else {
+			return Err(sp_runtime::TryRuntimeError::Other(
+				"Failed to decode length of next authorities",
+			));
+		}
+		Ok(())
+	}
+
+	/// # Invariants
+	///
+	/// `ValidatorSetId` must be present in `SetIdSession`
+	fn try_state_validators() -> Result<(), sp_runtime::TryRuntimeError> {
+		let validator_set_id = <ValidatorSetId<T>>::get();
+		ensure!(
+			SetIdSession::<T>::get(validator_set_id).is_some(),
+			"Validator set id must be present in SetIdSession"
+		);
+		Ok(())
+	}
+}
+
 impl<T: Config> Pallet<T> {
 	/// Return the current active BEEFY validator set.
 	pub fn validator_set() -> Option<ValidatorSet<T::BeefyId>> {

substrate/frame/beefy/src/mock.rs

@@ -27,7 +27,6 @@ use frame_support::{
 };
 use pallet_session::historical as pallet_session_historical;
 use sp_core::{crypto::KeyTypeId, ConstU128};
-use sp_io::TestExternalities;
 use sp_runtime::{
 	app_crypto::ecdsa::Public, curve::PiecewiseLinear, impl_opaque_keys, testing::TestXt,
 	traits::OpaqueKeys, BuildStorage, Perbill,
@@ -210,6 +209,73 @@ impl pallet_offences::Config for Test {
 	type OnOffenceHandler = Staking;
 }
 
+#[derive(Default)]
+pub struct ExtBuilder {
+	authorities: Vec<BeefyId>,
+}
+
+impl ExtBuilder {
+	/// Add some AccountIds to insert into `List`.
+	#[cfg(test)]
+	pub(crate) fn add_authorities(mut self, ids: Vec<BeefyId>) -> Self {
+		self.authorities = ids;
+		self
+	}
+
+	pub fn build(self) -> sp_io::TestExternalities {
+		let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
+
+		let balances: Vec<_> =
+			(0..self.authorities.len()).map(|i| (i as u64, 10_000_000)).collect();
+
+		pallet_balances::GenesisConfig::<Test> { balances }
+			.assimilate_storage(&mut t)
+			.unwrap();
+
+		let session_keys: Vec<_> = self
+			.authorities
+			.iter()
+			.enumerate()
+			.map(|(i, k)| (i as u64, i as u64, MockSessionKeys { dummy: k.clone() }))
+			.collect();
+
+		BasicExternalities::execute_with_storage(&mut t, || {
+			for (ref id, ..) in &session_keys {
+				frame_system::Pallet::<Test>::inc_providers(id);
+			}
+		});
+
+		pallet_session::GenesisConfig::<Test> { keys: session_keys }
+			.assimilate_storage(&mut t)
+			.unwrap();
+
+		// controllers are same as stash
+		let stakers: Vec<_> = (0..self.authorities.len())
+			.map(|i| (i as u64, i as u64, 10_000, pallet_staking::StakerStatus::<u64>::Validator))
+			.collect();
+
+		let staking_config = pallet_staking::GenesisConfig::<Test> {
+			stakers,
+			validator_count: 2,
+			force_era: pallet_staking::Forcing::ForceNew,
+			minimum_validator_count: 0,
+			invulnerables: vec![],
+			..Default::default()
+		};
+
+		staking_config.assimilate_storage(&mut t).unwrap();
+
+		t.into()
+	}
+
+	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
+		self.build().execute_with(|| {
+			test();
+			Beefy::do_try_state().expect("All invariants must hold after a test");
+		})
+	}
+}
+
 // Note, that we can't use `UintAuthorityId` here. Reason is that the implementation
 // of `to_public_key()` assumes, that a public key is 32 bytes long. This is true for
 // ed25519 and sr25519 but *not* for ecdsa. A compressed ecdsa public key is 33 bytes,
@@ -226,54 +292,6 @@ pub fn mock_authorities(vec: Vec<u8>) -> Vec<BeefyId> {
 	vec.into_iter().map(|id| mock_beefy_id(id)).collect()
 }
 
-pub fn new_test_ext(ids: Vec<u8>) -> TestExternalities {
-	new_test_ext_raw_authorities(mock_authorities(ids))
-}
-
-pub fn new_test_ext_raw_authorities(authorities: Vec<BeefyId>) -> TestExternalities {
-	let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
-
-	let balances: Vec<_> = (0..authorities.len()).map(|i| (i as u64, 10_000_000)).collect();
-
-	pallet_balances::GenesisConfig::<Test> { balances }
-		.assimilate_storage(&mut t)
-		.unwrap();
-
-	let session_keys: Vec<_> = authorities
-		.iter()
-		.enumerate()
-		.map(|(i, k)| (i as u64, i as u64, MockSessionKeys { dummy: k.clone() }))
-		.collect();
-
-	BasicExternalities::execute_with_storage(&mut t, || {
-		for (ref id, ..) in &session_keys {
-			frame_system::Pallet::<Test>::inc_providers(id);
-		}
-	});
-
-	pallet_session::GenesisConfig::<Test> { keys: session_keys }
-		.assimilate_storage(&mut t)
-		.unwrap();
-
-	// controllers are same as stash
-	let stakers: Vec<_> = (0..authorities.len())
-		.map(|i| (i as u64, i as u64, 10_000, pallet_staking::StakerStatus::<u64>::Validator))
-		.collect();
-
-	let staking_config = pallet_staking::GenesisConfig::<Test> {
-		stakers,
-		validator_count: 2,
-		force_era: pallet_staking::Forcing::ForceNew,
-		minimum_validator_count: 0,
-		invulnerables: vec![],
-		..Default::default()
-	};
-
-	staking_config.assimilate_storage(&mut t).unwrap();
-
-	t.into()
-}
-
 pub fn start_session(session_index: SessionIndex) {
 	for i in Session::current_index()..session_index {
 		System::on_finalize(System::block_number());

substrate/frame/beefy/src/tests.rs

@@ -47,7 +47,7 @@ fn genesis_session_initializes_authorities() {
 	let authorities = mock_authorities(vec![1, 2, 3, 4]);
 	let want = authorities.clone();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		let authorities = beefy::Authorities::<Test>::get();
 
 		assert_eq!(authorities.len(), 4);
@@ -69,130 +69,140 @@ fn session_change_updates_authorities() {
 	let authorities = mock_authorities(vec![1, 2, 3, 4]);
 	let want_validators = authorities.clone();
 
-	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
-		assert!(0 == beefy::ValidatorSetId::<Test>::get());
+	ExtBuilder::default()
+		.add_authorities(mock_authorities(vec![1, 2, 3, 4]))
+		.build_and_execute(|| {
+			assert!(0 == beefy::ValidatorSetId::<Test>::get());
 
-		init_block(1);
+			init_block(1);
 
-		assert!(1 == beefy::ValidatorSetId::<Test>::get());
+			assert!(1 == beefy::ValidatorSetId::<Test>::get());
 
-		let want = beefy_log(ConsensusLog::AuthoritiesChange(
-			ValidatorSet::new(want_validators, 1).unwrap(),
-		));
+			let want = beefy_log(ConsensusLog::AuthoritiesChange(
+				ValidatorSet::new(want_validators, 1).unwrap(),
+			));
 
-		let log = System::digest().logs[0].clone();
-		assert_eq!(want, log);
+			let log = System::digest().logs[0].clone();
+			assert_eq!(want, log);
 
-		init_block(2);
+			init_block(2);
 
-		assert!(2 == beefy::ValidatorSetId::<Test>::get());
+			assert!(2 == beefy::ValidatorSetId::<Test>::get());
 
-		let want = beefy_log(ConsensusLog::AuthoritiesChange(
-			ValidatorSet::new(vec![mock_beefy_id(2), mock_beefy_id(4)], 2).unwrap(),
-		));
+			let want = beefy_log(ConsensusLog::AuthoritiesChange(
+				ValidatorSet::new(vec![mock_beefy_id(2), mock_beefy_id(4)], 2).unwrap(),
+			));
 
-		let log = System::digest().logs[1].clone();
-		assert_eq!(want, log);
-	});
+			let log = System::digest().logs[1].clone();
+			assert_eq!(want, log);
+		});
 }
 
 #[test]
 fn session_change_updates_next_authorities() {
 	let want = vec![mock_beefy_id(1), mock_beefy_id(2), mock_beefy_id(3), mock_beefy_id(4)];
 
-	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
-		let next_authorities = beefy::NextAuthorities::<Test>::get();
+	ExtBuilder::default()
+		.add_authorities(mock_authorities(vec![1, 2, 3, 4]))
+		.build_and_execute(|| {
+			let next_authorities = beefy::NextAuthorities::<Test>::get();
 
-		assert_eq!(next_authorities.len(), 4);
-		assert_eq!(want[0], next_authorities[0]);
-		assert_eq!(want[1], next_authorities[1]);
-		assert_eq!(want[2], next_authorities[2]);
-		assert_eq!(want[3], next_authorities[3]);
+			assert_eq!(next_authorities.len(), 4);
+			assert_eq!(want[0], next_authorities[0]);
+			assert_eq!(want[1], next_authorities[1]);
+			assert_eq!(want[2], next_authorities[2]);
+			assert_eq!(want[3], next_authorities[3]);
 
-		init_block(1);
+			init_block(1);
 
-		let next_authorities = beefy::NextAuthorities::<Test>::get();
+			let next_authorities = beefy::NextAuthorities::<Test>::get();
 
-		assert_eq!(next_authorities.len(), 2);
-		assert_eq!(want[1], next_authorities[0]);
-		assert_eq!(want[3], next_authorities[1]);
-	});
+			assert_eq!(next_authorities.len(), 2);
+			assert_eq!(want[1], next_authorities[0]);
+			assert_eq!(want[3], next_authorities[1]);
+		});
 }
 
 #[test]
 fn validator_set_at_genesis() {
 	let want = vec![mock_beefy_id(1), mock_beefy_id(2)];
 
-	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
-		let vs = Beefy::validator_set().unwrap();
+	ExtBuilder::default()
+		.add_authorities(mock_authorities(vec![1, 2, 3, 4]))
+		.build_and_execute(|| {
+			let vs = Beefy::validator_set().unwrap();
 
-		assert_eq!(vs.id(), 0u64);
-		assert_eq!(vs.validators()[0], want[0]);
-		assert_eq!(vs.validators()[1], want[1]);
-	});
+			assert_eq!(vs.id(), 0u64);
+			assert_eq!(vs.validators()[0], want[0]);
+			assert_eq!(vs.validators()[1], want[1]);
+		});
 }
 
 #[test]
 fn validator_set_updates_work() {
 	let want = vec![mock_beefy_id(1), mock_beefy_id(2), mock_beefy_id(3), mock_beefy_id(4)];
 
-	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
-		let vs = Beefy::validator_set().unwrap();
-		assert_eq!(vs.id(), 0u64);
-		assert_eq!(want[0], vs.validators()[0]);
-		assert_eq!(want[1], vs.validators()[1]);
-		assert_eq!(want[2], vs.validators()[2]);
-		assert_eq!(want[3], vs.validators()[3]);
+	ExtBuilder::default()
+		.add_authorities(mock_authorities(vec![1, 2, 3, 4]))
+		.build_and_execute(|| {
+			let vs = Beefy::validator_set().unwrap();
+			assert_eq!(vs.id(), 0u64);
+			assert_eq!(want[0], vs.validators()[0]);
+			assert_eq!(want[1], vs.validators()[1]);
+			assert_eq!(want[2], vs.validators()[2]);
+			assert_eq!(want[3], vs.validators()[3]);
 
-		init_block(1);
+			init_block(1);
 
-		let vs = Beefy::validator_set().unwrap();
+			let vs = Beefy::validator_set().unwrap();
 
-		assert_eq!(vs.id(), 1u64);
-		assert_eq!(want[0], vs.validators()[0]);
-		assert_eq!(want[1], vs.validators()[1]);
+			assert_eq!(vs.id(), 1u64);
+			assert_eq!(want[0], vs.validators()[0]);
+			assert_eq!(want[1], vs.validators()[1]);
 
-		init_block(2);
+			init_block(2);
 
-		let vs = Beefy::validator_set().unwrap();
+			let vs = Beefy::validator_set().unwrap();
 
-		assert_eq!(vs.id(), 2u64);
-		assert_eq!(want[1], vs.validators()[0]);
-		assert_eq!(want[3], vs.validators()[1]);
-	});
+			assert_eq!(vs.id(), 2u64);
+			assert_eq!(want[1], vs.validators()[0]);
+			assert_eq!(want[3], vs.validators()[1]);
+		});
 }
 
 #[test]
 fn cleans_up_old_set_id_session_mappings() {
-	new_test_ext(vec![1, 2, 3, 4]).execute_with(|| {
-		let max_set_id_session_entries = MaxSetIdSessionEntries::get();
+	ExtBuilder::default()
+		.add_authorities(mock_authorities(vec![1, 2, 3, 4]))
+		.build_and_execute(|| {
+			let max_set_id_session_entries = MaxSetIdSessionEntries::get();
 
-		// we have 3 sessions per era
-		let era_limit = max_set_id_session_entries / 3;
-		// sanity check against division precision loss
-		assert_eq!(0, max_set_id_session_entries % 3);
-		// go through `max_set_id_session_entries` sessions
-		start_era(era_limit);
+			// we have 3 sessions per era
+			let era_limit = max_set_id_session_entries / 3;
+			// sanity check against division precision loss
+			assert_eq!(0, max_set_id_session_entries % 3);
+			// go through `max_set_id_session_entries` sessions
+			start_era(era_limit);
 
-		// we should have a session id mapping for all the set ids from
-		// `max_set_id_session_entries` eras we have observed
-		for i in 1..=max_set_id_session_entries {
-			assert!(beefy::SetIdSession::<Test>::get(i as u64).is_some());
-		}
+			// we should have a session id mapping for all the set ids from
+			// `max_set_id_session_entries` eras we have observed
+			for i in 1..=max_set_id_session_entries {
+				assert!(beefy::SetIdSession::<Test>::get(i as u64).is_some());
+			}
 
-		// go through another `max_set_id_session_entries` sessions
-		start_era(era_limit * 2);
+			// go through another `max_set_id_session_entries` sessions
+			start_era(era_limit * 2);
 
-		// we should keep tracking the new mappings for new sessions
-		for i in max_set_id_session_entries + 1..=max_set_id_session_entries * 2 {
-			assert!(beefy::SetIdSession::<Test>::get(i as u64).is_some());
-		}
+			// we should keep tracking the new mappings for new sessions
+			for i in max_set_id_session_entries + 1..=max_set_id_session_entries * 2 {
+				assert!(beefy::SetIdSession::<Test>::get(i as u64).is_some());
+			}
 
-		// but the old ones should have been pruned by now
-		for i in 1..=max_set_id_session_entries {
-			assert!(beefy::SetIdSession::<Test>::get(i as u64).is_none());
-		}
-	});
+			// but the old ones should have been pruned by now
+			for i in 1..=max_set_id_session_entries {
+				assert!(beefy::SetIdSession::<Test>::get(i as u64).is_none());
+			}
+		});
 }
 
 /// Returns a list with 3 authorities with known keys:
@@ -259,7 +269,7 @@ fn should_sign_and_verify() {
 fn report_equivocation_current_set_works() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		assert_eq!(Staking::current_era(), Some(0));
 		assert_eq!(Session::current_index(), 0);
 
@@ -339,7 +349,7 @@ fn report_equivocation_current_set_works() {
 fn report_equivocation_old_set_works() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -422,7 +432,7 @@ fn report_equivocation_old_set_works() {
 fn report_equivocation_invalid_set_id() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -460,7 +470,7 @@ fn report_equivocation_invalid_set_id() {
 fn report_equivocation_invalid_session() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -503,7 +513,7 @@ fn report_equivocation_invalid_session() {
 fn report_equivocation_invalid_key_owner_proof() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -551,7 +561,7 @@ fn report_equivocation_invalid_key_owner_proof() {
 fn report_equivocation_invalid_equivocation_proof() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -624,7 +634,7 @@ fn report_equivocation_validate_unsigned_prevents_duplicates() {
 
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -728,7 +738,7 @@ fn report_equivocation_has_valid_weight() {
 fn valid_equivocation_reports_dont_pay_fees() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let block_num = System::block_number();
@@ -796,7 +806,7 @@ fn valid_equivocation_reports_dont_pay_fees() {
 fn set_new_genesis_works() {
 	let authorities = test_authorities();
 
-	new_test_ext_raw_authorities(authorities).execute_with(|| {
+	ExtBuilder::default().add_authorities(authorities).build_and_execute(|| {
 		start_era(1);
 
 		let new_genesis_delay = 10u64;