avalanche_types/message/
ping.rs

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

avalanche_types/message/
ping.rs
