peat_lite/protocol/
header.rs1use super::constants::{HEADER_SIZE, MAGIC, PROTOCOL_VERSION};
13use super::error::MessageError;
14use super::message_type::MessageType;
15
16#[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
25pub 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
41pub 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 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}