Introduce submit_finality_proof_ex call to bridges GRANDPA pallet (#3225)

backport of
https://github.com/paritytech/parity-bridges-common/pull/2821 (see
detailed description there)

bridges/modules/grandpa/src/benchmarking.rs

@@ -16,8 +16,9 @@
 
 //! Benchmarks for the GRANDPA Pallet.
 //!
-//! The main dispatchable for the GRANDPA pallet is `submit_finality_proof`, so these benchmarks are
-//! based around that. There are to main factors which affect finality proof verification:
+//! The main dispatchable for the GRANDPA pallet is `submit_finality_proof_ex`. Our benchmarks
+//! are based around `submit_finality_proof`, though - from weight PoV they are the same calls.
+//! There are to main factors which affect finality proof verification:
 //!
 //! 1. The number of `votes-ancestries` in the justification
 //! 2. The number of `pre-commits` in the justification

bridges/modules/grandpa/src/call_ext.rs

@@ -14,7 +14,10 @@
 // You should have received a copy of the GNU General Public License
 // along with Parity Bridges Common.  If not, see <http://www.gnu.org/licenses/>.
 
-use crate::{weights::WeightInfo, BridgedBlockNumber, BridgedHeader, Config, Error, Pallet};
+use crate::{
+	weights::WeightInfo, BridgedBlockNumber, BridgedHeader, Config, CurrentAuthoritySet, Error,
+	Pallet,
+};
 use bp_header_chain::{
 	justification::GrandpaJustification, max_expected_submit_finality_proof_arguments_size,
 	ChainWithGrandpa, GrandpaConsensusLogReader,
@@ -22,6 +25,7 @@ use bp_header_chain::{
 use bp_runtime::{BlockNumberOf, OwnedBridgeModule};
 use codec::Encode;
 use frame_support::{dispatch::CallableCallFor, traits::IsSubType, weights::Weight};
+use sp_consensus_grandpa::SetId;
 use sp_runtime::{
 	traits::{Header, Zero},
 	transaction_validity::{InvalidTransaction, TransactionValidity, ValidTransaction},
@@ -33,6 +37,9 @@ use sp_runtime::{
 pub struct SubmitFinalityProofInfo<N> {
 	/// Number of the finality target.
 	pub block_number: N,
+	/// An identifier of the validators set that has signed the submitted justification.
+	/// It might be `None` if deprecated version of the `submit_finality_proof` is used.
+	pub current_set_id: Option<SetId>,
 	/// Extra weight that we assume is included in the call.
 	///
 	/// We have some assumptions about headers and justifications of the bridged chain.
@@ -61,9 +68,11 @@ pub struct SubmitFinalityProofHelper<T: Config<I>, I: 'static> {
 
 impl<T: Config<I>, I: 'static> SubmitFinalityProofHelper<T, I> {
 	/// Check that the GRANDPA head provided by the `SubmitFinalityProof` is better than the best
-	/// one we know.
+	/// one we know. Additionally, checks if `current_set_id` matches the current authority set
+	/// id, if specified.
 	pub fn check_obsolete(
 		finality_target: BlockNumberOf<T::BridgedChain>,
+		current_set_id: Option<SetId>,
 	) -> Result<(), Error<T, I>> {
 		let best_finalized = crate::BestFinalized::<T, I>::get().ok_or_else(|| {
 			log::trace!(
@@ -85,6 +94,20 @@ impl<T: Config<I>, I: 'static> SubmitFinalityProofHelper<T, I> {
 			return Err(Error::<T, I>::OldHeader)
 		}
 
+		if let Some(current_set_id) = current_set_id {
+			let actual_set_id = <CurrentAuthoritySet<T, I>>::get().set_id;
+			if current_set_id != actual_set_id {
+				log::trace!(
+					target: crate::LOG_TARGET,
+					"Cannot finalize header signed by unknown authority set: bundled {:?}, best {:?}",
+					current_set_id,
+					actual_set_id,
+				);
+
+				return Err(Error::<T, I>::InvalidAuthoritySetId)
+			}
+		}
+
 		Ok(())
 	}
 
@@ -111,6 +134,18 @@ pub trait CallSubType<T: Config<I, RuntimeCall = Self>, I: 'static>:
 			return Some(submit_finality_proof_info_from_args::<T, I>(
 				finality_target,
 				justification,
+				None,
+			))
+		} else if let Some(crate::Call::<T, I>::submit_finality_proof_ex {
+			finality_target,
+			justification,
+			current_set_id,
+		}) = self.is_sub_type()
+		{
+			return Some(submit_finality_proof_info_from_args::<T, I>(
+				finality_target,
+				justification,
+				Some(*current_set_id),
 			))
 		}
 
@@ -133,7 +168,10 @@ pub trait CallSubType<T: Config<I, RuntimeCall = Self>, I: 'static>:
 			return InvalidTransaction::Call.into()
 		}
 
-		match SubmitFinalityProofHelper::<T, I>::check_obsolete(finality_target.block_number) {
+		match SubmitFinalityProofHelper::<T, I>::check_obsolete(
+			finality_target.block_number,
+			finality_target.current_set_id,
+		) {
 			Ok(_) => Ok(ValidTransaction::default()),
 			Err(Error::<T, I>::OldHeader) => InvalidTransaction::Stale.into(),
 			Err(_) => InvalidTransaction::Call.into(),
@@ -150,6 +188,7 @@ impl<T: Config<I>, I: 'static> CallSubType<T, I> for T::RuntimeCall where
 pub(crate) fn submit_finality_proof_info_from_args<T: Config<I>, I: 'static>(
 	finality_target: &BridgedHeader<T, I>,
 	justification: &GrandpaJustification<BridgedHeader<T, I>>,
+	current_set_id: Option<SetId>,
 ) -> SubmitFinalityProofInfo<BridgedBlockNumber<T, I>> {
 	let block_number = *finality_target.number();
 
@@ -191,7 +230,7 @@ pub(crate) fn submit_finality_proof_info_from_args<T: Config<I>, I: 'static>(
 	);
 	let extra_size = actual_call_size.saturating_sub(max_expected_call_size);
 
-	SubmitFinalityProofInfo { block_number, extra_weight, extra_size }
+	SubmitFinalityProofInfo { block_number, current_set_id, extra_weight, extra_size }
 }
 
 #[cfg(test)]
@@ -199,20 +238,24 @@ mod tests {
 	use crate::{
 		call_ext::CallSubType,
 		mock::{run_test, test_header, RuntimeCall, TestBridgedChain, TestNumber, TestRuntime},
-		BestFinalized, Config, PalletOperatingMode, WeightInfo,
+		BestFinalized, Config, CurrentAuthoritySet, PalletOperatingMode, StoredAuthoritySet,
+		SubmitFinalityProofInfo, WeightInfo,
 	};
 	use bp_header_chain::ChainWithGrandpa;
 	use bp_runtime::{BasicOperatingMode, HeaderId};
 	use bp_test_utils::{
 		make_default_justification, make_justification_for_header, JustificationGeneratorParams,
+		TEST_GRANDPA_SET_ID,
 	};
 	use frame_support::weights::Weight;
 	use sp_runtime::{testing::DigestItem, traits::Header as _, SaturatedConversion};
 
 	fn validate_block_submit(num: TestNumber) -> bool {
-		let bridge_grandpa_call = crate::Call::<TestRuntime, ()>::submit_finality_proof {
+		let bridge_grandpa_call = crate::Call::<TestRuntime, ()>::submit_finality_proof_ex {
 			finality_target: Box::new(test_header(num)),
 			justification: make_default_justification(&test_header(num)),
+			// not initialized => zero
+			current_set_id: 0,
 		};
 		RuntimeCall::check_obsolete_submit_finality_proof(&RuntimeCall::Grandpa(
 			bridge_grandpa_call,
@@ -256,6 +299,18 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn extension_rejects_new_header_if_set_id_is_invalid() {
+		run_test(|| {
+			// when set id is different from the passed one => tx is rejected
+			sync_to_header_10();
+			let next_set = StoredAuthoritySet::<TestRuntime, ()>::try_new(vec![], 0x42).unwrap();
+			CurrentAuthoritySet::<TestRuntime, ()>::put(next_set);
+
+			assert!(!validate_block_submit(15));
+		});
+	}
+
 	#[test]
 	fn extension_accepts_new_header() {
 		run_test(|| {
@@ -266,6 +321,42 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn submit_finality_proof_info_is_parsed() {
+		// when `submit_finality_proof` is used, `current_set_id` is set to `None`
+		let deprecated_call =
+			RuntimeCall::Grandpa(crate::Call::<TestRuntime, ()>::submit_finality_proof {
+				finality_target: Box::new(test_header(42)),
+				justification: make_default_justification(&test_header(42)),
+			});
+		assert_eq!(
+			deprecated_call.submit_finality_proof_info(),
+			Some(SubmitFinalityProofInfo {
+				block_number: 42,
+				current_set_id: None,
+				extra_weight: Weight::zero(),
+				extra_size: 0,
+			})
+		);
+
+		// when `submit_finality_proof_ex` is used, `current_set_id` is set to `Some`
+		let deprecated_call =
+			RuntimeCall::Grandpa(crate::Call::<TestRuntime, ()>::submit_finality_proof_ex {
+				finality_target: Box::new(test_header(42)),
+				justification: make_default_justification(&test_header(42)),
+				current_set_id: 777,
+			});
+		assert_eq!(
+			deprecated_call.submit_finality_proof_info(),
+			Some(SubmitFinalityProofInfo {
+				block_number: 42,
+				current_set_id: Some(777),
+				extra_weight: Weight::zero(),
+				extra_size: 0,
+			})
+		);
+	}
+
 	#[test]
 	fn extension_returns_correct_extra_size_if_call_arguments_are_too_large() {
 		// when call arguments are below our limit => no refund
@@ -275,9 +366,10 @@ mod tests {
 			..Default::default()
 		};
 		let small_justification = make_justification_for_header(justification_params);
-		let small_call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof {
+		let small_call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof_ex {
 			finality_target: Box::new(small_finality_target),
 			justification: small_justification,
+			current_set_id: TEST_GRANDPA_SET_ID,
 		});
 		assert_eq!(small_call.submit_finality_proof_info().unwrap().extra_size, 0);
 
@@ -291,9 +383,10 @@ mod tests {
 			..Default::default()
 		};
 		let large_justification = make_justification_for_header(justification_params);
-		let large_call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof {
+		let large_call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof_ex {
 			finality_target: Box::new(large_finality_target),
 			justification: large_justification,
+			current_set_id: TEST_GRANDPA_SET_ID,
 		});
 		assert_ne!(large_call.submit_finality_proof_info().unwrap().extra_size, 0);
 	}
@@ -309,9 +402,10 @@ mod tests {
 
 		// when there are `REASONABLE_HEADERS_IN_JUSTIFICATON_ANCESTRY` headers => no refund
 		let justification = make_justification_for_header(justification_params.clone());
-		let call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof {
+		let call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof_ex {
 			finality_target: Box::new(finality_target.clone()),
 			justification,
+			current_set_id: TEST_GRANDPA_SET_ID,
 		});
 		assert_eq!(call.submit_finality_proof_info().unwrap().extra_weight, Weight::zero());
 
@@ -322,9 +416,10 @@ mod tests {
 			justification.commit.precommits.len().saturated_into(),
 			justification.votes_ancestries.len().saturated_into(),
 		);
-		let call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof {
+		let call = RuntimeCall::Grandpa(crate::Call::submit_finality_proof_ex {
 			finality_target: Box::new(finality_target),
 			justification,
+			current_set_id: TEST_GRANDPA_SET_ID,
 		});
 		assert_eq!(call.submit_finality_proof_info().unwrap().extra_weight, call_weight);
 	}

bridges/modules/grandpa/src/lib.rs

@@ -151,102 +151,31 @@ pub mod pallet {
 
 	#[pallet::call]
 	impl<T: Config<I>, I: 'static> Pallet<T, I> {
-		/// Verify a target header is finalized according to the given finality proof.
-		///
-		/// It will use the underlying storage pallet to fetch information about the current
-		/// authorities and best finalized header in order to verify that the header is finalized.
-		///
-		/// If successful in verification, it will write the target header to the underlying storage
-		/// pallet.
-		///
-		/// The call fails if:
-		///
-		/// - the pallet is halted;
-		///
-		/// - the pallet knows better header than the `finality_target`;
-		///
-		/// - verification is not optimized or invalid;
-		///
-		/// - header contains forced authorities set change or change with non-zero delay.
+		/// This call is deprecated and will be removed around May 2024. Use the
+		/// `submit_finality_proof_ex` instead. Semantically, this call is an equivalent of the
+		/// `submit_finality_proof_ex` call without current authority set id check.
 		#[pallet::call_index(0)]
 		#[pallet::weight(<T::WeightInfo as WeightInfo>::submit_finality_proof(
 			justification.commit.precommits.len().saturated_into(),
 			justification.votes_ancestries.len().saturated_into(),
 		))]
+		#[allow(deprecated)]
+		#[deprecated(
+			note = "`submit_finality_proof` will be removed in May 2024. Use `submit_finality_proof_ex` instead."
+		)]
 		pub fn submit_finality_proof(
 			origin: OriginFor<T>,
 			finality_target: Box<BridgedHeader<T, I>>,
 			justification: GrandpaJustification<BridgedHeader<T, I>>,
 		) -> DispatchResultWithPostInfo {
-			Self::ensure_not_halted().map_err(Error::<T, I>::BridgeModule)?;
-			ensure_signed(origin)?;
-
-			let (hash, number) = (finality_target.hash(), *finality_target.number());
-			log::trace!(
-				target: LOG_TARGET,
-				"Going to try and finalize header {:?}",
-				finality_target
-			);
-
-			SubmitFinalityProofHelper::<T, I>::check_obsolete(number)?;
-
-			let authority_set = <CurrentAuthoritySet<T, I>>::get();
-			let unused_proof_size = authority_set.unused_proof_size();
-			let set_id = authority_set.set_id;
-			let authority_set: AuthoritySet = authority_set.into();
-			verify_justification::<T, I>(&justification, hash, number, authority_set)?;
-
-			let maybe_new_authority_set =
-				try_enact_authority_change::<T, I>(&finality_target, set_id)?;
-			let may_refund_call_fee = maybe_new_authority_set.is_some() &&
-				// if we have seen too many mandatory headers in this block, we don't want to refund
-				Self::free_mandatory_headers_remaining() > 0 &&
-				// if arguments out of expected bounds, we don't want to refund
-				submit_finality_proof_info_from_args::<T, I>(&finality_target, &justification)
-					.fits_limits();
-			if may_refund_call_fee {
-				FreeMandatoryHeadersRemaining::<T, I>::mutate(|count| {
-					*count = count.saturating_sub(1)
-				});
-			}
-			insert_header::<T, I>(*finality_target, hash);
-			log::info!(
-				target: LOG_TARGET,
-				"Successfully imported finalized header with hash {:?}!",
-				hash
-			);
-
-			// mandatory header is a header that changes authorities set. The pallet can't go
-			// further without importing this header. So every bridge MUST import mandatory headers.
-			//
-			// We don't want to charge extra costs for mandatory operations. So relayer is not
-			// paying fee for mandatory headers import transactions.
-			//
-			// If size/weight of the call is exceeds our estimated limits, the relayer still needs
-			// to pay for the transaction.
-			let pays_fee = if may_refund_call_fee { Pays::No } else { Pays::Yes };
-
-			// the proof size component of the call weight assumes that there are
-			// `MaxBridgedAuthorities` in the `CurrentAuthoritySet` (we use `MaxEncodedLen`
-			// estimation). But if their number is lower, then we may "refund" some `proof_size`,
-			// making proof smaller and leaving block space to other useful transactions
-			let pre_dispatch_weight = T::WeightInfo::submit_finality_proof(
-				justification.commit.precommits.len().saturated_into(),
-				justification.votes_ancestries.len().saturated_into(),
-			);
-			let actual_weight = pre_dispatch_weight
-				.set_proof_size(pre_dispatch_weight.proof_size().saturating_sub(unused_proof_size));
-
-			Self::deposit_event(Event::UpdatedBestFinalizedHeader {
-				number,
-				hash,
-				grandpa_info: StoredHeaderGrandpaInfo {
-					finality_proof: justification,
-					new_verification_context: maybe_new_authority_set,
-				},
-			});
-
-			Ok(PostDispatchInfo { actual_weight: Some(actual_weight), pays_fee })
+			Self::submit_finality_proof_ex(
+				origin,
+				finality_target,
+				justification,
+				// the `submit_finality_proof_ex` also reads this value, but it is done from the
+				// cache, so we don't treat it as an additional db access
+				<CurrentAuthoritySet<T, I>>::get().set_id,
+			)
 		}
 
 		/// Bootstrap the bridge pallet with an initial header and authority set from which to sync.
@@ -299,6 +228,111 @@ pub mod pallet {
 		) -> DispatchResult {
 			<Self as OwnedBridgeModule<_>>::set_operating_mode(origin, operating_mode)
 		}
+
+		/// Verify a target header is finalized according to the given finality proof. The proof
+		/// is assumed to be signed by GRANDPA authorities set with `current_set_id` id.
+		///
+		/// It will use the underlying storage pallet to fetch information about the current
+		/// authorities and best finalized header in order to verify that the header is finalized.
+		///
+		/// If successful in verification, it will write the target header to the underlying storage
+		/// pallet.
+		///
+		/// The call fails if:
+		///
+		/// - the pallet is halted;
+		///
+		/// - the pallet knows better header than the `finality_target`;
+		///
+		/// - the id of best GRANDPA authority set, known to the pallet is not equal to the
+		///   `current_set_id`;
+		///
+		/// - verification is not optimized or invalid;
+		///
+		/// - header contains forced authorities set change or change with non-zero delay.
+		#[pallet::call_index(4)]
+		#[pallet::weight(<T::WeightInfo as WeightInfo>::submit_finality_proof(
+			justification.commit.precommits.len().saturated_into(),
+			justification.votes_ancestries.len().saturated_into(),
+		))]
+		pub fn submit_finality_proof_ex(
+			origin: OriginFor<T>,
+			finality_target: Box<BridgedHeader<T, I>>,
+			justification: GrandpaJustification<BridgedHeader<T, I>>,
+			current_set_id: sp_consensus_grandpa::SetId,
+		) -> DispatchResultWithPostInfo {
+			Self::ensure_not_halted().map_err(Error::<T, I>::BridgeModule)?;
+			ensure_signed(origin)?;
+
+			let (hash, number) = (finality_target.hash(), *finality_target.number());
+			log::trace!(
+				target: LOG_TARGET,
+				"Going to try and finalize header {:?}",
+				finality_target
+			);
+
+			// it checks whether the `number` is better than the current best block number
+			// and whether the `current_set_id` matches the best known set id
+			SubmitFinalityProofHelper::<T, I>::check_obsolete(number, Some(current_set_id))?;
+
+			let authority_set = <CurrentAuthoritySet<T, I>>::get();
+			let unused_proof_size = authority_set.unused_proof_size();
+			let set_id = authority_set.set_id;
+			let authority_set: AuthoritySet = authority_set.into();
+			verify_justification::<T, I>(&justification, hash, number, authority_set)?;
+
+			let maybe_new_authority_set =
+				try_enact_authority_change::<T, I>(&finality_target, set_id)?;
+			let may_refund_call_fee = maybe_new_authority_set.is_some() &&
+				// if we have seen too many mandatory headers in this block, we don't want to refund
+				Self::free_mandatory_headers_remaining() > 0 &&
+				// if arguments out of expected bounds, we don't want to refund
+				submit_finality_proof_info_from_args::<T, I>(&finality_target, &justification, Some(current_set_id))
+					.fits_limits();
+			if may_refund_call_fee {
+				FreeMandatoryHeadersRemaining::<T, I>::mutate(|count| {
+					*count = count.saturating_sub(1)
+				});
+			}
+			insert_header::<T, I>(*finality_target, hash);
+			log::info!(
+				target: LOG_TARGET,
+				"Successfully imported finalized header with hash {:?}!",
+				hash
+			);
+
+			// mandatory header is a header that changes authorities set. The pallet can't go
+			// further without importing this header. So every bridge MUST import mandatory headers.
+			//
+			// We don't want to charge extra costs for mandatory operations. So relayer is not
+			// paying fee for mandatory headers import transactions.
+			//
+			// If size/weight of the call is exceeds our estimated limits, the relayer still needs
+			// to pay for the transaction.
+			let pays_fee = if may_refund_call_fee { Pays::No } else { Pays::Yes };
+
+			// the proof size component of the call weight assumes that there are
+			// `MaxBridgedAuthorities` in the `CurrentAuthoritySet` (we use `MaxEncodedLen`
+			// estimation). But if their number is lower, then we may "refund" some `proof_size`,
+			// making proof smaller and leaving block space to other useful transactions
+			let pre_dispatch_weight = T::WeightInfo::submit_finality_proof(
+				justification.commit.precommits.len().saturated_into(),
+				justification.votes_ancestries.len().saturated_into(),
+			);
+			let actual_weight = pre_dispatch_weight
+				.set_proof_size(pre_dispatch_weight.proof_size().saturating_sub(unused_proof_size));
+
+			Self::deposit_event(Event::UpdatedBestFinalizedHeader {
+				number,
+				hash,
+				grandpa_info: StoredHeaderGrandpaInfo {
+					finality_proof: justification,
+					new_verification_context: maybe_new_authority_set,
+				},
+			});
+
+			Ok(PostDispatchInfo { actual_weight: Some(actual_weight), pays_fee })
+		}
 	}
 
 	/// Number mandatory headers that we may accept in the current block for free (returning
@@ -436,6 +470,9 @@ pub mod pallet {
 		TooManyAuthoritiesInSet,
 		/// Error generated by the `OwnedBridgeModule` trait.
 		BridgeModule(bp_runtime::OwnedBridgeModuleError),
+		/// The `current_set_id` argument of the `submit_finality_proof_ex` doesn't match
+		/// the id of the current set, known to the pallet.
+		InvalidAuthoritySetId,
 	}
 
 	/// Check the given header for a GRANDPA scheduled authority set change. If a change
@@ -663,6 +700,7 @@ mod tests {
 	use bp_test_utils::{
 		authority_list, generate_owned_bridge_module_tests, make_default_justification,
 		make_justification_for_header, JustificationGeneratorParams, ALICE, BOB,
+		TEST_GRANDPA_SET_ID,
 	};
 	use codec::Encode;
 	use frame_support::{
@@ -693,7 +731,7 @@ mod tests {
 		let init_data = InitializationData {
 			header: Box::new(genesis),
 			authority_list: authority_list(),
-			set_id: 1,
+			set_id: TEST_GRANDPA_SET_ID,
 			operating_mode: BasicOperatingMode::Normal,
 		};
 
@@ -704,10 +742,11 @@ mod tests {
 		let header = test_header(header.into());
 		let justification = make_default_justification(&header);
 
-		Pallet::<TestRuntime>::submit_finality_proof(
+		Pallet::<TestRuntime>::submit_finality_proof_ex(
 			RuntimeOrigin::signed(1),
 			Box::new(header),
 			justification,
+			TEST_GRANDPA_SET_ID,
 		)
 	}
 
@@ -722,10 +761,11 @@ mod tests {
 			..Default::default()
 		});
 
-		Pallet::<TestRuntime>::submit_finality_proof(
+		Pallet::<TestRuntime>::submit_finality_proof_ex(
 			RuntimeOrigin::signed(1),
 			Box::new(header),
 			justification,
+			set_id,
 		)
 	}
 
@@ -749,10 +789,11 @@ mod tests {
 			..Default::default()
 		});
 
-		Pallet::<TestRuntime>::submit_finality_proof(
+		Pallet::<TestRuntime>::submit_finality_proof_ex(
 			RuntimeOrigin::signed(1),
 			Box::new(header),
 			justification,
+			set_id,
 		)
 	}
 
@@ -955,17 +996,30 @@ mod tests {
 
 			let header = test_header(1);
 
-			let params =
-				JustificationGeneratorParams::<TestHeader> { set_id: 2, ..Default::default() };
+			let next_set_id = 2;
+			let params = JustificationGeneratorParams::<TestHeader> {
+				set_id: next_set_id,
+				..Default::default()
+			};
 			let justification = make_justification_for_header(params);
 
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
+					RuntimeOrigin::signed(1),
+					Box::new(header.clone()),
+					justification.clone(),
+					TEST_GRANDPA_SET_ID,
+				),
+				<Error<TestRuntime>>::InvalidJustification
+			);
+			assert_err!(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
 					justification,
+					next_set_id,
 				),
-				<Error<TestRuntime>>::InvalidJustification
+				<Error<TestRuntime>>::InvalidAuthoritySetId
 			);
 		})
 	}
@@ -980,10 +1034,11 @@ mod tests {
 			justification.round = 42;
 
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
 					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				<Error<TestRuntime>>::InvalidJustification
 			);
@@ -1009,10 +1064,11 @@ mod tests {
 			let justification = make_default_justification(&header);
 
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
 					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				<Error<TestRuntime>>::InvalidAuthoritySet
 			);
@@ -1047,10 +1103,11 @@ mod tests {
 			let justification = make_default_justification(&header);
 
 			// Let's import our test header
-			let result = Pallet::<TestRuntime>::submit_finality_proof(
+			let result = Pallet::<TestRuntime>::submit_finality_proof_ex(
 				RuntimeOrigin::signed(1),
 				Box::new(header.clone()),
 				justification.clone(),
+				TEST_GRANDPA_SET_ID,
 			);
 			assert_ok!(result);
 			assert_eq!(result.unwrap().pays_fee, frame_support::dispatch::Pays::No);
@@ -1109,10 +1166,11 @@ mod tests {
 
 			// without large digest item ^^^ the relayer would have paid zero transaction fee
 			// (`Pays::No`)
-			let result = Pallet::<TestRuntime>::submit_finality_proof(
+			let result = Pallet::<TestRuntime>::submit_finality_proof_ex(
 				RuntimeOrigin::signed(1),
 				Box::new(header.clone()),
 				justification,
+				TEST_GRANDPA_SET_ID,
 			);
 			assert_ok!(result);
 			assert_eq!(result.unwrap().pays_fee, frame_support::dispatch::Pays::Yes);
@@ -1140,10 +1198,11 @@ mod tests {
 
 			// without many headers in votes ancestries ^^^ the relayer would have paid zero
 			// transaction fee (`Pays::No`)
-			let result = Pallet::<TestRuntime>::submit_finality_proof(
+			let result = Pallet::<TestRuntime>::submit_finality_proof_ex(
 				RuntimeOrigin::signed(1),
 				Box::new(header.clone()),
 				justification,
+				TEST_GRANDPA_SET_ID,
 			);
 			assert_ok!(result);
 			assert_eq!(result.unwrap().pays_fee, frame_support::dispatch::Pays::Yes);
@@ -1169,10 +1228,11 @@ mod tests {
 
 			// Should not be allowed to import this header
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
-					justification
+					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				<Error<TestRuntime>>::UnsupportedScheduledChange
 			);
@@ -1194,10 +1254,11 @@ mod tests {
 
 			// Should not be allowed to import this header
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
-					justification
+					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				<Error<TestRuntime>>::UnsupportedScheduledChange
 			);
@@ -1219,10 +1280,11 @@ mod tests {
 
 			// Should not be allowed to import this header
 			assert_err!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
-					justification
+					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				<Error<TestRuntime>>::TooManyAuthoritiesInSet
 			);
@@ -1283,10 +1345,11 @@ mod tests {
 				let mut invalid_justification = make_default_justification(&header);
 				invalid_justification.round = 42;
 
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::signed(1),
 					Box::new(header),
 					invalid_justification,
+					TEST_GRANDPA_SET_ID,
 				)
 			};
 
@@ -1451,10 +1514,11 @@ mod tests {
 			let justification = make_default_justification(&header);
 
 			assert_noop!(
-				Pallet::<TestRuntime>::submit_finality_proof(
+				Pallet::<TestRuntime>::submit_finality_proof_ex(
 					RuntimeOrigin::root(),
 					Box::new(header),
 					justification,
+					TEST_GRANDPA_SET_ID,
 				),
 				DispatchError::BadOrigin,
 			);