Skip to main content

peat_lite/protocol/
header.rs

1//! Peat-Lite packet header codec.
2//!
3//! The header is a fixed 16-byte prefix on every packet:
4//!
5//! ```text
6//! ┌──────────┬─────────┬──────────┬──────────┬──────────┬──────────────┐
7//! │  MAGIC   │ Version │   Type   │  Flags   │  NodeID  │   SeqNum     │
8//! │  4 bytes │ 1 byte  │  1 byte  │  2 bytes │  4 bytes │   4 bytes    │
9//! └──────────┴─────────┴──────────┴──────────┴──────────┴──────────────┘
10//! ```
11
12use super::constants::{HEADER_SIZE, MAGIC, PROTOCOL_VERSION};
13use super::error::MessageError;
14use super::message_type::MessageType;
15
16/// Decoded header fields.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct Header {
19    pub msg_type: MessageType,
20    pub flags: u16,
21    pub node_id: u32,
22    pub seq_num: u32,
23}
24
25/// Encode a header into the first 16 bytes of `buf`.
26///
27/// Returns `Err(MessageError::BufferTooSmall)` if `buf.len() < HEADER_SIZE`.
28pub fn encode_header(header: &Header, buf: &mut [u8]) -> Result<(), MessageError> {
29    if buf.len() < HEADER_SIZE {
30        return Err(MessageError::BufferTooSmall);
31    }
32    buf[0..4].copy_from_slice(&MAGIC);
33    buf[4] = PROTOCOL_VERSION;
34    buf[5] = header.msg_type as u8;
35    buf[6..8].copy_from_slice(&header.flags.to_le_bytes());
36    buf[8..12].copy_from_slice(&header.node_id.to_le_bytes());
37    buf[12..16].copy_from_slice(&header.seq_num.to_le_bytes());
38    Ok(())
39}
40
41/// Decode a header from `buf`, returning the header and a slice of the
42/// remaining payload bytes.
43///
44/// Validates magic bytes, protocol version, and message type.
45pub fn decode_header(buf: &[u8]) -> Result<(Header, &[u8]), MessageError> {
46    if buf.len() < HEADER_SIZE {
47        return Err(MessageError::TooShort);
48    }
49    if buf[0..4] != MAGIC {
50        return Err(MessageError::InvalidMagic);
51    }
52    if buf[4] != PROTOCOL_VERSION {
53        return Err(MessageError::UnsupportedVersion);
54    }
55    let msg_type = MessageType::from_u8(buf[5]).ok_or(MessageError::InvalidMessageType)?;
56    let flags = u16::from_le_bytes([buf[6], buf[7]]);
57    let node_id = u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]);
58    let seq_num = u32::from_le_bytes([buf[12], buf[13], buf[14], buf[15]]);
59
60    let header = Header {
61        msg_type,
62        flags,
63        node_id,
64        seq_num,
65    };
66    Ok((header, &buf[HEADER_SIZE..]))
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn roundtrip() {
75        let hdr = Header {
76            msg_type: MessageType::Announce,
77            flags: 0x1234,
78            node_id: 0xDEADBEEF,
79            seq_num: 42,
80        };
81        let mut buf = [0u8; 32];
82        encode_header(&hdr, &mut buf).unwrap();
83        // Append some payload bytes
84        buf[16] = 0xAA;
85        buf[17] = 0xBB;
86
87        let (decoded, payload) = decode_header(&buf[..18]).unwrap();
88        assert_eq!(decoded, hdr);
89        assert_eq!(payload, &[0xAA, 0xBB]);
90    }
91
92    #[test]
93    fn too_short() {
94        let buf = [0u8; 10];
95        assert_eq!(decode_header(&buf), Err(MessageError::TooShort));
96    }
97
98    #[test]
99    fn bad_magic() {
100        let mut buf = [0u8; 16];
101        buf[0..4].copy_from_slice(&[0, 0, 0, 0]);
102        buf[4] = PROTOCOL_VERSION;
103        buf[5] = MessageType::Heartbeat as u8;
104        assert_eq!(decode_header(&buf), Err(MessageError::InvalidMagic));
105    }
106
107    #[test]
108    fn bad_version() {
109        let mut buf = [0u8; 16];
110        buf[0..4].copy_from_slice(&MAGIC);
111        buf[4] = 99;
112        buf[5] = MessageType::Heartbeat as u8;
113        assert_eq!(decode_header(&buf), Err(MessageError::UnsupportedVersion));
114    }
115
116    #[test]
117    fn bad_message_type() {
118        let mut buf = [0u8; 16];
119        buf[0..4].copy_from_slice(&MAGIC);
120        buf[4] = PROTOCOL_VERSION;
121        buf[5] = 0xFF;
122        assert_eq!(decode_header(&buf), Err(MessageError::InvalidMessageType));
123    }
124
125    #[test]
126    fn header_only_no_payload() {
127        let hdr = Header {
128            msg_type: MessageType::Leave,
129            flags: 0,
130            node_id: 1,
131            seq_num: 0,
132        };
133        let mut buf = [0u8; 16];
134        encode_header(&hdr, &mut buf).unwrap();
135        let (decoded, payload) = decode_header(&buf).unwrap();
136        assert_eq!(decoded, hdr);
137        assert!(payload.is_empty());
138    }
139
140    #[test]
141    fn encode_buffer_too_small() {
142        let hdr = Header {
143            msg_type: MessageType::Data,
144            flags: 0,
145            node_id: 0,
146            seq_num: 0,
147        };
148        let mut buf = [0u8; 10];
149        assert_eq!(
150            encode_header(&hdr, &mut buf),
151            Err(MessageError::BufferTooSmall)
152        );
153    }
154}