fix(gossipsub): bound IWANT/IDONTWANT message_ids at decode time

protocols/gossipsub/CHANGELOG.md

@@ -1,5 +1,10 @@
 ## 0.50.0
 
+- Unify gossipsub control-message limits under max_control_messages (replacing per-type control ID caps),
+  and truncate control vectors immediately after RPC decode.
+  rename `max_ihave_messages` to `max_ihave_messages_heartbeat`.
+  See PR XXXX (https://github.com/libp2p/rust-libp2p/pull/XXXX)
+
 - Rename metric `topic_msg_sent_bytes` to `topic_msg_last_sent_bytes` for accuracy.
   See [PR 6283](https://github.com/libp2p/rust-libp2p/pull/6283)
 

protocols/gossipsub/src/behaviour.rs

@@ -1207,7 +1207,7 @@ where
         // IHAVE flood protection
         let peer_have = self.count_received_ihave.entry(*peer_id).or_insert(0);
         *peer_have += 1;
-        if *peer_have > self.config.max_ihave_messages() {
+        if *peer_have > self.config.max_ihave_messages_heartbeat() {
             tracing::debug!(
                 peer=%peer_id,
                 "IHAVE: peer has advertised too many times ({}) within this heartbeat \
@@ -1218,7 +1218,7 @@ where
         }
 
         if let Some(iasked) = self.count_sent_iwant.get(peer_id) {
-            if *iasked >= self.config.max_ihave_length() {
+            if *iasked >= self.config.max_control_messages() {
                 tracing::debug!(
                     peer=%peer_id,
                     "IHAVE: peer has already advertised too many messages ({}); ignoring",
@@ -1263,8 +1263,8 @@ where
         if !iwant_ids.is_empty() {
             let iasked = self.count_sent_iwant.entry(*peer_id).or_insert(0);
             let mut iask = iwant_ids.len();
-            if *iasked + iask > self.config.max_ihave_length() {
-                iask = self.config.max_ihave_length().saturating_sub(*iasked);
+            if *iasked + iask > self.config.max_control_messages() {
+                iask = self.config.max_control_messages().saturating_sub(*iasked);
             }
 
             // Send the list of IWANT control messages
@@ -2561,7 +2561,7 @@ where
             }
 
             // if we are emitting more than GossipSubMaxIHaveLength message_ids, truncate the list
-            if message_ids.len() > self.config.max_ihave_length() {
+            if message_ids.len() > self.config.max_control_messages() {
                 // we do the truncation (with shuffling) per peer below
                 tracing::debug!(
                     "too many messages for gossip; will truncate IHAVE list ({} messages)",
@@ -2595,12 +2595,12 @@ where
             for peer_id in to_msg_peers {
                 let mut peer_message_ids = message_ids.clone();
 
-                if peer_message_ids.len() > self.config.max_ihave_length() {
+                if peer_message_ids.len() > self.config.max_control_messages() {
                     // We do this per peer so that we emit a different set for each peer.
                     // we have enough redundancy in the system that this will significantly increase
                     // the message coverage when we do truncate.
-                    peer_message_ids.partial_shuffle(&mut rng, self.config.max_ihave_length());
-                    peer_message_ids.truncate(self.config.max_ihave_length());
+                    peer_message_ids.partial_shuffle(&mut rng, self.config.max_control_messages());
+                    peer_message_ids.truncate(self.config.max_control_messages());
                 }
 
                 // send an IHAVE message
@@ -3289,16 +3289,7 @@ where
                 }
 
                 // Handle messages
-                for (count, raw_message) in rpc.messages.into_iter().enumerate() {
-                    // Only process the amount of messages the configuration allows.
-                    if self
-                        .config
-                        .max_messages_per_rpc()
-                        .is_some_and(|max_msg| count >= max_msg)
-                    {
-                        tracing::warn!("Received more messages than permitted. Ignoring further messages. Processed: {}", count);
-                        break;
-                    }
+                for raw_message in rpc.messages {
                     self.handle_received_message(raw_message, &propagation_source);
                 }
 
@@ -3308,26 +3299,15 @@ where
                 let mut ihave_msgs = vec![];
                 let mut graft_msgs = vec![];
                 let mut prune_msgs = vec![];
-                for (count, control_msg) in rpc.control_msgs.into_iter().enumerate() {
-                    // Only process the amount of messages the configuration allows.
-                    if self
-                        .config
-                        .max_messages_per_rpc()
-                        .is_some_and(|max_msg| count >= max_msg)
-                    {
-                        tracing::warn!("Received more control messages than permitted. Ignoring further messages. Processed: {}", count);
-                        break;
-                    }
-
+                for control_msg in rpc.control_msgs {
                     match control_msg {
                         ControlAction::IHave(IHave {
                             topic_hash,
                             message_ids,
                         }) => {
                             ihave_msgs.push((topic_hash, message_ids));
                         }
-                        ControlAction::IWant(IWant { mut message_ids }) => {
-                            message_ids.truncate(self.config.max_iwant_length());
+                        ControlAction::IWant(IWant { message_ids }) => {
                             self.handle_iwant(&propagation_source, message_ids)
                         }
                         ControlAction::Graft(Graft { topic_hash }) => graft_msgs.push(topic_hash),
@@ -3336,15 +3316,14 @@ where
                             peers,
                             backoff,
                         }) => prune_msgs.push((topic_hash, peers, backoff)),
-                        ControlAction::IDontWant(IDontWant { mut message_ids }) => {
+                        ControlAction::IDontWant(IDontWant { message_ids }) => {
                             let Some(peer) = self.connected_peers.get_mut(&propagation_source)
                             else {
                                 tracing::error!(peer = %propagation_source,
                                     "Could not handle IDONTWANT, peer doesn't exist in connected peer list");
                                 continue;
                             };
 
-                            message_ids.truncate(self.config.max_idontwant_messages());
                             // Remove messages from the queue.
                             #[allow(unused)]
                             let removed = peer.messages.remove_data_messages(&message_ids);

protocols/gossipsub/src/behaviour/tests/gossip.rs

@@ -451,7 +451,7 @@ fn test_ignore_too_many_iwants_from_same_peer_for_same_message() {
 #[test]
 fn test_ignore_too_many_ihaves() {
     let config = ConfigBuilder::default()
-        .max_ihave_messages(10)
+        .max_ihave_messages_heartbeat(10)
         .build()
         .unwrap();
     // build gossipsub with full mesh
@@ -528,8 +528,8 @@ fn test_ignore_too_many_ihaves() {
 #[test]
 fn test_ignore_too_many_messages_in_ihave() {
     let config = ConfigBuilder::default()
-        .max_ihave_messages(10)
-        .max_ihave_length(10)
+        .max_ihave_messages_heartbeat(10)
+        .max_control_messages(10)
         .build()
         .unwrap();
     // build gossipsub with full mesh
@@ -609,8 +609,8 @@ fn test_ignore_too_many_messages_in_ihave() {
 #[test]
 fn test_limit_number_of_message_ids_inside_ihave() {
     let config = ConfigBuilder::default()
-        .max_ihave_messages(10)
-        .max_ihave_length(100)
+        .max_ihave_messages_heartbeat(10)
+        .max_control_messages(100)
         .build()
         .unwrap();
     // build gossipsub with full mesh

protocols/gossipsub/src/behaviour/tests/topic_config.rs

@@ -667,6 +667,8 @@ fn test_validation_error_message_size_too_large_topic_specific() {
         Config::default_max_transmit_size() * 2,
         ValidationMode::None,
         max_transmit_size_map,
+        5000,
+        5000,
     );
     let mut buf = BytesMut::new();
     let rpc = proto::RPC {
@@ -771,6 +773,8 @@ fn test_validation_message_size_within_topic_specific() {
         Config::default_max_transmit_size() * 2,
         ValidationMode::None,
         max_transmit_size_map,
+        5000,
+        5000,
     );
     let mut buf = BytesMut::new();
     let rpc = proto::RPC {

protocols/gossipsub/src/config.rs

@@ -123,11 +123,9 @@ pub struct Config {
     opportunistic_graft_ticks: u64,
     opportunistic_graft_peers: usize,
     gossip_retransimission: u32,
-    max_messages_per_rpc: Option<usize>,
-    max_ihave_length: usize,
-    max_iwant_length: usize,
-    max_ihave_messages: usize,
-    max_idontwant_messages: usize,
+    max_publish_messages: usize,
+    max_control_messages: usize,
+    max_ihave_messages_heartbeat: usize,
     iwant_followup_time: Duration,
     connection_handler_queue_len: usize,
     connection_handler_publish_duration: Duration,
@@ -412,30 +410,16 @@ impl Config {
         self.opportunistic_graft_peers
     }
 
-    /// The maximum number of messages we will process in a given RPC. If this is unset, there is
+    /// The maximum number of publish messages we will process in a given RPC. If this is unset, there is
     /// no limit. The default is None.
-    pub fn max_messages_per_rpc(&self) -> Option<usize> {
-        self.max_messages_per_rpc
+    pub fn max_publish_messages(&self) -> usize {
+        self.max_publish_messages
     }
 
-    /// The maximum number of messages to include in an IHAVE message.
-    /// Also controls the maximum number of IHAVE ids we will accept and request with IWANT from a
-    /// peer within a heartbeat, to protect from IHAVE floods. You should adjust this value from the
-    /// default if your system is pushing more than 5000 messages in GossipSubHistoryGossip
-    /// heartbeats; with the defaults this is 1666 messages/s. The default is 5000.
-    pub fn max_ihave_length(&self) -> usize {
-        self.max_ihave_length
-    }
-
-    /// The maximum number of `message_ids` that we accept in a single IWANT message.
-    pub fn max_iwant_length(&self) -> usize {
-        self.max_iwant_length
-    }
-
-    /// GossipSubMaxIHaveMessages is the maximum number of IHAVE messages to accept from a peer
-    /// within a heartbeat.
-    pub fn max_ihave_messages(&self) -> usize {
-        self.max_ihave_messages
+    /// The maximum number of control messages per type we will process in a given RPC. If this is unset, there is
+    /// no limit. The default is None.
+    pub fn max_control_messages(&self) -> usize {
+        self.max_control_messages
     }
 
     /// Time to wait for a message requested through IWANT following an IHAVE advertisement.
@@ -484,9 +468,10 @@ impl Config {
         self.idontwant_on_publish
     }
 
-    /// The maximum number of `message_ids` that we accept in a single IDONTWANT message.
-    pub fn max_idontwant_messages(&self) -> usize {
-        self.max_idontwant_messages
+    /// GossipSubMaxIHaveMessages is the maximum number of IHAVE messages to accept from a peer
+    /// within a heartbeat.
+    pub fn max_ihave_messages_heartbeat(&self) -> usize {
+        self.max_ihave_messages_heartbeat
     }
 }
 
@@ -548,16 +533,14 @@ impl Default for ConfigBuilder {
                 opportunistic_graft_ticks: 60,
                 opportunistic_graft_peers: 2,
                 gossip_retransimission: 3,
-                max_messages_per_rpc: None,
-                max_ihave_length: 5000,
-                max_iwant_length: 5000,
-                max_ihave_messages: 10,
+                max_publish_messages: 5000,
+                max_control_messages: 5000,
+                max_ihave_messages_heartbeat: 10,
                 iwant_followup_time: Duration::from_secs(3),
                 connection_handler_queue_len: 5000,
                 connection_handler_publish_duration: Duration::from_secs(5),
                 connection_handler_forward_duration: Duration::from_secs(1),
                 idontwant_message_size_threshold: 1000,
-                max_idontwant_messages: 1000,
                 idontwant_on_publish: false,
                 topic_configuration: TopicConfigs::default(),
             },
@@ -969,27 +952,17 @@ impl ConfigBuilder {
         self
     }
 
-    /// The maximum number of messages we will process in a given RPC. If this is unset, there is
-    /// no limit. The default is None.
-    pub fn max_messages_per_rpc(&mut self, max: Option<usize>) -> &mut Self {
-        self.config.max_messages_per_rpc = max;
-        self
-    }
-
-    /// The maximum number of messages to include in an IHAVE message.
-    /// Also controls the maximum number of IHAVE ids we will accept and request with IWANT from a
-    /// peer within a heartbeat, to protect from IHAVE floods. You should adjust this value from the
-    /// default if your system is pushing more than 5000 messages in GossipSubHistoryGossip
-    /// heartbeats; with the defaults this is 1666 messages/s. The default is 5000.
-    pub fn max_ihave_length(&mut self, max_ihave_length: usize) -> &mut Self {
-        self.config.max_ihave_length = max_ihave_length;
+    /// The maximum number of messages we will process in a given RPC. The default is 5000.
+    pub fn max_messages_per_rpc(&mut self, max: usize) -> &mut Self {
+        self.config.max_publish_messages = max;
+        self.config.protocol.max_publish_messages = max;
         self
     }
 
     /// GossipSubMaxIHaveMessages is the maximum number of IHAVE messages to accept from a peer
     /// within a heartbeat.
-    pub fn max_ihave_messages(&mut self, max_ihave_messages: usize) -> &mut Self {
-        self.config.max_ihave_messages = max_ihave_messages;
+    pub fn max_ihave_messages_heartbeat(&mut self, max_ihave_messages: usize) -> &mut Self {
+        self.config.max_ihave_messages_heartbeat = max_ihave_messages;
         self
     }
 
@@ -1063,15 +1036,10 @@ impl ConfigBuilder {
         self
     }
 
-    /// The maximum number of `message_ids` that we accept in a single IDONTWANT message.
-    pub fn max_idontwant_messages(&mut self, size: usize) -> &mut Self {
-        self.config.max_idontwant_messages = size;
-        self
-    }
-
-    /// The maximum number of `message_ids` that we accept in a single IWANT message.
-    pub fn max_iwant_messages(&mut self, size: usize) -> &mut Self {
-        self.config.max_iwant_length = size;
+    /// The max number of control messages to decode in a single RPC. The default is 5000.
+    pub fn max_control_messages(&mut self, size: usize) -> &mut Self {
+        self.config.max_control_messages = size;
+        self.config.protocol.max_control_messages = size;
         self
     }
 
@@ -1183,9 +1151,12 @@ impl std::fmt::Debug for Config {
         );
         let _ = builder.field("opportunistic_graft_ticks", &self.opportunistic_graft_ticks);
         let _ = builder.field("opportunistic_graft_peers", &self.opportunistic_graft_peers);
-        let _ = builder.field("max_messages_per_rpc", &self.max_messages_per_rpc);
-        let _ = builder.field("max_ihave_length", &self.max_ihave_length);
-        let _ = builder.field("max_ihave_messages", &self.max_ihave_messages);
+        let _ = builder.field("max_messages_per_rpc", &self.max_publish_messages);
+        let _ = builder.field("max_control_messages", &self.max_control_messages);
+        let _ = builder.field(
+            "max_ihave_messages_heartbeat",
+            &self.max_ihave_messages_heartbeat,
+        );
         let _ = builder.field("iwant_followup_time", &self.iwant_followup_time);
         let _ = builder.field(
             "idontwant_message_size_threshold",