1use crate::error::{PwrError, Result};
20use crate::protocol::MessageType;
21
22pub const FRAME_MAGIC: [u8; 4] = [0x50, 0x57, 0x52, 0x46]; pub const PROTOCOL_VERSION: u8 = 0x01;
27
28pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;
30
31pub const HEADER_SIZE: usize = 10;
33
34pub fn encode_frame<T: serde::Serialize>(
43 msg: &T,
44 msg_type: MessageType,
45) -> Result<Vec<u8>> {
46 let body = serde_json::to_vec(msg)
47 .map_err(|e| PwrError::Framing(format!("encode error: {}", e)))?;
48
49 if body.len() > MAX_FRAME_SIZE {
50 return Err(PwrError::Framing(format!(
51 "payload size {} exceeds maximum {}",
52 body.len(),
53 MAX_FRAME_SIZE
54 )));
55 }
56
57 let payload_len = body.len() as u32;
58 let mut frame = Vec::with_capacity(HEADER_SIZE + body.len());
59
60 frame.extend_from_slice(&FRAME_MAGIC);
62 frame.extend_from_slice(&payload_len.to_be_bytes());
64 frame.push(PROTOCOL_VERSION);
66 frame.push(msg_type as u8);
68 frame.extend_from_slice(&body);
70
71 Ok(frame)
72}
73
74#[derive(Debug)]
76pub struct FrameHeader {
77 pub payload_len: u32,
79 pub version: u8,
81 pub msg_type: MessageType,
83}
84
85pub fn decode_frame(data: &[u8]) -> Result<Option<(FrameHeader, Vec<u8>)>> {
94 if data.len() < HEADER_SIZE {
96 return Ok(None);
97 }
98
99 if data[0..4] != FRAME_MAGIC {
101 return Err(PwrError::Framing(format!(
102 "invalid magic bytes: expected {:02X?}, got {:02X?}",
103 &FRAME_MAGIC[..],
104 &data[0..4]
105 )));
106 }
107
108 let payload_len = u32::from_be_bytes([data[4], data[5], data[6], data[7]]) as usize;
110
111 if payload_len > MAX_FRAME_SIZE {
112 return Err(PwrError::Framing(format!(
113 "frame payload size {} exceeds maximum {}",
114 payload_len, MAX_FRAME_SIZE
115 )));
116 }
117
118 let total_frame_size = HEADER_SIZE + payload_len;
120 if data.len() < total_frame_size {
121 return Ok(None);
122 }
123
124 let version = data[8];
126 if version != PROTOCOL_VERSION {
127 return Err(PwrError::Framing(format!(
128 "unsupported protocol version: got {}, expected {}",
129 version, PROTOCOL_VERSION
130 )));
131 }
132
133 let msg_type_byte = data[9];
135 let msg_type = MessageType::from_byte(msg_type_byte).ok_or_else(|| {
136 PwrError::Framing(format!("unknown message type: 0x{:02X}", msg_type_byte))
137 })?;
138
139 let payload = data[HEADER_SIZE..total_frame_size].to_vec();
141
142 Ok(Some((
143 FrameHeader {
144 payload_len: payload_len as u32,
145 version,
146 msg_type,
147 },
148 payload,
149 )))
150}
151
152#[derive(Debug, Default)]
158pub struct FrameDecoder {
159 buffer: Vec<u8>,
160}
161
162impl FrameDecoder {
163 pub fn new() -> Self {
165 Self { buffer: Vec::new() }
166 }
167
168 pub fn push_bytes(&mut self, data: &[u8]) {
170 self.buffer.extend_from_slice(data);
171 }
172
173 pub fn try_decode(&mut self) -> Result<Option<(FrameHeader, Vec<u8>)>> {
180 match decode_frame(&self.buffer)? {
181 Some((header, payload)) => {
182 let consumed = HEADER_SIZE + header.payload_len as usize;
183 self.buffer.drain(..consumed);
184 Ok(Some((header, payload)))
185 }
186 None => Ok(None),
187 }
188 }
189
190 pub fn buffered_len(&self) -> usize {
192 self.buffer.len()
193 }
194}
195
196pub fn encode_file_chunk(data: &[u8]) -> Vec<u8> {
205 let len = data.len() as u32;
206 let mut chunk = Vec::with_capacity(4 + data.len());
207 chunk.extend_from_slice(&len.to_be_bytes());
208 chunk.extend_from_slice(data);
209 chunk
210}
211
212pub fn encode_file_eof() -> Vec<u8> {
215 vec![0u8, 0, 0, 0]
216}
217
218pub fn decode_file_chunk(data: &[u8]) -> Result<Option<&[u8]>> {
225 if data.len() < 4 {
226 return Err(PwrError::Framing(
227 "incomplete chunk header: need 4 bytes".into(),
228 ));
229 }
230
231 let chunk_len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
232
233 if chunk_len == 0 {
234 return Ok(None); }
236
237 if data.len() < 4 + chunk_len {
238 return Err(PwrError::Framing(format!(
239 "incomplete chunk data: need {} bytes, have {}",
240 4 + chunk_len,
241 data.len()
242 )));
243 }
244
245 Ok(Some(&data[4..4 + chunk_len]))
246}
247
248#[cfg(test)]
249mod tests {
250 use super::*;
251 use crate::protocol::{Handshake, StatusRequest};
252
253 #[test]
254 fn test_encode_decode_round_trip() -> Result<()> {
255 let msg = Handshake {
256 version: 1,
257 client_id: "test-client".into(),
258 nonce: [0x42; 32],
259 proof: [0x99; 32],
260 };
261
262 let frame = encode_frame(&msg, MessageType::Handshake)?;
263
264 assert_eq!(&frame[0..4], &FRAME_MAGIC);
266
267 let (header, payload) = decode_frame(&frame)?.unwrap();
269
270 assert_eq!(header.version, PROTOCOL_VERSION);
271 assert_eq!(header.msg_type, MessageType::Handshake);
272
273 let decoded: Handshake = serde_json::from_slice(&payload).unwrap();
274 assert_eq!(decoded.client_id, "test-client");
275
276 Ok(())
277 }
278
279 #[test]
280 fn test_decode_needs_more_data() -> Result<()> {
281 assert!(decode_frame(&[0x50, 0x57, 0x52, 0x46, 0x00])?.is_none());
283 Ok(())
284 }
285
286 #[test]
287 fn test_decode_rejects_bad_magic() {
288 let mut bad = vec![0x00; 10];
289 bad[0] = 0xDE;
290 bad[1] = 0xAD;
291 assert!(decode_frame(&bad).is_err());
292 }
293
294 #[test]
295 fn test_decode_rejects_oversized_frame() {
296 let mut header = Vec::new();
297 header.extend_from_slice(&FRAME_MAGIC);
298 header.extend_from_slice(&((MAX_FRAME_SIZE as u32 + 1).to_be_bytes()));
299 header.push(PROTOCOL_VERSION);
300 header.push(MessageType::Error as u8);
301
302 assert!(decode_frame(&header).is_err());
303 }
304
305 #[test]
306 fn test_frame_decoder_buffering() -> Result<()> {
307 let mut decoder = FrameDecoder::new();
308
309 let msg = StatusRequest { project_uuid: None };
310 let frame = encode_frame(&msg, MessageType::StatusRequest)?;
311
312 let mid = frame.len() / 2;
314 decoder.push_bytes(&frame[..mid]);
315 assert!(decoder.try_decode()?.is_none()); decoder.push_bytes(&frame[mid..]);
319 let result = decoder.try_decode()?;
320 assert!(result.is_some());
321
322 let (header, _payload) = result.unwrap();
323 assert_eq!(header.msg_type, MessageType::StatusRequest);
324 assert_eq!(decoder.buffered_len(), 0);
325
326 Ok(())
327 }
328
329 #[test]
330 fn test_file_chunk_round_trip() {
331 let data = b"Hello, file streaming!";
332 let chunk = encode_file_chunk(data);
333
334 let decoded = decode_file_chunk(&chunk).unwrap();
335 assert_eq!(decoded, Some(&data[..]));
336 }
337
338 #[test]
339 fn test_file_eof() {
340 let eof = encode_file_eof();
341 assert_eq!(decode_file_chunk(&eof).unwrap(), None);
342 }
343}