Validate ancestry proofs generated at older blocks

substrate/frame/beefy-mmr/src/lib.rs

@@ -33,21 +33,22 @@
 //!
 //! and thanks to versioning can be easily updated in the future.
 
-use sp_runtime::traits::{Convert, Member};
+use sp_runtime::traits::{Convert, Header, Member};
 use sp_std::prelude::*;
 
 use codec::Decode;
 use pallet_mmr::{primitives::AncestryProof, LeafDataProvider, ParentNumberAndHash};
 use sp_consensus_beefy::{
 	known_payloads,
 	mmr::{BeefyAuthoritySet, BeefyDataProvider, BeefyNextAuthoritySet, MmrLeaf, MmrLeafVersion},
-	AncestryHelper, Commitment, ValidatorSet as BeefyValidatorSet,
+	AncestryHelper, Commitment, ConsensusLog, ValidatorSet as BeefyValidatorSet,
 };
 
 use frame_support::{crypto::ecdsa::ECDSAExt, traits::Get};
-use frame_system::pallet_prelude::BlockNumberFor;
+use frame_system::pallet_prelude::{BlockNumberFor, HeaderFor};
 
 pub use pallet::*;
+use sp_runtime::generic::OpaqueDigestItemId;
 
 #[cfg(test)]
 mod mock;
@@ -173,10 +174,35 @@ where
 	}
 }
 
