tower_lsp/
codec.rs

Help
1//! Encoder and decoder for Language Server Protocol messages.
2
3use std::error::Error;
4use std::fmt::{self, Display, Formatter};
5use std::io::{Error as IoError, Write};
6use std::marker::PhantomData;
7use std::num::ParseIntError;
8use std::str::Utf8Error;
9
10use bytes::buf::BufMut;
11use bytes::{Buf, BytesMut};
12use memchr::memmem;
13use serde::{de::DeserializeOwned, Serialize};
14use tracing::{trace, warn};
15
16#[cfg(feature = "runtime-agnostic")]
17use async_codec_lite::{Decoder, Encoder};
18#[cfg(feature = "runtime-tokio")]
19use tokio_util::codec::{Decoder, Encoder};
20
21/// Errors that can occur when processing an LSP message.
22#[derive(Debug)]
23pub enum ParseError {
24    /// Failed to parse the JSON body.
25    Body(serde_json::Error),
26    /// Failed to encode the response.
27    Encode(IoError),
28    /// Failed to parse headers.
29    Headers(httparse::Error),
30    /// The media type in the `Content-Type` header is invalid.
31    InvalidContentType,
32    /// The length value in the `Content-Length` header is invalid.
33    InvalidContentLength(ParseIntError),
34    /// Request lacks the required `Content-Length` header.
35    MissingContentLength,
36    /// Request contains invalid UTF8.
37    Utf8(Utf8Error),
38}
39
40impl Display for ParseError {
41    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
42        match *self {
43            ParseError::Body(ref e) => write!(f, "unable to parse JSON body: {e}"),
44            ParseError::Encode(ref e) => write!(f, "failed to encode response: {e}"),
45            ParseError::Headers(ref e) => write!(f, "failed to parse headers: {e}"),
46            ParseError::InvalidContentType => write!(f, "unable to parse content type"),
47            ParseError::InvalidContentLength(ref e) => {
48                write!(f, "unable to parse content length: {e}")
49            }
50            ParseError::MissingContentLength => {
51                write!(f, "missing required `Content-Length` header")
52            }
53            ParseError::Utf8(ref e) => write!(f, "request contains invalid UTF8: {e}"),
54        }
55    }
56}
57
58impl Error for ParseError {
59    fn source(&self) -> Option<&(dyn Error + 'static)> {
60        match *self {
61            ParseError::Body(ref e) => Some(e),
62            ParseError::Encode(ref e) => Some(e),
63            ParseError::InvalidContentLength(ref e) => Some(e),
64            ParseError::Utf8(ref e) => Some(e),
65            _ => None,
66        }
67    }
68}
69
70impl From<serde_json::Error> for ParseError {
71    fn from(error: serde_json::Error) -> Self {
72        ParseError::Body(error)
73    }
74}
75
76impl From<IoError> for ParseError {
77    fn from(error: IoError) -> Self {
78        ParseError::Encode(error)
79    }
80}
81
82impl From<httparse::Error> for ParseError {
83    fn from(error: httparse::Error) -> Self {
84        ParseError::Headers(error)
85    }
86}
87
88impl From<ParseIntError> for ParseError {
89    fn from(error: ParseIntError) -> Self {
90        ParseError::InvalidContentLength(error)
91    }
92}
93
94impl From<Utf8Error> for ParseError {
95    fn from(error: Utf8Error) -> Self {
96        ParseError::Utf8(error)
97    }
98}
99
100/// Encodes and decodes Language Server Protocol messages.
101pub struct LanguageServerCodec<T> {
102    content_len: Option<usize>,
103    _marker: PhantomData<T>,
104}
105
106impl<T> Default for LanguageServerCodec<T> {
107    fn default() -> Self {
108        LanguageServerCodec {
109            content_len: None,
110            _marker: PhantomData,
111        }
112    }
113}
114
115#[cfg(feature = "runtime-agnostic")]
116impl<T: Serialize> Encoder for LanguageServerCodec<T> {
117    type Item = T;
118    type Error = ParseError;
119
120    fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
121        let msg = serde_json::to_string(&item)?;
122        trace!("-> {}", msg);
123
124        // Reserve just enough space to hold the `Content-Length: ` and `\r\n\r\n` constants,
125        // the length of the message, and the message body.
126        dst.reserve(msg.len() + number_of_digits(msg.len()) + 20);
127        let mut writer = dst.writer();
128        write!(writer, "Content-Length: {}\r\n\r\n{}", msg.len(), msg)?;
129        writer.flush()?;
130
131        Ok(())
132    }
133}
134
135#[cfg(feature = "runtime-tokio")]
136impl<T: Serialize> Encoder<T> for LanguageServerCodec<T> {
137    type Error = ParseError;
138
139    fn encode(&mut self, item: T, dst: &mut BytesMut) -> Result<(), Self::Error> {
140        let msg = serde_json::to_string(&item)?;
141        trace!("-> {}", msg);
142
143        // Reserve just enough space to hold the `Content-Length: ` and `\r\n\r\n` constants,
144        // the length of the message, and the message body.
145        dst.reserve(msg.len() + number_of_digits(msg.len()) + 20);
146        let mut writer = dst.writer();
147        write!(writer, "Content-Length: {}\r\n\r\n{}", msg.len(), msg)?;
148        writer.flush()?;
149
150        Ok(())
151    }
152}
153
154fn number_of_digits(mut n: usize) -> usize {
155    let mut num_digits = 0;
156
157    while n > 0 {
158        n /= 10;
159        num_digits += 1;
160    }
161
162    num_digits
163}
164
165impl<T: DeserializeOwned> Decoder for LanguageServerCodec<T> {
166    type Item = T;
167    type Error = ParseError;
168
169    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
170        if let Some(content_len) = self.content_len {
171            if src.len() < content_len {
172                return Ok(None);
173            }
174
175            let bytes = &src[..content_len];
176            let message = std::str::from_utf8(bytes)?;
177
178            let result = if message.is_empty() {
179                Ok(None)
180            } else {
181                trace!("<- {}", message);
182                match serde_json::from_str(message) {
183                    Ok(parsed) => Ok(Some(parsed)),
184                    Err(err) => Err(err.into()),
185                }
186            };
187
188            src.advance(content_len);
189            self.content_len = None; // Reset state in preparation for parsing next message.
190
191            result
192        } else {
193            let mut dst = [httparse::EMPTY_HEADER; 2];
194
195            let (headers_len, headers) = match httparse::parse_headers(src, &mut dst)? {
196                httparse::Status::Complete(output) => output,
197                httparse::Status::Partial => return Ok(None),
198            };
199
200            match decode_headers(headers) {
201                Ok(content_len) => {
202                    src.advance(headers_len);
203                    self.content_len = Some(content_len);
204                    self.decode(src) // Recurse right back in, now that `Content-Length` is known.
205                }
206                Err(err) => {
207                    match err {
208                        ParseError::MissingContentLength => {}
209                        _ => src.advance(headers_len),
210                    }
211
212                    // Skip any garbage bytes by scanning ahead for another potential message.
213                    src.advance(memmem::find(src, b"Content-Length").unwrap_or_default());
214                    Err(err)
215                }
216            }
217        }
218    }
219}
220
221fn decode_headers(headers: &[httparse::Header<'_>]) -> Result<usize, ParseError> {
222    let mut content_len = None;
223
224    for header in headers {
225        match header.name {
226            "Content-Length" => {
227                let string = std::str::from_utf8(header.value)?;
228                let parsed_len = string.parse()?;
229                content_len = Some(parsed_len);
230            }
231            "Content-Type" => {
232                let string = std::str::from_utf8(header.value)?;
233                let charset = string
234                    .split(';')
235                    .skip(1)
236                    .map(|param| param.trim())
237                    .find_map(|param| param.strip_prefix("charset="));
238
239                match charset {
240                    Some("utf-8") | Some("utf8") => {}
241                    _ => return Err(ParseError::InvalidContentType),
242                }
243            }
244            other => warn!("encountered unsupported header: {:?}", other),
245        }
246    }
247
248    if let Some(content_len) = content_len {
249        Ok(content_len)
250    } else {
251        Err(ParseError::MissingContentLength)
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use bytes::BytesMut;
258    use serde_json::Value;
259
260    use super::*;
261
262    macro_rules! assert_err {
263        ($expression:expr, $($pattern:tt)+) => {
264            match $expression {
265                $($pattern)+ => (),
266                ref e => panic!("expected `{}` but got `{:?}`", stringify!($($pattern)+), e),
267            }
268        }
269    }
270
271    fn encode_message(content_type: Option<&str>, message: &str) -> String {
272        let content_type = content_type
273            .map(|ty| format!("\r\nContent-Type: {ty}"))
274            .unwrap_or_default();
275
276        format!(
277            "Content-Length: {}{}\r\n\r\n{}",
278            message.len(),
279            content_type,
280            message
281        )
282    }
283
284    #[test]
285    fn encode_and_decode() {
286        let decoded = r#"{"jsonrpc":"2.0","method":"exit"}"#;
287        let encoded = encode_message(None, decoded);
288
289        let mut codec = LanguageServerCodec::default();
290        let mut buffer = BytesMut::new();
291        let item: Value = serde_json::from_str(decoded).unwrap();
292        codec.encode(item, &mut buffer).unwrap();
293        assert_eq!(buffer, BytesMut::from(encoded.as_str()));
294
295        let mut buffer = BytesMut::from(encoded.as_str());
296        let message = codec.decode(&mut buffer).unwrap();
297        let decoded = serde_json::from_str(decoded).unwrap();
298        assert_eq!(message, Some(decoded));
299    }
300
301    #[test]
302    fn decodes_optional_content_type() {
303        let decoded = r#"{"jsonrpc":"2.0","method":"exit"}"#;
304        let content_type = "application/vscode-jsonrpc; charset=utf-8";
305        let encoded = encode_message(Some(content_type), decoded);
306
307        let mut codec = LanguageServerCodec::default();
308        let mut buffer = BytesMut::from(encoded.as_str());
309        let message = codec.decode(&mut buffer).unwrap();
310        let decoded_: Value = serde_json::from_str(decoded).unwrap();
311        assert_eq!(message, Some(decoded_));
312
313        let content_type = "application/vscode-jsonrpc; charset=utf8";
314        let encoded = encode_message(Some(content_type), decoded);
315
316        let mut buffer = BytesMut::from(encoded.as_str());
317        let message = codec.decode(&mut buffer).unwrap();
318        let decoded_: Value = serde_json::from_str(decoded).unwrap();
319        assert_eq!(message, Some(decoded_));
320
321        let content_type = "application/vscode-jsonrpc; charset=invalid";
322        let encoded = encode_message(Some(content_type), decoded);
323
324        let mut buffer = BytesMut::from(encoded.as_str());
325        assert_err!(
326            codec.decode(&mut buffer),
327            Err(ParseError::InvalidContentType)
328        );
329
330        let content_type = "application/vscode-jsonrpc";
331        let encoded = encode_message(Some(content_type), decoded);
332
333        let mut buffer = BytesMut::from(encoded.as_str());
334        assert_err!(
335            codec.decode(&mut buffer),
336            Err(ParseError::InvalidContentType)
337        );
338
339        let content_type = "this-mime-should-be-ignored; charset=utf8";
340        let encoded = encode_message(Some(content_type), decoded);
341
342        let mut buffer = BytesMut::from(encoded.as_str());
343        let message = codec.decode(&mut buffer).unwrap();
344        let decoded_: Value = serde_json::from_str(decoded).unwrap();
345        assert_eq!(message, Some(decoded_));
346    }
347
348    #[test]
349    fn decodes_zero_length_message() {
350        let content_type = "application/vscode-jsonrpc; charset=utf-8";
351        let encoded = encode_message(Some(content_type), "");
352
353        let mut codec = LanguageServerCodec::default();
354        let mut buffer = BytesMut::from(encoded.as_str());
355        let message: Option<Value> = codec.decode(&mut buffer).unwrap();
356        assert_eq!(message, None);
357    }
358
359    #[test]
360    fn recovers_from_parse_error() {
361        let decoded = r#"{"jsonrpc":"2.0","method":"exit"}"#;
362        let encoded = encode_message(None, decoded);
363        let mixed = format!("foobar{encoded}Content-Length: foobar\r\n\r\n{encoded}");
364
365        let mut codec = LanguageServerCodec::default();
366        let mut buffer = BytesMut::from(mixed.as_str());
367        assert_err!(
368            codec.decode(&mut buffer),
369            Err(ParseError::MissingContentLength)
370        );
371
372        let message: Option<Value> = codec.decode(&mut buffer).unwrap();
373        let first_valid = serde_json::from_str(decoded).unwrap();
374        assert_eq!(message, Some(first_valid));
375        assert_err!(
376            codec.decode(&mut buffer),
377            Err(ParseError::InvalidContentLength(_))
378        );
379
380        let message = codec.decode(&mut buffer).unwrap();
381        let second_valid = serde_json::from_str(decoded).unwrap();
382        assert_eq!(message, Some(second_valid));
383
384        let message = codec.decode(&mut buffer).unwrap();
385        assert_eq!(message, None);
386    }
387
388    #[test]
389    fn decodes_small_chunks() {
390        let decoded = r#"{"jsonrpc":"2.0","method":"exit"}"#;
391        let content_type = "application/vscode-jsonrpc; charset=utf-8";
392        let encoded = encode_message(Some(content_type), decoded);
393
394        let mut codec = LanguageServerCodec::default();
395        let mut buffer = BytesMut::from(encoded.as_str());
396
397        let rest = buffer.split_off(40);
398        let message = codec.decode(&mut buffer).unwrap();
399        assert_eq!(message, None);
400        buffer.unsplit(rest);
401
402        let rest = buffer.split_off(80);
403        let message = codec.decode(&mut buffer).unwrap();
404        assert_eq!(message, None);
405        buffer.unsplit(rest);
406
407        let rest = buffer.split_off(16);
408        let message = codec.decode(&mut buffer).unwrap();
409        assert_eq!(message, None);
410        buffer.unsplit(rest);
411
412        let decoded: Value = serde_json::from_str(decoded).unwrap();
413        let message = codec.decode(&mut buffer).unwrap();
414        assert_eq!(message, Some(decoded));
415    }
416}
tower_lsp/codec.rs

tower_lsp/
codec.rs
