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    /// Drain all buffered bytes from the decoder.
196    ///
197    /// This is used when switching from frame-decoded mode to raw
198    /// chunk-streaming mode (e.g., archive upload). Any bytes already
199    /// read from the transport but not yet consumed as a frame are
200    /// returned so they can be processed as chunk data.
201    pub fn drain_bytes(&mut self) -> Vec<u8> {
202        std::mem::take(&mut self.buffer)
203    }
204}
205
206// ---------------------------------------------------------------------------
207// File chunk streaming
208// ---------------------------------------------------------------------------
209
210/// Write a chunk of file data in the streaming format.
211///
212/// Format: 4-byte big-endian chunk length (0 = EOF marker),
213/// followed by that many bytes of file content.
214pub fn encode_file_chunk(data: &[u8]) -> Vec<u8> {
215    let len = data.len() as u32;
216    let mut chunk = Vec::with_capacity(4 + data.len());
217    chunk.extend_from_slice(&len.to_be_bytes());
218    chunk.extend_from_slice(data);
219    chunk
220}
221
222/// The EOF marker for file chunk streaming.
223/// A chunk with length 0 signals the end of the current file.
224pub fn encode_file_eof() -> Vec<u8> {
225    vec![0u8, 0, 0, 0]
226}
227
228/// Decode one chunk from a byte buffer.
229///
230/// Returns:
231/// - `Ok(Some(data))` for a data chunk
232/// - `Ok(None)` if the chunk is the EOF marker (length = 0)
233/// - `Err` if the buffer doesn't contain enough data for the claimed length
234pub fn decode_file_chunk(data: &[u8]) -> Result<Option<&[u8]>> {
235    if data.len() < 4 {
236        return Err(PwrError::Framing(
237            "incomplete chunk header: need 4 bytes".into(),
238        ));
239    }
240
241    let chunk_len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
242
243    if chunk_len == 0 {
244        return Ok(None); // EOF marker
245    }
246
247    if data.len() < 4 + chunk_len {
248        return Err(PwrError::Framing(format!(
249            "incomplete chunk data: need {} bytes, have {}",
250            4 + chunk_len,
251            data.len()
252        )));
253    }
254
255    Ok(Some(&data[4..4 + chunk_len]))
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261    use crate::protocol::{Handshake, StatusRequest};
262
263    #[test]
264    fn test_encode_decode_round_trip() -> Result<()> {
265        let msg = Handshake {
266            version: 1,
267            client_id: "test-client".into(),
268            nonce: [0x42; 32],
269            proof: [0x99; 32],
270        };
271
272        let frame = encode_frame(&msg, MessageType::Handshake)?;
273
274        // Verify magic bytes
275        assert_eq!(&frame[0..4], &FRAME_MAGIC);
276
277        // Decode
278        let (header, payload) = decode_frame(&frame)?.unwrap();
279
280        assert_eq!(header.version, PROTOCOL_VERSION);
281        assert_eq!(header.msg_type, MessageType::Handshake);
282
283        let decoded: Handshake = serde_json::from_slice(&payload).unwrap();
284        assert_eq!(decoded.client_id, "test-client");
285
286        Ok(())
287    }
288
289    #[test]
290    fn test_decode_needs_more_data() -> Result<()> {
291        // Only 5 bytes — incomplete header
292        assert!(decode_frame(&[0x50, 0x57, 0x52, 0x46, 0x00])?.is_none());
293        Ok(())
294    }
295
296    #[test]
297    fn test_decode_rejects_bad_magic() {
298        let mut bad = vec![0x00; 10];
299        bad[0] = 0xDE;
300        bad[1] = 0xAD;
301        assert!(decode_frame(&bad).is_err());
302    }
303
304    #[test]
305    fn test_decode_rejects_oversized_frame() {
306        let mut header = Vec::new();
307        header.extend_from_slice(&FRAME_MAGIC);
308        header.extend_from_slice(&((MAX_FRAME_SIZE as u32 + 1).to_be_bytes()));
309        header.push(PROTOCOL_VERSION);
310        header.push(MessageType::Error as u8);
311
312        assert!(decode_frame(&header).is_err());
313    }
314
315    #[test]
316    fn test_frame_decoder_buffering() -> Result<()> {
317        let mut decoder = FrameDecoder::new();
318
319        let msg = StatusRequest { project_uuid: None };
320        let frame = encode_frame(&msg, MessageType::StatusRequest)?;
321
322        // Feed half the frame
323        let mid = frame.len() / 2;
324        decoder.push_bytes(&frame[..mid]);
325        assert!(decoder.try_decode()?.is_none()); // Not enough data
326
327        // Feed the rest
328        decoder.push_bytes(&frame[mid..]);
329        let result = decoder.try_decode()?;
330        assert!(result.is_some());
331
332        let (header, _payload) = result.unwrap();
333        assert_eq!(header.msg_type, MessageType::StatusRequest);
334        assert_eq!(decoder.buffered_len(), 0);
335
336        Ok(())
337    }
338
339    #[test]
340    fn test_file_chunk_round_trip() {
341        let data = b"Hello, file streaming!";
342        let chunk = encode_file_chunk(data);
343
344        let decoded = decode_file_chunk(&chunk).unwrap();
345        assert_eq!(decoded, Some(&data[..]));
346    }
347
348    #[test]
349    fn test_file_eof() {
350        let eof = encode_file_eof();
351        assert_eq!(decode_file_chunk(&eof).unwrap(), None);
352    }
353}
pwr_core/frame.rs

pwr_core/
frame.rs
