Skip to main content

pwr_core/
frame.rs

1//! Frame encoding and decoding for the pwr wire protocol.
2//!
3//! Every message on the wire is wrapped in a 10-byte frame header:
4//!
5//! ```text
6//! Offset  Size  Field
7//! ------  ----  -----
8//! 0       4     Magic bytes "PWRF" (0x50 0x57 0x52 0x46)
9//! 4       4     Payload length in bytes (big-endian u32)
10//! 8       1     Protocol version (currently 0x01)
11//! 9       1     Message type (see protocol::MessageType)
12//! 10      N     Payload (bincode-encoded message body)
13//! ```
14//!
15//! The maximum payload size is 16 MiB. File content is NOT sent
16//! through the frame layer; it is streamed as raw bytes with a
17//! separate chunking protocol (4-byte length prefix, 0 = EOF).
18
19use crate::error::{PwrError, Result};
20use crate::protocol::MessageType;
21
22/// Magic bytes that prefix every frame.
23pub const FRAME_MAGIC: [u8; 4] = [0x50, 0x57, 0x52, 0x46]; // "PWRF"
24
25/// Current protocol version.
26pub const PROTOCOL_VERSION: u8 = 0x01;
27
28/// Maximum frame payload size: 16 MiB.
29pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;
30
31/// Total frame header size in bytes.
32pub const HEADER_SIZE: usize = 10;
33
34/// Encode a message into a framed byte vector suitable for writing
35/// to a TCP stream.
36///
37/// The caller provides the message (which must implement serde::Serialize),
38/// its MessageType discriminant, and optionally a specific protocol version
39/// (defaults to PROTOCOL_VERSION).
40///
41/// Returns the complete frame bytes: 10-byte header + JSON-encoded body.
42pub 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    // Magic bytes
61    frame.extend_from_slice(&FRAME_MAGIC);
62    // Payload length (big-endian)
63    frame.extend_from_slice(&payload_len.to_be_bytes());
64    // Protocol version
65    frame.push(PROTOCOL_VERSION);
66    // Message type
67    frame.push(msg_type as u8);
68    // Payload
69    frame.extend_from_slice(&body);
70
71    Ok(frame)
72}
73
74/// Decoded frame components returned by `decode_frame_header`.
75#[derive(Debug)]
76pub struct FrameHeader {
77    /// Length of the payload in bytes.
78    pub payload_len: u32,
79    /// Protocol version from the frame.
80    pub version: u8,
81    /// Message type discriminant.
82    pub msg_type: MessageType,
83}
84
85/// Decode a frame from raw bytes, returning the message type and
86/// the raw payload bytes.
87///
88/// The caller is responsible for deserializing the payload into
89/// the appropriate typed struct based on the message type.
90///
91/// Returns `None` if more bytes are needed to complete the frame
92/// (the caller should buffer and retry).
93pub fn decode_frame(data: &[u8]) -> Result<Option<(FrameHeader, Vec<u8>)>> {
94    // Need at least the 10-byte header
95    if data.len() < HEADER_SIZE {
96        return Ok(None);
97    }
98
99    // Verify magic bytes
100    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    // Read payload length
109    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    // Check if we have the full payload
119    let total_frame_size = HEADER_SIZE + payload_len;
120    if data.len() < total_frame_size {
121        return Ok(None);
122    }
123
124    // Read version
125    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    // Read message type
134    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    // Extract payload
140    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/// Streaming frame decoder that maintains an internal buffer.
153///
154/// Call `push_bytes` with data read from the socket, then call
155/// `try_decode` to extract complete frames. Incomplete frames
156/// are buffered until enough data arrives.
157#[derive(Debug, Default)]
158pub struct FrameDecoder {
159    buffer: Vec<u8>,
160}
161
162impl FrameDecoder {
163    /// Create a new frame decoder with an empty buffer.
164    pub fn new() -> Self {
165        Self { buffer: Vec::new() }
166    }
167
168    /// Append raw bytes received from the transport.
169    pub fn push_bytes(&mut self, data: &[u8]) {
170        self.buffer.extend_from_slice(data);
171    }
172
173    /// Attempt to decode one complete frame from the buffer.
174    ///
175    /// Returns `Ok(Some((header, payload)))` if a complete frame was decoded.
176    /// The decoded bytes are removed from the internal buffer.
177    /// Returns `Ok(None)` if more data is needed.
178    /// Returns `Err` if the frame is malformed.
179    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    /// Return the number of buffered bytes waiting to be decoded.
191    pub fn buffered_len(&self) -> usize {
192        self.buffer.len()
193    }
194}
195
196// ---------------------------------------------------------------------------
197// File chunk streaming
198// ---------------------------------------------------------------------------
199
200/// Write a chunk of file data in the streaming format.
201///
202/// Format: 4-byte big-endian chunk length (0 = EOF marker),
203/// followed by that many bytes of file content.
204pub 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
212/// The EOF marker for file chunk streaming.
213/// A chunk with length 0 signals the end of the current file.
214pub fn encode_file_eof() -> Vec<u8> {
215    vec![0u8, 0, 0, 0]
216}
217
218/// Decode one chunk from a byte buffer.
219///
220/// Returns:
221/// - `Ok(Some(data))` for a data chunk
222/// - `Ok(None)` if the chunk is the EOF marker (length = 0)
223/// - `Err` if the buffer doesn't contain enough data for the claimed length
224pub 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); // EOF marker
235    }
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        // Verify magic bytes
265        assert_eq!(&frame[0..4], &FRAME_MAGIC);
266
267        // Decode
268        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        // Only 5 bytes — incomplete header
282        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        // Feed half the frame
313        let mid = frame.len() / 2;
314        decoder.push_bytes(&frame[..mid]);
315        assert!(decoder.try_decode()?.is_none()); // Not enough data
316
317        // Feed the rest
318        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}