1use std::fmt;
6use std::io::{self, Read, Write};
7
8use super::types::{HEADER_SIZE, MAGIC, MAX_PAYLOAD};
9
10#[derive(Debug)]
12pub enum Error {
13 Io(io::Error),
14 BadMagic,
15 PayloadTooLarge,
16 ShortBuffer,
17 Truncated,
18 BadWire,
19 BadTag,
20 UnexpectedFrame {
23 want: &'static str,
24 got: u16,
25 },
26}
27
28impl fmt::Display for Error {
29 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30 match self {
31 Error::Io(e) => write!(f, "pxb: io: {e}"),
32 Error::BadMagic => f.write_str("pxb: bad magic"),
33 Error::PayloadTooLarge => f.write_str("pxb: payload too large"),
34 Error::ShortBuffer => f.write_str("pxb: short buffer"),
35 Error::Truncated => f.write_str("pxb: truncated payload"),
36 Error::BadWire => f.write_str("pxb: bad wire kind"),
37 Error::BadTag => f.write_str("pxb: bad field tag"),
38 Error::UnexpectedFrame { want, got } => {
39 write!(f, "pxb: expected {want} frame, got {got}")
40 }
41 }
42 }
43}
44
45impl std::error::Error for Error {
46 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
47 match self {
48 Error::Io(e) => Some(e),
49 _ => None,
50 }
51 }
52}
53
54impl From<io::Error> for Error {
55 fn from(e: io::Error) -> Self {
56 Error::Io(e)
57 }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub struct Header {
63 pub typ: u16,
64 pub flags: u16,
65 pub id: u32,
66 pub payload: u32,
67}
68
69pub fn encode_header(h: Header) -> [u8; HEADER_SIZE] {
71 let mut b = [0u8; HEADER_SIZE];
72 b[0..4].copy_from_slice(&MAGIC);
73 b[4..6].copy_from_slice(&h.typ.to_le_bytes());
74 b[6..8].copy_from_slice(&h.flags.to_le_bytes());
75 b[8..12].copy_from_slice(&h.id.to_le_bytes());
76 b[12..16].copy_from_slice(&h.payload.to_le_bytes());
77 b
78}
79
80pub fn decode_header(src: &[u8]) -> Result<Header, Error> {
82 if src.len() < HEADER_SIZE {
83 return Err(Error::ShortBuffer);
84 }
85 if src[0..4] != MAGIC {
86 return Err(Error::BadMagic);
87 }
88 let h = Header {
89 typ: u16::from_le_bytes([src[4], src[5]]),
90 flags: u16::from_le_bytes([src[6], src[7]]),
91 id: u32::from_le_bytes([src[8], src[9], src[10], src[11]]),
92 payload: u32::from_le_bytes([src[12], src[13], src[14], src[15]]),
93 };
94 if h.payload as usize > MAX_PAYLOAD {
95 return Err(Error::PayloadTooLarge);
96 }
97 Ok(h)
98}
99
100#[derive(Debug)]
102pub struct Frame {
103 pub header: Header,
104 pub body: Vec<u8>,
105}
106
107pub fn write_frame(
111 w: &mut impl Write,
112 typ: u16,
113 flags: u16,
114 id: u32,
115 body: &[u8],
116) -> Result<(), Error> {
117 if body.len() > MAX_PAYLOAD {
118 return Err(Error::PayloadTooLarge);
119 }
120 let hdr = encode_header(Header {
121 typ,
122 flags,
123 id,
124 payload: body.len() as u32,
125 });
126 w.write_all(&hdr)?;
127 w.write_all(body)?;
128 w.flush()?;
129 Ok(())
130}
131
132pub fn read_frame(r: &mut impl Read) -> Result<Frame, Error> {
134 let mut hdr = [0u8; HEADER_SIZE];
135 r.read_exact(&mut hdr)?;
136 let header = decode_header(&hdr)?;
137 let mut body = vec![0u8; header.payload as usize];
138 r.read_exact(&mut body)?;
139 Ok(Frame { header, body })
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145 use std::io::Cursor;
146
147 #[test]
148 fn header_roundtrip() {
149 let h = Header {
150 typ: 42,
151 flags: 2,
152 id: 7,
153 payload: 1024,
154 };
155 let b = encode_header(h);
156 assert_eq!(decode_header(&b).unwrap(), h);
157 }
158
159 #[test]
160 fn header_rejects_bad_input() {
161 assert!(matches!(
162 decode_header(&[0; 4]).unwrap_err(),
163 Error::ShortBuffer
164 ));
165 assert!(matches!(
166 decode_header(&[0; HEADER_SIZE]).unwrap_err(),
167 Error::BadMagic
168 ));
169
170 let mut b = encode_header(Header {
171 typ: 1,
172 flags: 0,
173 id: 0,
174 payload: 1 << 30,
175 });
176 let err = decode_header(&b).unwrap_err();
177 assert!(matches!(err, Error::PayloadTooLarge));
178 b[12..16].copy_from_slice(&(MAX_PAYLOAD as u32 + 1).to_le_bytes());
180 assert!(matches!(
181 decode_header(&b).unwrap_err(),
182 Error::PayloadTooLarge
183 ));
184 }
185
186 #[test]
187 fn frame_roundtrip_via_stream() {
188 let body = b"hello".to_vec();
189 let mut buf = Vec::new();
190 write_frame(&mut buf, 1, 0, 0, &body).unwrap();
191
192 let mut cur = Cursor::new(buf);
193 let f = read_frame(&mut cur).unwrap();
194 assert_eq!(f.header.typ, 1);
195 assert_eq!(f.body, body);
196
197 assert_eq!(cur.position(), (HEADER_SIZE + body.len()) as u64);
199 }
200
201 #[test]
202 fn empty_body_roundtrip() {
203 let mut buf = Vec::new();
204 write_frame(&mut buf, 9, 0, 0, &[]).unwrap();
205 let f = read_frame(&mut Cursor::new(buf)).unwrap();
206 assert_eq!(f.header.typ, 9);
207 assert_eq!(f.header.payload, 0);
208 assert!(f.body.is_empty());
209 }
210}