avalanche_types/message/
pong.rs

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

avalanche_types/message/
pong.rs