-impl<T: Config> AncestryHelper<BlockNumberFor<T>> for Pallet<T> {
+impl<T: Config> AncestryHelper<HeaderFor<T>> for Pallet<T>
+where
+	T: pallet_mmr::Config<Hashing = sp_consensus_beefy::MmrHashing>,
+{
 	type Proof = AncestryProof<MerkleRootOf<T>>;
+	type ValidationContext = MerkleRootOf<T>;
+
+	fn extract_validation_context(header: HeaderFor<T>) -> Option<Self::ValidationContext> {
+		// Check if the provided header is canonical.
+		let expected_hash = frame_system::Pallet::<T>::block_hash(header.number());
+		if expected_hash != header.hash() {
+			return None;
+		}
+
+		// Extract the MMR root from the header digest
+		header.digest().convert_first(|l| {
+			l.try_to(OpaqueDigestItemId::Consensus(&sp_consensus_beefy::BEEFY_ENGINE_ID))
+				.and_then(|log: ConsensusLog<<T as pallet_beefy::Config>::BeefyId>| match log {
+					ConsensusLog::MmrRoot(mmr_root) => Some(mmr_root),
+					_ => None,
+				})
+		})
+	}
 
-	fn is_non_canonical(commitment: &Commitment<BlockNumberFor<T>>, proof: Self::Proof) -> bool {
+	fn is_non_canonical(
+		commitment: &Commitment<BlockNumberFor<T>>,
+		proof: Self::Proof,
+		context: Self::ValidationContext,
+	) -> bool {
 		let commitment_leaf_count =
 			match pallet_mmr::Pallet::<T>::block_num_to_leaf_count(commitment.block_number) {
 				Ok(commitment_leaf_count) => commitment_leaf_count,
@@ -192,7 +218,7 @@ impl<T: Config> AncestryHelper<BlockNumberFor<T>> for Pallet<T> {
 			return false;
 		}
 
-		let canonical_mmr_root = pallet_mmr::Pallet::<T>::mmr_root();
+		let canonical_mmr_root = context;
 		let canonical_prev_root =
 			match pallet_mmr::Pallet::<T>::verify_ancestry_proof(canonical_mmr_root, proof) {
 				Ok(canonical_prev_root) => canonical_prev_root,

substrate/frame/beefy-mmr/src/tests.rs

@@ -35,9 +35,9 @@ use frame_support::traits::OnInitialize;
 
 use crate::mock::*;
 
-fn init_block(block: u64) {
-	let hash = H256::repeat_byte(block as u8);
-	System::initialize(&block, &hash, &Default::default());
+fn init_block(block: u64, maybe_parent_hash: Option<H256>) {
+	let parent_hash = maybe_parent_hash.unwrap_or(H256::repeat_byte(block as u8));
+	System::initialize(&block, &parent_hash, &Default::default());
 	Session::on_initialize(block);
 	Mmr::on_initialize(block);
 	Beefy::on_initialize(block);
@@ -66,7 +66,7 @@ fn read_mmr_leaf(ext: &mut TestExternalities, key: Vec<u8>) -> MmrLeaf {
 fn should_contain_mmr_digest() {
 	let mut ext = new_test_ext(vec![1, 2, 3, 4]);
 	ext.execute_with(|| {
-		init_block(1);
+		init_block(1, None);
 		assert_eq!(
 			System::digest().logs,
 			vec![
@@ -81,7 +81,7 @@ fn should_contain_mmr_digest() {
 		);
 
 		// unique every time
-		init_block(2);
+		init_block(2, None);
 		assert_eq!(
 			System::digest().logs,
 			vec![
@@ -105,7 +105,7 @@ fn should_contain_valid_leaf_data() {
 
 	let mut ext = new_test_ext(vec![1, 2, 3, 4]);
 	let parent_hash = ext.execute_with(|| {
-		init_block(1);
+		init_block(1, None);
 		frame_system::Pallet::<Test>::parent_hash()
 	});
 
@@ -130,7 +130,7 @@ fn should_contain_valid_leaf_data() {
 
 	// build second block on top
 	let parent_hash = ext.execute_with(|| {
-		init_block(2);
+		init_block(2, None);
 		frame_system::Pallet::<Test>::parent_hash()
 	});
 
@@ -174,7 +174,7 @@ fn should_update_authorities() {
 		assert_eq!(auth_set.keyset_commitment, next_auth_set.keyset_commitment);
 
 		let announced_set = next_auth_set;
-		init_block(1);
+		init_block(1, None);
 		let auth_set = BeefyMmr::authority_set_proof();
 		let next_auth_set = BeefyMmr::next_authority_set_proof();
 
@@ -190,7 +190,7 @@ fn should_update_authorities() {
 		assert_eq!(want, next_auth_set.keyset_commitment);
 
 		let announced_set = next_auth_set;
-		init_block(2);
+		init_block(2, None);
 		let auth_set = BeefyMmr::authority_set_proof();
 		let next_auth_set = BeefyMmr::next_authority_set_proof();
 
@@ -207,14 +207,56 @@ fn should_update_authorities() {
 	});
 }
 
+#[test]
+fn extract_validation_context_should_work_correctly() {
+	let mut ext = new_test_ext(vec![1, 2]);
+
+	// Register offchain ext.
+	let (offchain, _offchain_state) = TestOffchainExt::with_offchain_db(ext.offchain_db());
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
+	ext.register_extension(OffchainWorkerExt::new(offchain));
+
+	ext.execute_with(|| {
+		init_block(1, None);
+		let h1 = System::finalize();
+		init_block(2, Some(h1.hash()));
+		let h2 = System::finalize();
+
+		// Check the MMR root log
+		let expected_mmr_root: [u8; 32] = array_bytes::hex_n_into_unchecked(
+			"b2106eff9894288bc212b3a9389caa54efd37962c3a7b71b3b0b06a0911b88a5",
+		);
+		assert_eq!(
+			System::digest().logs,
+			vec![beefy_log(ConsensusLog::MmrRoot(H256::from_slice(&expected_mmr_root)))]
+		);
+
+		// Make sure that all the info about h2 was stored on-chain
+		init_block(3, Some(h2.hash()));
+
+		// `extract_validation_context` should return the MMR root when the provided header
+		// is part of the chain,
+		assert_eq!(
+			BeefyMmr::extract_validation_context(h2.clone()),
+			Some(H256::from_slice(&expected_mmr_root))
+		);
+
+		// `extract_validation_context` should return `None` when the provided header
+		// is not part of the chain.
+		let mut fork_h2 = h2;
+		fork_h2.state_root = H256::repeat_byte(0);
+		assert_eq!(BeefyMmr::extract_validation_context(fork_h2), None);
+	});
+}
+
 #[test]
 fn is_non_canonical_should_work_correctly() {
 	let mut ext = new_test_ext(vec![1, 2]);
 
 	let mut prev_roots = vec![];
 	ext.execute_with(|| {
 		for block_num in 1..=500 {
-			init_block(block_num);
+			init_block(block_num, None);
 			prev_roots.push(Mmr::mmr_root())
 		}
 	});
@@ -239,6 +281,7 @@ fn is_non_canonical_should_work_correctly() {
 					validator_set_id: 0
 				},
 				valid_proof.clone(),
+				Mmr::mmr_root(),
 			),
 			true
 		);
@@ -256,6 +299,7 @@ fn is_non_canonical_should_work_correctly() {
 					validator_set_id: 0,
 				},
 				valid_proof.clone(),
+				Mmr::mmr_root(),
 			),
 			true
 		);
@@ -272,6 +316,7 @@ fn is_non_canonical_should_work_correctly() {
 					validator_set_id: 0
 				},
 				invalid_proof,
+				Mmr::mmr_root(),
 			),
 			false
 		);
@@ -289,6 +334,7 @@ fn is_non_canonical_should_work_correctly() {
 					validator_set_id: 0,
 				},
 				valid_proof,
+				Mmr::mmr_root(),
 			),
 			false
 		);
@@ -310,6 +356,7 @@ fn is_non_canonical_should_work_correctly() {
 						validator_set_id: 0,
 					},
 					proof.clone(),
+					Mmr::mmr_root(),
 				),
 				false
 			);
@@ -325,6 +372,7 @@ fn is_non_canonical_should_work_correctly() {
 						validator_set_id: 0,
 					},
 					proof,
+					Mmr::mmr_root(),
 				),
 				true
 			)

substrate/frame/beefy/src/equivocation.rs

@@ -36,7 +36,7 @@
 
 use codec::{self as codec, Decode, Encode};
 use frame_support::traits::{Get, KeyOwnerProofSystem};
-use frame_system::pallet_prelude::BlockNumberFor;
+use frame_system::pallet_prelude::{BlockNumberFor, HeaderFor};
 use log::{error, info};
 use sp_consensus_beefy::{
 	check_commitment_signature, AncestryHelper, DoubleVotingProof, ForkVotingProof, ValidatorSetId,
@@ -136,9 +136,9 @@ pub enum EquivocationEvidenceFor<T: Config> {
 	),
 	ForkVotingProof(
 		ForkVotingProof<
-			BlockNumberFor<T>,
+			HeaderFor<T>,
 			T::BeefyId,
-			<T::AncestryHelper as AncestryHelper<BlockNumberFor<T>>>::Proof,
+			<T::AncestryHelper as AncestryHelper<HeaderFor<T>>>::Proof,
 		>,
 		T::KeyOwnerProof,
 	),
@@ -202,13 +202,23 @@ impl<T: Config> EquivocationEvidenceFor<T> {
 				return Ok(())
 			},
 			EquivocationEvidenceFor::ForkVotingProof(equivocation_proof, _) => {
-				let (vote, ancestry_proof) =
-					(equivocation_proof.vote, equivocation_proof.ancestry_proof);
+				let ForkVotingProof { vote, ancestry_proof, header } = equivocation_proof;
+
+				let maybe_validation_context = <T::AncestryHelper as AncestryHelper<
+					HeaderFor<T>,
+				>>::extract_validation_context(header);
+				let validation_context = match maybe_validation_context {
+					Some(validation_context) => validation_context,
+					None => {
+						return Err(Error::<T>::InvalidForkVotingProof);
+					},
+				};
 
 				let is_non_canonical =
-					<T::AncestryHelper as AncestryHelper<BlockNumberFor<T>>>::is_non_canonical(
+					<T::AncestryHelper as AncestryHelper<HeaderFor<T>>>::is_non_canonical(
 						&vote.commitment,
 						ancestry_proof,
+						validation_context,
 					);
 				if !is_non_canonical {
 					return Err(Error::<T>::InvalidForkVotingProof);

substrate/frame/beefy/src/lib.rs

@@ -28,7 +28,7 @@ use frame_support::{
 };
 use frame_system::{
 	ensure_none, ensure_signed,
-	pallet_prelude::{BlockNumberFor, OriginFor},
+	pallet_prelude::{BlockNumberFor, HeaderFor, OriginFor},
 };
 use log;
 use sp_runtime::{
@@ -99,7 +99,7 @@ pub mod pallet {
 		type OnNewValidatorSet: OnNewValidatorSet<<Self as Config>::BeefyId>;
 
 		/// Hook for checking commitment canonicity.
-		type AncestryHelper: AncestryHelper<BlockNumberFor<Self>>;
+		type AncestryHelper: AncestryHelper<HeaderFor<Self>>;
 
 		/// Weights for this pallet.
 		type WeightInfo: WeightInfo;
@@ -298,9 +298,9 @@ pub mod pallet {
 			origin: OriginFor<T>,
 			equivocation_proof: Box<
 				ForkVotingProof<
-					BlockNumberFor<T>,
+					HeaderFor<T>,
 					T::BeefyId,
-					<T::AncestryHelper as AncestryHelper<BlockNumberFor<T>>>::Proof,
+					<T::AncestryHelper as AncestryHelper<HeaderFor<T>>>::Proof,
 				>,
 			>,
 			key_owner_proof: T::KeyOwnerProof,
@@ -332,9 +332,9 @@ pub mod pallet {
 			origin: OriginFor<T>,
 			equivocation_proof: Box<
 				ForkVotingProof<
-					BlockNumberFor<T>,
+					HeaderFor<T>,
 					T::BeefyId,
-					<T::AncestryHelper as AncestryHelper<BlockNumberFor<T>>>::Proof,
+					<T::AncestryHelper as AncestryHelper<HeaderFor<T>>>::Proof,
 				>,
 			>,
 			key_owner_proof: T::KeyOwnerProof,