feat(gossipsub): improve message size validation

Pull-Request: #6468.

Author: jxs

misc/prost-codec/CHANGELOG.md

@@ -1,6 +1,9 @@
 ## 0.4.0
 
 - Raise MSRV to 1.88.0.
+- Add `decode_field_tag`, `consume_message`, and `consume_message_prefix` helper functions
+  for parsing protobuf wire format without allocating.
+  See [PR 6428](https://github.com/libp2p/rust-libp2p/pull/6428).
   See [PR 6273](https://github.com/libp2p/rust-libp2p/pull/6273).
 - Migrate from `prost` 0.11 to `prost` 0.14.
 - Replace `UviBytes`-based framing with manual unsigned-varint encoding/decoding for improved performance.

misc/prost-codec/src/lib.rs

@@ -4,7 +4,10 @@ use std::{io, marker::PhantomData};
 
 use asynchronous_codec::{Decoder, Encoder};
 use bytes::{Buf, BytesMut};
-use prost::Message;
+use prost::{
+    Message,
+    encoding::{DecodeContext, WireType},
+};
 
 mod generated;
 
@@ -101,7 +104,7 @@ where
 }
 
 #[derive(thiserror::Error, Debug)]
-#[error("Failed to encode/decode message")]
+#[error("Failed to encode/decode message: {0}")]
 pub struct Error(#[from] io::Error);
 
 impl From<Error> for io::Error {
@@ -110,6 +113,37 @@ impl From<Error> for io::Error {
     }
 }
 
+/// Consumes the varint length prefix from a length-prefixed buffer.
+/// Advances the buffer past the prefix, positioning it at the message bytes.
+pub fn consume_message_prefix(buf: &mut &[u8]) -> io::Result<bool> {
+    let (message_length, remaining) = match unsigned_varint::decode::usize(buf) {
+        Ok((len, remaining)) => (len, remaining),
+        Err(unsigned_varint::decode::Error::Insufficient) => return Ok(false),
+        Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidData, e)),
+    };
+    if remaining.len() < message_length {
+        return Ok(false);
+    }
+    // Advance buffer past the length prefix, capping at the end of the buffer.
+    *buf = &remaining[..message_length];
+    Ok(true)
+}
+
+/// Decodes a protobuf field key (tag + wire type) and returns them.
+/// Advances the buffer past the key only (does NOT consume the field value).
+pub fn decode_field_tag(buf: &mut &[u8]) -> io::Result<(u32, WireType)> {
+    use prost::encoding::decode_key;
+    decode_key(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+}
+
+/// Consumes a length-delimited protobuf message from the buffer.
+/// Advances the buffer past the message.
+pub fn consume_message(wire_type: WireType, tag: u32, buf: &mut &[u8]) -> io::Result<()> {
+    use prost::encoding::skip_field;
+    skip_field(wire_type, tag, buf, DecodeContext::default())
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+}
+
 #[cfg(test)]
 mod tests {
     use std::error::Error;

protocols/gossipsub/CHANGELOG.md

@@ -1,4 +1,12 @@
 ## 0.50.0
+- Simplify gossipsub control message validation by scanning protobuf bytes before decoding
+  to validate cumulative control message size, rather than decoding then counting individual IDs
+  to prevent memory amplification from decoding.
+  Introduce `max_control_message_size` (default 16KB) to limit total control bytes in an RPC.
+  Rename `max_ihave_messages` to `max_ihave_messages_heartbeat`.
+  Rename `max_control_messages` to `max_control_messages_sent`.
+  See [PR #6486](https://github.com/libp2p/rust-libp2p/pull/6486)
+
 - Send all topic subscriptions in a single hello RPC when connecting to a new peer, aligning with the GossipSub spec and other implementations (Go, Nim, JS).
   See [PR 6385](https://github.com/libp2p/rust-libp2p/pull/6385).
 
@@ -11,12 +19,6 @@
 - Optimize IDONTWANT sending by avoiding broadcasts for already-seen messages and deduplicating recipient peers.
   See [PR 6356](https://github.com/libp2p/rust-libp2p/pull/6356)
 
-- 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`.
-  Introduce `max_ids_per_control_message` to limit the number of message IDs per control message.
-  See [PR 6409](https://github.com/libp2p/rust-libp2p/pull/6409) and [PR 6428](https://github.com/libp2p/rust-libp2p/pull/6428)
-
 - 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

@@ -1352,7 +1352,7 @@ where
         }
 
         if let Some(iasked) = self.count_sent_iwant.get(peer_id)
-            && *iasked >= self.config.max_control_messages()
+            && *iasked >= self.config.max_control_messages_sent()
         {
             tracing::debug!(
                 peer=%peer_id,
@@ -1411,8 +1411,11 @@ 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_control_messages() {
-                iask = self.config.max_control_messages().saturating_sub(*iasked);
+            if *iasked + iask > self.config.max_control_messages_sent() {
+                iask = self
+                    .config
+                    .max_control_messages_sent()
+                    .saturating_sub(*iasked);
             }
 
             // Send the list of IWANT control messages
@@ -2797,7 +2800,7 @@ where
             }
 
             // if we are emitting more than GossipSubMaxIHaveLength message_ids, truncate the list
-            if message_ids.len() > self.config.max_control_messages() {
+            if message_ids.len() > self.config.max_control_messages_sent() {
                 // we do the truncation (with shuffling) per peer below
                 tracing::debug!(
                     "too many messages for gossip; will truncate IHAVE list ({} messages)",
@@ -2839,12 +2842,13 @@ where
             for peer_id in to_msg_peers {
                 let mut peer_message_ids = message_ids.clone();
 
-                if peer_message_ids.len() > self.config.max_control_messages() {
+                if peer_message_ids.len() > self.config.max_control_messages_sent() {
                     // 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_control_messages());
-                    peer_message_ids.truncate(self.config.max_control_messages());
+                    peer_message_ids
+                        .partial_shuffle(&mut rng, self.config.max_control_messages_sent());
+                    peer_message_ids.truncate(self.config.max_control_messages_sent());
                 }
 
                 // send an IHAVE message

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

@@ -530,7 +530,7 @@ fn test_ignore_too_many_ihaves() {
 fn test_ignore_too_many_messages_in_ihave() {
     let config = ConfigBuilder::default()
         .max_ihave_messages_heartbeat(10)
-        .max_control_messages(10)
+        .max_control_messages_sent(10)
         .build()
         .unwrap();
     // build gossipsub with full mesh
@@ -611,7 +611,7 @@ fn test_ignore_too_many_messages_in_ihave() {
 fn test_limit_number_of_message_ids_inside_ihave() {
     let config = ConfigBuilder::default()
         .max_ihave_messages_heartbeat(10)
-        .max_control_messages(100)
+        .max_control_messages_sent(100)
         .build()
         .unwrap();
     // build gossipsub with full mesh

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

@@ -671,7 +671,6 @@ fn test_validation_error_message_size_too_large_topic_specific() {
         max_transmit_size_map,
         5000,
         5000,
-        5000,
     );
     let mut buf = BytesMut::new();
     let rpc = proto::Rpc {
@@ -781,7 +780,6 @@ fn test_validation_message_size_within_topic_specific() {
         max_transmit_size_map,
         5000,
         5000,
-        5000,
     );
     let mut buf = BytesMut::new();
     let rpc = proto::Rpc {

protocols/gossipsub/src/config.rs

@@ -126,8 +126,8 @@ pub struct Config {
     #[cfg(feature = "partial-messages")]
     max_metadata_length: usize,
     max_publish_messages: usize,
-    max_control_messages: usize,
-    max_ids_per_control_message: usize,
+    max_control_message_size: usize,
+    max_control_messages_sent: usize,
     max_ihave_messages_heartbeat: usize,
     iwant_followup_time: Duration,
     connection_handler_queue_len: usize,
@@ -425,17 +425,18 @@ impl Config {
         self.max_publish_messages
     }
 
-    /// The maximum number of control messages by type we will process in a given RPC. The default
-    /// is 5000.
-    pub fn max_control_messages(&self) -> usize {
-        self.max_control_messages
+    /// The maximum number of control messages (IHAVE/IWANT) we will send/receive to/from a peer.
+    /// This limits the number of IHAVE messages sent during gossip and IWANT requests received.
+    /// The default is 5000.
+    pub fn max_control_messages_sent(&self) -> usize {
+        self.max_control_messages_sent
     }
 
-    /// The maximum number of message ids per IHAVE/IWANT/IDONTWANT control message we will
-    /// process in a given RPC. Other control message types (GRAFT, PRUNE) are not affected by
-    /// this limit as they do not contain message IDs. The default is 5000.
-    pub fn max_ids_per_control_message(&self) -> usize {
-        self.max_ids_per_control_message
+    /// The maximum total byte size of all control messages and subscriptions in an RPC.
+    /// Validates cumulative size by scanning protobuf bytes before decoding.
+    /// Messages exceeding this limit will be rejected. The default is 16KB.
+    pub fn max_control_message_size(&self) -> usize {
+        self.max_control_message_size
     }
 
     /// Time to wait for a message requested through IWANT following an IHAVE advertisement.
@@ -552,8 +553,8 @@ impl Default for ConfigBuilder {
                 #[cfg(feature = "partial-messages")]
                 max_metadata_length: 1000,
                 max_publish_messages: 5000,
-                max_control_messages: 5000,
-                max_ids_per_control_message: 5000,
+                max_control_messages_sent: 5000,
+                max_control_message_size: 16384, // 16KB
                 max_ihave_messages_heartbeat: 10,
                 iwant_followup_time: Duration::from_secs(3),
                 connection_handler_queue_len: 5000,
@@ -993,6 +994,14 @@ impl ConfigBuilder {
         self
     }
 
+    /// The maximum number of control messages (IHAVE/IWANT) we will send/receive to/from a peer.
+    /// This limits the number of IHAVE messages sent during gossip and IWANT requests received.
+    /// The default is 5000.
+    pub fn max_control_messages_sent(&mut self, max_control_messages: usize) -> &mut Self {
+        self.config.max_control_messages_sent = max_control_messages;
+        self
+    }
+
     /// By default, gossipsub will reject messages that are sent to us that has the same message
     /// source as we have specified locally. Enabling this, allows these messages and prevents
     /// penalizing the peer that sent us the message. Default is false.
@@ -1063,20 +1072,12 @@ impl ConfigBuilder {
         self
     }
 
-    /// The maximum number of control messages by type we will process 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
-    }
-
-    /// The maximum number of message ids per IHAVE/IWANT/IDONTWANT control message we will
-    /// process in a single RPC. Other control message types (GRAFT, PRUNE) are not affected by
-    /// this limit as they do not contain message IDs. The default is 5000.
-    pub fn max_ids_per_control_message(&mut self, size: usize) -> &mut Self {
-        self.config.max_ids_per_control_message = size;
-        self.config.protocol.max_ids_per_control_message = size;
+    /// The maximum total byte size of all control messages and subscriptions in an RPC.
+    /// Validates cumulative size by scanning protobuf bytes before decoding.
+    /// Messages exceeding this limit will be rejected. The default is 16KB.
+    pub fn max_control_message_size(&mut self, size: usize) -> &mut Self {
+        self.config.max_control_message_size = size;
+        self.config.protocol.max_control_message_size = size;
         self
     }
 
@@ -1189,11 +1190,8 @@ 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_publish_messages);
-        let _ = builder.field("max_control_messages", &self.max_control_messages);
-        let _ = builder.field(
-            "max_ids_per_control_message",
-            &self.max_ids_per_control_message,
-        );
+        let _ = builder.field("max_control_messages_sent", &self.max_control_messages_sent);
+        let _ = builder.field("max_control_message_size", &self.max_control_message_size);
         let _ = builder.field(
             "max_ihave_messages_heartbeat",
             &self.max_ihave_messages_heartbeat,