From abb39da13a060081973202ebccad5e927bd25104 Mon Sep 17 00:00:00 2001
From: Vincent Geddes <vincent.geddes@hey.com>
Date: Thu, 14 Dec 2023 21:36:40 +0200
Subject: [PATCH] Check if max nonce reached in outbound-queue pallet

---
 parachain/pallets/outbound-queue/src/lib.rs  | 25 ++++++++------
 parachain/pallets/outbound-queue/src/test.rs | 34 ++++++++++++++++++--
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/parachain/pallets/outbound-queue/src/lib.rs b/parachain/pallets/outbound-queue/src/lib.rs
index c8f2e2d6d0..14f5912a95 100644
--- a/parachain/pallets/outbound-queue/src/lib.rs
+++ b/parachain/pallets/outbound-queue/src/lib.rs
@@ -198,6 +198,8 @@ pub mod pallet {
 		InvalidFeeConfig,
 		/// Invalid Channel
 		InvalidChannel,
+		/// Maximum Nonce reached
+		MaxNonceReached,
 	}
 
 	/// Messages to be committed in the current block. This storage value is killed in
@@ -308,10 +310,19 @@ pub mod pallet {
 			let queued_message: QueuedMessage =
 				versioned_queued_message.try_into().map_err(|_| Unsupported)?;
 
-			let pricing_params = T::PricingParameters::get();
-
-			let next_nonce = Nonce::<T>::get(queued_message.channel_id).saturating_add(1);
+			// Obtain next nonce
+			let nonce = <Nonce<T>>::try_mutate(
+				queued_message.channel_id,
+				|nonce| -> Result<u64, ProcessMessageError> {
+					if *nonce == u64::MAX {
+						return Err(Unsupported)
+					}
+					*nonce = nonce.saturating_add(1);
+					Ok(*nonce)
+				},
+			)?;
 
+			let pricing_params = T::PricingParameters::get();
 			let command = queued_message.command.index();
 			let params = queued_message.command.abi_encode();
 			let max_dispatch_gas =
@@ -321,7 +332,7 @@ pub mod pallet {
 			// Construct the final committed message
 			let message = CommittedMessage {
 				channel_id: queued_message.channel_id,
-				nonce: next_nonce,
+				nonce,
 				command,
 				params,
 				max_dispatch_gas,
@@ -339,12 +350,8 @@ pub mod pallet {
 
 			Messages::<T>::append(Box::new(message));
 			MessageLeaves::<T>::append(message_abi_encoded_hash);
-			Nonce::<T>::set(queued_message.channel_id, next_nonce);
 
-			Self::deposit_event(Event::MessageAccepted {
-				id: queued_message.id,
-				nonce: next_nonce,
-			});
+			Self::deposit_event(Event::MessageAccepted { id: queued_message.id, nonce });
 
 			Ok(true)
 		}
diff --git a/parachain/pallets/outbound-queue/src/test.rs b/parachain/pallets/outbound-queue/src/test.rs
index 9bdc636495..c73d51fdca 100644
--- a/parachain/pallets/outbound-queue/src/test.rs
+++ b/parachain/pallets/outbound-queue/src/test.rs
@@ -99,16 +99,46 @@ fn process_message_yields_on_max_messages_per_block() {
 	})
 }
 
+#[test]
+fn process_message_fails_on_max_nonce_reached() {
+	new_tester().execute_with(|| {
+		let sibling_id = 1000;
+		let channel_id: ChannelId = ParaId::from(sibling_id).into();
+		let origin = AggregateMessageOrigin::Snowbridge(channel_id.into());
+		let message: QueuedMessage = QueuedMessage {
+			id: H256::zero(),
+			channel_id,
+			command: mock_message(sibling_id).command,
+		};
+		let versioned_queued_message: VersionedQueuedMessage = message.try_into().unwrap();
+		let encoded = versioned_queued_message.encode();
+		let mut meter = WeightMeter::with_limit(Weight::MAX);
+
+		Nonce::<Test>::set(channel_id, u64::MAX);
+
+		assert_noop!(
+			OutboundQueue::process_message(&encoded.as_slice(), origin, &mut meter, &mut [0u8; 32]),
+			ProcessMessageError::Unsupported
+		);
+	})
+}
+
 #[test]
 fn process_message_fails_on_overweight_message() {
 	new_tester().execute_with(|| {
 		let sibling_id = 1000;
 		let channel_id: ChannelId = ParaId::from(sibling_id).into();
 		let origin = AggregateMessageOrigin::Snowbridge(channel_id.into());
-		let message = mock_message(sibling_id).encode();
+		let message: QueuedMessage = QueuedMessage {
+			id: H256::zero(),
+			channel_id,
+			command: mock_message(sibling_id).command,
+		};
+		let versioned_queued_message: VersionedQueuedMessage = message.try_into().unwrap();
+		let encoded = versioned_queued_message.encode();
 		let mut meter = WeightMeter::with_limit(Weight::from_parts(1, 1));
 		assert_noop!(
-			OutboundQueue::process_message(&message.as_slice(), origin, &mut meter, &mut [0u8; 32]),
+			OutboundQueue::process_message(&encoded.as_slice(), origin, &mut meter, &mut [0u8; 32]),
 			ProcessMessageError::Overweight(<Test as Config>::WeightInfo::do_process_message())
 		);
 	})