Skip to main content

thunder/wire/
frame.rs

1//! Length-prefixed MessagePack frame codec.
2//!
3//! ```text
4//! ┌───────────────────┬──────────────────────────┐
5//! │  length: u32 (LE) │  body: MessagePack bytes │
6//! └───────────────────┴──────────────────────────┘
7//!     4 bytes              length bytes
8//! ```
9//!
10//! The cap is validated against the length prefix **before** the body
11//! buffer is allocated (WIRE-020/021), so a hostile prefix cannot exhaust
12//! memory. Decode reports how many bytes one frame consumed, which is also
13//! the frame size consumers feed to metrics — nothing ever re-encodes a
14//! frame just to measure it (SRV-007).
15
16use serde::{Deserialize, Serialize};
17
18use crate::wire::DEFAULT_MAX_FRAME_BYTES;
19
20/// Errors from the sync decoder.
21#[derive(Debug, thiserror::Error)]
22pub enum DecodeError {
23    /// The length prefix declared a body larger than the caller's cap.
24    /// Raised before any body allocation (WIRE-021).
25    #[error("frame body {body} bytes exceeds limit {max} bytes")]
26    FrameTooLarge { body: usize, max: usize },
27    /// Well-formed frame, malformed MessagePack payload (WIRE-023).
28    #[error("decode error: {0}")]
29    Rmp(#[from] rmp_serde::decode::Error),
30}
31
32/// Encode a message into one complete frame (`u32 LE length` + body).
33pub fn encode_frame<T: Serialize>(msg: &T) -> Result<Vec<u8>, rmp_serde::encode::Error> {
34    let body = rmp_serde::to_vec(msg)?;
35    let len = body.len() as u32;
36    let mut frame = Vec::with_capacity(4 + body.len());
37    frame.extend_from_slice(&len.to_le_bytes());
38    frame.extend_from_slice(&body);
39    Ok(frame)
40}
41
42/// Decode one frame from a byte slice using [`DEFAULT_MAX_FRAME_BYTES`].
43///
44/// Returns `Ok(None)` when the buffer does not yet hold a complete frame
45/// (read more and retry — WIRE-022). On success returns the value and the
46/// total bytes consumed (`4 + body`), which is the frame size for metrics.
47pub fn decode_frame<T: for<'de> Deserialize<'de>>(
48    buf: &[u8],
49) -> Result<Option<(T, usize)>, DecodeError> {
50    decode_frame_with_limit(buf, DEFAULT_MAX_FRAME_BYTES)
51}
52
53/// Decode one frame, rejecting bodies larger than `max` before the body
54/// is even inspected (WIRE-020/021).
55pub fn decode_frame_with_limit<T: for<'de> Deserialize<'de>>(
56    buf: &[u8],
57    max: usize,
58) -> Result<Option<(T, usize)>, DecodeError> {
59    if buf.len() < 4 {
60        return Ok(None);
61    }
62    let len = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]) as usize;
63    if len > max {
64        return Err(DecodeError::FrameTooLarge { body: len, max });
65    }
66    let total = 4 + len;
67    if buf.len() < total {
68        return Ok(None);
69    }
70    let value = rmp_serde::from_slice(&buf[4..total])?;
71    Ok(Some((value, total)))
72}
73
74// ── Async helpers (feature = "tokio") ───────────────────────────────────────
75
76#[cfg(feature = "tokio")]
77mod tokio_io {
78    use std::io;
79
80    use serde::{Deserialize, Serialize};
81    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
82
83    use crate::wire::value::{Request, Response};
84    use crate::wire::DEFAULT_MAX_FRAME_BYTES;
85
86    /// Read one frame; returns the decoded value and the frame size in
87    /// bytes (`4 + body` — the metrics input, SRV-007). The cap is checked
88    /// between reading the prefix and allocating the body (WIRE-020).
89    pub async fn read_frame<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
90        reader: &mut R,
91        max: usize,
92    ) -> io::Result<(T, usize)> {
93        let mut len_buf = [0u8; 4];
94        reader.read_exact(&mut len_buf).await?;
95        let len = u32::from_le_bytes(len_buf) as usize;
96        if len > max {
97            return Err(io::Error::new(
98                io::ErrorKind::InvalidData,
99                format!("frame body {len} bytes exceeds limit {max} bytes"),
100            ));
101        }
102        let mut body = vec![0u8; len];
103        reader.read_exact(&mut body).await?;
104        let value = rmp_serde::from_slice(&body)
105            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
106        Ok((value, 4 + len))
107    }
108
109    /// Read one [`Request`] with the default cap.
110    pub async fn read_request<R: AsyncRead + Unpin>(
111        reader: &mut R,
112    ) -> io::Result<(Request, usize)> {
113        read_frame(reader, DEFAULT_MAX_FRAME_BYTES).await
114    }
115
116    /// Read one [`Request`] with a caller-supplied cap (server hot path).
117    pub async fn read_request_with_limit<R: AsyncRead + Unpin>(
118        reader: &mut R,
119        max: usize,
120    ) -> io::Result<(Request, usize)> {
121        read_frame(reader, max).await
122    }
123
124    /// Read one [`Response`] with the default cap.
125    pub async fn read_response<R: AsyncRead + Unpin>(
126        reader: &mut R,
127    ) -> io::Result<(Response, usize)> {
128        read_frame(reader, DEFAULT_MAX_FRAME_BYTES).await
129    }
130
131    /// Read one [`Response`] with a caller-supplied cap.
132    pub async fn read_response_with_limit<R: AsyncRead + Unpin>(
133        reader: &mut R,
134        max: usize,
135    ) -> io::Result<(Response, usize)> {
136        read_frame(reader, max).await
137    }
138
139    /// Encode and write one frame; returns the frame size written.
140    pub async fn write_frame<T: Serialize, W: AsyncWrite + Unpin>(
141        writer: &mut W,
142        msg: &T,
143    ) -> io::Result<usize> {
144        let frame = crate::wire::frame::encode_frame(msg)
145            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
146        writer.write_all(&frame).await?;
147        Ok(frame.len())
148    }
149
150    /// Write one [`Request`] frame.
151    pub async fn write_request<W: AsyncWrite + Unpin>(
152        writer: &mut W,
153        req: &Request,
154    ) -> io::Result<usize> {
155        write_frame(writer, req).await
156    }
157
158    /// Write one [`Response`] frame.
159    pub async fn write_response<W: AsyncWrite + Unpin>(
160        writer: &mut W,
161        resp: &Response,
162    ) -> io::Result<usize> {
163        write_frame(writer, resp).await
164    }
165}
166
167#[cfg(feature = "tokio")]
168pub use tokio_io::{
169    read_frame, read_request, read_request_with_limit, read_response, read_response_with_limit,
170    write_frame, write_request, write_response,
171};
172
173#[cfg(test)]
174#[allow(clippy::unwrap_used, clippy::expect_used)]
175mod tests {
176    use super::*;
177    use crate::wire::value::{Request, Response, Value};
178    use crate::wire::PUSH_ID;
179
180    fn hex(bytes: &[u8]) -> String {
181        bytes
182            .iter()
183            .map(|b| format!("{b:02x}"))
184            .collect::<Vec<_>>()
185            .join(" ")
186    }
187
188    // ── Golden vectors (family-pinned bytes, corpus canonical group) ───────
189
190    #[test]
191    fn ping_request_matches_family_golden_vector() {
192        let req = Request {
193            id: 1,
194            command: "PING".to_owned(),
195            args: vec![],
196        };
197        let frame = encode_frame(&req).unwrap();
198        assert_eq!(
199            hex(&frame),
200            "08 00 00 00 93 01 a4 50 49 4e 47 90",
201            "frame must match the corpus request-ping vector"
202        );
203        let (decoded, consumed): (Request, usize) = decode_frame(&frame).unwrap().unwrap();
204        assert_eq!(decoded, req);
205        assert_eq!(consumed, frame.len());
206    }
207
208    #[test]
209    fn pong_response_matches_nested_ok_golden_vector() {
210        let resp = Response::ok(1, Value::Str("PONG".to_owned()));
211        let frame = encode_frame(&resp).unwrap();
212        // Result<Value, String> nests two one-key maps: {"Ok": {"Str": "PONG"}}.
213        assert_eq!(
214            hex(&frame),
215            "10 00 00 00 92 01 81 a2 4f 6b 81 a3 53 74 72 a4 50 4f 4e 47"
216        );
217        let (decoded, _): (Response, usize) = decode_frame(&frame).unwrap().unwrap();
218        assert_eq!(decoded, resp);
219    }
220
221    #[test]
222    fn null_is_bare_string_and_int_is_single_key_map() {
223        assert_eq!(
224            hex(&rmp_serde::to_vec(&Value::Null).unwrap()),
225            "a4 4e 75 6c 6c"
226        );
227        assert_eq!(
228            hex(&rmp_serde::to_vec(&Value::Int(42)).unwrap()),
229            "81 a3 49 6e 74 2a"
230        );
231    }
232
233    // ── Bytes canonicalization (WIRE-010/011, probe T-029) ────────────────
234
235    #[test]
236    fn bytes_emit_as_bin_canonical() {
237        let encoded = rmp_serde::to_vec(&Value::Bytes(vec![1, 2, 3, 255])).unwrap();
238        // {"Bytes": bin8(4)} — c4 04, never the int-array form (94 ... cc ff).
239        assert_eq!(hex(&encoded), "81 a5 42 79 74 65 73 c4 04 01 02 03 ff");
240    }
241
242    #[test]
243    fn bytes_decode_legacy_int_array_form() {
244        // The seq-of-u8 form every pre-Thunder Rust implementation emits.
245        let legacy: Vec<u8> = vec![
246            0x81, 0xa5, 0x42, 0x79, 0x74, 0x65, 0x73, // {"Bytes":
247            0x94, 0x01, 0x02, 0x03, 0xcc, 0xff, // [1, 2, 3, 255] as ints
248        ];
249        let decoded: Value = rmp_serde::from_slice(&legacy).unwrap();
250        assert_eq!(decoded, Value::Bytes(vec![1, 2, 3, 255]));
251    }
252
253    // ── Request shape tolerance (WIRE-012/013) ────────────────────────────
254
255    #[test]
256    fn map_shaped_request_decodes() {
257        // Some pre-v1 encoders emit Request as a named map.
258        let req = Request {
259            id: 7,
260            command: "GET".to_owned(),
261            args: vec![Value::Str("key".to_owned())],
262        };
263        let map_shaped = rmp_serde::to_vec_named(&req).unwrap();
264        let array_shaped = rmp_serde::to_vec(&req).unwrap();
265        assert_ne!(map_shaped, array_shaped, "shapes must actually differ");
266        let decoded: Request = rmp_serde::from_slice(&map_shaped).unwrap();
267        assert_eq!(decoded, req);
268    }
269
270    // ── Round-trip matrix (donor test suite, WIRE-002/014/015) ────────────
271
272    #[test]
273    fn round_trip_all_variants() {
274        let all = Value::Array(vec![
275            Value::Null,
276            Value::Bool(true),
277            Value::Bool(false),
278            Value::Int(0),
279            Value::Int(i64::MIN),
280            Value::Int(i64::MAX),
281            Value::Int(-32),
282            Value::Int(127),
283            Value::Int(255),
284            Value::Int(65535),
285            Value::Float(0.0),
286            Value::Float(-0.0),
287            Value::Float(f64::INFINITY),
288            Value::Float(f64::NEG_INFINITY),
289            Value::Bytes(vec![]),
290            Value::Bytes(vec![0, 1, 2, 255]),
291            Value::Str(String::new()),
292            Value::Str("héllo wörld".to_owned()),
293            Value::Array(vec![]),
294            Value::Map(vec![]),
295            Value::Map(vec![
296                (Value::Str("k".to_owned()), Value::Int(1)),
297                (Value::Int(2), Value::Str("non-string key".to_owned())),
298            ]),
299        ]);
300        let frame = encode_frame(&all).unwrap();
301        let (decoded, consumed): (Value, usize) = decode_frame(&frame).unwrap().unwrap();
302        assert_eq!(decoded, all);
303        assert_eq!(consumed, frame.len());
304    }
305
306    #[test]
307    fn nan_bit_pattern_survives() {
308        let frame = encode_frame(&Value::Float(f64::NAN)).unwrap();
309        let (decoded, _): (Value, usize) = decode_frame(&frame).unwrap().unwrap();
310        match decoded {
311            Value::Float(f) => assert_eq!(f.to_bits(), f64::NAN.to_bits()),
312            other => panic!("expected Float, got {other:?}"),
313        }
314    }
315
316    #[test]
317    fn error_response_round_trips_with_prefix_conventions() {
318        for msg in [
319            "ERR unknown command",
320            "NOAUTH Authentication required.",
321            "WRONGPASS invalid username-password pair or user is disabled.",
322            "[collection_not_found] no such collection: docs",
323        ] {
324            let resp = Response::err(9, msg);
325            let frame = encode_frame(&resp).unwrap();
326            let (decoded, _): (Response, usize) = decode_frame(&frame).unwrap().unwrap();
327            assert_eq!(decoded.result, Err(msg.to_owned()));
328        }
329    }
330
331    // ── Framing edges (WIRE-020..023) ─────────────────────────────────────
332
333    #[test]
334    fn partial_header_and_partial_body_return_none() {
335        let frame = encode_frame(&Request {
336            id: 1,
337            command: "PING".to_owned(),
338            args: vec![],
339        })
340        .unwrap();
341        for cut in [0, 1, 3, 4, frame.len() - 1] {
342            let out: Option<(Request, usize)> = decode_frame(&frame[..cut]).unwrap();
343            assert!(out.is_none(), "cut at {cut} must ask for more bytes");
344        }
345    }
346
347    #[test]
348    fn two_frames_in_one_buffer_consume_exactly_one_each() {
349        let a = encode_frame(&Response::ok(1, Value::Int(1))).unwrap();
350        let b = encode_frame(&Response::ok(2, Value::Int(2))).unwrap();
351        let mut buf = a.clone();
352        buf.extend_from_slice(&b);
353        let (first, used): (Response, usize) = decode_frame(&buf).unwrap().unwrap();
354        assert_eq!(first.id, 1);
355        assert_eq!(used, a.len());
356        let (second, used2): (Response, usize) = decode_frame(&buf[used..]).unwrap().unwrap();
357        assert_eq!(second.id, 2);
358        assert_eq!(used2, b.len());
359    }
360
361    #[test]
362    fn oversized_prefix_rejected_before_body_arrives() {
363        // Only the 4-byte prefix claiming cap+1: the check fires without the
364        // body being present at all — allocation cannot have happened.
365        let over = (DEFAULT_MAX_FRAME_BYTES + 1) as u32;
366        let buf = over.to_le_bytes();
367        let err = decode_frame::<Request>(&buf).unwrap_err();
368        match err {
369            DecodeError::FrameTooLarge { body, max } => {
370                assert_eq!(body, DEFAULT_MAX_FRAME_BYTES + 1);
371                assert_eq!(max, DEFAULT_MAX_FRAME_BYTES);
372            }
373            other => panic!("expected FrameTooLarge, got {other:?}"),
374        }
375    }
376
377    #[test]
378    fn custom_limit_is_honored() {
379        let frame = encode_frame(&Value::Str("x".repeat(100))).unwrap();
380        let err = decode_frame_with_limit::<Value>(&frame, 8).unwrap_err();
381        assert!(matches!(err, DecodeError::FrameTooLarge { .. }));
382    }
383
384    #[test]
385    fn garbage_body_is_a_typed_error_not_a_panic() {
386        let mut buf = 4u32.to_le_bytes().to_vec();
387        buf.extend_from_slice(&[0xc1, 0xc1, 0xc1, 0xc1]); // 0xc1 is never valid
388        let err = decode_frame::<Request>(&buf).unwrap_err();
389        assert!(matches!(err, DecodeError::Rmp(_)));
390    }
391
392    #[test]
393    fn zero_length_body_is_a_decode_error() {
394        let buf = 0u32.to_le_bytes();
395        let err = decode_frame::<Request>(&buf).unwrap_err();
396        assert!(matches!(err, DecodeError::Rmp(_)));
397    }
398
399    #[test]
400    fn push_id_is_reserved_u32_max() {
401        assert_eq!(PUSH_ID, u32::MAX);
402    }
403
404    // ── Async path (feature = "tokio") ────────────────────────────────────
405
406    #[cfg(feature = "tokio")]
407    #[tokio::test]
408    async fn async_write_then_read_reports_frame_size() {
409        let req = Request {
410            id: 3,
411            command: "PING".to_owned(),
412            args: vec![],
413        };
414        let mut buf = Vec::new();
415        let written = write_request(&mut buf, &req).await.unwrap();
416        assert_eq!(written, buf.len());
417        let mut cursor = std::io::Cursor::new(buf);
418        let (decoded, size) = read_request(&mut cursor).await.unwrap();
419        assert_eq!(decoded, req);
420        assert_eq!(size, written);
421    }
422
423    #[cfg(feature = "tokio")]
424    #[tokio::test]
425    async fn async_read_rejects_oversized_prefix_without_reading_body() {
426        let over = ((DEFAULT_MAX_FRAME_BYTES + 1) as u32).to_le_bytes();
427        let mut cursor = std::io::Cursor::new(over.to_vec());
428        let err = read_request(&mut cursor).await.unwrap_err();
429        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
430    }
431}