bitcoin_embed/
message.rs

1//! # Message Encoding
2
3use crate::varint;
4
5use std::fmt;
6
7/// Errors that can occur during encoding/decoding
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum Error {
10    /// Zero tag and repeating tags are invalid
11    InvalidTag,
12    /// Invalid LEB128 encoding
13    InvalidVarInt,
14    /// Byte count exceeds 2^32 - 1
15    InvalidByteCount,
16    /// Final size byte must be zero
17    InvalidFinalSizeByte,
18    /// Variable-length encoding indicates bytes are missing
19    MissingBytes,
20}
21
22/// Type representing a protocol tag
23pub type Tag = u128;
24
25/// Standard tag definitions
26pub mod tags {
27    use super::Tag;
28
29    /// Repeat
30    pub const REPEAT: Tag = 0;
31}
32
33/// A struct containing a tag and associated bytes
34#[derive(Debug, Clone, PartialEq)]
35pub struct Message {
36    /// A tag
37    pub tag: Tag,
38    /// A vector of bytes
39    pub body: Vec<u8>,
40    /// Private field to prevent direct construction
41    _private: bool,
42}
43
44impl Message {
45    /// Constructs a Message struct. Throws an error if tag is zero, tag
46    /// exceeds 2^127 - 1, or byte count exceeds 2^32 - 1.
47    pub fn new(tag: u128, body: Vec<u8>) -> Result<Self, Error> {
48        if tag == tags::REPEAT || tag > ((1 << 127) - 1) {
49            return Err(Error::InvalidTag);
50        }
51
52        if body.len() > u32::MAX as usize {
53            return Err(Error::InvalidByteCount);
54        }
55
56        Ok(Self {
57            tag,
58            body,
59            _private: false,
60        })
61    }
62
63    /// Encodes messages as raw bytes.
64    ///
65    /// Details:
66    /// - Tags are LEB128-encoded as 2 * tag + (1 if terminal tag else 0)
67    /// - Repeating tags are encoded using a tag of zero
68    /// - Chunk lengths are LEB32-encoded, except for the final chunk
69    pub fn encode(messsages: Vec<Self>) -> Vec<u8> {
70        let mut bytes = Vec::new();
71        let len = messsages.len();
72        let mut last_tag = 0;
73
74        for (i, message) in messsages.into_iter().enumerate() {
75            let is_last = i == len - 1;
76
77            if message.tag == last_tag {
78                bytes.push(is_last as u8);
79            } else {
80                bytes.extend(varint::encode(2 * message.tag + (is_last as u128)));
81                last_tag = message.tag;
82            }
83
84            if !is_last {
85                bytes.extend(varint::encode(message.body.len() as u128));
86            }
87
88            bytes.extend(message.body);
89        }
90
91        bytes
92    }
93
94    /// Decodes messages from raw bytes.
95    ///
96    /// Returns an empty array if an invalid varint encoding or a chunk length that exceeds
97    /// the remaining array length is encountered.
98    pub fn decode(bytes: &[u8]) -> Result<Vec<Self>, Error> {
99        let mut messages = Vec::new();
100        let mut index = 0;
101        let mut last_tag = 0;
102
103        while index < bytes.len() {
104            let (value, size) =
105                varint::decode(&bytes[index..]).map_err(|_| Error::InvalidVarInt)?;
106            index += size;
107
108            let is_last = value % 2 == 1;
109            let tag = value / 2;
110
111            if tag == last_tag {
112                return Err(Error::InvalidTag);
113            }
114
115            let tag = if tag == tags::REPEAT {
116                last_tag
117            } else {
118                last_tag = tag;
119                tag
120            };
121
122            if is_last {
123                messages.push(Self {
124                    tag,
125                    body: bytes[index..].to_vec(),
126                    _private: false,
127                });
128                break;
129            }
130
131            if index >= bytes.len() {
132                return Err(Error::MissingBytes);
133            }
134
135            let (n, size) = varint::decode(&bytes[index..]).map_err(|_| Error::InvalidVarInt)?;
136            index += size;
137
138            let length: usize = if n == 0 {
139                bytes.len() - index
140            } else if n > u32::MAX.into() {
141                return Err(Error::InvalidByteCount);
142            } else {
143                n.try_into().unwrap()
144            };
145
146            if index + length > bytes.len() {
147                return Err(Error::MissingBytes);
148            }
149
150            if n > 0 && index + length == bytes.len() {
151                return Err(Error::InvalidFinalSizeByte);
152            }
153
154            messages.push(Self {
155                tag,
156                body: bytes[index..(index + length)].to_vec(),
157                _private: false,
158            });
159            index += length;
160        }
161
162        Ok(messages)
163    }
164}
165
166impl std::error::Error for Error {}
167
168impl fmt::Display for Error {
169    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
170        match self {
171            Error::InvalidTag => write!(f, "Invalid tag"),
172            Error::InvalidVarInt => write!(f, "Invalid variable integer encoding"),
173            Error::InvalidByteCount => write!(f, "Byte count exceeds 2^32 - 1"),
174            Error::InvalidFinalSizeByte => write!(f, "Final size byte must be zero"),
175            Error::MissingBytes => {
176                write!(f, "Variable-length encoding indicates bytes are missing")
177            }
178        }
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn test_new_valid() {
188        let data = Message::new(123, vec![1, 2, 3]).unwrap();
189        assert_eq!(data.tag, 123);
190        assert_eq!(data.body, vec![1, 2, 3]);
191    }
192
193    #[test]
194    fn test_new_invalid_tag() {
195        let result = Message::new(tags::REPEAT, vec![1, 2, 3]);
196        assert_eq!(result.err(), Some(Error::InvalidTag));
197
198        let result = Message::new(1 << 127, vec![1, 2, 3]);
199        assert_eq!(result.err(), Some(Error::InvalidTag));
200    }
201
202    #[test]
203    fn test_new_invalid_byte_count() {
204        if u32::MAX as usize == usize::MAX {
205            // Skip test on platforms where we can't exceed u32::MAX
206            return;
207        }
208
209        let result = Message::new(1, vec![0; (u32::MAX as usize) + 1]);
210        assert_eq!(result.err(), Some(Error::InvalidByteCount));
211    }
212
213    #[test]
214    fn test_encode_single_chunk() {
215        let chunk = Message::new(1, vec![5, 6, 7]).unwrap();
216        let encoded = Message::encode(vec![chunk]);
217
218        // Tag (2*1+1=3 terminal) + body [5,6,7]
219        assert_eq!(encoded, vec![3, 5, 6, 7]);
220    }
221
222    #[test]
223    fn test_encode_multiple_chunks() {
224        let chunk1 = Message::new(1, vec![1, 2]).unwrap();
225        let chunk2 = Message::new(2, vec![3, 4, 5]).unwrap();
226        let encoded = Message::encode(vec![chunk1, chunk2]);
227
228        // Tag (2*1+0=2 non-terminal) + Size(2) + [1,2] + Tag (2*2+1=5 terminal) + [3,4,5]
229        assert_eq!(encoded, vec![2, 2, 1, 2, 5, 3, 4, 5]);
230    }
231
232    #[test]
233    fn test_encode_repeated_tag() {
234        let chunk1 = Message::new(1, vec![1, 2]).unwrap();
235        let chunk2 = Message::new(1, vec![3, 4]).unwrap();
236        let chunk3 = Message::new(2, vec![5, 6]).unwrap();
237        let encoded = Message::encode(vec![chunk1, chunk2, chunk3]);
238
239        // Tag (2*1+0=2 non-terminal) + Size(2) + [1,2] +
240        // Repeat tag (0) + Size(2) + [3,4] +
241        // Tag (2*2+1=5 terminal) + [5,6]
242        assert_eq!(encoded, vec![2, 2, 1, 2, 0, 2, 3, 4, 5, 5, 6]);
243    }
244
245    #[test]
246    fn test_decode_valid() {
247        // Tag (2*1+0=2 non-terminal) + Size(2) + [1,2] + Tag (2*2+1=5 terminal) + [3,4,5]
248        let encoded = vec![2, 2, 1, 2, 5, 3, 4, 5];
249        let decoded = Message::decode(&encoded).unwrap();
250
251        assert_eq!(decoded.len(), 2);
252        assert_eq!(decoded[0].tag, 1);
253        assert_eq!(decoded[0].body, vec![1, 2]);
254        assert_eq!(decoded[1].tag, 2);
255        assert_eq!(decoded[1].body, vec![3, 4, 5]);
256    }
257
258    #[test]
259    fn test_decode_repeated_tag() {
260        // Tag (2*1+0=2 non-terminal) + Size(2) + [1,2] +
261        // Repeat tag (0) + Size(2) + [3,4] +
262        // Tag (2*2+1=5 terminal) + [5,6]
263        let encoded = vec![2, 2, 1, 2, 0, 2, 3, 4, 5, 5, 6];
264        let decoded = Message::decode(&encoded).unwrap();
265
266        assert_eq!(decoded.len(), 3);
267        assert_eq!(decoded[0].tag, 1);
268        assert_eq!(decoded[0].body, vec![1, 2]);
269        assert_eq!(decoded[1].tag, 1); // Repeated tag should be the same as previous
270        assert_eq!(decoded[1].body, vec![3, 4]);
271        assert_eq!(decoded[2].tag, 2);
272        assert_eq!(decoded[2].body, vec![5, 6]);
273    }
274
275    #[test]
276    fn test_decode_invalid_varint() {
277        // Invalid VarInt (0xFF is not a complete encoding)
278        let encoded = vec![0xFF];
279        let result = Message::decode(&encoded);
280
281        assert_eq!(result.err(), Some(Error::InvalidVarInt));
282    }
283
284    #[test]
285    fn test_decode_missing_bytes() {
286        // Tag (2*1+0=2 non-terminal) + Size(10) but not enough bytes follow
287        let encoded = vec![2, 10, 1, 2, 3];
288        let result = Message::decode(&encoded);
289
290        assert_eq!(result.err(), Some(Error::MissingBytes));
291    }
292
293    #[test]
294    fn test_decode_invalid_final_size_byte() {
295        // Tag (2*1+0=2 non-terminal) + Size(3) + [1,2,3] (should have terminal tag)
296        let encoded = vec![2, 3, 1, 2, 3];
297        let result = Message::decode(&encoded);
298
299        assert_eq!(result.err(), Some(Error::InvalidFinalSizeByte));
300    }
301
302    #[test]
303    fn test_empty_final_body() {
304        // Test with a terminal tag and empty body
305        let chunk = Message::new(3, vec![]).unwrap();
306        let encoded = Message::encode(vec![chunk]);
307
308        // Just the terminal tag (2*3+1=7)
309        assert_eq!(encoded, vec![7]);
310
311        let decoded = Message::decode(&encoded).unwrap();
312        assert_eq!(decoded.len(), 1);
313        assert_eq!(decoded[0].tag, 3);
314        assert_eq!(decoded[0].body, Vec::<u8>::new());
315    }
316
317    #[test]
318    fn test_decode_with_termination_only() {
319        // Just a terminal tag with no body (2*5+1=11)
320        let encoded = vec![11];
321        let decoded = Message::decode(&encoded).unwrap();
322
323        assert_eq!(decoded.len(), 1);
324        assert_eq!(decoded[0].tag, 5);
325        assert_eq!(decoded[0].body, Vec::<u8>::new());
326    }
327
328    #[test]
329    fn test_roundtrip_encode_decode() {
330        let original = vec![
331            Message::new(1, vec![1, 2]).unwrap(),
332            Message::new(2, vec![3, 4, 5]).unwrap(),
333        ];
334
335        let encoded = Message::encode(original.clone());
336        let decoded = Message::decode(&encoded).unwrap();
337
338        assert_eq!(original, decoded);
339    }
340
341    #[test]
342    fn test_large_tag_values() {
343        // Test encoding/decoding with large tag values that use multiple bytes in LEB128
344        let large_tag = (1 << 127) - 1;
345        let data = Message::new(large_tag, vec![9, 8, 7]).unwrap();
346        let encoded = Message::encode(vec![data]);
347        let decoded = Message::decode(&encoded).unwrap();
348
349        assert_eq!(decoded.len(), 1);
350        assert_eq!(decoded[0].tag, large_tag);
351        assert_eq!(decoded[0].body, vec![9, 8, 7]);
352    }
353
354    #[test]
355    fn test_multi_chunk_with_repeated_tag() {
356        let chunks = vec![
357            Message::new(10, vec![1, 2]).unwrap(),
358            Message::new(10, vec![3, 4]).unwrap(),
359            Message::new(10, vec![5, 6]).unwrap(),
360            Message::new(20, vec![7, 8, 9]).unwrap(),
361        ];
362
363        let encoded = Message::encode(chunks.clone());
364        let decoded = Message::decode(&encoded).unwrap();
365
366        assert_eq!(chunks, decoded);
367    }
368}
bitcoin_embed/message.rs

bitcoin_embed/
message.rs
