Update weight for im-online (#5771)

* implementation using Keys::decode_len

* add heartbeat new param

* address issues

* improve doc

* fix test

* fix overflow

substrate/frame/im-online/src/benchmarking.rs

@@ -23,7 +23,7 @@ use super::*;
 use frame_system::RawOrigin;
 use frame_benchmarking::benchmarks;
 use sp_core::offchain::{OpaquePeerId, OpaqueMultiaddr};
-use sp_runtime::traits::{ValidateUnsigned, Zero};
+use sp_runtime::traits::{ValidateUnsigned, Zero, Dispatchable};
 use sp_runtime::transaction_validity::TransactionSource;
 
 use crate::Module as ImOnline;
@@ -49,6 +49,7 @@ pub fn create_heartbeat<T: Trait>(k: u32, e: u32) ->
 		network_state,
 		session_index: 0,
 		authority_index: k-1,
+		validators_len: keys.len() as u32,
 	};
 
 	let encoded_heartbeat = input_heartbeat.encode();
@@ -75,6 +76,16 @@ benchmarks! {
 	}: {
 		ImOnline::<T>::validate_unsigned(TransactionSource::InBlock, &call)?;
 	}
+
+	validate_unsigned_and_then_heartbeat {
+		let k in 1 .. MAX_KEYS;
+		let e in 1 .. MAX_EXTERNAL_ADDRESSES;
+		let (input_heartbeat, signature) = create_heartbeat::<T>(k, e)?;
+		let call = Call::heartbeat(input_heartbeat, signature);
+	}: {
+		ImOnline::<T>::validate_unsigned(TransactionSource::InBlock, &call)?;
+		call.dispatch(RawOrigin::None.into())?;
+	}
 }
 
 #[cfg(test)]
@@ -88,6 +99,7 @@ mod tests {
 		new_test_ext().execute_with(|| {
 			assert_ok!(test_benchmark_heartbeat::<Runtime>());
 			assert_ok!(test_benchmark_validate_unsigned::<Runtime>());
+			assert_ok!(test_benchmark_validate_unsigned_and_then_heartbeat::<Runtime>());
 		});
 	}
 }

substrate/frame/im-online/src/lib.rs

