protocols/gossipsub: Convert message_id_fn's to closures (#2103)

This will allow capturing variables in these closures so that we can
make these functions aware of the forkId (necessary for altair).

Co-authored-by: Max Inden <mail@max-inden.de>

protocols/gossipsub/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 - Update dependencies.
 
+- Allow `message_id_fn`s to accept closures that capture variables.
+  [PR 2103](https://github.com/libp2p/rust-libp2p/pull/2103)
+
 # 0.32.0 [2021-07-12]
 
 - Update dependencies.

protocols/gossipsub/src/config.rs

@@ -19,6 +19,7 @@
 // DEALINGS IN THE SOFTWARE.
 
 use std::borrow::Cow;
+use std::sync::Arc;
 use std::time::Duration;
 
 use libp2p_core::PeerId;
@@ -69,8 +70,9 @@ pub struct GossipsubConfig {
     duplicate_cache_time: Duration,
     validate_messages: bool,
     validation_mode: ValidationMode,
-    message_id_fn: fn(&GossipsubMessage) -> MessageId,
-    fast_message_id_fn: Option<fn(&RawGossipsubMessage) -> FastMessageId>,
+    message_id_fn: Arc<dyn Fn(&GossipsubMessage) -> MessageId + Send + Sync + 'static>,
+    fast_message_id_fn:
+        Option<Arc<dyn Fn(&RawGossipsubMessage) -> FastMessageId + Send + Sync + 'static>>,
     allow_self_origin: bool,
     do_px: bool,
     prune_peers: usize,
@@ -234,6 +236,7 @@ impl GossipsubConfig {
     /// interpreted as the fast message id. Default is None.
     pub fn fast_message_id(&self, message: &RawGossipsubMessage) -> Option<FastMessageId> {
         self.fast_message_id_fn
+            .as_ref()
             .map(|fast_message_id_fn| fast_message_id_fn(message))
     }
 
@@ -396,7 +399,7 @@ impl Default for GossipsubConfigBuilder {
                 duplicate_cache_time: Duration::from_secs(60),
                 validate_messages: false,
                 validation_mode: ValidationMode::Strict,
-                message_id_fn: |message| {
+                message_id_fn: Arc::new(|message| {
                     // default message id is: source + sequence number
                     // NOTE: If either the peer_id or source is not provided, we set to 0;
                     let mut source_string = if let Some(peer_id) = message.source.as_ref() {
@@ -409,7 +412,7 @@ impl Default for GossipsubConfigBuilder {
                     source_string
                         .push_str(&message.sequence_number.unwrap_or_default().to_string());
                     MessageId::from(source_string)
-                },
+                }),
                 fast_message_id_fn: None,
                 allow_self_origin: false,
                 do_px: false,
@@ -574,8 +577,11 @@ impl GossipsubConfigBuilder {
     ///
     /// The function takes a [`GossipsubMessage`] as input and outputs a String to be
     /// interpreted as the message id.
-    pub fn message_id_fn(&mut self, id_fn: fn(&GossipsubMessage) -> MessageId) -> &mut Self {
-        self.config.message_id_fn = id_fn;
+    pub fn message_id_fn<F>(&mut self, id_fn: F) -> &mut Self
+    where
+        F: Fn(&GossipsubMessage) -> MessageId + Send + Sync + 'static,
+    {
+        self.config.message_id_fn = Arc::new(id_fn);
         self
     }
 
@@ -587,11 +593,11 @@ impl GossipsubConfigBuilder {
     ///
     /// The function takes a [`RawGossipsubMessage`] as input and outputs a String to be interpreted
     /// as the fast message id. Default is None.
-    pub fn fast_message_id_fn(
-        &mut self,
-        fast_id_fn: fn(&RawGossipsubMessage) -> FastMessageId,
-    ) -> &mut Self {
-        self.config.fast_message_id_fn = Some(fast_id_fn);
+    pub fn fast_message_id_fn<F>(&mut self, fast_id_fn: F) -> &mut Self
+    where
+        F: Fn(&RawGossipsubMessage) -> FastMessageId + Send + Sync + 'static,
+    {
+        self.config.fast_message_id_fn = Some(Arc::new(fast_id_fn));
         self
     }
 
@@ -815,6 +821,10 @@ impl std::fmt::Debug for GossipsubConfig {
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::topic::IdentityHash;
+    use crate::Topic;
+    use std::collections::hash_map::DefaultHasher;
+    use std::hash::{Hash, Hasher};
 
     #[test]
     fn create_thing() {
@@ -825,4 +835,80 @@ mod test {
 
         dbg!(builder);
     }
+
+    fn get_gossipsub_message() -> GossipsubMessage {
+        GossipsubMessage {
+            source: None,
+            data: vec![12, 34, 56],
+            sequence_number: None,
+            topic: Topic::<IdentityHash>::new("test").hash(),
+        }
+    }
+
+    fn get_expected_message_id() -> MessageId {
+        MessageId::from([
+            49, 55, 56, 51, 56, 52, 49, 51, 52, 51, 52, 55, 51, 51, 53, 52, 54, 54, 52, 49, 101,
+        ])
+    }
+
+    fn message_id_plain_function(message: &GossipsubMessage) -> MessageId {
+        let mut s = DefaultHasher::new();
+        message.data.hash(&mut s);
+        let mut v = s.finish().to_string();
+        v.push('e');
+        MessageId::from(v)
+    }
+
+    #[test]
+    fn create_config_with_message_id_as_plain_function() {
+        let builder: GossipsubConfig = GossipsubConfigBuilder::default()
+            .protocol_id_prefix("purple")
+            .message_id_fn(message_id_plain_function)
+            .build()
+            .unwrap();
+
+        let result = builder.message_id(&get_gossipsub_message());
+        assert_eq!(result, get_expected_message_id());
+    }
+
+    #[test]
+    fn create_config_with_message_id_as_closure() {
+        let closure = |message: &GossipsubMessage| {
+            let mut s = DefaultHasher::new();
+            message.data.hash(&mut s);
+            let mut v = s.finish().to_string();
+            v.push('e');
+            MessageId::from(v)
+        };
+
+        let builder: GossipsubConfig = GossipsubConfigBuilder::default()
+            .protocol_id_prefix("purple")
+            .message_id_fn(closure)
+            .build()
+            .unwrap();
+
+        let result = builder.message_id(&get_gossipsub_message());
+        assert_eq!(result, get_expected_message_id());
+    }
+
+    #[test]
+    fn create_config_with_message_id_as_closure_with_variable_capture() {
+        let captured: char = 'e';
+        let closure = move |message: &GossipsubMessage| {
+            let mut s = DefaultHasher::new();
+            message.data.hash(&mut s);
+            let mut v = s.finish().to_string();
+            v.push(captured);
+            MessageId::from(v)
+        };
+
+        let builder: GossipsubConfig = GossipsubConfigBuilder::default()
+            .protocol_id_prefix("purple")
+            .message_id_fn(closure)
+            .build()
+            .unwrap();
+
+        let result = builder.message_id(&get_gossipsub_message());
+        assert_eq!(result, get_expected_message_id());
+    }
 }