shep_core/protocol/
wire.rs1use core::fmt;
7
8use bytes::Bytes;
9use serde::Serialize;
10use serde::de::DeserializeOwned;
11use tokio_util::codec::LengthDelimitedCodec;
12
13pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
15
16#[must_use]
18pub fn codec() -> LengthDelimitedCodec {
19 LengthDelimitedCodec::builder()
20 .length_field_type::<u32>()
21 .max_frame_length(MAX_FRAME_BYTES)
22 .new_codec()
23}
24
25pub fn encode_frame<T: Serialize>(value: &T) -> Result<Bytes, WireError> {
32 let vec = serde_json::to_vec(value).map_err(|e| WireError::Json(e.to_string()))?;
33 if vec.len() > MAX_FRAME_BYTES {
34 return Err(WireError::FrameTooLarge(vec.len()));
35 }
36 Ok(Bytes::from(vec)) }
38
39pub fn decode_frame<T: DeserializeOwned>(frame: &[u8]) -> Result<T, WireError> {
45 serde_json::from_slice(frame).map_err(|e| WireError::Json(e.to_string()))
46}
47
48#[non_exhaustive]
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum WireError {
58 Json(String),
60 FrameTooLarge(usize),
62}
63
64impl fmt::Display for WireError {
65 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66 match self {
67 Self::Json(m) => write!(f, "wire frame JSON error: {m}"),
68 Self::FrameTooLarge(n) => {
69 write!(
70 f,
71 "frame of {n} bytes exceeds the {MAX_FRAME_BYTES}-byte limit"
72 )
73 }
74 }
75 }
76}
77
78impl core::error::Error for WireError {}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83 use crate::protocol::{Envelope, Request};
84
85 #[test]
86 fn encode_decode_round_trip() {
87 let env = Envelope {
88 id: 9,
89 deadline_ms: Some(5000),
90 body: Request::Ping,
91 };
92 let bytes = encode_frame(&env).unwrap();
93 let back: Envelope = decode_frame(&bytes).unwrap();
94 assert_eq!(back, env);
95 }
96
97 #[test]
98 fn decode_rejects_garbage_with_json_error() {
99 assert!(matches!(
100 decode_frame::<Envelope>(b"not json"),
101 Err(WireError::Json(_))
102 ));
103 }
104
105 #[test]
106 fn codec_uses_u32_prefix_and_max_frame() {
107 let c = codec();
108 assert_eq!(c.max_frame_length(), MAX_FRAME_BYTES);
111 }
112
113 #[tokio::test]
114 async fn framed_stream_round_trip() {
115 use futures_util::{SinkExt, StreamExt};
116 use tokio_util::codec::{FramedRead, FramedWrite};
117
118 let (client, server) = tokio::io::duplex(64 * 1024);
119 let mut writer = FramedWrite::new(client, codec());
120 let mut reader = FramedRead::new(server, codec());
121
122 let env = Envelope {
123 id: 1,
124 deadline_ms: None,
125 body: Request::ListFlock,
126 };
127 writer.send(encode_frame(&env).unwrap()).await.unwrap();
128
129 let frame = reader.next().await.unwrap().unwrap();
130 let back: Envelope = decode_frame(&frame).unwrap();
131 assert_eq!(back, env);
132 }
133}