update trait return types and remove non used impl

Author: jxs

protocols/gossipsub/src/behaviour.rs

@@ -806,9 +806,11 @@ where
     ) -> Result<(), PublishError> {
         let topic_id = topic.into();
 
-        let available_parts = partial_message.available_parts().map(|p| p.to_vec());
-        let missing_parts = partial_message.missing_parts().map(|p| p.to_vec());
-        let group_id = partial_message.group_id().to_vec();
+        let available_parts = partial_message
+            .available_parts()
+            .map(|p| p.as_ref().to_vec());
+        let missing_parts = partial_message.missing_parts().map(|p| p.as_ref().to_vec());
+        let group_id = partial_message.group_id().as_ref().to_vec();
 
         // TODO: should we construct a recipient list just for partials?
         let recipient_peers = self.get_publish_peers(&topic_id);
@@ -824,8 +826,9 @@ where
             let peer_partials = peer.partial_messages.entry(topic_id.clone()).or_default();
             let peer_partial = peer_partials.entry(group_id.clone()).or_default();
 
-            let Ok((message_data, rest_wanted)) =
-                partial_message.partial_message_bytes_from_metadata(&peer_partial.wanted)
+            let Ok((message_data, rest_wanted)) = partial_message
+                .partial_message_bytes_from_metadata(&peer_partial.wanted)
+                .map(|(m, r)| (m.as_ref().to_vec(), r.map(|r| r.as_ref().to_vec())))
             else {
                 tracing::error!(peer = %peer_id, group_id = ?group_id,
                     "Could not reconstruct message bytes for peer metadata");

protocols/gossipsub/src/error.rs

@@ -180,10 +180,8 @@ pub enum PartialMessageError {
 
     /// The partial data doesn't belong to this message group.
     WrongGroup {
-        /// Expected minimum number of bytes.
-        expected: usize,
-        /// Actual number of bytes received.
-        received: usize,
+        /// Group Id of the received message.
+        received: Vec<u8>,
     },
 
     /// The partial data is a duplicate of already received data.
@@ -214,12 +212,8 @@ impl std::fmt::Display for PartialMessageError {
             Self::InvalidFormat => {
                 write!(f, "Invalid data format")
             }
-            Self::WrongGroup { expected, received } => {
-                write!(
-                    f,
-                    "Wrong group ID: expected {:?}, got {:?}",
-                    expected, received
-                )
+            Self::WrongGroup { received } => {
+                write!(f, "Wrong group ID: got {:?}", received)
             }
             Self::DuplicateData(part_id) => {
                 write!(f, "Duplicate data for part {:?}", part_id)

protocols/gossipsub/src/partial.rs

@@ -38,7 +38,7 @@ pub trait Partial {
     /// All partial messages belonging to the same logical message should return
     /// the same group ID. This is used to associate partial messages together
     /// during reconstruction.
-    fn group_id(&self) -> &[u8];
+    fn group_id(&self) -> impl AsRef<[u8]>;
 
     /// Returns metadata describing which parts of the message are missing.
     ///
@@ -48,7 +48,7 @@ pub trait Partial {
     ///
     /// The returned bytes will be sent in PartialIWANT messages to request
     /// missing parts from peers.
-    fn missing_parts(&self) -> Option<&[u8]>;
+    fn missing_parts(&self) -> Option<impl AsRef<[u8]>>;
 
     /// Returns metadata describing which parts of the message are available.
     ///
@@ -58,7 +58,7 @@ pub trait Partial {
     ///
     /// The returned bytes will be sent in PartialIHAVE messages to advertise
     /// available parts to peers.
-    fn available_parts(&self) -> Option<&[u8]>;
+    fn available_parts(&self) -> Option<impl AsRef<[u8]>>;
 
     /// Generates partial message bytes from the given metadata.
     ///
@@ -71,8 +71,8 @@ pub trait Partial {
     /// - Optional remaining metadata if more parts are still available after this one
     fn partial_message_bytes_from_metadata(
         &self,
-        metadata: &[u8],
-    ) -> Result<(Vec<u8>, Option<Vec<u8>>), PartialMessageError>;
+        metadata: impl AsRef<[u8]>,
+    ) -> Result<(impl AsRef<[u8]>, Option<impl AsRef<[u8]>>), PartialMessageError>;
 
     /// Extends this message with received partial message data.
     ///
@@ -87,34 +87,3 @@ pub trait Partial {
         data: &[u8],
     ) -> Result<(), PartialMessageError>;
 }
-
-/// Default implementation that disables partial messages.
-impl Partial for () {
-    fn group_id(&self) -> &[u8] {
-        &[]
-    }
-
-    fn missing_parts(&self) -> Option<&[u8]> {
-        None
-    }
-
-    fn available_parts(&self) -> Option<&[u8]> {
-        None
-    }
-
-    fn partial_message_bytes_from_metadata(
-        &self,
-        _metadata: &[u8],
-    ) -> Result<(Vec<u8>, Option<Vec<u8>>), PartialMessageError> {
-        Ok((vec![], None))
-    }
-
-    fn extend_from_encoded_partial_message(
-        &mut self,
-        _data: &[u8],
-    ) -> Result<(), PartialMessageError> {
-        // This should never be called since we never advertise having or wanting parts,
-        // but if it is called, just ignore the data silently
-        Ok(())
-    }
-}

protocols/gossipsub/src/types.rs

@@ -31,6 +31,8 @@ use libp2p_swarm::ConnectionId;
 use quick_protobuf::MessageWrite;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
+#[cfg(feature = "partial_messages")]
+use std::collections::HashMap;
 use web_time::Instant;
 
 use crate::{queue::Queue, rpc_proto::proto, TopicHash};
@@ -103,8 +105,7 @@ pub(crate) struct PeerDetails {
     pub(crate) dont_send: LinkedHashMap<MessageId, Instant>,
     /// Peer Partial messages.
     #[cfg(feature = "partial_messages")]
-    pub(crate) partial_messages:
-        std::collections::HashMap<TopicHash, HashMap<Vec<u8>, PartialData>>,
+    pub(crate) partial_messages: HashMap<TopicHash, HashMap<Vec<u8>, PartialData>>,
     /// Message queue consumed by the connection handler.
     pub(crate) messages: Queue,
 }