1use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
31
32pub const MAX_BODY: usize = 128 << 20;
38
39pub const MAX_HEADER_LINE: u64 = 8 << 10;
42
43#[derive(Debug, PartialEq, Eq)]
45pub enum Frame {
46 Forward,
48 Repaired(Vec<u8>),
50 Skip(String),
52}
53
54pub fn classify(body: &[u8]) -> Frame {
56 let error = match serde_json::from_slice::<serde_json::Value>(body) {
57 Ok(_) => return Frame::Forward,
58 Err(e) => e,
59 };
60 if let Ok(text) = std::str::from_utf8(body)
61 && let Some(repaired) = repair_lone_surrogates(text)
62 && serde_json::from_str::<serde_json::Value>(&repaired).is_ok()
63 {
64 return Frame::Repaired(repaired.into_bytes());
65 }
66 Frame::Skip(error.to_string())
67}
68
69fn repair_lone_surrogates(text: &str) -> Option<String> {
76 let bytes = text.as_bytes();
77 let mut out = String::with_capacity(text.len());
78 let mut changed = false;
79 let mut in_string = false;
80 let mut i = 0;
81 while i < bytes.len() {
82 let b = bytes[i];
83 if !in_string {
84 if b == b'"' {
85 in_string = true;
86 }
87 out.push(b as char);
88 i += 1;
89 continue;
90 }
91 match b {
92 b'"' => {
93 in_string = false;
94 out.push('"');
95 i += 1;
96 }
97 b'\\' if bytes.get(i + 1) == Some(&b'u') => {
98 let unit = hex4(bytes, i + 2);
99 match unit {
100 Some(0xD800..=0xDBFF) => {
101 let low = (bytes.get(i + 6) == Some(&b'\\')
103 && bytes.get(i + 7) == Some(&b'u'))
104 .then(|| hex4(bytes, i + 8))
105 .flatten();
106 if matches!(low, Some(0xDC00..=0xDFFF)) {
107 out.push_str(&text[i..i + 12]);
108 i += 12;
109 } else {
110 out.push_str("\\uFFFD");
111 changed = true;
112 i += 6;
113 }
114 }
115 Some(0xDC00..=0xDFFF) => {
116 out.push_str("\\uFFFD");
117 changed = true;
118 i += 6;
119 }
120 _ => {
121 out.push_str("\\u");
122 i += 2;
123 }
124 }
125 }
126 b'\\' => {
127 out.push('\\');
130 if let Some(next) = text[i + 1..].chars().next() {
131 out.push(next);
132 i += 1 + next.len_utf8();
133 } else {
134 i += 1;
135 }
136 }
137 _ => {
138 let ch = text[i..].chars().next().expect("in bounds");
139 out.push(ch);
140 i += ch.len_utf8();
141 }
142 }
143 }
144 changed.then_some(out)
145}
146
147const MAX_EXIT_BODY: usize = 1024;
151
152pub fn is_exit(body: &[u8]) -> bool {
155 body.len() <= MAX_EXIT_BODY
156 && serde_json::from_slice::<serde_json::Value>(body)
157 .is_ok_and(|message| message.get("method").and_then(|m| m.as_str()) == Some("exit"))
158}
159
160fn hex4(bytes: &[u8], at: usize) -> Option<u32> {
162 let digits = std::str::from_utf8(bytes.get(at..at + 4)?).ok()?;
163 u32::from_str_radix(digits, 16).ok()
164}
165
166pub async fn pump<R, W>(input: R, mut output: W)
170where
171 R: tokio::io::AsyncRead + Unpin,
172 W: AsyncWrite + Unpin,
173{
174 let mut input = BufReader::new(input);
175 loop {
176 let mut length: Option<usize> = None;
178 let mut saw_header = false;
179 loop {
180 let mut line = String::new();
181 let read = (&mut input)
182 .take(MAX_HEADER_LINE)
183 .read_line(&mut line)
184 .await;
185 if read.as_ref().is_ok_and(|n| *n as u64 == MAX_HEADER_LINE) && !line.ends_with('\n') {
186 tracing::error!(
187 "a client header line is over {MAX_HEADER_LINE} bytes; \
188 the input can't be resynchronised, so it is closed"
189 );
190 return;
191 }
192 match read {
193 Ok(0) => {
194 if saw_header {
195 tracing::error!("client input ended inside a message's headers");
196 }
197 return;
198 }
199 Ok(_) => {}
200 Err(e) => {
201 tracing::error!("reading client input failed: {e}");
202 return;
203 }
204 }
205 let line = line.trim_end_matches(['\r', '\n']);
206 if line.is_empty() {
207 if saw_header {
208 break;
209 }
210 continue;
211 }
212 saw_header = true;
213 if let Some((name, value)) = line.split_once(':')
214 && name.trim().eq_ignore_ascii_case("content-length")
215 {
216 length = value.trim().parse().ok();
217 }
218 }
219 let Some(length) = length else {
220 tracing::error!(
221 "a client message had no usable Content-Length header; \
222 the input can't be resynchronised, so it is closed"
223 );
224 return;
225 };
226 if length > MAX_BODY {
227 tracing::error!(
228 "a client message claims a {length}-byte body, over the {MAX_BODY}-byte limit; \
229 the input can't be resynchronised, so it is closed"
230 );
231 return;
232 }
233 let mut body = vec![0u8; length];
234 if let Err(e) = input.read_exact(&mut body).await {
235 tracing::error!(
236 "client input ended inside a message body ({length} bytes expected): {e}"
237 );
238 return;
239 }
240 let body = match classify(&body) {
241 Frame::Forward => body,
242 Frame::Repaired(repaired) => {
243 tracing::warn!(
244 "replaced lone UTF-16 surrogate escapes in a client message with U+FFFD"
245 );
246 repaired
247 }
248 Frame::Skip(reason) => {
249 let shown = String::from_utf8_lossy(&body[..body.len().min(200)]).into_owned();
250 tracing::error!(
251 "skipped a client message that isn't valid JSON ({reason}): {shown}"
252 );
253 continue;
254 }
255 };
256 let header = format!("Content-Length: {}\r\n\r\n", body.len());
257 if output.write_all(header.as_bytes()).await.is_err()
258 || output.write_all(&body).await.is_err()
259 || output.flush().await.is_err()
260 {
261 return;
263 }
264 if is_exit(&body) {
265 return;
266 }
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273
274 #[test]
275 fn valid_json_is_forwarded() {
276 assert_eq!(classify(br#"{"a":"\ud83d\ude00"}"#), Frame::Forward);
277 }
278
279 #[test]
280 fn a_lone_surrogate_is_repaired_to_the_replacement_character() {
281 let Frame::Repaired(body) = classify(br#"{"text":"a\udcffb"}"#) else {
282 panic!("repaired");
283 };
284 let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
285 assert_eq!(value["text"], "a\u{FFFD}b");
286 }
287
288 #[test]
289 fn a_lone_high_surrogate_is_repaired_and_a_pair_is_kept() {
290 let Frame::Repaired(body) = classify(br#"{"t":"\ud83d x \ud83d\ude00"}"#) else {
291 panic!("repaired");
292 };
293 let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
294 assert_eq!(value["t"], "\u{FFFD} x \u{1F600}");
295 }
296
297 #[test]
298 fn an_escaped_backslash_before_u_is_not_a_surrogate_escape() {
299 assert_eq!(classify(br#"{"t":"\\udcff"}"#), Frame::Forward);
301 }
302
303 #[test]
304 fn other_invalid_json_is_skipped() {
305 assert!(matches!(classify(b"{not json"), Frame::Skip(_)));
306 }
307
308 #[test]
309 fn a_surrogate_escape_cut_off_at_the_end_is_skipped() {
310 assert!(matches!(classify(br#"{"t":"\ud8"#), Frame::Skip(_)));
311 }
312
313 async fn pumped(input: &[u8]) -> Vec<u8> {
315 let (mut server_side, client_side) = tokio::io::duplex(64 * 1024);
316 pump(input, client_side).await;
317 let mut out = Vec::new();
318 server_side.read_to_end(&mut out).await.unwrap();
319 out
320 }
321
322 #[tokio::test]
325 async fn unrecoverable_framing_ends_the_pump() {
326 assert!(pumped(b"Content-Type: x\r\n\r\n{}").await.is_empty());
328 assert!(pumped(b"Content-Length: 50\r\n\r\nhello").await.is_empty());
330 assert!(pumped(b"\r\n\r\n\r\n").await.is_empty());
332 assert!(
334 pumped(b"Content-Length: 99999999999\r\n\r\n{}")
335 .await
336 .is_empty()
337 );
338 assert!(
339 pumped(b"Content-Length: 18446744073709551615\r\n\r\n{}")
340 .await
341 .is_empty()
342 );
343 let endless = vec![b'x'; (MAX_HEADER_LINE as usize) * 2];
345 assert!(pumped(&endless).await.is_empty());
346 }
347
348 #[tokio::test]
352 async fn the_pump_skips_a_bad_frame_and_keeps_going() {
353 let frame = |body: &str| format!("Content-Length: {}\r\n\r\n{body}", body.len());
354 let input = [
355 frame(r#"{"n":1}"#),
356 frame("{not json"),
357 frame(r#"{"t":"\udcff"}"#),
358 frame(r#"{"n":2}"#),
359 ]
360 .concat();
361 let (server_side, client_side) = tokio::io::duplex(64 * 1024);
362 pump(input.as_bytes(), client_side).await;
363 let mut out = String::new();
364 let mut server_side = server_side;
365 server_side.read_to_string(&mut out).await.unwrap();
366 assert_eq!(out.matches("Content-Length").count(), 3, "{out}");
367 assert!(
368 out.contains(r#"{"n":1}"#) && out.contains(r#"{"n":2}"#),
369 "{out}"
370 );
371 assert!(out.contains(r#"\uFFFD"#), "{out}");
372 assert!(!out.contains("not json"), "{out}");
373 }
374
375 #[test]
376 fn only_the_exit_notification_is_exit() {
377 assert!(is_exit(br#"{"jsonrpc":"2.0","method":"exit"}"#));
378 assert!(!is_exit(br#"{"jsonrpc":"2.0","id":2,"method":"shutdown"}"#));
379 assert!(!is_exit(
380 br#"{"jsonrpc":"2.0","method":"initialized","params":{}}"#
381 ));
382 assert!(!is_exit(
384 br#"{"jsonrpc":"2.0","method":"x","params":{"t":"exit"}}"#
385 ));
386 }
387
388 #[tokio::test]
392 async fn the_pump_stops_after_exit() {
393 let frame = |body: &str| format!("Content-Length: {}\r\n\r\n{body}", body.len());
394 let (mut client, input) = tokio::io::duplex(64 * 1024);
395 let written = [
396 frame(r#"{"jsonrpc":"2.0","method":"exit"}"#),
397 frame(r#"{"jsonrpc":"2.0","method":"initialized","params":{}}"#),
398 ]
399 .concat();
400 client.write_all(written.as_bytes()).await.unwrap();
401 let (mut server_side, output) = tokio::io::duplex(64 * 1024);
403 tokio::time::timeout(std::time::Duration::from_secs(5), pump(input, output))
404 .await
405 .expect("the pump returns after `exit`");
406 let mut out = String::new();
407 server_side.read_to_string(&mut out).await.unwrap();
408 assert!(out.contains(r#""method":"exit""#), "{out}");
409 assert!(!out.contains("initialized"), "{out}");
410 drop(client);
411 }
412}