Fixed messages count check (#659)

* fixed messages count check

* explicit check of `messages_count` in the receive_messages_proof

* change messages_count to be u32

* Update modules/message-lane/src/lib.rs

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

bridges/bin/millau/runtime/src/rialto_messages.rs

@@ -192,7 +192,7 @@ impl SourceHeaderChain<bp_rialto::Balance> for Rialto {
 
 	fn verify_messages_proof(
 		proof: Self::MessagesProof,
-		messages_count: MessageNonce,
+		messages_count: u32,
 	) -> Result<ProvedMessages<Message<bp_rialto::Balance>>, Self::Error> {
 		messages::target::verify_messages_proof::<WithRialtoMessageBridge, Runtime>(proof, messages_count)
 	}

bridges/bin/rialto/runtime/src/millau_messages.rs

@@ -193,7 +193,7 @@ impl SourceHeaderChain<bp_millau::Balance> for Millau {
 
 	fn verify_messages_proof(
 		proof: Self::MessagesProof,
-		messages_count: MessageNonce,
+		messages_count: u32,
 	) -> Result<ProvedMessages<Message<bp_millau::Balance>>, Self::Error> {
 		messages::target::verify_messages_proof::<WithMillauMessageBridge, Runtime>(proof, messages_count)
 	}

bridges/bin/runtime-common/src/messages.rs

@@ -396,9 +396,13 @@ pub mod target {
 	}
 
 	/// Verify proof of Bridged -> This chain messages.
+	///
+	/// The `messages_count` argument verification (sane limits) is supposed to be made
+	/// outside of this function. This function only verifies that the proof declares exactly
+	/// `messages_count` messages.
 	pub fn verify_messages_proof<B: MessageBridge, ThisRuntime>(
 		proof: FromBridgedChainMessagesProof<B>,
-		messages_count: MessageNonce,
+		messages_count: u32,
 	) -> Result<ProvedMessages<Message<BalanceOf<BridgedChain<B>>>>, &'static str>
 	where
 		ThisRuntime: pallet_substrate_bridge::Config,
@@ -487,7 +491,7 @@ pub mod target {
 	/// Verify proof of Bridged -> This chain messages using given message proof parser.
 	pub(crate) fn verify_messages_proof_with_parser<B: MessageBridge, BuildParser, Parser>(
 		proof: FromBridgedChainMessagesProof<B>,
-		messages_count: MessageNonce,
+		messages_count: u32,
 		build_parser: BuildParser,
 	) -> Result<ProvedMessages<Message<BalanceOf<BridgedChain<B>>>>, MessageProofError>
 	where
@@ -497,11 +501,18 @@ pub mod target {
 		let (bridged_header_hash, bridged_storage_proof, lane_id, begin, end) = proof;
 
 		// receiving proofs where end < begin is ok (if proof includes outbound lane state)
-		// => hence unwrap_or(0)
+		let messages_in_the_proof = if let Some(nonces_difference) = end.checked_sub(begin) {
-		let messages_in_the_proof = end.checked_sub(begin).and_then(|diff| diff.checked_add(1)).unwrap_or(0);
+			// let's check that the user (relayer) has passed correct `messages_count`
-		if messages_in_the_proof != messages_count {
+			// (this bounds maximal capacity of messages vec below)
-			return Err(MessageProofError::MessagesCountMismatch);
+			let messages_in_the_proof = nonces_difference.saturating_add(1);
-		}
+			if messages_in_the_proof != MessageNonce::from(messages_count) {
+				return Err(MessageProofError::MessagesCountMismatch);
+			}
+
+			messages_in_the_proof
+		} else {
+			0
+		};
 
 		let parser = build_parser(bridged_header_hash, bridged_storage_proof)?;
 
@@ -509,7 +520,7 @@ pub mod target {
 		// be in the proof. So any error in `read_value`, or even missing value is fatal.
 		//
 		// Mind that we allow proofs with no messages if outbound lane state is proved.
-		let mut messages = Vec::with_capacity(end.saturating_sub(begin) as _);
+		let mut messages = Vec::with_capacity(messages_in_the_proof as _);
 		for nonce in begin..=end {
 			let message_key = MessageKey { lane_id, nonce };
 			let raw_message_data = parser
@@ -1183,4 +1194,30 @@ mod tests {
 			.collect()),
 		);
 	}
+
+	#[test]
+	fn verify_messages_proof_with_parser_does_not_panic_if_messages_count_mismatches() {
+		assert_eq!(
+			target::verify_messages_proof_with_parser::<OnThisChainBridge, _, _>(
+				(
+					Default::default(),
+					StorageProof::new(vec![]),
+					Default::default(),
+					0,
+					u64::MAX
+				),
+				0,
+				|_, _| Ok(TestMessageProofParser {
+					failing: false,
+					messages: 0..=u64::MAX,
+					outbound_lane_data: Some(OutboundLaneData {
+						oldest_unpruned_nonce: 1,
+						latest_received_nonce: 1,
+						latest_generated_nonce: 1,
+					}),
+				}),
+			),
+			Err(target::MessageProofError::MessagesCountMismatch),
+		);
+	}
 }

bridges/modules/message-lane/src/lib.rs

@@ -101,6 +101,10 @@ pub trait Config<I = DefaultInstance>: frame_system::Config {
 	///
 	/// There is no point of making this parameter lesser than MaxUnrewardedRelayerEntriesAtInboundLane,
 	/// because then maximal number of relayer entries will be limited by maximal number of messages.
+	///
+	/// This value also represents maximal number of messages in single delivery transaction. Transaction
+	/// that is declaring more messages than this value, will be rejected. Even if these messages are
+	/// from different lanes.
 	type MaxUnconfirmedMessagesAtInboundLane: Get<MessageNonce>;
 
 	/// Payload type of outbound messages. This payload is dispatched on the bridged chain.
@@ -156,6 +160,8 @@ decl_error! {
 		MessageRejectedByLaneVerifier,
 		/// Submitter has failed to pay fee for delivering and dispatching messages.
 		FailedToWithdrawMessageFee,
+		/// The transaction brings too many messages.
+		TooManyMessagesInTheProof,
 		/// Invalid messages has been submitted.
 		InvalidMessagesProof,
 		/// Invalid messages dispatch weight has been declared by the relayer.
@@ -344,19 +350,25 @@ decl_module! {
 		/// this data in the transaction, so reward confirmations lags should be minimal.
 		#[weight = T::WeightInfo::receive_messages_proof_overhead()
 			.saturating_add(T::WeightInfo::receive_messages_proof_outbound_lane_state_overhead())
-			.saturating_add(T::WeightInfo::receive_messages_proof_messages_overhead(*messages_count))
+			.saturating_add(T::WeightInfo::receive_messages_proof_messages_overhead(MessageNonce::from(*messages_count)))
 			.saturating_add(*dispatch_weight)
 		]
 		pub fn receive_messages_proof(
 			origin,
 			relayer_id: T::InboundRelayer,
 			proof: MessagesProofOf<T, I>,
-			messages_count: MessageNonce,
+			messages_count: u32,
 			dispatch_weight: Weight,
 		) -> DispatchResult {
 			ensure_operational::<T, I>()?;
 			let _ = ensure_signed(origin)?;
 
+			// reject transactions that are declaring too many messages
+			ensure!(
+				MessageNonce::from(messages_count) <= T::MaxUnconfirmedMessagesAtInboundLane::get(),
+				Error::<T, I>::TooManyMessagesInTheProof
+			);
+
 			// verify messages proof && convert proof into messages
 			let messages = verify_and_decode_messages_proof::<
 				T::SourceHeaderChain,
@@ -457,7 +469,8 @@ decl_module! {
 			// verify that the relayer has declared correct `lane_data::relayers` state
 			// (we only care about total number of entries and messages, because this affects call weight)
 			ensure!(
-				total_unrewarded_messages(&lane_data.relayers) == relayers_state.total_messages
+				total_unrewarded_messages(&lane_data.relayers)
+					.unwrap_or(MessageNonce::MAX) == relayers_state.total_messages
 					&& lane_data.relayers.len() as MessageNonce == relayers_state.unrewarded_relayer_entries,
 				Error::<T, I>::InvalidUnrewardedRelayersState
 			);
@@ -545,7 +558,7 @@ impl<T: Config<I>, I: Instance> Module<T, I> {
 		bp_message_lane::UnrewardedRelayersState {
 			unrewarded_relayer_entries: relayers.len() as _,
 			messages_in_oldest_entry: relayers.front().map(|(begin, end, _)| 1 + end - begin).unwrap_or(0),
-			total_messages: total_unrewarded_messages(&relayers),
+			total_messages: total_unrewarded_messages(&relayers).unwrap_or(MessageNonce::MAX),
 		}
 	}
 
@@ -733,8 +746,11 @@ impl<T: Config<I>, I: Instance> OutboundLaneStorage for RuntimeOutboundLaneStora
 /// Verify messages proof and return proved messages with decoded payload.
 fn verify_and_decode_messages_proof<Chain: SourceHeaderChain<Fee>, Fee, DispatchPayload: Decode>(
 	proof: Chain::MessagesProof,
-	messages_count: MessageNonce,
+	messages_count: u32,
 ) -> Result<ProvedMessages<DispatchMessage<DispatchPayload, Fee>>, Chain::Error> {
+	// `receive_messages_proof` weight formula and `MaxUnconfirmedMessagesAtInboundLane` check
+	// guarantees that the `message_count` is sane and Vec<Message> may be allocated.
+	// (tx with too many messages will either be rejected from the pool, or will fail earlier)
 	Chain::verify_messages_proof(proof, messages_count).map(|messages_by_lane| {
 		messages_by_lane
 			.into_iter()
@@ -1073,6 +1089,22 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn receive_messages_proof_rejects_proof_with_too_many_messages() {
+		run_test(|| {
+			assert_noop!(
+				Module::<TestRuntime, DefaultInstance>::receive_messages_proof(
+					Origin::signed(1),
+					TEST_RELAYER_A,
+					Ok(vec![message(1, REGULAR_PAYLOAD)]).into(),
+					u32::MAX,
+					0,
+				),
+				Error::<TestRuntime, DefaultInstance>::TooManyMessagesInTheProof,
+			);
+		});
+	}
+
 	#[test]
 	fn receive_messages_delivery_proof_works() {
 		run_test(|| {

bridges/modules/message-lane/src/mock.rs

@@ -309,7 +309,7 @@ impl SourceHeaderChain<TestMessageFee> for TestSourceHeaderChain {
 
 	fn verify_messages_proof(
 		proof: Self::MessagesProof,
-		_messages_count: MessageNonce,
+		_messages_count: u32,
 	) -> Result<ProvedMessages<Message<TestMessageFee>>, Self::Error> {
 		proof
 			.result

bridges/primitives/message-lane/src/lib.rs

@@ -158,11 +158,36 @@ impl Default for OutboundLaneData {
 }
 
 /// Returns total number of messages in the `InboundLaneData::relayers` vector.
+///
+/// Returns `None` if there are more messages that `MessageNonce` may fit (i.e. `MessageNonce + 1`).
 pub fn total_unrewarded_messages<RelayerId>(
 	relayers: &VecDeque<(MessageNonce, MessageNonce, RelayerId)>,
-) -> MessageNonce {
+) -> Option<MessageNonce> {
 	match (relayers.front(), relayers.back()) {
-		(Some((begin, _, _)), Some((_, end, _))) => end.checked_sub(*begin).and_then(|d| d.checked_add(1)).unwrap_or(0),
+		(Some((begin, _, _)), Some((_, end, _))) => {
-		_ => 0,
+			if let Some(difference) = end.checked_sub(*begin) {
+				difference.checked_add(1)
+			} else {
+				Some(0)
+			}
+		}
+		_ => Some(0),
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn total_unrewarded_messages_does_not_overflow() {
+		assert_eq!(
+			total_unrewarded_messages(
+				&vec![(0, 0, 1), (MessageNonce::MAX, MessageNonce::MAX, 2)]
+					.into_iter()
+					.collect()
+			),
+			None,
+		);
 	}
 }

bridges/primitives/message-lane/src/target_chain.rs

@@ -16,7 +16,7 @@
 
 //! Primitives of message lane module, that are used on the target chain.
 
-use crate::{LaneId, Message, MessageData, MessageKey, MessageNonce, OutboundLaneData};
+use crate::{LaneId, Message, MessageData, MessageKey, OutboundLaneData};
 
 use codec::{Decode, Encode, Error as CodecError};
 use frame_support::{weights::Weight, Parameter, RuntimeDebug};
@@ -72,9 +72,13 @@ pub trait SourceHeaderChain<Fee> {
 	///
 	/// Messages vector is required to be sorted by nonce within each lane. Out-of-order
 	/// messages will be rejected.
+	///
+	/// The `messages_count` argument verification (sane limits) is supposed to be made
+	/// outside of this function. This function only verifies that the proof declares exactly
+	/// `messages_count` messages.
 	fn verify_messages_proof(
 		proof: Self::MessagesProof,
-		messages_count: MessageNonce,
+		messages_count: u32,
 	) -> Result<ProvedMessages<Message<Fee>>, Self::Error>;
 }
 

bridges/relays/substrate/src/millau_messages_to_rialto.rs

@@ -81,7 +81,7 @@ impl SubstrateMessageLane for MillauMessagesToRialto {
 		let call = rialto_runtime::MessageLaneCall::receive_messages_proof(
 			self.relayer_id_at_source.clone(),
 			proof,
-			messages_count,
+			messages_count as _,
 			dispatch_weight,
 		)
 		.into();

bridges/relays/substrate/src/rialto_messages_to_millau.rs

@@ -81,7 +81,7 @@ impl SubstrateMessageLane for RialtoMessagesToMillau {
 		let call = millau_runtime::MessageLaneCall::receive_messages_proof(
 			self.relayer_id_at_source.clone(),
 			proof,
-			messages_count,
+			messages_count as _,
 			dispatch_weight,
 		)
 		.into();