@@ -94,6 +94,7 @@ use sp_staking::{
 use frame_support::{
 	decl_module, decl_event, decl_storage, Parameter, debug, decl_error,
 	traits::Get,
+	weights::Weight,
 };
 use frame_system::{self as system, ensure_none};
 use frame_system::offchain::{
@@ -220,6 +221,8 @@ pub struct Heartbeat<BlockNumber>
 	pub session_index: SessionIndex,
 	/// An index of the authority on the list of validators.
 	pub authority_index: AuthIndex,
+	/// The length of session validator set
+	pub validators_len: u32,
 }
 
 pub trait Trait: SendTransactionTypes<Call<Self>> + pallet_session::historical::Trait {
@@ -313,13 +316,30 @@ decl_module! {
 
 		fn deposit_event() = default;
 
-		#[weight = 0]
+		/// # <weight>
+		/// - Complexity: `O(K + E)` where K is length of `Keys` and E is length of
+		///   `Heartbeat.network_state.external_address`
+		///
+		///   - `O(K)`: decoding of length `K`
+		///   - `O(E)`: decoding/encoding of length `E`
+		/// - DbReads: pallet_session `Validators`, pallet_session `CurrentIndex`, `Keys`,
+		///   `ReceivedHeartbeats`
+		/// - DbWrites: `ReceivedHeartbeats`
+		/// # </weight>
+		// NOTE: the weight include cost of validate_unsigned as it is part of the cost to import
+		// block with such an extrinsic.
+		#[weight = (310_000_000 + T::DbWeight::get().reads_writes(4, 1))
+			.saturating_add(750_000.saturating_mul(heartbeat.validators_len as Weight))
+			.saturating_add(
+				1_200_000.saturating_mul(heartbeat.network_state.external_addresses.len() as Weight)
+			)
+		]
 		fn heartbeat(
 			origin,
 			heartbeat: Heartbeat<T::BlockNumber>,
 			// since signature verification is done in `validate_unsigned`
 			// we can skip doing it here again.
-			_signature: <T::AuthorityId as RuntimeAppPublic>::Signature
+			_signature: <T::AuthorityId as RuntimeAppPublic>::Signature,
 		) {
 			ensure_none(origin)?;
 
@@ -439,10 +459,17 @@ impl<T: Trait> Module<T> {
 		}
 
 		let session_index = <pallet_session::Module<T>>::current_index();
+		let validators_len = <pallet_session::Module<T>>::validators().len() as u32;
 
 		Ok(Self::local_authority_keys()
 			.map(move |(authority_index, key)|
-				Self::send_single_heartbeat(authority_index, key, session_index, block_number)
+				Self::send_single_heartbeat(
+					authority_index,
+					key,
+					session_index,
+					block_number,
+					validators_len,
+				)
 			))
 	}
 
@@ -451,7 +478,8 @@ impl<T: Trait> Module<T> {
 		authority_index: u32,
 		key: T::AuthorityId,
 		session_index: SessionIndex,
-		block_number: T::BlockNumber
+		block_number: T::BlockNumber,
+		validators_len: u32,
 	) -> OffchainResult<T, ()> {
 		// A helper function to prepare heartbeat call.
 		let prepare_heartbeat = || -> OffchainResult<T, Call<T>> {
@@ -462,6 +490,7 @@ impl<T: Trait> Module<T> {
 				network_state,
 				session_index,
 				authority_index,
+				validators_len,
 			};
 
 			let signature = key.sign(&heartbeat_data.encode()).ok_or(OffchainErr::FailedSigning)?;
@@ -643,6 +672,9 @@ impl<T: Trait> pallet_session::OneSessionHandler<T::AccountId> for Module<T> {
 	}
 }
 
+/// Invalid transaction custom error. Returned when validators_len field in heartbeat is incorrect.
+const INVALID_VALIDATORS_LEN: u8 = 10;
+
 impl<T: Trait> frame_support::unsigned::ValidateUnsigned for Module<T> {
 	type Call = Call<T>;
 
@@ -664,6 +696,9 @@ impl<T: Trait> frame_support::unsigned::ValidateUnsigned for Module<T> {
 
 			// verify that the incoming (unverified) pubkey is actually an authority id
 			let keys = Keys::<T>::get();
+			if keys.len() as u32 != heartbeat.validators_len {
+				return InvalidTransaction::Custom(INVALID_VALIDATORS_LEN).into();
+			}
 			let authority_id = match keys.get(heartbeat.authority_index as usize) {
 				Some(id) => id,
 				None => return InvalidTransaction::BadProof.into(),

substrate/frame/im-online/src/tests.rs

@@ -27,7 +27,7 @@ use sp_core::offchain::{
 	testing::{TestOffchainExt, TestTransactionPoolExt},
 };
 use frame_support::{dispatch, assert_noop};
-use sp_runtime::testing::UintAuthorityId;
+use sp_runtime::{testing::UintAuthorityId, transaction_validity::TransactionValidityError};
 
 #[test]
 fn test_unresponsiveness_slash_fraction() {
@@ -87,7 +87,7 @@ fn should_report_offline_validators() {
 
 		// should not report when heartbeat is sent
 		for (idx, v) in validators.into_iter().take(4).enumerate() {
-			let _ = heartbeat(block, 3, idx as u32, v.into()).unwrap();
+			let _ = heartbeat(block, 3, idx as u32, v.into(), Session::validators()).unwrap();
 		}
 		advance_session();
 
@@ -111,6 +111,7 @@ fn heartbeat(
 	session_index: u32,
 	authority_index: u32,
 	id: UintAuthorityId,
+	validators: Vec<u64>,
 ) -> dispatch::DispatchResult {
 	use frame_support::unsigned::ValidateUnsigned;
 
@@ -122,15 +123,20 @@ fn heartbeat(
 		},
 		session_index,
 		authority_index,
+		validators_len: validators.len() as u32,
 	};
 	let signature = id.sign(&heartbeat.encode()).unwrap();
 
 	ImOnline::pre_dispatch(&crate::Call::heartbeat(heartbeat.clone(), signature.clone()))
-		.map_err(|e| <&'static str>::from(e))?;
+		.map_err(|e| match e {
+			TransactionValidityError::Invalid(InvalidTransaction::Custom(INVALID_VALIDATORS_LEN)) =>
+				"invalid validators len",
+			e @ _ => <&'static str>::from(e),
+		})?;
 	ImOnline::heartbeat(
 		Origin::system(frame_system::RawOrigin::None),
 		heartbeat,
-		signature
+		signature,
 	)
 }
 
@@ -152,7 +158,7 @@ fn should_mark_online_validator_when_heartbeat_is_received() {
 		assert!(!ImOnline::is_online(2));
 
 		// when
-		let _ = heartbeat(1, 2, 0, 1.into()).unwrap();
+		let _ = heartbeat(1, 2, 0, 1.into(), Session::validators()).unwrap();
 
 		// then
 		assert!(ImOnline::is_online(0));
@@ -160,7 +166,7 @@ fn should_mark_online_validator_when_heartbeat_is_received() {
 		assert!(!ImOnline::is_online(2));
 
 		// and when
-		let _ = heartbeat(1, 2, 2, 3.into()).unwrap();
+		let _ = heartbeat(1, 2, 2, 3.into(), Session::validators()).unwrap();
 
 		// then
 		assert!(ImOnline::is_online(0));
@@ -170,7 +176,7 @@ fn should_mark_online_validator_when_heartbeat_is_received() {
 }
 
 #[test]
-fn late_heartbeat_should_fail() {
+fn late_heartbeat_and_invalid_keys_len_should_fail() {
 	new_test_ext().execute_with(|| {
 		advance_session();
 		// given
@@ -183,8 +189,11 @@ fn late_heartbeat_should_fail() {
 		assert_eq!(Session::validators(), vec![1, 2, 3]);
 
 		// when
-		assert_noop!(heartbeat(1, 3, 0, 1.into()), "Transaction is outdated");
-		assert_noop!(heartbeat(1, 1, 0, 1.into()), "Transaction is outdated");
+		assert_noop!(heartbeat(1, 3, 0, 1.into(), Session::validators()), "Transaction is outdated");
+		assert_noop!(heartbeat(1, 1, 0, 1.into(), Session::validators()), "Transaction is outdated");
+
+		// invalid validators_len
+		assert_noop!(heartbeat(1, 2, 0, 1.into(), vec![]), "invalid validators len");
 	});
 }
 
@@ -220,7 +229,7 @@ fn should_generate_heartbeats() {
 		// check stuff about the transaction.
 		let ex: Extrinsic = Decode::decode(&mut &*transaction).unwrap();
 		let heartbeat = match ex.call {
-			crate::mock::Call::ImOnline(crate::Call::heartbeat(h, _)) => h,
+			crate::mock::Call::ImOnline(crate::Call::heartbeat(h, ..)) => h,
 			e => panic!("Unexpected call: {:?}", e),
 		};
 
@@ -229,6 +238,7 @@ fn should_generate_heartbeats() {
 			network_state: sp_io::offchain::network_state().unwrap(),
 			session_index: 2,
 			authority_index: 2,
+			validators_len: 3,
 		});
 	});
 }
@@ -248,7 +258,7 @@ fn should_cleanup_received_heartbeats_on_session_end() {
 		assert_eq!(Session::validators(), vec![1, 2, 3]);
 
 		// send an heartbeat from authority id 0 at session 2
-		let _ = heartbeat(1, 2, 0, 1.into()).unwrap();
+		let _ = heartbeat(1, 2, 0, 1.into(), Session::validators()).unwrap();
 
 		// the heartbeat is stored
 		assert!(!ImOnline::received_heartbeats(&2, &0).is_none());
@@ -330,7 +340,7 @@ fn should_not_send_a_report_if_already_online() {
 		// check stuff about the transaction.
 		let ex: Extrinsic = Decode::decode(&mut &*transaction).unwrap();
 		let heartbeat = match ex.call {
-			crate::mock::Call::ImOnline(crate::Call::heartbeat(h, _)) => h,
+			crate::mock::Call::ImOnline(crate::Call::heartbeat(h, ..)) => h,
 			e => panic!("Unexpected call: {:?}", e),
 		};
 
@@ -339,6 +349,7 @@ fn should_not_send_a_report_if_already_online() {
 			network_state: sp_io::offchain::network_state().unwrap(),
 			session_index: 2,
 			authority_index: 0,
+			validators_len: 3,
 		});
 	});
 }