avalanche_types/message/
get_accepted.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::GetAccepted,
9    pub gzip_compress: bool,
10}
11
12impl Default for Message {
13    fn default() -> Self {
14        Message {
15            msg: p2p::GetAccepted {
16                chain_id: prost::bytes::Bytes::new(),
17                request_id: 0,
18                deadline: 0,
19                container_ids: Vec::new(),
20                engine_type: p2p::EngineType::Unspecified.into(),
21            },
22            gzip_compress: false,
23        }
24    }
25}
26
27impl Message {
28    #[must_use]
29    pub fn chain_id(mut self, chain_id: ids::Id) -> Self {
30        self.msg.chain_id = prost::bytes::Bytes::from(chain_id.to_vec());
31        self
32    }
33
34    #[must_use]
35    pub fn request_id(mut self, request_id: u32) -> Self {
36        self.msg.request_id = request_id;
37        self
38    }
39
40    #[must_use]
41    pub fn deadline(mut self, deadline: u64) -> Self {
42        self.msg.deadline = deadline;
43        self
44    }
45
46    #[must_use]
47    pub fn container_ids(mut self, container_ids: Vec<ids::Id>) -> Self {
48        let mut container_ids_bytes: Vec<prost::bytes::Bytes> =
49            Vec::with_capacity(container_ids.len());
50        for id in container_ids.iter() {
51            container_ids_bytes.push(prost::bytes::Bytes::from(id.to_vec()));
52        }
53        self.msg.container_ids = container_ids_bytes;
54        self
55    }
56
57    #[must_use]
58    pub fn gzip_compress(mut self, gzip_compress: bool) -> Self {
59        self.gzip_compress = gzip_compress;
60        self
61    }
62
63    pub fn serialize(&self) -> io::Result<Vec<u8>> {
64        let msg = p2p::Message {
65            message: Some(p2p::message::Message::GetAccepted(self.msg.clone())),
66        };
67        let encoded = ProstMessage::encode_to_vec(&msg);
68        if !self.gzip_compress {
69            return Ok(encoded);
70        }
71
72        let uncompressed_len = encoded.len();
73        let compressed = message::compress::pack_gzip(&encoded)?;
74        let msg = p2p::Message {
75            message: Some(p2p::message::Message::CompressedGzip(
76                prost::bytes::Bytes::from(compressed),
77            )),
78        };
79
80        let compressed_len = msg.encoded_len();
81        if uncompressed_len > compressed_len {
82            log::debug!(
83                "get_accepted compression saved {} bytes",
84                uncompressed_len - compressed_len
85            );
86        } else {
87            log::debug!(
88                "get_accepted compression added {} byte(s)",
89                compressed_len - uncompressed_len
90            );
91        }
92
93        Ok(ProstMessage::encode_to_vec(&msg))
94    }
95
96    pub fn deserialize(d: impl AsRef<[u8]>) -> io::Result<Self> {
97        let buf = bytes::Bytes::from(d.as_ref().to_vec());
98        let p2p_msg: p2p::Message = ProstMessage::decode(buf).map_err(|e| {
99            Error::new(
100                ErrorKind::InvalidData,
101                format!("failed prost::Message::decode '{}'", e),
102            )
103        })?;
104
105        match p2p_msg.message.unwrap() {
106            // was not compressed
107            p2p::message::Message::GetAccepted(msg) => Ok(Message {
108                msg,
109                gzip_compress: false,
110            }),
111
112            // was compressed, so need decompress first
113            p2p::message::Message::CompressedGzip(msg) => {
114                let decompressed = message::compress::unpack_gzip(msg.as_ref())?;
115                let decompressed_msg: p2p::Message =
116                    ProstMessage::decode(prost::bytes::Bytes::from(decompressed)).map_err(|e| {
117                        Error::new(
118                            ErrorKind::InvalidData,
119                            format!("failed prost::Message::decode '{}'", e),
120                        )
121                    })?;
122                match decompressed_msg.message.unwrap() {
123                    p2p::message::Message::GetAccepted(msg) => Ok(Message {
124                        msg,
125                        gzip_compress: false,
126                    }),
127                    _ => Err(Error::new(
128                        ErrorKind::InvalidInput,
129                        "unknown message type after decompress",
130                    )),
131                }
132            }
133
134            // unknown message enum
135            _ => Err(Error::new(ErrorKind::InvalidInput, "unknown message type")),
136        }
137    }
138}
139
140/// RUST_LOG=debug cargo test --package avalanche-types --lib -- message::get_accepted::test_message --exact --show-output
141#[test]
142fn test_message() {
143    let _ = env_logger::builder()
144        .filter_level(log::LevelFilter::Debug)
145        .is_test(true)
146        .try_init();
147
148    let msg1_with_no_compression = Message::default()
149        .chain_id(ids::Id::from_slice(
150            &random_manager::secure_bytes(32).unwrap(),
151        ))
152        .request_id(random_manager::u32())
153        .deadline(random_manager::u64())
154        .container_ids(vec![
155            ids::Id::empty(),
156            ids::Id::empty(),
157            ids::Id::empty(),
158            ids::Id::empty(),
159            ids::Id::empty(),
160            ids::Id::empty(),
161            ids::Id::empty(),
162            ids::Id::empty(),
163            ids::Id::empty(),
164            ids::Id::empty(),
165            ids::Id::from_slice(&random_manager::secure_bytes(32).unwrap()),
166            ids::Id::from_slice(&random_manager::secure_bytes(32).unwrap()),
167        ]);
168
169    let data1 = msg1_with_no_compression.serialize().unwrap();
170    let msg1_with_no_compression_deserialized = Message::deserialize(data1).unwrap();
171    assert_eq!(
172        msg1_with_no_compression,
173        msg1_with_no_compression_deserialized
174    );
175
176    let msg2_with_compression = msg1_with_no_compression.clone().gzip_compress(true);
177    assert_ne!(msg1_with_no_compression, msg2_with_compression);
178
179    let data2 = msg2_with_compression.serialize().unwrap();
180    let msg2_with_compression_deserialized = Message::deserialize(data2).unwrap();
181    assert_eq!(msg1_with_no_compression, msg2_with_compression_deserialized);
182}
avalanche_types/message/get_accepted.rs

avalanche_types/message/
get_accepted.rs
