beefy: put not only lease parachain heads into mmr (#4751)

Short-term addresses
#4737.

- [x] Resolve benchmarking
I've digged into benchmarking mentioned
#4737 (comment),
but it seemed to me that this code is different proof/path. @acatangiu
could you confirm? (btw, in this
[bench](https://github.com/paritytech/polkadot-sdk/blob/b65313e81465dd730e48d4ce00deb76922618375/bridges/modules/parachains/src/benchmarking.rs#L57),
where do you actually set the `fn parachains()` to a reasonable number?
i've only seen 1)
- [ ] Communicate to Snowfork team:
This seems to be the relevant code:
https://github.com/Snowfork/snowbridge/blob/1e18e010331777042aa7e8fff3c118094af856ba/relayer/cmd/parachain_head_proof.go#L95-L120
- [x] Is it preferred to iter() in some random order as suggested in
#4737 (comment)
or take lowest para ids instead as implemented here currently?
- [x] PRDoc

## Updating Polkadot and Kusama runtimes:

New weights need to be generated (`pallet_mmr`) and configs updated
similar to Rococo/Westend:
```patch
diff --git a/polkadot/runtime/rococo/src/lib.rs b/polkadot/runtime/rococo/src/lib.rs
index 5adffbd7422..c7da339b981 100644
--- a/polkadot/runtime/rococo/src/lib.rs
+++ b/polkadot/runtime/rococo/src/lib.rs
@@ -1307,9 +1307,11 @@ impl pallet_mmr::Config for Runtime {
        const INDEXING_PREFIX: &'static [u8] = mmr::INDEXING_PREFIX;
        type Hashing = Keccak256;
        type OnNewRoot = pallet_beefy_mmr::DepositBeefyDigest<Runtime>;
-       type WeightInfo = ();
        type LeafData = pallet_beefy_mmr::Pallet<Runtime>;
        type BlockHashProvider = pallet_mmr::DefaultBlockHashProvider<Runtime>;
+       type WeightInfo = weights::pallet_mmr::WeightInfo<Runtime>;
+       #[cfg(feature = "runtime-benchmarks")]
+       type BenchmarkHelper = parachains_paras::benchmarking::mmr_setup::MmrSetup<Runtime>;
 }

 parameter_types! {
@@ -1319,13 +1321,8 @@ parameter_types! {
 pub struct ParaHeadsRootProvider;
 impl BeefyDataProvider<H256> for ParaHeadsRootProvider {
        fn extra_data() -> H256 {
-               let mut para_heads: Vec<(u32, Vec<u8>)> = parachains_paras::Parachains::<Runtime>::get()
-                       .into_iter()
-                       .filter_map(|id| {
-                               parachains_paras::Heads::<Runtime>::get(&id).map(|head| (id.into(), head.0))
-                       })
-                       .collect();
-               para_heads.sort();
+               let para_heads: Vec<(u32, Vec<u8>)> =
+                       parachains_paras::Pallet::<Runtime>::sorted_para_heads();
                binary_merkle_tree::merkle_root::<mmr::Hashing, _>(
                        para_heads.into_iter().map(|pair| pair.encode()),
                )
@@ -1746,6 +1743,7 @@ mod benches {
                [pallet_identity, Identity]
                [pallet_indices, Indices]
                [pallet_message_queue, MessageQueue]
+               [pallet_mmr, Mmr]
                [pallet_multisig, Multisig]
                [pallet_parameters, Parameters]
                [pallet_preimage, Preimage]
```

---------

Co-authored-by: Adrian Catangiu <adrian@parity.io>

polkadot/runtime/parachains/Cargo.toml

@@ -30,13 +30,15 @@ sp-keystore = { optional = true, workspace = true }
 sp-application-crypto = { optional = true, workspace = true }
 sp-tracing = { optional = true, workspace = true }
 sp-arithmetic = { workspace = true }
+sp-std = { workspace = true, optional = true }
 
 pallet-authority-discovery = { workspace = true }
 pallet-authorship = { workspace = true }
 pallet-balances = { workspace = true }
 pallet-babe = { workspace = true }
 pallet-broker = { workspace = true }
 pallet-message-queue = { workspace = true }
+pallet-mmr = { workspace = true, optional = true }
 pallet-session = { workspace = true }
 pallet-staking = { workspace = true }
 pallet-timestamp = { workspace = true }
@@ -86,6 +88,7 @@ std = [
 	"pallet-balances/std",
 	"pallet-broker/std",
 	"pallet-message-queue/std",
+	"pallet-mmr?/std",
 	"pallet-session/std",
 	"pallet-staking/std",
 	"pallet-timestamp/std",
@@ -109,6 +112,7 @@ std = [
 	"sp-runtime/std",
 	"sp-session/std",
 	"sp-staking/std",
+	"sp-std?/std",
 	"xcm-executor/std",
 	"xcm/std",
 ]
@@ -120,6 +124,7 @@ runtime-benchmarks = [
 	"pallet-balances/runtime-benchmarks",
 	"pallet-broker/runtime-benchmarks",
 	"pallet-message-queue/runtime-benchmarks",
+	"pallet-mmr/runtime-benchmarks",
 	"pallet-staking/runtime-benchmarks",
 	"pallet-timestamp/runtime-benchmarks",
 	"pallet-vesting/runtime-benchmarks",
@@ -128,6 +133,7 @@ runtime-benchmarks = [
 	"sp-application-crypto",
 	"sp-runtime/runtime-benchmarks",
 	"sp-staking/runtime-benchmarks",
+	"sp-std",
 	"static_assertions",
 	"xcm-executor/runtime-benchmarks",
 ]
@@ -141,6 +147,7 @@ try-runtime = [
 	"pallet-balances/try-runtime",
 	"pallet-broker/try-runtime",
 	"pallet-message-queue/try-runtime",
+	"pallet-mmr/try-runtime",
 	"pallet-session/try-runtime",
 	"pallet-staking/try-runtime",
 	"pallet-timestamp/try-runtime",

polkadot/runtime/parachains/src/paras/benchmarking.rs

@@ -24,6 +24,7 @@ use polkadot_primitives::{
 };
 use sp_runtime::traits::{One, Saturating};
 
+pub mod mmr_setup;
 mod pvf_check;
 
 use self::pvf_check::{VoteCause, VoteOutcome};

polkadot/runtime/parachains/src/paras/benchmarking/mmr_setup.rs

@@ -0,0 +1,40 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Implements benchmarking setup for the `merkle-mountain-range` pallet.
+
+use crate::paras::*;
+use pallet_mmr::BenchmarkHelper;
+use sp_std::vec;
+
+/// Struct to setup benchmarks for the `merkle-mountain-range` pallet.
+pub struct MmrSetup<T>(core::marker::PhantomData<T>);
+
+impl<T> BenchmarkHelper for MmrSetup<T>
+where
+	T: Config,
+{
+	fn setup() {
+		// Create a head with 1024 bytes of data.
+		let head = vec![42u8; 1024];
+
+		for para in 0..MAX_PARA_HEADS {
+			let id = (para as u32).into();
+			let h = head.clone().into();
+			Pallet::<T>::heads_insert(&id, h);
+		}
+	}
+}

polkadot/runtime/parachains/src/paras/mod.rs

@@ -135,7 +135,7 @@ use serde::{Deserialize, Serialize};
 pub use crate::Origin as ParachainOrigin;
 
 #[cfg(feature = "runtime-benchmarks")]
-pub(crate) mod benchmarking;
+pub mod benchmarking;
 
 #[cfg(test)]
 pub(crate) mod tests;
@@ -1222,6 +1222,15 @@ const INVALID_TX_BAD_VALIDATOR_IDX: u8 = 1;
 const INVALID_TX_BAD_SUBJECT: u8 = 2;
 const INVALID_TX_DOUBLE_VOTE: u8 = 3;
 
+/// This is intermediate "fix" for this issue:
+/// <https://github.com/paritytech/polkadot-sdk/issues/4737>
+///
+/// It does not actually fix it, but makes the worst case better. Without that limit someone
+/// could completely DoS the relay chain by registering a ridiculously high amount of paras.
+/// With this limit the same attack could lead to some parachains ceasing to being able to
+/// communicate via offchain XCMP. Snowbridge will still work as it only cares about `BridgeHub`.
+pub const MAX_PARA_HEADS: usize = 1024;
+
 impl<T: Config> Pallet<T> {
 	/// This is a call to schedule code upgrades for parachains which is safe to be called
 	/// outside of this module. That means this function does all checks necessary to ensure
@@ -1291,6 +1300,16 @@ impl<T: Config> Pallet<T> {
 		})
 	}
 
+	/// Get a list of the first [`MAX_PARA_HEADS`] para heads sorted by para_id.
+	/// This method is likely to be removed in the future.
+	pub fn sorted_para_heads() -> Vec<(u32, Vec<u8>)> {
+		let mut heads: Vec<(u32, Vec<u8>)> =
+			Heads::<T>::iter().map(|(id, head)| (id.into(), head.0)).collect();
+		heads.sort_by_key(|(id, _)| *id);
+		heads.truncate(MAX_PARA_HEADS);
+		heads
+	}
+
 	// Apply all para actions queued for the given session index.
 	//
 	// The actions to take are based on the lifecycle of of the paras.

polkadot/runtime/rococo/Cargo.toml

@@ -226,6 +226,7 @@ runtime-benchmarks = [
 	"pallet-asset-rate/runtime-benchmarks",
 	"pallet-babe/runtime-benchmarks",
 	"pallet-balances/runtime-benchmarks",
+	"pallet-beefy-mmr/runtime-benchmarks",
 	"pallet-bounties/runtime-benchmarks",
 	"pallet-child-bounties/runtime-benchmarks",
 	"pallet-collective/runtime-benchmarks",
@@ -330,7 +331,10 @@ metadata-hash = ["substrate-wasm-builder/metadata-hash"]
 # Set timing constants (e.g. session period) to faster versions to speed up testing.
 fast-runtime = ["rococo-runtime-constants/fast-runtime"]
 
-runtime-metrics = ["polkadot-runtime-parachains/runtime-metrics", "sp-io/with-tracing"]
+runtime-metrics = [
+	"polkadot-runtime-parachains/runtime-metrics",
+	"sp-io/with-tracing",
+]
 
 # A feature that should be enabled when the runtime should be built for on-chain
 # deployment. This will disable stuff that shouldn't be part of the on-chain wasm

polkadot/runtime/rococo/src/lib.rs

@@ -1307,9 +1307,11 @@ impl pallet_mmr::Config for Runtime {
 	const INDEXING_PREFIX: &'static [u8] = mmr::INDEXING_PREFIX;
 	type Hashing = Keccak256;
 	type OnNewRoot = pallet_beefy_mmr::DepositBeefyDigest<Runtime>;
-	type WeightInfo = ();
 	type LeafData = pallet_beefy_mmr::Pallet<Runtime>;
 	type BlockHashProvider = pallet_mmr::DefaultBlockHashProvider<Runtime>;
+	type WeightInfo = weights::pallet_mmr::WeightInfo<Runtime>;
+	#[cfg(feature = "runtime-benchmarks")]
+	type BenchmarkHelper = parachains_paras::benchmarking::mmr_setup::MmrSetup<Runtime>;
 }
 
 parameter_types! {
@@ -1319,13 +1321,8 @@ parameter_types! {
 pub struct ParaHeadsRootProvider;
 impl BeefyDataProvider<H256> for ParaHeadsRootProvider {
 	fn extra_data() -> H256 {
-		let mut para_heads: Vec<(u32, Vec<u8>)> = parachains_paras::Parachains::<Runtime>::get()
-			.into_iter()
-			.filter_map(|id| {
-				parachains_paras::Heads::<Runtime>::get(&id).map(|head| (id.into(), head.0))
-			})
-			.collect();
-		para_heads.sort();
+		let para_heads: Vec<(u32, Vec<u8>)> =
+			parachains_paras::Pallet::<Runtime>::sorted_para_heads();
 		binary_merkle_tree::merkle_root::<mmr::Hashing, _>(
 			para_heads.into_iter().map(|pair| pair.encode()),
 		)
@@ -1746,6 +1743,7 @@ mod benches {
 		[pallet_identity, Identity]
 		[pallet_indices, Indices]
 		[pallet_message_queue, MessageQueue]
+		[pallet_mmr, Mmr]
 		[pallet_multisig, Multisig]
 		[pallet_parameters, Parameters]
 		[pallet_preimage, Preimage]

polkadot/runtime/rococo/src/weights/mod.rs

@@ -25,6 +25,7 @@ pub mod pallet_conviction_voting;
 pub mod pallet_identity;
 pub mod pallet_indices;
 pub mod pallet_message_queue;
+pub mod pallet_mmr;
 pub mod pallet_multisig;
 pub mod pallet_nis;
 pub mod pallet_parameters;

polkadot/runtime/rococo/src/weights/pallet_mmr.rs

@@ -0,0 +1,77 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Autogenerated weights for `pallet_mmr`
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE SUBSTRATE BENCHMARK CLI VERSION 32.0.0
+//! DATE: 2024-07-15, STEPS: `5`, REPEAT: `1`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! WASM-EXECUTION: `Compiled`, CHAIN: `Some("westend-dev")`, DB CACHE: 1024
+
+// Executed Command:
+// target/testnet/polkadot
+// benchmark
+// pallet
+// --steps=5
+// --repeat=1
+// --extrinsic=*
+// --wasm-execution=compiled
+// --heap-pages=4096
+// --pallet=pallet_mmr
+// --chain=westend-dev
+// --header=./polkadot/file_header.txt
+// --output=./polkadot/runtime/westend/src/weights/
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+#![allow(missing_docs)]
+
+use frame_support::{traits::Get, weights::Weight};
+use core::marker::PhantomData;
+
+/// Weight functions for `pallet_mmr`.
+pub struct WeightInfo<T>(PhantomData<T>);
+impl<T: frame_system::Config> pallet_mmr::WeightInfo for WeightInfo<T> {
+	/// Storage: `Mmr::NumberOfLeaves` (r:1 w:1)
+	/// Proof: `Mmr::NumberOfLeaves` (`max_values`: Some(1), `max_size`: Some(8), added: 503, mode: `MaxEncodedLen`)
+	/// Storage: `System::ParentHash` (r:1 w:0)
+	/// Proof: `System::ParentHash` (`max_values`: Some(1), `max_size`: Some(32), added: 527, mode: `MaxEncodedLen`)
+	/// Storage: `Paras::Heads` (r:2049 w:0)
+	/// Proof: `Paras::Heads` (`max_values`: None, `max_size`: None, mode: `Measured`)
+	/// Storage: `BeefyMmrLeaf::BeefyNextAuthorities` (r:1 w:0)
+	/// Proof: `BeefyMmrLeaf::BeefyNextAuthorities` (`max_values`: Some(1), `max_size`: Some(44), added: 539, mode: `MaxEncodedLen`)
+	/// Storage: `System::Digest` (r:1 w:1)
+	/// Proof: `System::Digest` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `Mmr::Nodes` (r:0 w:1000)
+	/// Proof: `Mmr::Nodes` (`max_values`: None, `max_size`: Some(40), added: 2515, mode: `MaxEncodedLen`)
+	/// Storage: `Mmr::RootHash` (r:0 w:1)
+	/// Proof: `Mmr::RootHash` (`max_values`: Some(1), `max_size`: Some(32), added: 527, mode: `MaxEncodedLen`)
+	/// The range of component `x` is `[1, 1000]`.
+	fn on_initialize(x: u32) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `2140817`
+		//  Estimated: `7213082`
+		// Minimum execution time: 20_387_000_000 picoseconds.
+		Weight::from_parts(223_625_477_528, 0)
+			.saturating_add(Weight::from_parts(0, 7213082))
+			// Standard Error: 310_550_970
+			.saturating_add(Weight::from_parts(16_906_397_286, 0).saturating_mul(x.into()))
+			.saturating_add(T::DbWeight::get().reads(2053))
+			.saturating_add(T::DbWeight::get().writes(3))
+			.saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(x.into())))
+	}
+}

polkadot/runtime/westend/Cargo.toml

@@ -240,6 +240,7 @@ runtime-benchmarks = [
 	"pallet-babe/runtime-benchmarks",
 	"pallet-bags-list/runtime-benchmarks",
 	"pallet-balances/runtime-benchmarks",
+	"pallet-beefy-mmr/runtime-benchmarks",
 	"pallet-collective/runtime-benchmarks",
 	"pallet-conviction-voting/runtime-benchmarks",
 	"pallet-delegated-staking/runtime-benchmarks",
@@ -346,7 +347,10 @@ metadata-hash = ["substrate-wasm-builder/metadata-hash"]
 # Set timing constants (e.g. session period) to faster versions to speed up testing.
 fast-runtime = []
 
-runtime-metrics = ["polkadot-runtime-parachains/runtime-metrics", "sp-io/with-tracing"]
+runtime-metrics = [
+	"polkadot-runtime-parachains/runtime-metrics",
+	"sp-io/with-tracing",
+]
 
 # A feature that should be enabled when the runtime should be built for on-chain
 # deployment. This will disable stuff that shouldn't be part of the on-chain wasm

polkadot/runtime/westend/src/lib.rs

@@ -342,9 +342,11 @@ impl pallet_mmr::Config for Runtime {
 	const INDEXING_PREFIX: &'static [u8] = mmr::INDEXING_PREFIX;
 	type Hashing = Keccak256;
 	type OnNewRoot = pallet_beefy_mmr::DepositBeefyDigest<Runtime>;
-	type WeightInfo = ();
 	type LeafData = pallet_beefy_mmr::Pallet<Runtime>;
 	type BlockHashProvider = pallet_mmr::DefaultBlockHashProvider<Runtime>;
+	type WeightInfo = weights::pallet_mmr::WeightInfo<Runtime>;
+	#[cfg(feature = "runtime-benchmarks")]
+	type BenchmarkHelper = parachains_paras::benchmarking::mmr_setup::MmrSetup<Runtime>;
 }
 
 /// MMR helper types.
@@ -366,13 +368,8 @@ parameter_types! {
 pub struct ParaHeadsRootProvider;
 impl BeefyDataProvider<H256> for ParaHeadsRootProvider {
 	fn extra_data() -> H256 {
-		let mut para_heads: Vec<(u32, Vec<u8>)> = parachains_paras::Parachains::<Runtime>::get()
-			.into_iter()
-			.filter_map(|id| {
-				parachains_paras::Heads::<Runtime>::get(&id).map(|head| (id.into(), head.0))
-			})
-			.collect();
-		para_heads.sort_by_key(|k| k.0);
+		let para_heads: Vec<(u32, Vec<u8>)> =
+			parachains_paras::Pallet::<Runtime>::sorted_para_heads();
 		binary_merkle_tree::merkle_root::<mmr::Hashing, _>(
 			para_heads.into_iter().map(|pair| pair.encode()),
 		)
@@ -1755,6 +1752,7 @@ mod benches {
 		[pallet_identity, Identity]
 		[pallet_indices, Indices]
 		[pallet_message_queue, MessageQueue]
+		[pallet_mmr, Mmr]
 		[pallet_multisig, Multisig]
 		[pallet_nomination_pools, NominationPoolsBench::<Runtime>]
 		[pallet_offences, OffencesBench::<Runtime>]

polkadot/runtime/westend/src/weights/mod.rs

@@ -26,6 +26,7 @@ pub mod pallet_fast_unstake;
 pub mod pallet_identity;
 pub mod pallet_indices;
 pub mod pallet_message_queue;
+pub mod pallet_mmr;
 pub mod pallet_multisig;
 pub mod pallet_nomination_pools;
 pub mod pallet_preimage;