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, MessageType},
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/// Decode only the message kind from an encoded control body, without
279/// copying its payload. Returns `None` when the body is not a CBOR envelope.
280pub fn peek_message_type(body: &[u8]) -> Option<MessageType> {
281    #[derive(serde::Deserialize)]
282    struct Kind {
283        t: MessageType,
284    }
285
286    ciborium::from_reader::<Kind, _>(body)
287        .ok()
288        .map(|kind| kind.t)
289}
290
291/// Encodes a message to a byte buffer using the length-prefixed frame format.
292///
293/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
294pub fn encode_to_buf(msg: &Message, buf: &mut Vec<u8>) -> ProtocolResult<()> {
295    let mut body = Vec::new();
296    ciborium::into_writer(msg, &mut body)?;
297    encode_raw_to_buf(
298        &RawFrame {
299            id: msg.id,
300            flags: msg.flags,
301            body,
302        },
303        buf,
304    )
305}
306
307/// Tries to decode a complete message from a byte buffer.
308///
309/// Returns `Some(Message)` if a complete frame is available, consuming
310/// the bytes. Returns `None` if more data is needed.
311///
312/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
313pub fn try_decode_from_buf(buf: &mut Vec<u8>) -> ProtocolResult<Option<Message>> {
314    if buf.len() < 4 {
315        return Ok(None);
316    }
317
318    let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
319
320    if frame_len > MAX_FRAME_SIZE {
321        return Err(ProtocolError::FrameTooLarge {
322            size: frame_len,
323            max: MAX_FRAME_SIZE,
324        });
325    }
326
327    let frame_len = frame_len as usize;
328    let total = 4 + frame_len;
329
330    if buf.len() < total {
331        return Ok(None);
332    }
333
334    let msg = decode_message_frame(&buf[..total])?;
335    buf.drain(..total);
336    Ok(Some(msg))
337}
338
339/// Reads a length-prefixed message from the given reader.
340///
341/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
342pub async fn read_message<R: AsyncRead + Unpin>(reader: &mut R) -> ProtocolResult<Message> {
343    let frame = read_raw_frame(reader).await?;
344    raw_frame_to_message(frame)
345}
346
347/// Writes a length-prefixed message to the given writer.
348///
349/// Frame format: `[len: u32 BE][id: u32 BE][flags: u8][CBOR(v, t, p)]`
350pub async fn write_message<W: AsyncWrite + Unpin>(
351    writer: &mut W,
352    message: &Message,
353) -> ProtocolResult<()> {
354    let mut body = Vec::new();
355    ciborium::into_writer(message, &mut body)?;
356    write_raw_frame(
357        writer,
358        &RawFrame {
359            id: message.id,
360            flags: message.flags,
361            body,
362        },
363    )
364    .await
365}
366
367/// Decodes a [`RawFrame`] into a typed [`Message`] by CBOR-deserializing the body.
368pub fn raw_frame_to_message(frame: RawFrame) -> ProtocolResult<Message> {
369    let mut msg: Message = ciborium::from_reader(&frame.body[..])?;
370    msg.id = frame.id;
371    msg.flags = frame.flags;
372    Ok(msg)
373}
374
375/// Decodes one complete length-prefixed frame from a borrowed byte slice.
376///
377/// The input must include the 4-byte length prefix, frame header, and CBOR body.
378/// The slice is not consumed or copied.
379pub fn decode_message_frame(frame: &[u8]) -> ProtocolResult<Message> {
380    if frame.len() < 4 {
381        return Err(ProtocolError::UnexpectedEof);
382    }
383
384    let frame_len = u32::from_be_bytes([frame[0], frame[1], frame[2], frame[3]]);
385    if frame_len > MAX_FRAME_SIZE {
386        return Err(ProtocolError::FrameTooLarge {
387            size: frame_len,
388            max: MAX_FRAME_SIZE,
389        });
390    }
391
392    let frame_len = frame_len as usize;
393    let total = 4 + frame_len;
394    if frame.len() < total {
395        return Err(ProtocolError::UnexpectedEof);
396    }
397
398    if frame_len < FRAME_HEADER_SIZE {
399        return Err(ProtocolError::FrameTooShort {
400            size: frame_len as u32,
401            min: FRAME_HEADER_SIZE as u32,
402        });
403    }
404
405    let mut msg: Message = ciborium::from_reader(&frame[4 + FRAME_HEADER_SIZE..total])?;
406    msg.id = u32::from_be_bytes([frame[4], frame[5], frame[6], frame[7]]);
407    msg.flags = frame[8];
408    Ok(msg)
409}
410
411/// Encodes the fixed outer and generation-8 bulk headers, leaving the payload separate.
412pub fn encode_bulk_header(
413    record: &BulkRecord,
414) -> ProtocolResult<[u8; 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE]> {
415    let payload_len = record.payload.len();
416    if payload_len == 0 || payload_len > MAX_BULK_RECORD_PAYLOAD as usize {
417        return Err(invalid_bulk(format!(
418            "payload length {payload_len} is outside 1..={MAX_BULK_RECORD_PAYLOAD}"
419        )));
420    }
421    record
422        .offset
423        .checked_add(payload_len as u64)
424        .ok_or_else(|| invalid_bulk("record end offset overflows u64"))?;
425
426    let frame_len = FRAME_HEADER_SIZE + BULK_HEADER_SIZE + payload_len;
427    let frame_len = u32::try_from(frame_len).map_err(|_| ProtocolError::FrameTooLarge {
428        size: u32::MAX,
429        max: MAX_FRAME_SIZE,
430    })?;
431    if frame_len > MAX_FRAME_SIZE {
432        return Err(ProtocolError::FrameTooLarge {
433            size: frame_len,
434            max: MAX_FRAME_SIZE,
435        });
436    }
437
438    let mut header = [0u8; 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE];
439    header[..4].copy_from_slice(&frame_len.to_be_bytes());
440    header[4..8].copy_from_slice(&record.id.to_be_bytes());
441    header[8] = FLAG_BULK;
442    header[9] = record.kind as u8;
443    header[10] = record.flow as u8;
444    // Bytes 11..13 are reserved and deliberately remain zero.
445    header[13..21].copy_from_slice(&record.offset.to_be_bytes());
446    Ok(header)
447}
448
449pub(crate) fn decode_bulk_body(
450    id: u32,
451    body: Bytes,
452    max_payload: u32,
453) -> ProtocolResult<BulkRecord> {
454    if max_payload == 0 || max_payload > MAX_BULK_RECORD_PAYLOAD {
455        return Err(invalid_bulk(format!(
456            "invalid decoder payload limit {max_payload}"
457        )));
458    }
459    if body.len() < BULK_HEADER_SIZE + 1 {
460        return Err(invalid_bulk("bulk record has no payload"));
461    }
462    if body[2] != 0 || body[3] != 0 {
463        return Err(invalid_bulk("reserved bulk header bytes must be zero"));
464    }
465
466    let kind = BulkKind::from_wire(body[0])
467        .ok_or_else(|| invalid_bulk(format!("unknown bulk kind {}", body[0])))?;
468    let flow = BulkFlow::from_wire(body[1])
469        .ok_or_else(|| invalid_bulk(format!("unknown bulk flow {}", body[1])))?;
470    let offset = u64::from_be_bytes(body[4..12].try_into().unwrap());
471    let payload = body.slice(BULK_HEADER_SIZE..);
472    if payload.len() > max_payload as usize {
473        return Err(invalid_bulk(format!(
474            "payload length {} exceeds negotiated maximum {max_payload}",
475            payload.len()
476        )));
477    }
478    offset
479        .checked_add(payload.len() as u64)
480        .ok_or_else(|| invalid_bulk("record end offset overflows u64"))?;
481
482    Ok(BulkRecord {
483        id,
484        kind,
485        flow,
486        offset,
487        payload,
488    })
489}
490
491fn complete_frame_len(buf: &BytesMut) -> ProtocolResult<Option<(usize, usize)>> {
492    if buf.len() < 4 {
493        return Ok(None);
494    }
495
496    let frame_len = u32::from_be_bytes(buf[..4].try_into().unwrap());
497    if frame_len > MAX_FRAME_SIZE {
498        return Err(ProtocolError::FrameTooLarge {
499            size: frame_len,
500            max: MAX_FRAME_SIZE,
501        });
502    }
503    if frame_len < FRAME_HEADER_SIZE as u32 {
504        return Err(ProtocolError::FrameTooShort {
505            size: frame_len,
506            min: FRAME_HEADER_SIZE as u32,
507        });
508    }
509
510    let frame_len = frame_len as usize;
511    let total = 4 + frame_len;
512    if buf.len() < total {
513        return Ok(None);
514    }
515    Ok(Some((frame_len, total)))
516}
517
518fn invalid_bulk(message: impl Into<String>) -> ProtocolError {
519    ProtocolError::InvalidBulkFrame(message.into())
520}
521
522async fn write_vectored_all<W: AsyncWrite + Unpin>(
523    writer: &mut W,
524    header: &[u8],
525    body: &[u8],
526) -> std::io::Result<()> {
527    let mut header_offset = 0;
528    let mut body_offset = 0;
529
530    while header_offset < header.len() || body_offset < body.len() {
531        let written = if header_offset < header.len() {
532            let slices = [
533                IoSlice::new(&header[header_offset..]),
534                IoSlice::new(&body[body_offset..]),
535            ];
536            let slice_count = if body_offset < body.len() { 2 } else { 1 };
537            writer.write_vectored(&slices[..slice_count]).await?
538        } else {
539            // Some AsyncWrite implementations stop at an empty first IoSlice, so never leave the
540            // exhausted header in front of a non-empty body.
541            writer.write(&body[body_offset..]).await?
542        };
543        if written == 0 {
544            return Err(std::io::ErrorKind::WriteZero.into());
545        }
546
547        let header_remaining = header.len() - header_offset;
548        if written < header_remaining {
549            header_offset += written;
550        } else {
551            header_offset = header.len();
552            body_offset += written - header_remaining;
553        }
554    }
555
556    Ok(())
557}
558
559//--------------------------------------------------------------------------------------------------
560// Tests
561//--------------------------------------------------------------------------------------------------
562
563#[cfg(test)]
564mod tests {
565    use std::pin::Pin;
566    use std::task::{Context, Poll};
567
568    use super::*;
569    use crate::message::{FLAG_SESSION_START, FLAG_TERMINAL, MessageType, PROTOCOL_VERSION};
570
571    #[tokio::test]
572    async fn test_codec_roundtrip_empty_payload() {
573        let msg = Message::new(MessageType::Ready, 0, Vec::new());
574
575        let mut buf = Vec::new();
576        write_message(&mut buf, &msg).await.unwrap();
577
578        let mut cursor = &buf[..];
579        let decoded = read_message(&mut cursor).await.unwrap();
580
581        assert_eq!(decoded.v, msg.v);
582        assert_eq!(decoded.t, msg.t);
583        assert_eq!(decoded.id, msg.id);
584        assert_eq!(decoded.flags, 0);
585    }
586
587    #[tokio::test]
588    async fn test_codec_roundtrip_with_payload() {
589        use crate::exec::ExecExited;
590
591        let msg =
592            Message::with_payload(MessageType::ExecExited, 7, &ExecExited { code: 42 }).unwrap();
593
594        let mut buf = Vec::new();
595        write_message(&mut buf, &msg).await.unwrap();
596
597        let mut cursor = &buf[..];
598        let decoded = read_message(&mut cursor).await.unwrap();
599
600        assert_eq!(decoded.v, PROTOCOL_VERSION);
601        assert_eq!(decoded.t, MessageType::ExecExited);
602        assert_eq!(decoded.id, 7);
603        assert_eq!(decoded.flags, FLAG_TERMINAL);
604
605        let payload: ExecExited = decoded.payload().unwrap();
606        assert_eq!(payload.code, 42);
607    }
608
609    #[tokio::test]
610    async fn test_codec_multiple_messages() {
611        let messages = vec![
612            Message::new(MessageType::Ready, 0, Vec::new()),
613            Message::new(MessageType::ExecExited, 1, Vec::new()),
614            Message::new(MessageType::Shutdown, 2, Vec::new()),
615        ];
616
617        let mut buf = Vec::new();
618        for msg in &messages {
619            write_message(&mut buf, msg).await.unwrap();
620        }
621
622        let mut cursor = &buf[..];
623        for expected in &messages {
624            let decoded = read_message(&mut cursor).await.unwrap();
625            assert_eq!(decoded.t, expected.t);
626            assert_eq!(decoded.id, expected.id);
627            assert_eq!(decoded.flags, expected.flags);
628        }
629    }
630
631    #[tokio::test]
632    async fn test_codec_unexpected_eof() {
633        let mut cursor: &[u8] = &[];
634        let result = read_message(&mut cursor).await;
635        assert!(matches!(result, Err(ProtocolError::UnexpectedEof)));
636    }
637
638    #[test]
639    fn test_sync_encode_decode_roundtrip() {
640        use crate::exec::ExecExited;
641
642        let msg =
643            Message::with_payload(MessageType::ExecExited, 5, &ExecExited { code: 0 }).unwrap();
644
645        let mut buf = Vec::new();
646        encode_to_buf(&msg, &mut buf).unwrap();
647
648        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
649        assert_eq!(decoded.t, MessageType::ExecExited);
650        assert_eq!(decoded.id, 5);
651        assert_eq!(decoded.flags, FLAG_TERMINAL);
652
653        let payload: ExecExited = decoded.payload().unwrap();
654        assert_eq!(payload.code, 0);
655        assert!(buf.is_empty());
656    }
657
658    #[test]
659    fn test_borrowed_decode_message_frame_roundtrip() {
660        use crate::exec::ExecExited;
661
662        let msg =
663            Message::with_payload(MessageType::ExecExited, 5, &ExecExited { code: 0 }).unwrap();
664
665        let mut buf = Vec::new();
666        encode_to_buf(&msg, &mut buf).unwrap();
667
668        let decoded = decode_message_frame(&buf).unwrap();
669        assert_eq!(decoded.t, MessageType::ExecExited);
670        assert_eq!(decoded.id, 5);
671        assert_eq!(decoded.flags, FLAG_TERMINAL);
672
673        let payload: ExecExited = decoded.payload().unwrap();
674        assert_eq!(payload.code, 0);
675        assert!(!buf.is_empty(), "borrowed decode must not consume input");
676    }
677
678    #[test]
679    fn test_borrowed_decode_message_frame_rejects_incomplete() {
680        let buf = vec![0, 0, 0, 10];
681        assert!(matches!(
682            decode_message_frame(&buf),
683            Err(ProtocolError::UnexpectedEof)
684        ));
685    }
686
687    #[test]
688    fn test_sync_decode_incomplete() {
689        let mut buf = vec![0, 0, 0, 10]; // Length 10 but no payload bytes.
690        assert!(try_decode_from_buf(&mut buf).unwrap().is_none());
691    }
692
693    #[test]
694    fn test_sync_decode_frame_too_large() {
695        let huge_len: u32 = MAX_FRAME_SIZE + 1;
696        let mut buf = Vec::new();
697        buf.extend_from_slice(&huge_len.to_be_bytes());
698        let result = try_decode_from_buf(&mut buf);
699        assert!(matches!(result, Err(ProtocolError::FrameTooLarge { .. })));
700    }
701
702    #[test]
703    fn test_frame_header_wire_format() {
704        let msg = Message::new(MessageType::ExecRequest, 0x12345678, Vec::new());
705
706        let mut buf = Vec::new();
707        encode_to_buf(&msg, &mut buf).unwrap();
708
709        // Bytes 0–3: length prefix (u32 BE).
710        let len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
711        assert_eq!(len as usize + 4, buf.len());
712
713        // Bytes 4–7: correlation ID (u32 BE).
714        let id = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
715        assert_eq!(id, 0x12345678);
716
717        // Byte 8: flags.
718        assert_eq!(buf[8], FLAG_SESSION_START);
719
720        // Bytes 9..: CBOR body (v, t, p — no id or flags).
721    }
722
723    #[test]
724    fn test_flags_roundtrip_terminal() {
725        let msg = Message::new(MessageType::ExecExited, 99, Vec::new());
726
727        let mut buf = Vec::new();
728        encode_to_buf(&msg, &mut buf).unwrap();
729
730        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
731        assert_ne!(decoded.flags & FLAG_TERMINAL, 0);
732        assert_eq!(decoded.flags & FLAG_SESSION_START, 0);
733    }
734
735    #[test]
736    fn test_flags_roundtrip_session_start() {
737        let msg = Message::new(MessageType::FsRequest, 42, Vec::new());
738
739        let mut buf = Vec::new();
740        encode_to_buf(&msg, &mut buf).unwrap();
741
742        let decoded = try_decode_from_buf(&mut buf).unwrap().unwrap();
743        assert_ne!(decoded.flags & FLAG_SESSION_START, 0);
744        assert_eq!(decoded.flags & FLAG_TERMINAL, 0);
745    }
746
747    #[test]
748    fn test_sync_decode_frame_too_short() {
749        // Frame with len=3 (too short for id+flags header).
750        let mut buf = Vec::new();
751        buf.extend_from_slice(&3u32.to_be_bytes());
752        buf.extend_from_slice(&[0, 0, 0]); // 3 bytes of payload.
753
754        let result = try_decode_from_buf(&mut buf);
755        assert!(matches!(result, Err(ProtocolError::FrameTooShort { .. })));
756    }
757
758    #[tokio::test]
759    async fn test_raw_frame_roundtrip() {
760        let frame = RawFrame {
761            id: 0xDEADBEEF,
762            flags: FLAG_TERMINAL,
763            body: vec![1, 2, 3, 4, 5],
764        };
765
766        let mut buf = Vec::new();
767        write_raw_frame(&mut buf, &frame).await.unwrap();
768
769        let mut cursor = &buf[..];
770        let decoded = read_raw_frame(&mut cursor).await.unwrap();
771
772        assert_eq!(decoded.id, frame.id);
773        assert_eq!(decoded.flags, frame.flags);
774        assert_eq!(decoded.body, frame.body);
775    }
776
777    #[test]
778    fn smallest_bulk_record_has_exact_wire_layout() {
779        let record = BulkRecord {
780            id: 0x0102_0304,
781            kind: BulkKind::Filesystem,
782            flow: BulkFlow::HostToGuest,
783            offset: 0x0102_0304_0506_0708,
784            payload: Bytes::from_static(&[0xFF]),
785        };
786        let mut encoded = Vec::new();
787        encode_bulk_to_buf(&record, &mut encoded).unwrap();
788
789        assert_eq!(
790            encoded,
791            vec![
792                0x00, 0x00, 0x00, 0x12, // frame length: 5 + 12 + 1
793                0x01, 0x02, 0x03, 0x04, // correlation ID
794                0x08, // FLAG_BULK
795                0x01, // filesystem
796                0x01, // host to guest
797                0x00, 0x00, // reserved
798                0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, // offset
799                0xFF, // opaque payload
800            ]
801        );
802    }
803
804    #[test]
805    fn largest_bulk_record_roundtrips_without_payload_copy() {
806        let record = BulkRecord {
807            id: 81,
808            kind: BulkKind::Tcp,
809            flow: BulkFlow::GuestToHost,
810            offset: 7,
811            payload: Bytes::from(vec![0xA5; MAX_BULK_RECORD_PAYLOAD as usize]),
812        };
813        let mut encoded = Vec::new();
814        encode_bulk_to_buf(&record, &mut encoded).unwrap();
815        let mut cursor = BytesMut::from(encoded.as_slice());
816
817        let Some(DecodedFrame::Bulk(decoded)) = try_decode_frame_from_bytes(&mut cursor).unwrap()
818        else {
819            panic!("expected bulk record");
820        };
821        assert_eq!(decoded, record);
822        assert!(cursor.is_empty());
823    }
824
825    #[test]
826    fn bulk_payload_is_opaque_even_when_it_looks_like_cbor() {
827        let payload = Bytes::from_static(&[0xA3, 0x61, b'v', 0x07, 0x61, b't', 0xF6]);
828        let record = BulkRecord {
829            id: 3,
830            kind: BulkKind::Filesystem,
831            flow: BulkFlow::GuestToHost,
832            offset: 0,
833            payload: payload.clone(),
834        };
835        let mut encoded = Vec::new();
836        encode_bulk_to_buf(&record, &mut encoded).unwrap();
837        let mut cursor = BytesMut::from(encoded.as_slice());
838
839        let Some(DecodedFrame::Bulk(decoded)) = try_decode_frame_from_bytes(&mut cursor).unwrap()
840        else {
841            panic!("expected bulk record");
842        };
843        assert_eq!(decoded.payload, payload);
844    }
845
846    #[test]
847    fn bulk_decoder_accepts_every_frame_fragment_boundary() {
848        let record = BulkRecord {
849            id: 34,
850            kind: BulkKind::Tcp,
851            flow: BulkFlow::HostToGuest,
852            offset: 99,
853            payload: Bytes::from(vec![0x73; 128]),
854        };
855        let mut encoded = Vec::new();
856        encode_bulk_to_buf(&record, &mut encoded).unwrap();
857
858        for split in 0..=encoded.len() {
859            let mut cursor = BytesMut::from(&encoded[..split]);
860            let first = try_decode_frame_from_bytes(&mut cursor).unwrap();
861            if split < encoded.len() {
862                assert!(first.is_none(), "decoded incomplete frame at split {split}");
863                cursor.extend_from_slice(&encoded[split..]);
864            }
865            let decoded = first
866                .or_else(|| try_decode_frame_from_bytes(&mut cursor).unwrap())
867                .unwrap();
868            assert!(matches!(decoded, DecodedFrame::Bulk(ref value) if value == &record));
869            assert!(cursor.is_empty());
870        }
871    }
872
873    #[test]
874    fn bulk_decoder_rejects_malformed_wire_shapes() {
875        let record = BulkRecord {
876            id: 5,
877            kind: BulkKind::Filesystem,
878            flow: BulkFlow::HostToGuest,
879            offset: 0,
880            payload: Bytes::from_static(b"x"),
881        };
882        let mut valid = Vec::new();
883        encode_bulk_to_buf(&record, &mut valid).unwrap();
884
885        for (index, value) in [(8, FLAG_BULK | FLAG_TERMINAL), (9, 0), (10, 3), (11, 1)] {
886            let mut malformed = valid.clone();
887            malformed[index] = value;
888            let error =
889                try_decode_frame_from_bytes(&mut BytesMut::from(malformed.as_slice())).unwrap_err();
890            assert!(matches!(error, ProtocolError::InvalidBulkFrame(_)));
891        }
892
893        let mut no_payload = valid;
894        no_payload.truncate(4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE);
895        no_payload[..4]
896            .copy_from_slice(&((FRAME_HEADER_SIZE + BULK_HEADER_SIZE) as u32).to_be_bytes());
897        let error =
898            try_decode_frame_from_bytes(&mut BytesMut::from(no_payload.as_slice())).unwrap_err();
899        assert!(matches!(error, ProtocolError::InvalidBulkFrame(_)));
900    }
901
902    #[tokio::test]
903    async fn bulk_vectored_writer_handles_one_byte_short_writes() {
904        let record = BulkRecord {
905            id: 91,
906            kind: BulkKind::Tcp,
907            flow: BulkFlow::GuestToHost,
908            offset: 42,
909            payload: Bytes::from(vec![0xCD; 257]),
910        };
911        let mut expected = Vec::new();
912        encode_bulk_to_buf(&record, &mut expected).unwrap();
913
914        let mut writer = OneByteWriter::default();
915        write_bulk_record(&mut writer, &record).await.unwrap();
916
917        assert_eq!(writer.bytes, expected);
918    }
919
920    #[test]
921    fn cursor_decoder_accepts_every_frame_fragment_boundary() {
922        let msg = Message::new(MessageType::Ready, 77, vec![0xAB; 1024]);
923        let mut encoded = Vec::new();
924        encode_to_buf(&msg, &mut encoded).unwrap();
925
926        for split in 0..=encoded.len() {
927            let mut buf = BytesMut::new();
928            buf.extend_from_slice(&encoded[..split]);
929            let first = try_decode_from_bytes(&mut buf).unwrap();
930            if split < encoded.len() {
931                assert!(first.is_none(), "decoded incomplete frame at split {split}");
932                buf.extend_from_slice(&encoded[split..]);
933            }
934
935            let decoded = first
936                .or_else(|| try_decode_from_bytes(&mut buf).unwrap())
937                .unwrap();
938            assert_eq!(decoded.id, msg.id);
939            assert_eq!(decoded.t, msg.t);
940            assert!(buf.is_empty());
941        }
942    }
943
944    #[tokio::test]
945    async fn vectored_writer_handles_one_byte_short_writes() {
946        let frame = RawFrame {
947            id: 91,
948            flags: FLAG_TERMINAL,
949            body: vec![0xCD; 257],
950        };
951        let mut expected = Vec::new();
952        encode_raw_to_buf(&frame, &mut expected).unwrap();
953
954        let mut writer = OneByteWriter::default();
955        write_raw_frame(&mut writer, &frame).await.unwrap();
956
957        assert_eq!(writer.bytes, expected);
958    }
959
960    #[test]
961    fn test_raw_frame_sync_roundtrip() {
962        let frame = RawFrame {
963            id: 42,
964            flags: FLAG_SESSION_START,
965            body: vec![0xAA; 100],
966        };
967
968        let mut buf = Vec::new();
969        encode_raw_to_buf(&frame, &mut buf).unwrap();
970
971        let decoded = try_decode_raw_from_buf(&mut buf).unwrap().unwrap();
972        assert_eq!(decoded.id, frame.id);
973        assert_eq!(decoded.flags, frame.flags);
974        assert_eq!(decoded.body, frame.body);
975        assert!(buf.is_empty());
976    }
977
978    #[test]
979    fn test_raw_frame_to_message() {
980        use crate::exec::ExecExited;
981
982        let msg =
983            Message::with_payload(MessageType::ExecExited, 13, &ExecExited { code: 7 }).unwrap();
984
985        let mut buf = Vec::new();
986        encode_to_buf(&msg, &mut buf).unwrap();
987
988        let frame = try_decode_raw_from_buf(&mut buf).unwrap().unwrap();
989        let decoded = raw_frame_to_message(frame).unwrap();
990
991        assert_eq!(decoded.id, 13);
992        assert_eq!(decoded.t, MessageType::ExecExited);
993        let payload: ExecExited = decoded.payload().unwrap();
994        assert_eq!(payload.code, 7);
995    }
996
997    #[derive(Default)]
998    struct OneByteWriter {
999        bytes: Vec<u8>,
1000    }
1001
1002    impl AsyncWrite for OneByteWriter {
1003        fn poll_write(
1004            mut self: Pin<&mut Self>,
1005            _cx: &mut Context<'_>,
1006            buf: &[u8],
1007        ) -> Poll<std::io::Result<usize>> {
1008            if let Some(byte) = buf.first() {
1009                self.bytes.push(*byte);
1010                Poll::Ready(Ok(1))
1011            } else {
1012                Poll::Ready(Ok(0))
1013            }
1014        }
1015
1016        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1017            Poll::Ready(Ok(()))
1018        }
1019
1020        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1021            Poll::Ready(Ok(()))
1022        }
1023
1024        fn is_write_vectored(&self) -> bool {
1025            true
1026        }
1027
1028        fn poll_write_vectored(
1029            mut self: Pin<&mut Self>,
1030            _cx: &mut Context<'_>,
1031            bufs: &[IoSlice<'_>],
1032        ) -> Poll<std::io::Result<usize>> {
1033            if let Some(byte) = bufs.iter().find_map(|buf| buf.first()) {
1034                self.bytes.push(*byte);
1035                Poll::Ready(Ok(1))
1036            } else {
1037                Poll::Ready(Ok(0))
1038            }
1039        }
1040    }
1041}
microsandbox_protocol/codec.rs

microsandbox_protocol/
codec.rs
