microsandbox_protocol/
codec.rs

1//! Length-prefixed frame codec for reading and writing protocol messages.
2//!
3//! Wire format: `[len: u32 BE][id: u32 BE][flags: u8][body]`.
4//! Control bodies are CBOR; generation-8 bulk bodies use a fixed raw header.
5//!
6//! The correlation ID and flags sit in a fixed-position binary header so that
7//! relay intermediaries can route frames without CBOR parsing.
8
9use std::io::IoSlice;
10
11use bytes::{Buf, Bytes, BytesMut};
12use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
13
14use crate::{
15    bulk::{BULK_HEADER_SIZE, BulkFlow, BulkKind, BulkRecord, MAX_BULK_RECORD_PAYLOAD},
16    error::{ProtocolError, ProtocolResult},
17    message::{FLAG_BULK, FRAME_HEADER_SIZE, Message},
18};
19
20//--------------------------------------------------------------------------------------------------
21// Constants
22//--------------------------------------------------------------------------------------------------
23
24/// Maximum allowed frame size (4 MiB).
25///
26/// This covers everything after the 4-byte length prefix:
27/// `id (4) + flags (1) + control or raw body`.
28pub const MAX_FRAME_SIZE: u32 = 4 * 1024 * 1024;
29
30/// Maximum complete encoded frame size, including the four-byte length prefix.
31pub const MAX_WIRE_FRAME: usize = 4 + MAX_FRAME_SIZE as usize;
32
33//--------------------------------------------------------------------------------------------------
34// Types
35//--------------------------------------------------------------------------------------------------
36
37/// A frame with the binary header parsed but the body left untouched.
38///
39/// Used by routers, relays, and FFI consumers that want to handle framing
40/// without interpreting the control or raw data body. The [`body`](Self::body)
41/// field contains the exact bytes that follow the binary header on the wire.
42#[derive(Debug, Clone)]
43pub struct RawFrame {
44    /// Correlation ID. Same as [`Message::id`].
45    pub id: u32,
46
47    /// Frame flags. Same as [`Message::flags`].
48    pub flags: u8,
49
50    /// Raw body bytes. `flags` determines whether these are CBOR control bytes or raw bulk bytes.
51    pub body: Vec<u8>,
52}
53
54/// One fully validated control message or raw bulk record.
55#[derive(Debug, Clone)]
56pub enum DecodedFrame {
57    /// CBOR control-plane message.
58    Control(Message),
59
60    /// Generation-8 data-plane record.
61    Bulk(BulkRecord),
62}
63
64//--------------------------------------------------------------------------------------------------
65// Functions: Raw frame codec (CBOR-blind)
66//--------------------------------------------------------------------------------------------------
67
68/// Encodes a raw frame to a byte buffer using the length-prefixed format.
69///
70/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][body...]`
71pub fn encode_raw_to_buf(frame: &RawFrame, buf: &mut Vec<u8>) -> ProtocolResult<()> {
72    let frame_len = u32::try_from(FRAME_HEADER_SIZE + frame.body.len()).map_err(|_| {
73        ProtocolError::FrameTooLarge {
74            size: u32::MAX,
75            max: MAX_FRAME_SIZE,
76        }
77    })?;
78
79    if frame_len > MAX_FRAME_SIZE {
80        return Err(ProtocolError::FrameTooLarge {
81            size: frame_len,
82            max: MAX_FRAME_SIZE,
83        });
84    }
85
86    buf.reserve(4 + frame_len as usize);
87    buf.extend_from_slice(&frame_len.to_be_bytes());
88    buf.extend_from_slice(&frame.id.to_be_bytes());
89    buf.push(frame.flags);
90    buf.extend_from_slice(&frame.body);
91    Ok(())
92}
93
94/// Tries to decode a complete raw frame from a byte buffer.
95///
96/// Returns `Some(RawFrame)` if a complete frame is available, consuming
97/// the bytes. Returns `None` if more data is needed.
98///
99/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][body...]`
100pub fn try_decode_raw_from_buf(buf: &mut Vec<u8>) -> ProtocolResult<Option<RawFrame>> {
101    if buf.len() < 4 {
102        return Ok(None);
103    }
104
105    let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
106
107    if frame_len > MAX_FRAME_SIZE {
108        return Err(ProtocolError::FrameTooLarge {
109            size: frame_len,
110            max: MAX_FRAME_SIZE,
111        });
112    }
113
114    let frame_len = frame_len as usize;
115    let total = 4 + frame_len;
116
117    if buf.len() < total {
118        return Ok(None);
119    }
120
121    if frame_len < FRAME_HEADER_SIZE {
122        return Err(ProtocolError::FrameTooShort {
123            size: frame_len as u32,
124            min: FRAME_HEADER_SIZE as u32,
125        });
126    }
127
128    let id = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
129    let flags = buf[8];
130    let body = buf[4 + FRAME_HEADER_SIZE..total].to_vec();
131
132    buf.drain(..total);
133    Ok(Some(RawFrame { id, flags, body }))
134}
135
136/// Reads a length-prefixed raw frame from the given reader.
137///
138/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][body...]`
139pub async fn read_raw_frame<R: AsyncRead + Unpin>(reader: &mut R) -> ProtocolResult<RawFrame> {
140    let mut len_buf = [0u8; 4];
141    match reader.read_exact(&mut len_buf).await {
142        Ok(_) => {}
143        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
144            return Err(ProtocolError::UnexpectedEof);
145        }
146        Err(e) => return Err(e.into()),
147    }
148
149    let frame_len = u32::from_be_bytes(len_buf);
150
151    if frame_len > MAX_FRAME_SIZE {
152        return Err(ProtocolError::FrameTooLarge {
153            size: frame_len,
154            max: MAX_FRAME_SIZE,
155        });
156    }
157
158    let frame_len = frame_len as usize;
159
160    if frame_len < FRAME_HEADER_SIZE {
161        return Err(ProtocolError::FrameTooShort {
162            size: frame_len as u32,
163            min: FRAME_HEADER_SIZE as u32,
164        });
165    }
166
167    // Read the fixed header separately so the body is allocated exactly once.
168    let mut header = [0u8; FRAME_HEADER_SIZE];
169    reader.read_exact(&mut header).await?;
170    let id = u32::from_be_bytes(header[..4].try_into().unwrap());
171    let flags = header[4];
172    let mut body = vec![0u8; frame_len - FRAME_HEADER_SIZE];
173    reader.read_exact(&mut body).await?;
174
175    Ok(RawFrame { id, flags, body })
176}
177
178/// Writes a length-prefixed raw frame to the given writer.
179///
180/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][body...]`
181pub async fn write_raw_frame<W: AsyncWrite + Unpin>(
182    writer: &mut W,
183    frame: &RawFrame,
184) -> ProtocolResult<()> {
185    let frame_len = u32::try_from(FRAME_HEADER_SIZE + frame.body.len()).map_err(|_| {
186        ProtocolError::FrameTooLarge {
187            size: u32::MAX,
188            max: MAX_FRAME_SIZE,
189        }
190    })?;
191    if frame_len > MAX_FRAME_SIZE {
192        return Err(ProtocolError::FrameTooLarge {
193            size: frame_len,
194            max: MAX_FRAME_SIZE,
195        });
196    }
197
198    let mut header = [0u8; 4 + FRAME_HEADER_SIZE];
199    header[..4].copy_from_slice(&frame_len.to_be_bytes());
200    header[4..8].copy_from_slice(&frame.id.to_be_bytes());
201    header[8] = frame.flags;
202    write_vectored_all(writer, &header, &frame.body).await?;
203    writer.flush().await?;
204    Ok(())
205}
206
207/// Writes a generation-8 raw bulk record without copying its payload.
208pub async fn write_bulk_record<W: AsyncWrite + Unpin>(
209    writer: &mut W,
210    record: &BulkRecord,
211) -> ProtocolResult<()> {
212    let header = encode_bulk_header(record)?;
213    write_vectored_all(writer, &header, &record.payload).await?;
214    writer.flush().await?;
215    Ok(())
216}
217
218/// Encodes a generation-8 raw bulk record into a contiguous byte buffer.
219pub fn encode_bulk_to_buf(record: &BulkRecord, buf: &mut Vec<u8>) -> ProtocolResult<()> {
220    let header = encode_bulk_header(record)?;
221    buf.reserve(header.len() + record.payload.len());
222    buf.extend_from_slice(&header);
223    buf.extend_from_slice(&record.payload);
224    Ok(())
225}
226
227/// Decodes and validates a raw frame as a generation-8 bulk record.
228pub fn raw_frame_to_bulk(frame: RawFrame, max_payload: u32) -> ProtocolResult<BulkRecord> {
229    if frame.flags != FLAG_BULK {
230        return Err(invalid_bulk("bulk flag must be set exclusively"));
231    }
232    decode_bulk_body(frame.id, Bytes::from(frame.body), max_payload)
233}
234
235/// Decodes one complete control or generation-8 bulk frame from a cursor-based buffer.
236pub fn try_decode_frame_from_bytes(buf: &mut BytesMut) -> ProtocolResult<Option<DecodedFrame>> {
237    let Some((frame_len, total)) = complete_frame_len(buf)? else {
238        return Ok(None);
239    };
240
241    let flags = buf[8];
242    if flags & FLAG_BULK == 0 {
243        let message = decode_message_frame(&buf[..total])?;
244        buf.advance(total);
245        return Ok(Some(DecodedFrame::Control(message)));
246    }
247    if flags != FLAG_BULK {
248        return Err(invalid_bulk(
249            "bulk flag cannot be combined with control flags",
250        ));
251    }
252    if frame_len < FRAME_HEADER_SIZE + BULK_HEADER_SIZE + 1 {
253        return Err(invalid_bulk("bulk record has no payload"));
254    }
255
256    let frame = buf.split_to(total).freeze();
257    let id = u32::from_be_bytes(frame[4..8].try_into().unwrap());
258    let body = frame.slice(4 + FRAME_HEADER_SIZE..);
259    let record = decode_bulk_body(id, body, MAX_BULK_RECORD_PAYLOAD)?;
260    Ok(Some(DecodedFrame::Bulk(record)))
261}
262
263/// Decode one complete typed frame from a cursor-based buffer without front-draining it.
264pub fn try_decode_from_bytes(buf: &mut BytesMut) -> ProtocolResult<Option<Message>> {
265    match try_decode_frame_from_bytes(buf)? {
266        Some(DecodedFrame::Control(message)) => Ok(Some(message)),
267        Some(DecodedFrame::Bulk(_)) => Err(invalid_bulk(
268            "raw bulk record passed to the control-message decoder",
269        )),
270        None => Ok(None),
271    }
272}
273
274//--------------------------------------------------------------------------------------------------
275// Functions: Typed message codec (CBOR-aware)
276//--------------------------------------------------------------------------------------------------
277
278/// Encodes a message to a byte buffer using the length-prefixed frame format.
279///
280/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
281pub fn encode_to_buf(msg: &Message, buf: &mut Vec<u8>) -> ProtocolResult<()> {
282    let mut body = Vec::new();
283    ciborium::into_writer(msg, &mut body)?;
284    encode_raw_to_buf(
285        &RawFrame {
286            id: msg.id,
287            flags: msg.flags,
288            body,
289        },
290        buf,
291    )
292}
293
294/// Tries to decode a complete message from a byte buffer.
295///
296/// Returns `Some(Message)` if a complete frame is available, consuming
297/// the bytes. Returns `None` if more data is needed.
298///
299/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
300pub fn try_decode_from_buf(buf: &mut Vec<u8>) -> ProtocolResult<Option<Message>> {
301    if buf.len() < 4 {
302        return Ok(None);
303    }
304
305    let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
306
307    if frame_len > MAX_FRAME_SIZE {
308        return Err(ProtocolError::FrameTooLarge {
309            size: frame_len,
310            max: MAX_FRAME_SIZE,
311        });
312    }
313
314    let frame_len = frame_len as usize;
315    let total = 4 + frame_len;
316
317    if buf.len() < total {
318        return Ok(None);
319    }
320
321    let msg = decode_message_frame(&buf[..total])?;
322    buf.drain(..total);
323    Ok(Some(msg))
324}
325
326/// Reads a length-prefixed message from the given reader.
327///
328/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
329pub async fn read_message<R: AsyncRead + Unpin>(reader: &mut R) -> ProtocolResult<Message> {
330    let frame = read_raw_frame(reader).await?;
331    raw_frame_to_message(frame)
332}
333
334/// Writes a length-prefixed message to the given writer.
335///
336/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
337pub async fn write_message<W: AsyncWrite + Unpin>(
338    writer: &mut W,
339    message: &Message,
340) -> ProtocolResult<()> {
341    let mut body = Vec::new();
342    ciborium::into_writer(message, &mut body)?;
343    write_raw_frame(
344        writer,
345        &RawFrame {
346            id: message.id,
347            flags: message.flags,
348            body,
349        },
350    )
351    .await
352}
353
354/// Decodes a [`RawFrame`] into a typed [`Message`] by CBOR-deserializing the body.
355pub fn raw_frame_to_message(frame: RawFrame) -> ProtocolResult<Message> {
356    let mut msg: Message = ciborium::from_reader(&frame.body[..])?;
357    msg.id = frame.id;
358    msg.flags = frame.flags;
359    Ok(msg)
360}
361
362/// Decodes one complete length-prefixed frame from a borrowed byte slice.
363///
364/// The input must include the 4-byte length prefix, frame header, and CBOR body.
365/// The slice is not consumed or copied.
366pub fn decode_message_frame(frame: &[u8]) -> ProtocolResult<Message> {
367    if frame.len() < 4 {
368        return Err(ProtocolError::UnexpectedEof);
369    }
370
371    let frame_len = u32::from_be_bytes([frame[0], frame[1], frame[2], frame[3]]);
372    if frame_len > MAX_FRAME_SIZE {
373        return Err(ProtocolError::FrameTooLarge {
374            size: frame_len,
375            max: MAX_FRAME_SIZE,
376        });
377    }
378
379    let frame_len = frame_len as usize;
380    let total = 4 + frame_len;
381    if frame.len() < total {
382        return Err(ProtocolError::UnexpectedEof);
383    }
384
385    if frame_len < FRAME_HEADER_SIZE {
386        return Err(ProtocolError::FrameTooShort {
387            size: frame_len as u32,
388            min: FRAME_HEADER_SIZE as u32,
389        });
390    }
391
392    let mut msg: Message = ciborium::from_reader(&frame[4 + FRAME_HEADER_SIZE..total])?;
393    msg.id = u32::from_be_bytes([frame[4], frame[5], frame[6], frame[7]]);
394    msg.flags = frame[8];
395    Ok(msg)
396}
397
398/// Encodes the fixed outer and generation-8 bulk headers, leaving the payload separate.
399pub fn encode_bulk_header(
400    record: &BulkRecord,
401) -> ProtocolResult<[u8; 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE]> {
402    let payload_len = record.payload.len();
403    if payload_len == 0 || payload_len > MAX_BULK_RECORD_PAYLOAD as usize {
404        return Err(invalid_bulk(format!(
405            "payload length {payload_len} is outside 1..={MAX_BULK_RECORD_PAYLOAD}"
406        )));
407    }
408    record
409        .offset
410        .checked_add(payload_len as u64)
411        .ok_or_else(|| invalid_bulk("record end offset overflows u64"))?;
412
413    let frame_len = FRAME_HEADER_SIZE + BULK_HEADER_SIZE + payload_len;
414    let frame_len = u32::try_from(frame_len).map_err(|_| ProtocolError::FrameTooLarge {
415        size: u32::MAX,
416        max: MAX_FRAME_SIZE,
417    })?;
418    if frame_len > MAX_FRAME_SIZE {
419        return Err(ProtocolError::FrameTooLarge {
420            size: frame_len,
421            max: MAX_FRAME_SIZE,
422        });
423    }
424
425    let mut header = [0u8; 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE];
426    header[..4].copy_from_slice(&frame_len.to_be_bytes());
427    header[4..8].copy_from_slice(&record.id.to_be_bytes());
428    header[8] = FLAG_BULK;
429    header[9] = record.kind as u8;
430    header[10] = record.flow as u8;
431    // Bytes 11..13 are reserved and deliberately remain zero.
432    header[13..21].copy_from_slice(&record.offset.to_be_bytes());
433    Ok(header)
434}
435
436pub(crate) fn decode_bulk_body(
437    id: u32,
438    body: Bytes,
439    max_payload: u32,
440) -> ProtocolResult<BulkRecord> {
441    if max_payload == 0 || max_payload > MAX_BULK_RECORD_PAYLOAD {
442        return Err(invalid_bulk(format!(
443            "invalid decoder payload limit {max_payload}"
444        )));
445    }
446    if body.len() < BULK_HEADER_SIZE + 1 {
447        return Err(invalid_bulk("bulk record has no payload"));
448    }
449    if body[2] != 0 || body[3] != 0 {
450        return Err(invalid_bulk("reserved bulk header bytes must be zero"));
451    }
452
453    let kind = BulkKind::from_wire(body[0])
454        .ok_or_else(|| invalid_bulk(format!("unknown bulk kind {}", body[0])))?;
455    let flow = BulkFlow::from_wire(body[1])
456        .ok_or_else(|| invalid_bulk(format!("unknown bulk flow {}", body[1])))?;
457    let offset = u64::from_be_bytes(body[4..12].try_into().unwrap());
458    let payload = body.slice(BULK_HEADER_SIZE..);
459    if payload.len() > max_payload as usize {
460        return Err(invalid_bulk(format!(
461            "payload length {} exceeds negotiated maximum {max_payload}",
462            payload.len()
463        )));
464    }
465    offset
466        .checked_add(payload.len() as u64)
467        .ok_or_else(|| invalid_bulk("record end offset overflows u64"))?;
468
469    Ok(BulkRecord {
470        id,
471        kind,
472        flow,
473        offset,
474        payload,
475    })
476}
477
478fn complete_frame_len(buf: &BytesMut) -> ProtocolResult<Option<(usize, usize)>> {
479    if buf.len() < 4 {
480        return Ok(None);
481    }
482
483    let frame_len = u32::from_be_bytes(buf[..4].try_into().unwrap());
484    if frame_len > MAX_FRAME_SIZE {
485        return Err(ProtocolError::FrameTooLarge {
486            size: frame_len,
487            max: MAX_FRAME_SIZE,
488        });
489    }
490    if frame_len < FRAME_HEADER_SIZE as u32 {
491        return Err(ProtocolError::FrameTooShort {
492            size: frame_len,
493            min: FRAME_HEADER_SIZE as u32,
494        });
495    }
496
497    let frame_len = frame_len as usize;
498    let total = 4 + frame_len;
499    if buf.len() < total {
500        return Ok(None);
501    }
502    Ok(Some((frame_len, total)))
503}
504
505fn invalid_bulk(message: impl Into<String>) -> ProtocolError {
506    ProtocolError::InvalidBulkFrame(message.into())
507}
508
509async fn write_vectored_all<W: AsyncWrite + Unpin>(
510    writer: &mut W,
511    header: &[u8],
512    body: &[u8],
513) -> std::io::Result<()> {
514    let mut header_offset = 0;
515    let mut body_offset = 0;
516
517    while header_offset < header.len() || body_offset < body.len() {
518        let written = if header_offset < header.len() {
519            let slices = [
520                IoSlice::new(&header[header_offset..]),
521                IoSlice::new(&body[body_offset..]),
522            ];
523            let slice_count = if body_offset < body.len() { 2 } else { 1 };
524            writer.write_vectored(&slices[..slice_count]).await?
525        } else {
526            // Some AsyncWrite implementations stop at an empty first IoSlice, so never leave the
527            // exhausted header in front of a non-empty body.
528            writer.write(&body[body_offset..]).await?
529        };
530        if written == 0 {
531            return Err(std::io::ErrorKind::WriteZero.into());
532        }
533
534        let header_remaining = header.len() - header_offset;
535        if written < header_remaining {
536            header_offset += written;
537        } else {
538            header_offset = header.len();
539            body_offset += written - header_remaining;
540        }
541    }
542
543    Ok(())
544}
545
546//--------------------------------------------------------------------------------------------------
547// Tests
548//--------------------------------------------------------------------------------------------------
549
550#[cfg(test)]
551mod tests {
552    use std::pin::Pin;
553    use std::task::{Context, Poll};
554
555    use super::*;
556    use crate::message::{FLAG_SESSION_START, FLAG_TERMINAL, MessageType, PROTOCOL_VERSION};
557
558    #[tokio::test]
559    async fn test_codec_roundtrip_empty_payload() {
560        let msg = Message::new(MessageType::Ready, 0, Vec::new());
561
562        let mut buf = Vec::new();
563        write_message(&mut buf, &msg).await.unwrap();
564
565        let mut cursor = &buf[..];
566        let decoded = read_message(&mut cursor).await.unwrap();
567
568        assert_eq!(decoded.v, msg.v);
569        assert_eq!(decoded.t, msg.t);
570        assert_eq!(decoded.id, msg.id);
571        assert_eq!(decoded.flags, 0);
572    }
573
574    #[tokio::test]
575    async fn test_codec_roundtrip_with_payload() {
576        use crate::exec::ExecExited;
577
578        let msg =
579            Message::with_payload(MessageType::ExecExited, 7, &ExecExited { code: 42 }).unwrap();
580
581        let mut buf = Vec::new();
582        write_message(&mut buf, &msg).await.unwrap();
583
584        let mut cursor = &buf[..];
585        let decoded = read_message(&mut cursor).await.unwrap();
586
587        assert_eq!(decoded.v, PROTOCOL_VERSION);
588        assert_eq!(decoded.t, MessageType::ExecExited);
589        assert_eq!(decoded.id, 7);
590        assert_eq!(decoded.flags, FLAG_TERMINAL);
591
592        let payload: ExecExited = decoded.payload().unwrap();
593        assert_eq!(payload.code, 42);
594    }
595
596    #[tokio::test]
597    async fn test_codec_multiple_messages() {
598        let messages = vec![
599            Message::new(MessageType::Ready, 0, Vec::new()),
600            Message::new(MessageType::ExecExited, 1, Vec::new()),
601            Message::new(MessageType::Shutdown, 2, Vec::new()),
602        ];
603
604        let mut buf = Vec::new();
605        for msg in &messages {
606            write_message(&mut buf, msg).await.unwrap();
607        }
608
609        let mut cursor = &buf[..];
610        for expected in &messages {
611            let decoded = read_message(&mut cursor).await.unwrap();
612            assert_eq!(decoded.t, expected.t);
613            assert_eq!(decoded.id, expected.id);
614            assert_eq!(decoded.flags, expected.flags);
615        }
616    }
617
618    #[tokio::test]
619    async fn test_codec_unexpected_eof() {
620        let mut cursor: &[u8] = &[];
621        let result = read_message(&mut cursor).await;
622        assert!(matches!(result, Err(ProtocolError::UnexpectedEof)));
623    }
624
625    #[test]
626    fn test_sync_encode_decode_roundtrip() {
627        use crate::exec::ExecExited;
628
629        let msg =
630            Message::with_payload(MessageType::ExecExited, 5, &ExecExited { code: 0 }).unwrap();
631
632        let mut buf = Vec::new();
633        encode_to_buf(&msg, &mut buf).unwrap();
634
635        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
636        assert_eq!(decoded.t, MessageType::ExecExited);
637        assert_eq!(decoded.id, 5);
638        assert_eq!(decoded.flags, FLAG_TERMINAL);
639
640        let payload: ExecExited = decoded.payload().unwrap();
641        assert_eq!(payload.code, 0);
642        assert!(buf.is_empty());
643    }
644
645    #[test]
646    fn test_borrowed_decode_message_frame_roundtrip() {
647        use crate::exec::ExecExited;
648
649        let msg =
650            Message::with_payload(MessageType::ExecExited, 5, &ExecExited { code: 0 }).unwrap();
651
652        let mut buf = Vec::new();
653        encode_to_buf(&msg, &mut buf).unwrap();
654
655        let decoded = decode_message_frame(&buf).unwrap();
656        assert_eq!(decoded.t, MessageType::ExecExited);
657        assert_eq!(decoded.id, 5);
658        assert_eq!(decoded.flags, FLAG_TERMINAL);
659
660        let payload: ExecExited = decoded.payload().unwrap();
661        assert_eq!(payload.code, 0);
662        assert!(!buf.is_empty(), "borrowed decode must not consume input");
663    }
664
665    #[test]
666    fn test_borrowed_decode_message_frame_rejects_incomplete() {
667        let buf = vec![0, 0, 0, 10];
668        assert!(matches!(
669            decode_message_frame(&buf),
670            Err(ProtocolError::UnexpectedEof)
671        ));
672    }
673
674    #[test]
675    fn test_sync_decode_incomplete() {
676        let mut buf = vec![0, 0, 0, 10]; // Length 10 but no payload bytes.
677        assert!(try_decode_from_buf(&mut buf).unwrap().is_none());
678    }
679
680    #[test]
681    fn test_sync_decode_frame_too_large() {
682        let huge_len: u32 = MAX_FRAME_SIZE + 1;
683        let mut buf = Vec::new();
684        buf.extend_from_slice(&huge_len.to_be_bytes());
685        let result = try_decode_from_buf(&mut buf);
686        assert!(matches!(result, Err(ProtocolError::FrameTooLarge { .. })));
687    }
688
689    #[test]
690    fn test_frame_header_wire_format() {
691        let msg = Message::new(MessageType::ExecRequest, 0x12345678, Vec::new());
692
693        let mut buf = Vec::new();
694        encode_to_buf(&msg, &mut buf).unwrap();
695
696        // Bytes 0–3: length prefix (u32 BE).
697        let len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
698        assert_eq!(len as usize + 4, buf.len());
699
700        // Bytes 4–7: correlation ID (u32 BE).
701        let id = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
702        assert_eq!(id, 0x12345678);
703
704        // Byte 8: flags.
705        assert_eq!(buf[8], FLAG_SESSION_START);
706
707        // Bytes 9..: CBOR body (v, t, p — no id or flags).
708    }
709
710    #[test]
711    fn test_flags_roundtrip_terminal() {
712        let msg = Message::new(MessageType::ExecExited, 99, Vec::new());
713
714        let mut buf = Vec::new();
715        encode_to_buf(&msg, &mut buf).unwrap();
716
717        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
718        assert_ne!(decoded.flags & FLAG_TERMINAL, 0);
719        assert_eq!(decoded.flags & FLAG_SESSION_START, 0);
720    }
721
722    #[test]
723    fn test_flags_roundtrip_session_start() {
724        let msg = Message::new(MessageType::FsRequest, 42, Vec::new());
725
726        let mut buf = Vec::new();
727        encode_to_buf(&msg, &mut buf).unwrap();
728
729        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
730        assert_ne!(decoded.flags & FLAG_SESSION_START, 0);
731        assert_eq!(decoded.flags & FLAG_TERMINAL, 0);
732    }
733
734    #[test]
735    fn test_sync_decode_frame_too_short() {
736        // Frame with len=3 (too short for id+flags header).
737        let mut buf = Vec::new();
738        buf.extend_from_slice(&3u32.to_be_bytes());
739        buf.extend_from_slice(&[0, 0, 0]); // 3 bytes of payload.
740
741        let result = try_decode_from_buf(&mut buf);
742        assert!(matches!(result, Err(ProtocolError::FrameTooShort { .. })));
743    }
744
745    #[tokio::test]
746    async fn test_raw_frame_roundtrip() {
747        let frame = RawFrame {
748            id: 0xDEADBEEF,
749            flags: FLAG_TERMINAL,
750            body: vec![1, 2, 3, 4, 5],
751        };
752
753        let mut buf = Vec::new();
754        write_raw_frame(&mut buf, &frame).await.unwrap();
755
756        let mut cursor = &buf[..];
757        let decoded = read_raw_frame(&mut cursor).await.unwrap();
758
759        assert_eq!(decoded.id, frame.id);
760        assert_eq!(decoded.flags, frame.flags);
761        assert_eq!(decoded.body, frame.body);
762    }
763
764    #[test]
765    fn smallest_bulk_record_has_exact_wire_layout() {
766        let record = BulkRecord {
767            id: 0x0102_0304,
768            kind: BulkKind::Filesystem,
769            flow: BulkFlow::HostToGuest,
770            offset: 0x0102_0304_0506_0708,
771            payload: Bytes::from_static(&[0xFF]),
772        };
773        let mut encoded = Vec::new();
774        encode_bulk_to_buf(&record, &mut encoded).unwrap();
775
776        assert_eq!(
777            encoded,
778            vec![
779                0x00, 0x00, 0x00, 0x12, // frame length: 5 + 12 + 1
780                0x01, 0x02, 0x03, 0x04, // correlation ID
781                0x08, // FLAG_BULK
782                0x01, // filesystem
783                0x01, // host to guest
784                0x00, 0x00, // reserved
785                0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, // offset
786                0xFF, // opaque payload
787            ]
788        );
789    }
790
791    #[test]
792    fn largest_bulk_record_roundtrips_without_payload_copy() {
793        let record = BulkRecord {
794            id: 81,
795            kind: BulkKind::Tcp,
796            flow: BulkFlow::GuestToHost,
797            offset: 7,
798            payload: Bytes::from(vec![0xA5; MAX_BULK_RECORD_PAYLOAD as usize]),
799        };
800        let mut encoded = Vec::new();
801        encode_bulk_to_buf(&record, &mut encoded).unwrap();
802        let mut cursor = BytesMut::from(encoded.as_slice());
803
804        let Some(DecodedFrame::Bulk(decoded)) = try_decode_frame_from_bytes(&mut cursor).unwrap()
805        else {
806            panic!("expected bulk record");
807        };
808        assert_eq!(decoded, record);
809        assert!(cursor.is_empty());
810    }
811
812    #[test]
813    fn bulk_payload_is_opaque_even_when_it_looks_like_cbor() {
814        let payload = Bytes::from_static(&[0xA3, 0x61, b'v', 0x07, 0x61, b't', 0xF6]);
815        let record = BulkRecord {
816            id: 3,
817            kind: BulkKind::Filesystem,
818            flow: BulkFlow::GuestToHost,
819            offset: 0,
820            payload: payload.clone(),
821        };
822        let mut encoded = Vec::new();
823        encode_bulk_to_buf(&record, &mut encoded).unwrap();
824        let mut cursor = BytesMut::from(encoded.as_slice());
825
826        let Some(DecodedFrame::Bulk(decoded)) = try_decode_frame_from_bytes(&mut cursor).unwrap()
827        else {
828            panic!("expected bulk record");
829        };
830        assert_eq!(decoded.payload, payload);
831    }
832
833    #[test]
834    fn bulk_decoder_accepts_every_frame_fragment_boundary() {
835        let record = BulkRecord {
836            id: 34,
837            kind: BulkKind::Tcp,
838            flow: BulkFlow::HostToGuest,
839            offset: 99,
840            payload: Bytes::from(vec![0x73; 128]),
841        };
842        let mut encoded = Vec::new();
843        encode_bulk_to_buf(&record, &mut encoded).unwrap();
844
845        for split in 0..=encoded.len() {
846            let mut cursor = BytesMut::from(&encoded[..split]);
847            let first = try_decode_frame_from_bytes(&mut cursor).unwrap();
848            if split < encoded.len() {
849                assert!(first.is_none(), "decoded incomplete frame at split {split}");
850                cursor.extend_from_slice(&encoded[split..]);
851            }
852            let decoded = first
853                .or_else(|| try_decode_frame_from_bytes(&mut cursor).unwrap())
854                .unwrap();
855            assert!(matches!(decoded, DecodedFrame::Bulk(ref value) if value == &record));
856            assert!(cursor.is_empty());
857        }
858    }
859
860    #[test]
861    fn bulk_decoder_rejects_malformed_wire_shapes() {
862        let record = BulkRecord {
863            id: 5,
864            kind: BulkKind::Filesystem,
865            flow: BulkFlow::HostToGuest,
866            offset: 0,
867            payload: Bytes::from_static(b"x"),
868        };
869        let mut valid = Vec::new();
870        encode_bulk_to_buf(&record, &mut valid).unwrap();
871
872        for (index, value) in [(8, FLAG_BULK | FLAG_TERMINAL), (9, 0), (10, 3), (11, 1)] {
873            let mut malformed = valid.clone();
874            malformed[index] = value;
875            let error =
876                try_decode_frame_from_bytes(&mut BytesMut::from(malformed.as_slice())).unwrap_err();
877            assert!(matches!(error, ProtocolError::InvalidBulkFrame(_)));
878        }
879
880        let mut no_payload = valid;
881        no_payload.truncate(4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE);
882        no_payload[..4]
883            .copy_from_slice(&((FRAME_HEADER_SIZE + BULK_HEADER_SIZE) as u32).to_be_bytes());
884        let error =
885            try_decode_frame_from_bytes(&mut BytesMut::from(no_payload.as_slice())).unwrap_err();
886        assert!(matches!(error, ProtocolError::InvalidBulkFrame(_)));
887    }
888
889    #[tokio::test]
890    async fn bulk_vectored_writer_handles_one_byte_short_writes() {
891        let record = BulkRecord {
892            id: 91,
893            kind: BulkKind::Tcp,
894            flow: BulkFlow::GuestToHost,
895            offset: 42,
896            payload: Bytes::from(vec![0xCD; 257]),
897        };
898        let mut expected = Vec::new();
899        encode_bulk_to_buf(&record, &mut expected).unwrap();
900
901        let mut writer = OneByteWriter::default();
902        write_bulk_record(&mut writer, &record).await.unwrap();
903
904        assert_eq!(writer.bytes, expected);
905    }
906
907    #[test]
908    fn cursor_decoder_accepts_every_frame_fragment_boundary() {
909        let msg = Message::new(MessageType::Ready, 77, vec![0xAB; 1024]);
910        let mut encoded = Vec::new();
911        encode_to_buf(&msg, &mut encoded).unwrap();
912
913        for split in 0..=encoded.len() {
914            let mut buf = BytesMut::new();
915            buf.extend_from_slice(&encoded[..split]);
916            let first = try_decode_from_bytes(&mut buf).unwrap();
917            if split < encoded.len() {
918                assert!(first.is_none(), "decoded incomplete frame at split {split}");
919                buf.extend_from_slice(&encoded[split..]);
920            }
921
922            let decoded = first
923                .or_else(|| try_decode_from_bytes(&mut buf).unwrap())
924                .unwrap();
925            assert_eq!(decoded.id, msg.id);
926            assert_eq!(decoded.t, msg.t);
927            assert!(buf.is_empty());
928        }
929    }
930
931    #[tokio::test]
932    async fn vectored_writer_handles_one_byte_short_writes() {
933        let frame = RawFrame {
934            id: 91,
935            flags: FLAG_TERMINAL,
936            body: vec![0xCD; 257],
937        };
938        let mut expected = Vec::new();
939        encode_raw_to_buf(&frame, &mut expected).unwrap();
940
941        let mut writer = OneByteWriter::default();
942        write_raw_frame(&mut writer, &frame).await.unwrap();
943
944        assert_eq!(writer.bytes, expected);
945    }
946
947    #[test]
948    fn test_raw_frame_sync_roundtrip() {
949        let frame = RawFrame {
950            id: 42,
951            flags: FLAG_SESSION_START,
952            body: vec![0xAA; 100],
953        };
954
955        let mut buf = Vec::new();
956        encode_raw_to_buf(&frame, &mut buf).unwrap();
957
958        let decoded = try_decode_raw_from_buf(&mut buf).unwrap().unwrap();
959        assert_eq!(decoded.id, frame.id);
960        assert_eq!(decoded.flags, frame.flags);
961        assert_eq!(decoded.body, frame.body);
962        assert!(buf.is_empty());
963    }
964
965    #[test]
966    fn test_raw_frame_to_message() {
967        use crate::exec::ExecExited;
968
969        let msg =
970            Message::with_payload(MessageType::ExecExited, 13, &ExecExited { code: 7 }).unwrap();
971
972        let mut buf = Vec::new();
973        encode_to_buf(&msg, &mut buf).unwrap();
974
975        let frame = try_decode_raw_from_buf(&mut buf).unwrap().unwrap();
976        let decoded = raw_frame_to_message(frame).unwrap();
977
978        assert_eq!(decoded.id, 13);
979        assert_eq!(decoded.t, MessageType::ExecExited);
980        let payload: ExecExited = decoded.payload().unwrap();
981        assert_eq!(payload.code, 7);
982    }
983
984    #[derive(Default)]
985    struct OneByteWriter {
986        bytes: Vec<u8>,
987    }
988
989    impl AsyncWrite for OneByteWriter {
990        fn poll_write(
991            mut self: Pin<&mut Self>,
992            _cx: &mut Context<'_>,
993            buf: &[u8],
994        ) -> Poll<std::io::Result<usize>> {
995            if let Some(byte) = buf.first() {
996                self.bytes.push(*byte);
997                Poll::Ready(Ok(1))
998            } else {
999                Poll::Ready(Ok(0))
1000            }
1001        }
1002
1003        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1004            Poll::Ready(Ok(()))
1005        }
1006
1007        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1008            Poll::Ready(Ok(()))
1009        }
1010
1011        fn is_write_vectored(&self) -> bool {
1012            true
1013        }
1014
1015        fn poll_write_vectored(
1016            mut self: Pin<&mut Self>,
1017            _cx: &mut Context<'_>,
1018            bufs: &[IoSlice<'_>],
1019        ) -> Poll<std::io::Result<usize>> {
1020            if let Some(byte) = bufs.iter().find_map(|buf| buf.first()) {
1021                self.bytes.push(*byte);
1022                Poll::Ready(Ok(1))
1023            } else {
1024                Poll::Ready(Ok(0))
1025            }
1026        }
1027    }
1028}
microsandbox_protocol/codec.rs

microsandbox_protocol/
codec.rs
