Generate storage info for pallet im_online (paritytech#9654)

* Integrating WrapperOpaque from PR paritytech#9738

* Adding storage_info to pallet im-online
Changing some `Vec` to `WeakBoundedVec`
Adding the following bounds:
* `MaxKeys
* `MaxPeerInHeartbeats`
* `MaxPeerDataEncodingSize`
to limit the size of `WeakBoundedVec`

* Fix syntax

* Need to clone keys

* Changes in formatting

bin/node/runtime/src/lib.rs

@@ -932,6 +932,9 @@ parameter_types! {
 	/// We prioritize im-online heartbeats over election solution submission.
 	pub const StakingUnsignedPriority: TransactionPriority = TransactionPriority::max_value() / 2;
 	pub const MaxAuthorities: u32 = 100;
+	pub const MaxKeys: u32 = 10_000;
+	pub const MaxPeerInHeartbeats: u32 = 10_000;
+	pub const MaxPeerDataEncodingSize: u32 = 1_000;
 }
 
 impl<LocalCall> frame_system::offchain::CreateSignedTransaction<LocalCall> for Runtime
@@ -996,6 +999,9 @@ impl pallet_im_online::Config for Runtime {
 	type ReportUnresponsiveness = Offences;
 	type UnsignedPriority = ImOnlineUnsignedPriority;
 	type WeightInfo = pallet_im_online::weights::SubstrateWeight<Runtime>;
+	type MaxKeys = MaxKeys;
+	type MaxPeerInHeartbeats = MaxPeerInHeartbeats;
+	type MaxPeerDataEncodingSize = MaxPeerDataEncodingSize;
 }
 
 impl pallet_offences::Config for Runtime {

frame/im-online/src/benchmarking.rs

@@ -22,7 +22,7 @@
 use super::*;
 
 use frame_benchmarking::{benchmarks, impl_benchmark_test_suite};
-use frame_support::traits::UnfilteredDispatchable;
+use frame_support::{traits::UnfilteredDispatchable, WeakBoundedVec};
 use frame_system::RawOrigin;
 use sp_core::{offchain::OpaqueMultiaddr, OpaquePeerId};
 use sp_runtime::{
@@ -46,7 +46,9 @@ pub fn create_heartbeat<T: Config>(
 	for _ in 0..k {
 		keys.push(T::AuthorityId::generate_pair(None));
 	}
-	Keys::<T>::put(keys.clone());
+	let bounded_keys = WeakBoundedVec::<_, T::MaxKeys>::try_from(keys.clone())
+		.map_err(|()| "More than the maximum number of keys provided")?;
+	Keys::<T>::put(bounded_keys);
 
 	let network_state = OpaqueNetworkState {
 		peer_id: OpaquePeerId::default(),

frame/im-online/src/lib.rs

@@ -74,9 +74,14 @@ mod mock;
 mod tests;
 pub mod weights;
 
-use codec::{Decode, Encode};
-use frame_support::traits::{
-	EstimateNextSessionRotation, OneSessionHandler, ValidatorSet, ValidatorSetWithIdentification,
+use codec::{Decode, Encode, MaxEncodedLen};
+use core::convert::TryFrom;
+use frame_support::{
+	traits::{
+		EstimateNextSessionRotation, Get, OneSessionHandler, ValidatorSet,
+		ValidatorSetWithIdentification, WrapperOpaque,
+	},
+	BoundedSlice, WeakBoundedVec,
 };
 use frame_system::offchain::{SendTransactionTypes, SubmitTransaction};
 pub use pallet::*;
@@ -220,6 +225,65 @@ where
 	pub validators_len: u32,
 }
 
+/// A type that is the same as [`OpaqueNetworkState`] but with [`Vec`] replaced with
+/// [`WeakBoundedVec<Limit>`] where Limit is the respective size limit
+/// `PeerIdEncodingLimit` represents the size limit of the encoding of `PeerId`
+/// `MultiAddrEncodingLimit` represents the size limit of the encoding of `MultiAddr`
+/// `AddressesLimit` represents the size limit of the vector of peers connected
+#[derive(Clone, Eq, PartialEq, Encode, Decode, MaxEncodedLen, TypeInfo)]
+#[codec(mel_bound(PeerIdEncodingLimit: Get<u32>,
+  	MultiAddrEncodingLimit: Get<u32>, AddressesLimit: Get<u32>))]
+#[scale_info(skip_type_params(PeerIdEncodingLimit, MultiAddrEncodingLimit, AddressesLimit))]
+pub struct BoundedOpaqueNetworkState<PeerIdEncodingLimit, MultiAddrEncodingLimit, AddressesLimit>
+where
+	PeerIdEncodingLimit: Get<u32>,
+	MultiAddrEncodingLimit: Get<u32>,
+	AddressesLimit: Get<u32>,
+{
+	/// PeerId of the local node in SCALE encoded.
+	pub peer_id: WeakBoundedVec<u8, PeerIdEncodingLimit>,
+	/// List of addresses the node knows it can be reached as.
+	pub external_addresses:
+		WeakBoundedVec<WeakBoundedVec<u8, MultiAddrEncodingLimit>, AddressesLimit>,
+}
+
+impl<PeerIdEncodingLimit: Get<u32>, MultiAddrEncodingLimit: Get<u32>, AddressesLimit: Get<u32>>
+	BoundedOpaqueNetworkState<PeerIdEncodingLimit, MultiAddrEncodingLimit, AddressesLimit>
+{
+	fn force_from(ons: &OpaqueNetworkState) -> Self {
+		let peer_id = WeakBoundedVec::<_, PeerIdEncodingLimit>::force_from(
+			ons.peer_id.0.clone(),
+			Some(
+				"Warning: The size of the encoding of PeerId \
+  				is bigger than expected. A runtime configuration \
+  				adjustment may be needed.",
+			),
+		);
+
+		let external_addresses = WeakBoundedVec::<_, AddressesLimit>::force_from(
+			ons.external_addresses
+				.iter()
+				.map(|x| {
+					WeakBoundedVec::<_, MultiAddrEncodingLimit>::force_from(
+						x.0.clone(),
+						Some(
+							"Warning: The size of the encoding of MultiAddr \
+  							is bigger than expected. A runtime configuration \
+  							adjustment may be needed.",
+						),
+					)
+				})
+				.collect(),
+			Some(
+				"Warning: The network has more peers than expected \
+  				A runtime configuration adjustment may be needed.",
+			),
+		);
+
+		Self { peer_id, external_addresses }
+	}
+}
+
 /// A type for representing the validator id in a session.
 pub type ValidatorId<T> = <<T as Config>::ValidatorSet as ValidatorSet<
 	<T as frame_system::Config>::AccountId,
@@ -251,6 +315,7 @@ pub mod pallet {
 
 	#[pallet::pallet]
 	#[pallet::generate_store(pub(super) trait Store)]
+	#[pallet::generate_storage_info]
 	pub struct Pallet<T>(_);
 
 	#[pallet::config]
@@ -261,7 +326,18 @@ pub mod pallet {
 			+ RuntimeAppPublic
 			+ Default
 			+ Ord
-			+ MaybeSerializeDeserialize;
+			+ MaybeSerializeDeserialize
+			+ MaxEncodedLen;
+
+		/// The maximum number of keys that can be added.
+		type MaxKeys: Get<u32>;
+
+		/// The maximum number of peers to be stored in `ReceivedHeartbeats`
+		type MaxPeerInHeartbeats: Get<u32>;
+
+		/// The maximum size of the encoding of `PeerId` and `MultiAddr` that are coming
+		/// from the hearbeat
+		type MaxPeerDataEncodingSize: Get<u32>;
 
 		/// The overarching event type.
 		type Event: From<Event<Self>> + IsType<<Self as frame_system::Config>::Event>;
@@ -333,14 +409,27 @@ pub mod pallet {
 	/// The current set of keys that may issue a heartbeat.
 	#[pallet::storage]
 	#[pallet::getter(fn keys)]
-	pub(crate) type Keys<T: Config> = StorageValue<_, Vec<T::AuthorityId>, ValueQuery>;
+	pub(crate) type Keys<T: Config> =
+		StorageValue<_, WeakBoundedVec<T::AuthorityId, T::MaxKeys>, ValueQuery>;
 
-	/// For each session index, we keep a mapping of `AuthIndex` to
-	/// `offchain::OpaqueNetworkState`.
+	/// For each session index, we keep a mapping of 'SessionIndex` and `AuthIndex` to
+	/// `WrapperOpaque<BoundedOpaqueNetworkState>`.
 	#[pallet::storage]
 	#[pallet::getter(fn received_heartbeats)]
-	pub(crate) type ReceivedHeartbeats<T> =
-		StorageDoubleMap<_, Twox64Concat, SessionIndex, Twox64Concat, AuthIndex, Vec<u8>>;
+	pub(crate) type ReceivedHeartbeats<T: Config> = StorageDoubleMap<
+		_,
+		Twox64Concat,
+		SessionIndex,
+		Twox64Concat,
+		AuthIndex,
+		WrapperOpaque<
+			BoundedOpaqueNetworkState<
+				T::MaxPeerDataEncodingSize,
+				T::MaxPeerDataEncodingSize,
+				T::MaxPeerInHeartbeats,
+			>,
+		>,
+	>;
 
 	/// For each session index, we keep a mapping of `ValidatorId<T>` to the
 	/// number of blocks authored by the given authority.
@@ -409,11 +498,15 @@ pub mod pallet {
 			if let (false, Some(public)) = (exists, public) {
 				Self::deposit_event(Event::<T>::HeartbeatReceived(public.clone()));
 
-				let network_state = heartbeat.network_state.encode();
+				let network_state_bounded = BoundedOpaqueNetworkState::<
+					T::MaxPeerDataEncodingSize,
+					T::MaxPeerDataEncodingSize,
+					T::MaxPeerInHeartbeats,
+				>::force_from(&heartbeat.network_state);
 				ReceivedHeartbeats::<T>::insert(
 					&current_session,
 					&heartbeat.authority_index,
-					&network_state,
+					WrapperOpaque::from(network_state_bounded),
 				);
 
 				Ok(())
@@ -739,13 +832,17 @@ impl<T: Config> Pallet<T> {
 	fn initialize_keys(keys: &[T::AuthorityId]) {
 		if !keys.is_empty() {
 			assert!(Keys::<T>::get().is_empty(), "Keys are already initialized!");
-			Keys::<T>::put(keys);
+			let bounded_keys = <BoundedSlice<'_, _, T::MaxKeys>>::try_from(keys)
+				.expect("More than the maximum number of keys provided");
+			Keys::<T>::put(bounded_keys);
 		}
 	}
 
 	#[cfg(test)]
 	fn set_keys(keys: Vec<T::AuthorityId>) {
-		Keys::<T>::put(&keys)
+		let bounded_keys = WeakBoundedVec::<_, T::MaxKeys>::try_from(keys)
+			.expect("More than the maximum number of keys provided");
+		Keys::<T>::put(bounded_keys);
 	}
 }
 
@@ -776,7 +873,15 @@ impl<T: Config> OneSessionHandler<T::AccountId> for Pallet<T> {
 		<HeartbeatAfter<T>>::put(block_number + half_session);
 
 		// Remember who the authorities are for the new session.
-		Keys::<T>::put(validators.map(|x| x.1).collect::<Vec<_>>());
+		let keys = validators.map(|x| x.1).collect::<Vec<_>>();
+		let bounded_keys = WeakBoundedVec::<_, T::MaxKeys>::force_from(
+			keys,
+			Some(
+				"Warning: The session has more keys than expected. \
+  				A runtime configuration adjustment may be needed.",
+			),
+		);
+		Keys::<T>::put(bounded_keys);
 	}
 
 	fn on_before_session_ending() {

frame/im-online/src/mock.rs

@@ -217,6 +217,9 @@ impl frame_support::traits::EstimateNextSessionRotation<u64> for TestNextSession
 
 parameter_types! {
 	pub const UnsignedPriority: u64 = 1 << 20;
+	pub const MaxKeys: u32 = 10_000;
+	pub const MaxPeerInHeartbeats: u32 = 10_000;
+	pub const MaxPeerDataEncodingSize: u32 = 1_000;
 }
 
 impl Config for Runtime {
@@ -227,6 +230,9 @@ impl Config for Runtime {
 	type ReportUnresponsiveness = OffenceHandler;
 	type UnsignedPriority = UnsignedPriority;
 	type WeightInfo = ();
+	type MaxKeys = MaxKeys;
+	type MaxPeerInHeartbeats = MaxPeerInHeartbeats;
+	type MaxPeerDataEncodingSize = MaxPeerDataEncodingSize;
 }
 
 impl<LocalCall> frame_system::offchain::SendTransactionTypes<LocalCall> for Runtime

frame/offences/benchmarking/src/mock.rs

@@ -146,6 +146,9 @@ pallet_staking_reward_curve::build! {
 parameter_types! {
 	pub const RewardCurve: &'static sp_runtime::curve::PiecewiseLinear<'static> = &I_NPOS;
 	pub const MaxNominatorRewardedPerValidator: u32 = 64;
+	pub const MaxKeys: u32 = 10_000;
+	  pub const MaxPeerInHeartbeats: u32 = 10_000;
+	  pub const MaxPeerDataEncodingSize: u32 = 1_000;
 }
 
 pub type Extrinsic = sp_runtime::testing::TestXt<Call, ()>;
@@ -186,6 +189,9 @@ impl pallet_im_online::Config for Test {
 	type ReportUnresponsiveness = Offences;
 	type UnsignedPriority = ();
 	type WeightInfo = ();
+	type MaxKeys = MaxKeys;
+	type MaxPeerInHeartbeats = MaxPeerInHeartbeats;
+	type MaxPeerDataEncodingSize = MaxPeerDataEncodingSize;
 }
 
 impl pallet_offences::Config for Test {

frame/session/src/lib.rs

@@ -114,7 +114,7 @@ mod mock;
 mod tests;
 pub mod weights;
 
-use codec::Decode;
+use codec::{Decode, MaxEncodedLen};
 use frame_support::{
 	decl_error, decl_event, decl_module, decl_storage,
 	dispatch::{self, DispatchError, DispatchResult},
@@ -367,7 +367,7 @@ pub trait Config: frame_system::Config {
 	type Event: From<Event> + Into<<Self as frame_system::Config>::Event>;
 
 	/// A stable ID for a validator.
-	type ValidatorId: Member + Parameter;
+	type ValidatorId: Member + Parameter + MaxEncodedLen;
 
 	/// A conversion from account ID to validator ID.
 	///

frame/support/src/traits.rs

@@ -52,7 +52,7 @@ mod misc;
 pub use misc::{
 	Backing, ConstU32, EnsureInherentsAreFirst, EstimateCallFee, ExecuteBlock, ExtrinsicCall, Get,
 	GetBacking, GetDefault, HandleLifetime, IsSubType, IsType, Len, OffchainWorker,
-	OnKilledAccount, OnNewAccount, SameOrOther, Time, TryDrop, UnixTime,
+	OnKilledAccount, OnNewAccount, SameOrOther, Time, TryDrop, UnixTime, WrapperOpaque,
 };
 
 mod stored_map;

frame/support/src/traits/misc.rs

@@ -17,8 +17,10 @@
 
 //! Smaller traits used in FRAME which don't need their own file.
 
-use crate::dispatch::Parameter;
+use crate::{dispatch::Parameter, TypeInfo};
+use codec::{Decode, Encode, EncodeLike, Input, MaxEncodedLen};
 use sp_runtime::{traits::Block as BlockT, DispatchError};
+use sp_std::vec::Vec;
 
 /// Anything that can have a `::len()` method.
 pub trait Len {
@@ -377,3 +379,48 @@ impl<Call, Balance: From<u32>, const T: u32> EstimateCallFee<Call, Balance> for
 		T.into()
 	}
 }
+
+/// A wrapper for any type `T` which implement encode/decode in a way compatible with `Vec<u8>`.
+///
+/// The encoding is the encoding of `T` prepended with the compact encoding of its size in bytes.
+/// Thus the encoded value can be decoded as a `Vec<u8>`.
+#[derive(Debug, Eq, PartialEq, Default, Clone, MaxEncodedLen, TypeInfo)]
+#[cfg_attr(feature = "std", derive(serde::Serialize, serde::Deserialize))]
+pub struct WrapperOpaque<T>(pub T);
+
+impl<T: Encode> EncodeLike for WrapperOpaque<T> {}
+
+impl<T: Encode> Encode for WrapperOpaque<T> {
+	fn size_hint(&self) -> usize {
+		// Compact<u32> usually takes at most 4 bytes
+		self.0.size_hint().saturating_add(4)
+	}
+
+	fn encode_to<O: codec::Output + ?Sized>(&self, dest: &mut O) {
+		self.0.encode().encode_to(dest);
+	}
+
+	fn encode(&self) -> Vec<u8> {
+		self.0.encode().encode()
+	}
+
+	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
+		self.0.encode().using_encoded(f)
+	}
+}
+
+impl<T: Decode> Decode for WrapperOpaque<T> {
+	fn decode<I: Input>(input: &mut I) -> Result<Self, codec::Error> {
+		Ok(Self(T::decode(&mut &<Vec<u8>>::decode(input)?[..])?))
+	}
+
+	fn skip<I: Input>(input: &mut I) -> Result<(), codec::Error> {
+		<Vec<u8>>::skip(input)
+	}
+}
+
+impl<T> From<T> for WrapperOpaque<T> {
+	fn from(t: T) -> Self {
+		Self(t)
+	}
+}

frame/support/src/traits/validation.rs

@@ -18,7 +18,7 @@
 //! Traits for dealing with validation and validators.
 
 use crate::{dispatch::Parameter, weights::Weight};
-use codec::{Codec, Decode};
+use codec::{Codec, Decode, MaxEncodedLen};
 use sp_runtime::{
 	traits::{Convert, Zero},
 	BoundToRuntimeAppPublic, ConsensusEngineId, Permill, RuntimeAppPublic,
@@ -31,7 +31,7 @@ use sp_std::prelude::*;
 /// Something that can give information about the current validator set.
 pub trait ValidatorSet<AccountId> {
 	/// Type for representing validator id in a session.
-	type ValidatorId: Parameter;
+	type ValidatorId: Parameter + MaxEncodedLen;
 	/// A type for converting `AccountId` to `ValidatorId`.
 	type ValidatorIdOf: Convert<AccountId, Option<Self::ValidatorId>>;