1use serde_json::Value;
11use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
12
13pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
15
16#[derive(Debug)]
22pub enum Inbound {
23 Frame(Value),
24 Violation(Violation),
25}
26
27#[derive(Debug, PartialEq, Eq)]
30#[non_exhaustive]
31pub enum Violation {
32 TooLarge,
33 InvalidJson,
34}
35
36#[derive(Debug)]
43pub struct FrameReader<R> {
44 reader: BufReader<R>,
45 max_frame: usize,
46 buf: Vec<u8>,
47 discarding: bool,
49}
50
51impl<R: AsyncRead + Unpin> FrameReader<R> {
52 pub fn new(reader: R, max_frame: usize) -> Self {
53 debug_assert!(max_frame > 0);
54 Self {
55 reader: BufReader::new(reader),
56 max_frame,
57 buf: Vec::new(),
58 discarding: false,
59 }
60 }
61
62 pub fn into_inner(self) -> BufReader<R> {
69 self.reader
70 }
71
72 pub async fn next(&mut self) -> std::io::Result<Option<Inbound>> {
78 let mut byte = [0u8; 1];
79 loop {
80 let n = self.reader.read(&mut byte).await?;
81 if n == 0 {
82 if std::mem::take(&mut self.discarding) {
83 return Ok(Some(Inbound::Violation(Violation::TooLarge)));
84 }
85 return Ok(None);
86 }
87 if byte[0] == b'\n' {
88 if std::mem::take(&mut self.discarding) {
89 return Ok(Some(Inbound::Violation(Violation::TooLarge)));
90 }
91 let line = std::mem::take(&mut self.buf);
94 return Ok(Some(match serde_json::from_slice::<Value>(&line) {
96 Ok(v) => Inbound::Frame(v),
97 Err(_) => Inbound::Violation(Violation::InvalidJson),
98 }));
99 }
100 if self.discarding {
101 continue;
102 }
103 if self.buf.len() >= self.max_frame {
104 self.discarding = true;
105 self.buf.clear();
106 continue;
107 }
108 self.buf.push(byte[0]);
109 }
110 }
111}
112
113pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, frame: &Value) -> std::io::Result<()> {
116 let mut line = serde_json::to_vec(frame)?;
117 line.push(b'\n');
118 w.write_all(&line).await?;
119 w.flush().await
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125 use serde_json::json;
126
127 fn reader_over(bytes: &[u8], cap: usize) -> FrameReader<&[u8]> {
128 FrameReader::new(bytes, cap)
129 }
130
131 #[tokio::test]
132 async fn valid_frames_roundtrip() {
133 let mut buf = Vec::new();
134 write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":1,"method":"ping"}))
135 .await
136 .unwrap();
137 write_frame(&mut buf, &json!({"jsonrpc":"2.0","id":2,"method":"pong"}))
138 .await
139 .unwrap();
140 let mut r = reader_over(&buf, 1024);
141 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
142 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
143 assert!(r.next().await.unwrap().is_none()); }
145
146 #[tokio::test]
147 async fn oversized_frame_is_discarded_and_reported_and_stream_continues() {
148 let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
149 let mut bytes = big.into_bytes();
150 bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"ok\"}\n");
151 let mut r = reader_over(&bytes, 64);
152 assert!(matches!(
153 r.next().await.unwrap().unwrap(),
154 Inbound::Violation(Violation::TooLarge)
155 ));
156 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 7));
157 }
158
159 #[tokio::test]
160 async fn invalid_json_is_a_violation() {
161 let mut r = reader_over(b"not json at all\n", 1024);
162 assert!(matches!(
163 r.next().await.unwrap().unwrap(),
164 Inbound::Violation(Violation::InvalidJson)
165 ));
166 }
167
168 #[tokio::test]
169 async fn oversized_frame_truncated_by_eof_is_still_reported() {
170 let bytes = format!("{{\"pad\":\"{}\"}}", "x".repeat(100)).into_bytes(); let mut r = reader_over(&bytes, 64);
172 assert!(matches!(
173 r.next().await.unwrap().unwrap(),
174 Inbound::Violation(Violation::TooLarge)
175 ));
176 assert!(r.next().await.unwrap().is_none());
177 }
178
179 #[tokio::test]
180 async fn back_to_back_oversized_frames_then_valid_frame() {
181 let big = format!("{{\"pad\":\"{}\"}}\n", "x".repeat(100));
182 let mut bytes = big.clone().into_bytes();
183 bytes.extend_from_slice(big.as_bytes());
184 bytes.extend_from_slice(b"{\"jsonrpc\":\"2.0\",\"id\":9,\"method\":\"ok\"}\n");
185 let mut r = reader_over(&bytes, 64);
186 assert!(matches!(
187 r.next().await.unwrap().unwrap(),
188 Inbound::Violation(Violation::TooLarge)
189 ));
190 assert!(matches!(
191 r.next().await.unwrap().unwrap(),
192 Inbound::Violation(Violation::TooLarge)
193 ));
194 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 9));
195 }
196
197 #[tokio::test]
198 async fn frame_of_exactly_max_frame_bytes_passes_one_more_is_too_large() {
199 let payload = "x".repeat(62);
202 let exact = format!("\"{payload}\"\n");
203 let mut r = reader_over(exact.as_bytes(), 64);
204 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v == payload));
205
206 let over = format!("\"{}\"\n", "x".repeat(63)); let mut r = reader_over(over.as_bytes(), 64);
208 assert!(matches!(
209 r.next().await.unwrap().unwrap(),
210 Inbound::Violation(Violation::TooLarge)
211 ));
212 }
213
214 #[tokio::test]
219 async fn into_inner_preserves_pipelined_read_ahead() {
220 let bytes = b"{\"id\":1}\n{\"id\":2}\n";
221 let mut r = reader_over(bytes, 1024);
222 assert!(matches!(r.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 1));
225 let boxed: Box<dyn AsyncRead + Unpin + Send> = Box::new(r.into_inner());
227 let mut r2 = FrameReader::new(boxed, 1024);
228 assert!(matches!(r2.next().await.unwrap().unwrap(), Inbound::Frame(v) if v["id"] == 2));
229 assert!(r2.next().await.unwrap().is_none());
230 }
231}