avalanche_types/message/
app_gossip.rs

1use std::io::{self, Error, ErrorKind};
2
3use crate::{ids, message, proto::pb::p2p};
4use prost::Message as ProstMessage;
5
6#[derive(Debug, PartialEq, Clone)]
7pub struct Message {
8    pub msg: p2p::AppGossip,
9    pub gzip_compress: bool,
10}
11
12impl Default for Message {
13    fn default() -> Self {
14        Message {
15            msg: p2p::AppGossip {
16                chain_id: prost::bytes::Bytes::new(),
17                app_bytes: prost::bytes::Bytes::new(),
18            },
19            gzip_compress: false,
20        }
21    }
22}
23
24impl Message {
25    #[must_use]
26    pub fn chain_id(mut self, chain_id: ids::Id) -> Self {
27        self.msg.chain_id = prost::bytes::Bytes::from(chain_id.to_vec());
28        self
29    }
30
31    #[must_use]
32    pub fn app_bytes(mut self, app_bytes: Vec<u8>) -> Self {
33        self.msg.app_bytes = prost::bytes::Bytes::from(app_bytes);
34        self
35    }
36
37    #[must_use]
38    pub fn gzip_compress(mut self, gzip_compress: bool) -> Self {
39        self.gzip_compress = gzip_compress;
40        self
41    }
42
43    pub fn serialize(&self) -> io::Result<Vec<u8>> {
44        let msg = p2p::Message {
45            message: Some(p2p::message::Message::AppGossip(self.msg.clone())),
46        };
47        let encoded = ProstMessage::encode_to_vec(&msg);
48        if !self.gzip_compress {
49            return Ok(encoded);
50        }
51
52        let uncompressed_len = encoded.len();
53        let compressed = message::compress::pack_gzip(&encoded)?;
54        let msg = p2p::Message {
55            message: Some(p2p::message::Message::CompressedGzip(
56                prost::bytes::Bytes::from(compressed),
57            )),
58        };
59
60        let compressed_len = msg.encoded_len();
61        if uncompressed_len > compressed_len {
62            log::debug!(
63                "app_gossip compression saved {} bytes",
64                uncompressed_len - compressed_len
65            );
66        } else {
67            log::debug!(
68                "app_gossip compression added {} byte(s)",
69                compressed_len - uncompressed_len
70            );
71        }
72
73        Ok(ProstMessage::encode_to_vec(&msg))
74    }
75
76    pub fn deserialize(d: impl AsRef<[u8]>) -> io::Result<Self> {
77        let buf = bytes::Bytes::from(d.as_ref().to_vec());
78        let p2p_msg: p2p::Message = ProstMessage::decode(buf).map_err(|e| {
79            Error::new(
80                ErrorKind::InvalidData,
81                format!("failed prost::Message::decode '{}'", e),
82            )
83        })?;
84
85        match p2p_msg.message.unwrap() {
86            // was not compressed
87            p2p::message::Message::AppGossip(msg) => Ok(Message {
88                msg,
89                gzip_compress: false,
90            }),
91
92            // was compressed, so need decompress first
93            p2p::message::Message::CompressedGzip(msg) => {
94                let decompressed = message::compress::unpack_gzip(msg.as_ref())?;
95                let decompressed_msg: p2p::Message =
96                    ProstMessage::decode(prost::bytes::Bytes::from(decompressed)).map_err(|e| {
97                        Error::new(
98                            ErrorKind::InvalidData,
99                            format!("failed prost::Message::decode '{}'", e),
100                        )
101                    })?;
102                match decompressed_msg.message.unwrap() {
103                    p2p::message::Message::AppGossip(msg) => Ok(Message {
104                        msg,
105                        gzip_compress: false,
106                    }),
107                    _ => Err(Error::new(
108                        ErrorKind::InvalidInput,
109                        "unknown message type after decompress",
110                    )),
111                }
112            }
113
114            // unknown message enum
115            _ => Err(Error::new(ErrorKind::InvalidInput, "unknown message type")),
116        }
117    }
118}
119
120/// RUST_LOG=debug cargo test --package avalanche-types --lib -- message::app_gossip::test_message --exact --show-output
121#[test]
122fn test_message() {
123    let _ = env_logger::builder()
124        .filter_level(log::LevelFilter::Debug)
125        .is_test(true)
126        .try_init();
127
128    let msg1_with_no_compression = Message::default()
129        .chain_id(ids::Id::from_slice(
130            &random_manager::secure_bytes(32).unwrap(),
131        ))
132        .app_bytes(vec![0u8; 100]);
133
134    let data1 = msg1_with_no_compression.serialize().unwrap();
135    let msg1_with_no_compression_deserialized = Message::deserialize(data1).unwrap();
136    assert_eq!(
137        msg1_with_no_compression,
138        msg1_with_no_compression_deserialized
139    );
140
141    let msg2_with_compression = msg1_with_no_compression.clone().gzip_compress(true);
142    assert_ne!(msg1_with_no_compression, msg2_with_compression);
143
144    let data2 = msg2_with_compression.serialize().unwrap();
145    let msg2_with_compression_deserialized = Message::deserialize(data2).unwrap();
146    assert_eq!(msg1_with_no_compression, msg2_with_compression_deserialized);
147}
avalanche_types/message/app_gossip.rs

avalanche_types/message/
app_gossip.rs
