1use serde::{Deserialize, Serialize};
17
18use crate::wire::DEFAULT_MAX_FRAME_BYTES;
19
20#[derive(Debug, thiserror::Error)]
22pub enum DecodeError {
23 #[error("frame body {body} bytes exceeds limit {max} bytes")]
26 FrameTooLarge { body: usize, max: usize },
27 #[error("decode error: {0}")]
29 Rmp(#[from] rmp_serde::decode::Error),
30}
31
32pub 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
42pub 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
53pub 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#[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 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 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 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 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 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 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 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 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 #[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 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 #[test]
236 fn bytes_emit_as_bin_canonical() {
237 let encoded = rmp_serde::to_vec(&Value::Bytes(vec![1, 2, 3, 255])).unwrap();
238 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 let legacy: Vec<u8> = vec![
246 0x81, 0xa5, 0x42, 0x79, 0x74, 0x65, 0x73, 0x94, 0x01, 0x02, 0x03, 0xcc, 0xff, ];
249 let decoded: Value = rmp_serde::from_slice(&legacy).unwrap();
250 assert_eq!(decoded, Value::Bytes(vec![1, 2, 3, 255]));
251 }
252
253 #[test]
256 fn map_shaped_request_decodes() {
257 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 #[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 #[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 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]); 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 #[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}