libp2p_pubsub_core/
message_id.rs

1use bytes::Bytes;
2use libp2p::identity::PeerId;
3
4use crate::framing::Message as FrameMessage;
5
6/// Macro for declaring message id types
7macro_rules! declare_message_id_type {
8    ($name: ident) => {
9        #[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
10        pub struct $name(Vec<u8>);
11
12        impl $name {
13            pub fn new<T: Into<Vec<u8>>>(value: T) -> Self {
14                Self(value.into())
15            }
16
17            pub fn new_from_slice(value: &[u8]) -> Self {
18                Self(value.to_vec())
19            }
20
21            fn into_vec(self) -> Vec<u8> {
22                self.0
23            }
24        }
25
26        impl From<Vec<u8>> for $name {
27            fn from(value: Vec<u8>) -> Self {
28                Self(value)
29            }
30        }
31
32        impl From<Bytes> for $name {
33            fn from(value: Bytes) -> Self {
34                Self(value.to_vec())
35            }
36        }
37
38        impl From<$name> for Bytes {
39            fn from(value: $name) -> Self {
40                Bytes::from(value.into_vec())
41            }
42        }
43
44        impl std::fmt::Display for $name {
45            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46                write!(f, "{}", hex_fmt::HexFmt(&self.0))
47            }
48        }
49
50        impl std::fmt::Debug for $name {
51            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52                write!(f, "{}({})", stringify!($name), hex_fmt::HexFmt(&self.0))
53            }
54        }
55    };
56}
57
58// A type for pubsub message IDs.
59declare_message_id_type!(MessageId);
60
61pub type MessageIdFn = dyn Fn(&FrameMessage) -> MessageId + Send + Sync + 'static;
62
63pub fn default_message_id_fn(msg: &FrameMessage) -> MessageId {
64    // default message id is: source + sequence number
65    // NOTE: If either the peer_id or source is not provided, we set to 0;
66    let mut source_string = if let Some(peer_id) = msg.source().as_ref() {
67        peer_id.to_base58().into_bytes()
68    } else {
69        PeerId::from_bytes(&[0, 1, 0])
70            .unwrap()
71            .to_base58()
72            .into_bytes()
73    };
74    source_string.extend(msg.sequence_number().unwrap_or_default());
75    MessageId::new(source_string)
76}
77
78#[cfg(test)]
79mod tests {
80    use rand::random;
81
82    use crate::topic::IdentTopic;
83
84    use super::*;
85
86    /// Helper function to create a random topic.
87    fn new_test_topic() -> IdentTopic {
88        IdentTopic::new(format!("/test-{}/0.1.0", random::<u32>()))
89    }
90
91    /// Helper function to create a random sequence number.
92    fn new_test_seqno() -> Bytes {
93        Bytes::from(random::<u32>().to_be_bytes().to_vec())
94    }
95
96    fn new_test_message(source: Option<PeerId>, seqno: Option<Bytes>) -> FrameMessage {
97        let mut message = FrameMessage::new(new_test_topic(), b"test-data".to_vec());
98        message.set_source(source);
99        message.set_sequence_number(seqno);
100        message
101    }
102
103    #[test]
104    fn default_message_id_fn_should_return_same_id_for_same_message() {
105        //// Given
106        let source = PeerId::random();
107        let message = new_test_message(Some(source), Some(new_test_seqno()));
108
109        let id_fn: Box<MessageIdFn> = Box::new(default_message_id_fn);
110
111        //// When
112        let message_id = id_fn(&message);
113        let message_id2 = id_fn(&message);
114
115        //// Then
116        assert_eq!(message_id, message_id2);
117    }
118
119    #[test]
120    fn default_message_id_fn_should_return_same_id_for_same_message_no_source() {
121        //// Given
122        let message = new_test_message(None, Some(new_test_seqno()));
123
124        let id_fn: Box<MessageIdFn> = Box::new(default_message_id_fn);
125
126        //// When
127        let message_id = id_fn(&message);
128        let message_id2 = id_fn(&message);
129
130        //// Then
131        assert_eq!(message_id, message_id2);
132    }
133}