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

avalanche_types/message/
app_response.rs
