1use serde_json::Value;
8
9use super::cbor::{self, CborError, CborOptions};
10use super::framing::{self, FrameDecoder, FrameError, DEFAULT_MAX_FRAME_LENGTH};
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum ProtocolError {
15 Frame(String),
16 Cbor(String),
17 Encode(String),
18}
19
20impl std::fmt::Display for ProtocolError {
21 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22 match self {
23 Self::Frame(m) => write!(f, "frame: {m}"),
24 Self::Cbor(m) => write!(f, "cbor: {m}"),
25 Self::Encode(m) => write!(f, "encode: {m}"),
26 }
27 }
28}
29
30impl std::error::Error for ProtocolError {}
31
32impl From<FrameError> for ProtocolError {
33 fn from(e: FrameError) -> Self {
34 Self::Frame(e.0)
35 }
36}
37
38impl From<CborError> for ProtocolError {
39 fn from(e: CborError) -> Self {
40 Self::Cbor(e.0)
41 }
42}
43
44pub fn encode_message(value: &Value) -> Result<Vec<u8>, ProtocolError> {
46 encode_message_with(value, &CborOptions::default(), DEFAULT_MAX_FRAME_LENGTH)
47}
48
49pub fn encode_message_with(
51 value: &Value,
52 cbor_options: &CborOptions,
53 max_frame_length: usize,
54) -> Result<Vec<u8>, ProtocolError> {
55 let payload = cbor::encode(value, cbor_options)?;
56 if payload.len() > max_frame_length {
57 return Err(ProtocolError::Encode(format!(
58 "encoded frame {} exceeds max frame length {max_frame_length}",
59 payload.len()
60 )));
61 }
62 Ok(framing::encode_frame(&payload)?)
63}
64
65pub struct MessageDecoder {
67 frames: FrameDecoder,
68 cbor_options: CborOptions,
69}
70
71impl MessageDecoder {
72 pub fn new(max_frame_length: usize, cbor_options: CborOptions) -> Self {
73 Self {
74 frames: FrameDecoder::new(max_frame_length),
75 cbor_options,
76 }
77 }
78
79 pub fn with_defaults() -> Self {
80 Self::new(DEFAULT_MAX_FRAME_LENGTH, CborOptions::default())
81 }
82
83 pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<Value>, ProtocolError> {
85 let payloads = self.frames.push(chunk)?;
86 let mut messages = Vec::with_capacity(payloads.len());
87 for payload in payloads {
88 messages.push(cbor::decode(&payload, &self.cbor_options)?);
89 }
90 Ok(messages)
91 }
92
93 pub fn end(&mut self) -> Result<(), ProtocolError> {
95 self.frames.end()?;
96 Ok(())
97 }
98}
99
100pub fn decode_message(
102 frame: &[u8],
103 cbor_options: &CborOptions,
104 max_frame_length: usize,
105) -> Result<Value, ProtocolError> {
106 framing::assert_complete_frame(frame, max_frame_length)?;
107 let payload = &frame[framing::FRAME_HEADER_LENGTH..];
108 Ok(cbor::decode(payload, cbor_options)?)
109}
110
111#[cfg(test)]
112mod tests {
113 use super::*;
114 use serde_json::json;
115
116 #[test]
117 fn round_trip_message() {
118 let message = json!({ "type": "hello", "version": 1, "tools": ["read", "write"] });
119 let frame = encode_message(&message).unwrap();
120 let decoded =
121 decode_message(&frame, &CborOptions::default(), DEFAULT_MAX_FRAME_LENGTH).unwrap();
122 assert_eq!(decoded, message);
123 }
124
125 #[test]
126 fn decoder_stream() {
127 let a = encode_message(&json!({ "type": "a" })).unwrap();
128 let b = encode_message(&json!({ "type": "b" })).unwrap();
129 let mut decoder = MessageDecoder::with_defaults();
130 assert!(decoder.push(&a[..3]).unwrap().is_empty());
131 let mut all = Vec::new();
132 all.extend(decoder.push(&a[3..]).unwrap());
133 all.extend(decoder.push(&b).unwrap());
134 assert_eq!(all, vec![json!({ "type": "a" }), json!({ "type": "b" })]);
135 decoder.end().unwrap();
136 }
137
138 #[test]
139 fn rejects_oversized_encoding() {
140 let big = json!({ "type": "x", "data": "a".repeat(64) });
141 assert!(encode_message_with(&big, &CborOptions::default(), 8).is_err());
142 }
143}