1use std::io::{self, Write};
2
3use eoka_protocol::sidecar::{
4 ErrorCode, ErrorData, Request, Response, ServerError, MAX_MESSAGE_BYTES,
5};
6use serde_json::Value;
7use tokio::io::{AsyncBufRead, AsyncBufReadExt};
8
9pub enum Frame {
10 Request(Request),
11 Error { id: Value, error: ServerError },
12 Empty,
13}
14
15pub struct FrameReader<R> {
16 input: R,
17 bytes: Vec<u8>,
18 oversized: bool,
19}
20
21impl<R: AsyncBufRead + Unpin> FrameReader<R> {
22 pub fn new(input: R) -> Self {
23 Self {
24 input,
25 bytes: Vec::new(),
26 oversized: false,
27 }
28 }
29
30 pub async fn next(&mut self) -> io::Result<Option<Frame>> {
31 loop {
32 let available = self.input.fill_buf().await?;
33 if available.is_empty() {
34 if self.oversized || !self.bytes.is_empty() {
35 return Ok(Some(self.finish()));
36 }
37 return Ok(None);
38 }
39 let newline = available.iter().position(|byte| *byte == b'\n');
40 let length = newline.unwrap_or(available.len());
41 if self.bytes.is_empty()
42 && !self.oversized
43 && newline.is_some()
44 && (length <= MAX_MESSAGE_BYTES
45 || (length == MAX_MESSAGE_BYTES + 1 && available[length - 1] == b'\r'))
46 {
47 let bytes = &available[..length];
48 let frame = parse(bytes.strip_suffix(b"\r").unwrap_or(bytes));
49 self.input.consume(length + 1);
50 return Ok(Some(frame));
51 }
52 if !self.oversized {
53 if length > (MAX_MESSAGE_BYTES + 1).saturating_sub(self.bytes.len()) {
54 self.bytes.clear();
55 self.oversized = true;
56 } else {
57 append_bounded(&mut self.bytes, &available[..length], MAX_MESSAGE_BYTES + 1);
58 }
59 }
60 self.input.consume(length + usize::from(newline.is_some()));
61 if newline.is_some() {
62 return Ok(Some(self.finish()));
63 }
64 }
65 }
66
67 fn finish(&mut self) -> Frame {
68 if self.bytes.last() == Some(&b'\r') {
69 self.bytes.pop();
70 }
71 if self.oversized || self.bytes.len() > MAX_MESSAGE_BYTES {
72 self.oversized = false;
73 self.bytes.clear();
74 return Frame::Error {
75 id: Value::Null,
76 error: ServerError::new(
77 ErrorCode::MessageTooLarge,
78 "request exceeds maxMessageBytes",
79 ),
80 };
81 }
82 let frame = parse(&self.bytes);
83 self.bytes.clear();
84 frame
85 }
86}
87
88fn append_bounded(buffer: &mut Vec<u8>, bytes: &[u8], limit: usize) {
89 if bytes.len().saturating_add(1) > buffer.capacity() - buffer.len() {
90 let capacity = buffer
91 .capacity()
92 .max(8192)
93 .saturating_mul(2)
94 .max(buffer.len() + bytes.len() + 1)
95 .min(limit);
96 buffer.reserve_exact(capacity - buffer.len());
97 }
98 buffer.extend_from_slice(bytes);
99}
100
101#[derive(serde::Deserialize)]
102struct RecoverId {
103 id: u64,
104}
105
106fn parse(bytes: &[u8]) -> Frame {
107 if bytes.iter().all(u8::is_ascii_whitespace) {
108 return Frame::Empty;
109 }
110 match serde_json::from_slice::<Request>(bytes) {
111 Ok(request) => Frame::Request(request),
112 Err(error) => {
113 let id = serde_json::from_slice::<RecoverId>(bytes)
114 .ok()
115 .map(|envelope| envelope.id)
116 .filter(|id| (1..=9_007_199_254_740_991).contains(id))
117 .map(Value::from)
118 .unwrap_or(Value::Null);
119 Frame::Error {
120 id,
121 error: ServerError::invalid_params(error.to_string()),
122 }
123 }
124 }
125}
126
127struct LimitedBuffer {
128 bytes: Vec<u8>,
129 exceeded: bool,
130}
131
132impl Write for LimitedBuffer {
133 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
134 if bytes.len() > MAX_MESSAGE_BYTES.saturating_sub(self.bytes.len()) {
135 self.exceeded = true;
136 return Err(io::Error::other("response exceeds maxMessageBytes"));
137 }
138 append_bounded(&mut self.bytes, bytes, MAX_MESSAGE_BYTES + 1);
139 Ok(bytes.len())
140 }
141
142 fn flush(&mut self) -> io::Result<()> {
143 Ok(())
144 }
145}
146
147pub fn encode(response: &Response, id: &Value, operation_id: Option<&str>) -> io::Result<Vec<u8>> {
148 let mut output = LimitedBuffer {
149 bytes: Vec::new(),
150 exceeded: false,
151 };
152 if let Err(error) = serde_json::to_writer(&mut output, response) {
153 if !output.exceeded {
154 return Err(io::Error::other(error));
155 }
156 output.bytes.clear();
157 output.exceeded = false;
158 let error = ServerError::new(
159 ErrorCode::MessageTooLarge,
160 "response exceeds maxMessageBytes; reduce the response size",
161 )
162 .with_data(ErrorData {
163 operation_id: operation_id.map(str::to_owned),
164 ..Default::default()
165 });
166 serde_json::to_writer(&mut output, &Response::err(id.clone(), error))
167 .map_err(io::Error::other)?;
168 }
169 output.bytes.push(b'\n');
170 Ok(output.bytes)
171}
172
173#[cfg(test)]
174mod tests {
175 use super::*;
176 use serde_json::json;
177 use tokio::io::BufReader;
178
179 #[tokio::test]
180 async fn rejects_oversized_input_and_recovers_next_frame() {
181 let mut input = vec![b'x'; MAX_MESSAGE_BYTES + 1];
182 input.extend_from_slice(b"\n{\"id\":7,\"method\":\"server.hello\"}\n");
183 let mut reader = FrameReader::new(BufReader::with_capacity(31, &input[..]));
184 assert!(matches!(
185 reader.next().await.unwrap(),
186 Some(Frame::Error {
187 error: ServerError {
188 code: ErrorCode::MessageTooLarge,
189 ..
190 },
191 ..
192 })
193 ));
194 assert!(
195 matches!(reader.next().await.unwrap(), Some(Frame::Request(Request { id, .. })) if id == json!(7))
196 );
197 assert!(reader.next().await.unwrap().is_none());
198 }
199
200 #[tokio::test]
201 async fn handles_split_utf8_and_partial_final_frame() {
202 let input = "{\"id\":1,\"method\":\"é\"}\n{\"id\":2,\"method\":\"server.shutdown\"}";
203 let mut reader = FrameReader::new(BufReader::with_capacity(1, input.as_bytes()));
204 assert!(
205 matches!(reader.next().await.unwrap(), Some(Frame::Request(Request { method, .. })) if method == "é")
206 );
207 assert!(
208 matches!(reader.next().await.unwrap(), Some(Frame::Request(Request { id, .. })) if id == json!(2))
209 );
210 assert!(reader.next().await.unwrap().is_none());
211 }
212
213 #[tokio::test]
214 async fn crlf_message_at_ceiling_is_accepted() {
215 let mut input = b"{\"id\":1,\"method\":\"server.hello\"}".to_vec();
216 input.resize(MAX_MESSAGE_BYTES, b' ');
217 input.extend_from_slice(b"\r\n");
218 for capacity in [8192, input.len()] {
219 let mut reader = FrameReader::new(BufReader::with_capacity(capacity, &input[..]));
220 assert!(
221 matches!(reader.next().await.unwrap(), Some(Frame::Request(Request { id, .. })) if id == json!(1))
222 );
223 assert!(reader.bytes.capacity() <= MAX_MESSAGE_BYTES + 1);
224 }
225 }
226
227 #[test]
228 fn output_at_ceiling_has_only_one_extra_terminator_byte() {
229 let overhead = serde_json::to_vec(&Response::ok(json!(1), json!("")))
230 .unwrap()
231 .len();
232 let response = Response::ok(
233 json!(1),
234 Value::String("x".repeat(MAX_MESSAGE_BYTES - overhead)),
235 );
236 let bytes = encode(&response, &json!(1), None).unwrap();
237 assert_eq!(bytes.len(), MAX_MESSAGE_BYTES + 1);
238 assert_eq!(bytes.last(), Some(&b'\n'));
239 assert!(bytes.capacity() <= MAX_MESSAGE_BYTES + 1);
240 }
241
242 #[test]
243 fn rejects_oversized_output_without_unbounded_serialized_buffer() {
244 let id = json!(9);
245 let response = Response::ok(
246 id.clone(),
247 json!({"result":"x".repeat(MAX_MESSAGE_BYTES + 1)}),
248 );
249 let bytes = encode(&response, &id, Some("operation")).unwrap();
250 let value: Value = serde_json::from_slice(&bytes).unwrap();
251 assert_eq!(value["id"], 9);
252 assert_eq!(value["error"]["code"], "MessageTooLarge");
253 assert_eq!(value["error"]["data"]["operation_id"], "operation");
254 assert!(bytes.len() < 1024);
255 }
256}