From 73c12b3f4f44e52e39610ccaef77826709c2fd0d Mon Sep 17 00:00:00 2001 From: vedhavyas Date: Tue, 21 Jan 2025 16:33:11 +0530 Subject: [PATCH] ensure max outgoing messages per channel is atleast one and ensure to return open channel that can still accept messages --- domains/pallets/messenger/src/benchmarking.rs | 3 ++ domains/pallets/messenger/src/lib.rs | 44 ++++++++++------ domains/pallets/messenger/src/messages.rs | 23 ++++++++- domains/pallets/messenger/src/tests.rs | 51 +++++++++++++++---- 4 files changed, 95 insertions(+), 26 deletions(-) diff --git a/domains/pallets/messenger/src/benchmarking.rs b/domains/pallets/messenger/src/benchmarking.rs index 8616d2b3c4..9422459e19 100644 --- a/domains/pallets/messenger/src/benchmarking.rs +++ b/domains/pallets/messenger/src/benchmarking.rs @@ -203,6 +203,9 @@ mod benchmarks { last_delivered_message_response_nonce: None, }; Outbox::::insert((dst_chain_id, channel_id, next_outbox_nonce), req_msg); + OutboxMessageCount::::mutate((dst_chain_id, channel_id), |count| { + *count = count.saturating_add(1u32); + }); // Insert a dummy response message which will be handled during the `relay_message_response` call let resp_msg: Message> = Message { src_chain_id: dst_chain_id, diff --git a/domains/pallets/messenger/src/lib.rs b/domains/pallets/messenger/src/lib.rs index 4fa397b3c1..ebc1492c6f 100644 --- a/domains/pallets/messenger/src/lib.rs +++ b/domains/pallets/messenger/src/lib.rs @@ -243,25 +243,21 @@ mod pallet { /// Used by the dst_chains to verify the message response. #[pallet::storage] #[pallet::getter(fn inbox_responses)] - pub(super) type InboxResponses = CountedStorageMap< - _, - Identity, - (ChainId, ChannelId, Nonce), - Message>, - OptionQuery, - >; + pub(super) type InboxResponses = + StorageMap<_, Identity, (ChainId, ChannelId, Nonce), Message>, OptionQuery>; /// Stores the outgoing messages that are awaiting message responses from the dst_chain. /// Messages are processed in the outbox nonce order of chain's channel. #[pallet::storage] #[pallet::getter(fn outbox)] - pub(super) type Outbox = CountedStorageMap< - _, - Identity, - (ChainId, ChannelId, Nonce), - Message>, - OptionQuery, - >; + pub(super) type Outbox = + StorageMap<_, Identity, (ChainId, ChannelId, Nonce), Message>, OptionQuery>; + + /// Stores the outgoing messages count that are awaiting message responses from the dst_chain. + #[pallet::storage] + #[pallet::getter(fn outbox_message_count)] + pub(super) type OutboxMessageCount = + StorageMap<_, Identity, (ChainId, ChannelId), u32, ValueQuery>; /// A temporary storage for storing decoded outbox response message between `pre_dispatch_relay_message_response` /// and `relay_message_response`. @@ -545,6 +541,15 @@ mod pallet { /// Invalid channel reserve fee InvalidChannelReserveFee, + + /// Invalid max outgoing messages + InvalidMaxOutgoingMessages, + + /// Message count overflow + MessageCountOverflow, + + /// Message count underflow + MessageCountUnderflow, } #[pallet::call] @@ -907,8 +912,11 @@ mod pallet { // loop through channels in descending order until open channel is found. // we always prefer latest opened channel. while let Some(channel_id) = next_channel_id.checked_sub(ChannelId::one()) { + let message_count = OutboxMessageCount::::get((dst_chain_id, channel_id)); if let Some(channel) = Channels::::get(dst_chain_id, channel_id) { - if channel.state == ChannelState::Open { + if channel.state == ChannelState::Open + && message_count < channel.max_outgoing_messages + { return Some((channel_id, channel.fee)); } } @@ -1006,6 +1014,12 @@ mod pallet { Error::::InvalidChain, ); + // ensure max outgoing messages is atleast 1 + ensure!( + init_params.max_outgoing_messages >= 1u32, + Error::::InvalidMaxOutgoingMessages + ); + // If the channel owner is in this chain then the channel reserve fee // must not be empty ensure!( diff --git a/domains/pallets/messenger/src/messages.rs b/domains/pallets/messenger/src/messages.rs index 26e9172068..65782096dc 100644 --- a/domains/pallets/messenger/src/messages.rs +++ b/domains/pallets/messenger/src/messages.rs @@ -1,7 +1,7 @@ #[cfg(not(feature = "std"))] extern crate alloc; -use crate::pallet::{ChainAllowlist, UpdatedChannels}; +use crate::pallet::{ChainAllowlist, OutboxMessageCount, UpdatedChannels}; use crate::{ BalanceOf, ChannelId, ChannelState, Channels, CloseChannelBy, Config, Error, Event, InboxResponses, MessageWeightTags as MessageWeightTagStore, Nonce, Outbox, OutboxMessageResult, @@ -52,7 +52,7 @@ impl Pallet { |maybe_channel| -> Result { let channel = maybe_channel.as_mut().ok_or(Error::::MissingChannel)?; // check if the outbox is full - let count = Outbox::::count(); + let count = OutboxMessageCount::::get((dst_chain_id, channel_id)); ensure!( count < channel.max_outgoing_messages, Error::::OutboxFull @@ -72,6 +72,15 @@ impl Pallet { .latest_response_received_message_nonce, }; Outbox::::insert((dst_chain_id, channel_id, next_outbox_nonce), msg); + OutboxMessageCount::::try_mutate( + (dst_chain_id, channel_id), + |count| -> Result<(), DispatchError> { + *count = count + .checked_add(1u32) + .ok_or(Error::::MessageCountOverflow)?; + Ok(()) + }, + )?; // update channel state channel.next_outbox_nonce = next_outbox_nonce @@ -340,6 +349,16 @@ impl Pallet { let req_msg = Outbox::::take((dst_chain_id, channel_id, nonce)) .ok_or(Error::::MissingMessage)?; + OutboxMessageCount::::try_mutate( + (dst_chain_id, channel_id), + |count| -> Result<(), DispatchError> { + *count = count + .checked_sub(1u32) + .ok_or(Error::::MessageCountUnderflow)?; + Ok(()) + }, + )?; + // clear out box message weight tag MessageWeightTagStore::::mutate(|maybe_messages| { let mut messages = maybe_messages.as_mut().cloned().unwrap_or_default(); diff --git a/domains/pallets/messenger/src/tests.rs b/domains/pallets/messenger/src/tests.rs index 8559c4db3a..f58ae7a9f0 100644 --- a/domains/pallets/messenger/src/tests.rs +++ b/domains/pallets/messenger/src/tests.rs @@ -6,6 +6,7 @@ use crate::mock::{ chain_a, chain_b, consensus_chain, storage_proof_of_inbox_message_responses, storage_proof_of_outbox_messages, AccountId, Balance, TestExternalities, }; +use crate::pallet::OutboxMessageCount; use crate::{ ChainAllowlist, ChainAllowlistUpdate, Channel, ChannelId, ChannelState, Channels, CloseChannelBy, Error, FeeModel, Inbox, InboxResponses, InitiateChannelParams, Nonce, Outbox, @@ -59,7 +60,10 @@ fn create_channel(chain_id: ChainId, channel_id: ChannelId) { assert_eq!(channel.next_inbox_nonce, Nonce::zero()); assert_eq!(channel.next_outbox_nonce, Nonce::one()); assert_eq!(channel.latest_response_received_message_nonce, None); - assert_eq!(Outbox::::count(), 1); + assert_eq!( + OutboxMessageCount::::get((chain_id, channel_id)), + 1 + ); let msg = Outbox::::get((chain_id, channel_id, Nonce::zero())).unwrap(); assert_eq!(msg.dst_chain_id, chain_id); assert_eq!(msg.channel_id, channel_id); @@ -338,6 +342,7 @@ fn send_message_between_chains( msg: EndpointPayload, channel_id: ChannelId, ) { + let chain_a_id = chain_a::SelfChainId::get(); let chain_b_id = chain_b::SelfChainId::get(); // send message form outbox @@ -368,21 +373,35 @@ fn send_message_between_chains( // check state on chain_b chain_b_test_ext.execute_with(|| { // Outbox, Outbox responses, Inbox, InboxResponses must be empty - assert_eq!(Outbox::::count(), 0); + assert_eq!( + OutboxMessageCount::::get((chain_a_id, channel_id)), + 0 + ); assert!(OutboxResponses::::get().is_none()); assert!(Inbox::::get().is_none()); // latest inbox message response is cleared on next message - assert_eq!(InboxResponses::::count(), 1); + assert!(InboxResponses::::contains_key(( + chain_a_id, + channel_id, + Nonce::one() + )),); }); // check state on chain_a chain_a_test_ext.execute_with(|| { // Outbox, Outbox responses, Inbox, InboxResponses must be empty - assert_eq!(Outbox::::count(), 0); + assert_eq!( + OutboxMessageCount::::get((chain_b_id, channel_id)), + 0 + ); assert!(OutboxResponses::::get().is_none()); assert!(Inbox::::get().is_none()); - assert_eq!(InboxResponses::::count(), 0); + assert!(!InboxResponses::::contains_key(( + chain_b_id, + channel_id, + Nonce::one() + ))); let channel = chain_a::Messenger::channels(chain_b_id, channel_id).unwrap(); assert_eq!( @@ -435,12 +454,19 @@ fn close_channel_between_chains( assert_eq!(channel.next_outbox_nonce, Nonce::zero()); // Outbox, Outbox responses, Inbox, InboxResponses must be empty - assert_eq!(Outbox::::count(), 0); + assert_eq!( + OutboxMessageCount::::get((chain_a_id, channel_id)), + 0 + ); assert!(OutboxResponses::::get().is_none()); assert!(Inbox::::get().is_none()); // latest inbox message response is cleared on next message - assert_eq!(InboxResponses::::count(), 1); + assert!(InboxResponses::::contains_key(( + chain_a_id, + channel_id, + Nonce::one() + ))); }); // check channel state be closed on chain_a @@ -464,10 +490,17 @@ fn close_channel_between_chains( })); // Outbox, Outbox responses, Inbox, InboxResponses must be empty - assert_eq!(Outbox::::count(), 0); + assert_eq!( + OutboxMessageCount::::get((chain_b_id, channel_id)), + 0 + ); assert!(OutboxResponses::::get().is_none()); assert!(Inbox::::get().is_none()); - assert_eq!(InboxResponses::::count(), 0); + assert!(!InboxResponses::::contains_key(( + chain_b_id, + channel_id, + Nonce::one() + ))); }) }