Add try-state hook to pallet aura

frame/aura/src/lib.rs

@@ -132,6 +132,11 @@ pub mod pallet {
 				T::DbWeight::get().reads(1)
 			}
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_: T::BlockNumber) -> Result<(), sp_runtime::TryRuntimeError> {
+			Self::do_try_state()
+		}
 	}
 
 	/// The current authority set.
@@ -217,6 +222,55 @@ impl<T: Config> Pallet<T> {
 		// the majority of its slot.
 		<T as pallet_timestamp::Config>::MinimumPeriod::get().saturating_mul(2u32.into())
 	}
+
+	/// Ensure the correctness of the state of this pallet.
+	///
+	/// This should be valid before or after each state transition of this pallet.
+	///
+	/// # Invariants
+	///
+	/// ## `CurrentSlot`
+	///
+	/// If we don't allow for multiple blocks per slot, then the current slot must be less than the
+	/// maximal slot number. Otherwise, it can be arbitrary.
+	///
+	/// ## `Authorities`
+	///
+	/// * The authorities must be non-empty.
+	/// * The current authority cannot be disabled.
+	/// * The number of authorities must be less than or equal to `T::MaxAuthorities`. This however,
+	///   is guarded by the type system.
+	#[cfg(any(test, feature = "try-runtime"))]
+	pub fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
+		// We don't have any guarantee that we are already after `on_initialize` and thus we have to
+		// check the current slot from the digest or take the last known slot.
+		let current_slot =
+			Self::current_slot_from_digests().unwrap_or_else(|| CurrentSlot::<T>::get());
+
+		// Check that the current slot is less than the maximal slot number, unless we allow for
+		// multiple blocks per slot.
+		if !T::AllowMultipleBlocksPerSlot::get() {
+			frame_support::ensure!(
+				current_slot < u64::MAX,
+				"Current slot has reached maximum value and cannot be incremented further.",
+			);
+		}
+
+		let authorities_len =
+			<Authorities<T>>::decode_len().ok_or("Failed to decode authorities length")?;
+
+		// Check that the authorities are non-empty.
+		frame_support::ensure!(!authorities_len.is_zero(), "Authorities must be non-empty.");
+
+		// Check that the current authority is not disabled.
+		let authority_index = *current_slot % authorities_len as u64;
+		frame_support::ensure!(
+			!T::DisabledValidators::is_disabled(authority_index as u32),
+			"Current validator is disabled and should not be attempting to author blocks.",
+		);
+
+		Ok(())
+	}
 }
 
 impl<T: Config> sp_runtime::BoundToRuntimeAppPublic for Pallet<T> {

frame/aura/src/mock.rs

@@ -110,12 +110,20 @@ impl pallet_aura::Config for Test {
 	type AllowMultipleBlocksPerSlot = AllowMultipleBlocksPerSlot;
 }
 
-pub fn new_test_ext(authorities: Vec<u64>) -> sp_io::TestExternalities {
-	let mut t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
+fn build_ext(authorities: Vec<u64>) -> sp_io::TestExternalities {
+	let mut storage = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
 	pallet_aura::GenesisConfig::<Test> {
 		authorities: authorities.into_iter().map(|a| UintAuthorityId(a).to_public_key()).collect(),
 	}
-	.assimilate_storage(&mut t)
+	.assimilate_storage(&mut storage)
 	.unwrap();
-	t.into()
+	storage.into()
+}
+
+pub fn build_ext_and_execute_test(authorities: Vec<u64>, test: impl FnOnce() -> ()) {
+	let mut ext = build_ext(authorities);
+	ext.execute_with(|| {
+		test();
+		Aura::do_try_state().expect("Storage invariants should hold")
+	});
 }

frame/aura/src/tests.rs

@@ -19,15 +19,15 @@
 
 #![cfg(test)]
 
-use crate::mock::{new_test_ext, Aura, MockDisabledValidators, System};
+use crate::mock::{build_ext_and_execute_test, Aura, MockDisabledValidators, System};
 use codec::Encode;
 use frame_support::traits::OnInitialize;
 use sp_consensus_aura::{Slot, AURA_ENGINE_ID};
 use sp_runtime::{Digest, DigestItem};
 
 #[test]
 fn initial_values() {
-	new_test_ext(vec![0, 1, 2, 3]).execute_with(|| {
+	build_ext_and_execute_test(vec![0, 1, 2, 3], || {
 		assert_eq!(Aura::current_slot(), 0u64);
 		assert_eq!(Aura::authorities().len(), 4);
 	});
@@ -38,7 +38,7 @@ fn initial_values() {
 	expected = "Validator with index 1 is disabled and should not be attempting to author blocks."
 )]
 fn disabled_validators_cannot_author_blocks() {
-	new_test_ext(vec![0, 1, 2, 3]).execute_with(|| {
+	build_ext_and_execute_test(vec![0, 1, 2, 3], || {
 		// slot 1 should be authored by validator at index 1
 		let slot = Slot::from(1);
 		let pre_digest =
@@ -58,7 +58,7 @@ fn disabled_validators_cannot_author_blocks() {
 #[test]
 #[should_panic(expected = "Slot must increase")]
 fn pallet_requires_slot_to_increase_unless_allowed() {
-	new_test_ext(vec![0, 1, 2, 3]).execute_with(|| {
+	build_ext_and_execute_test(vec![0, 1, 2, 3], || {
 		crate::mock::AllowMultipleBlocksPerSlot::set(false);
 
 		let slot = Slot::from(1);
@@ -76,7 +76,7 @@ fn pallet_requires_slot_to_increase_unless_allowed() {
 
 #[test]
 fn pallet_can_allow_unchanged_slot() {
-	new_test_ext(vec![0, 1, 2, 3]).execute_with(|| {
+	build_ext_and_execute_test(vec![0, 1, 2, 3], || {
 		let slot = Slot::from(1);
 		let pre_digest =
 			Digest { logs: vec![DigestItem::PreRuntime(AURA_ENGINE_ID, slot.encode())] };
@@ -95,7 +95,7 @@ fn pallet_can_allow_unchanged_slot() {
 #[test]
 #[should_panic(expected = "Slot must not decrease")]
 fn pallet_always_rejects_decreasing_slot() {
-	new_test_ext(vec![0, 1, 2, 3]).execute_with(|| {
+	build_ext_and_execute_test(vec![0, 1, 2, 3], || {
 		let slot = Slot::from(2);
 		let pre_digest =
 			Digest { logs: vec![DigestItem::PreRuntime(AURA_ENGINE_ID, slot.encode())] };