Skip to main content

bynk_lsp/
transport.rs

1//! #1667: a framing pump in front of `tower-lsp`'s transport.
2//!
3//! `tower-lsp` 0.20 reads stdin through a `FramedRead`, and a framed stream
4//! ends at its first decode error. So one message whose body isn't valid JSON
5//! stopped the whole server, with exit 0 and nothing on stderr. The trigger
6//! is reachable from a well-behaved client: `JSON.stringify` writes a lone
7//! UTF-16 surrogate in a string as `\udcff`, which `serde_json` rejects.
8//!
9//! [`pump`] reads the client's frames itself and forwards to the server only
10//! bodies that parse, so a bad one can't end the stream:
11//!
12//! - a body that parses is forwarded unchanged;
13//! - a body whose only fault is a lone surrogate escape is repaired (each
14//!   lone `\uD800`–`\uDFFF` becomes `\uFFFD`) and forwarded. U+FFFD is one
15//!   UTF-16 code unit, like the surrogate it replaces, so the client's
16//!   positions into the text stay aligned with the server's;
17//! - any other body is logged and skipped. Its `id` can't be recovered from
18//!   unparseable JSON, so no error response is possible.
19//!
20//! A header block without `Content-Length`, or a stream that ends mid-frame
21//! (a truncated body), leaves nothing to resynchronise on: the pump stops,
22//! with the reason logged, and the server sees end of input.
23//!
24//! #2239: the pump also stops once it has forwarded `exit`. `tower-lsp` 0.20
25//! handles a message that arrives after `exit` by leaving its read loop early,
26//! without the clean-up it does at end of input, so `Server::serve` never
27//! returned and the process never exited. Ending the server's input after
28//! `exit` means it only ever sees end of input next, whatever the client sends.
29
30use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
31
32/// #1771 review: the largest body [`pump`] accepts. A `Content-Length` is
33/// otherwise trusted and allocated in full before the body arrives, so one
34/// corrupted digit could abort the process. A frame over this is unrecoverable
35/// framing, like a missing `Content-Length`. 128 MiB is far above any real
36/// message (a `didOpen` carries one file's text).
37pub const MAX_BODY: usize = 128 << 20;
38
39/// The longest header line [`pump`] reads, so a client that never sends a
40/// newline can't grow the line without limit.
41pub const MAX_HEADER_LINE: u64 = 8 << 10;
42
43/// What [`pump`] does with one frame's body.
44#[derive(Debug, PartialEq, Eq)]
45pub enum Frame {
46    /// The body parses as JSON: forward it unchanged.
47    Forward,
48    /// The body parses once its lone surrogate escapes are replaced.
49    Repaired(Vec<u8>),
50    /// The body can't be parsed. The string says why.
51    Skip(String),
52}
53
54/// Decide what to do with one frame's body.
55pub 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
69/// Replace each lone surrogate escape inside a JSON string with `\uFFFD`.
70/// `None` when there is none to replace.
71///
72/// The scan tracks string and escape state, so an escaped backslash
73/// (`\\udcff`, a literal backslash then `udcff`) is left alone, and a valid
74/// pair (`\ud83d\ude00`) is kept as written.
75fn 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                        // A high surrogate is kept only with a low one after it.
102                        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                // Any other escape: copy it and the escaped byte together, so
128                // an escaped quote or backslash can't change the string state.
129                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
147/// #2239: the longest body [`is_exit`] parses. An `exit` notification is a
148/// few dozen bytes, so a larger body (a `didOpen`'s text) is never parsed
149/// twice.
150const MAX_EXIT_BODY: usize = 1024;
151
152/// Whether `body` is the `exit` notification, after which the server reads
153/// nothing more.
154pub 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
160/// The four hex digits at `at`, as a UTF-16 code unit.
161fn 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
166/// Read framed messages from `input` and write the usable ones, re-framed, to
167/// `output`. Returns when `input` ends or loses its framing, or once `exit`
168/// is forwarded (#2239). Dropping `output` then ends the server's input.
169pub 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        // Headers, up to the blank line. Only `Content-Length` matters.
177        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            // The server side has gone; nothing more to deliver.
262            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        // `\\udcff` is a literal backslash then `udcff`: valid JSON as is.
300        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    /// Run the pump over `input` and return everything it forwarded.
314    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    /// #1771 review: each unrecoverable-framing path returns (the server
323    /// then sees end of input) rather than hanging, spinning or allocating.
324    #[tokio::test]
325    async fn unrecoverable_framing_ends_the_pump() {
326        // A header block with no Content-Length.
327        assert!(pumped(b"Content-Type: x\r\n\r\n{}").await.is_empty());
328        // A truncated body: 50 bytes promised, 5 sent, then end of input.
329        assert!(pumped(b"Content-Length: 50\r\n\r\nhello").await.is_empty());
330        // Only blank lines, then end of input.
331        assert!(pumped(b"\r\n\r\n\r\n").await.is_empty());
332        // A Content-Length over the limit (or past usize) allocates nothing.
333        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        // A header line with no end.
344        let endless = vec![b'x'; (MAX_HEADER_LINE as usize) * 2];
345        assert!(pumped(&endless).await.is_empty());
346    }
347
348    /// The pump forwards good frames, repairs a lone surrogate, skips an
349    /// unparseable one, and keeps going: one bad message no longer ends the
350    /// stream the server reads.
351    #[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        // `exit` inside a string is not the method.
383        assert!(!is_exit(
384            br#"{"jsonrpc":"2.0","method":"x","params":{"t":"exit"}}"#
385        ));
386    }
387
388    /// #2239: the pump ends the server's input once it has forwarded `exit`,
389    /// without waiting for the client's input to end, and forwards nothing
390    /// after it.
391    #[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        // `client` stays open: the pump must return on `exit`, not end of input.
402        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}