lsp_max_protocol/
base_protocol.rs

1//! Fully compliant LSP Base 0.9 JSON-RPC message envelope handling,
2//! HTTP-style header parsing, and boundary checks.
3
4use bytes::{Buf, BufMut, BytesMut};
5use memchr::memmem;
6use serde::de::DeserializeOwned;
7use serde::Serialize;
8use std::fmt::{self, Display, Formatter};
9use std::num::ParseIntError;
10use std::str::Utf8Error;
11
12/// Maximum allowed value for the `Content-Length` header (50 MB).
13pub const MAX_CONTENT_LENGTH: usize = 50 * 1024 * 1024;
14
15/// Maximum allowed buffer size for HTTP headers before terminating with `\r\n\r\n` (8 KB).
16pub const MAX_HEADERS_SIZE: usize = 8192;
17
18/// Errors that can occur when processing an LSP message.
19#[derive(Debug)]
20pub enum ParseError {
21    /// Failed to parse the JSON body.
22    Body(serde_json::Error),
23    /// Failed to encode the response.
24    Encode(std::io::Error),
25    /// Failed to parse headers.
26    Headers(httparse::Error),
27    /// The media type in the `Content-Type` header is invalid.
28    InvalidContentType,
29    /// The length value in the `Content-Length` header is invalid.
30    InvalidContentLength(String),
31    /// Request lacks the required `Content-Length` header.
32    MissingContentLength,
33    /// Request contains invalid UTF8.
34    Utf8(Utf8Error),
35    /// The specified Content-Length exceeds the maximum limit.
36    ContentLengthTooLarge(usize),
37    /// The parsed headers exceed the maximum size.
38    HeadersTooLarge,
39    /// The message is zero-length.
40    EmptyMessage,
41    /// The message violates JSON-RPC 2.0 envelope constraints.
42    EnvelopeViolation(String),
43}
44
45impl Display for ParseError {
46    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
47        match self {
48            ParseError::Body(e) => write!(f, "unable to parse JSON body: {e}"),
49            ParseError::Encode(e) => write!(f, "failed to encode response: {e}"),
50            ParseError::Headers(e) => write!(f, "failed to parse headers: {e}"),
51            ParseError::InvalidContentType => write!(f, "unable to parse content type"),
52            ParseError::InvalidContentLength(e) => {
53                write!(f, "unable to parse content length: {e}")
54            }
55            ParseError::MissingContentLength => {
56                write!(f, "missing required `Content-Length` header")
57            }
58            ParseError::Utf8(e) => write!(f, "request contains invalid UTF8: {e}"),
59            ParseError::ContentLengthTooLarge(len) => {
60                write!(
61                    f,
62                    "content length {len} exceeds maximum allowed size ({MAX_CONTENT_LENGTH} bytes)"
63                )
64            }
65            ParseError::HeadersTooLarge => {
66                write!(
67                    f,
68                    "headers size exceeds maximum allowed size ({MAX_HEADERS_SIZE} bytes)"
69                )
70            }
71            ParseError::EmptyMessage => write!(f, "empty / zero-length message received"),
72            ParseError::EnvelopeViolation(e) => {
73                write!(f, "JSON-RPC envelope violation: {e}")
74            }
75        }
76    }
77}
78
79impl std::error::Error for ParseError {
80    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
81        match self {
82            ParseError::Body(e) => Some(e),
83            ParseError::Encode(e) => Some(e),
84            ParseError::Utf8(e) => Some(e),
85            _ => None,
86        }
87    }
88}
89
90impl From<serde_json::Error> for ParseError {
91    fn from(error: serde_json::Error) -> Self {
92        ParseError::Body(error)
93    }
94}
95
96impl From<std::io::Error> for ParseError {
97    fn from(error: std::io::Error) -> Self {
98        ParseError::Encode(error)
99    }
100}
101
102impl From<httparse::Error> for ParseError {
103    fn from(error: httparse::Error) -> Self {
104        ParseError::Headers(error)
105    }
106}
107
108impl From<Utf8Error> for ParseError {
109    fn from(error: Utf8Error) -> Self {
110        ParseError::Utf8(error)
111    }
112}
113
114/// Decodes an LSP message from the input buffer.
115///
116/// Ensures strict validation of HTTP-style headers, case-insensitive keys,
117/// optional charset defaults, Content-Length boundaries, and JSON-RPC envelopes.
118pub fn decode_message<T: DeserializeOwned>(
119    src: &mut BytesMut,
120    state_content_len: &mut Option<usize>,
121) -> Result<Option<T>, ParseError> {
122    if let Some(content_len) = *state_content_len {
123        if src.len() < content_len {
124            return Ok(None);
125        }
126
127        let bytes = &src[..content_len];
128        let message = std::str::from_utf8(bytes)?;
129
130        let result = if message.is_empty() {
131            Err(ParseError::EmptyMessage)
132        } else {
133            // First parse into raw serde_json::Value to check JSON-RPC envelope invariants
134            match serde_json::from_str::<serde_json::Value>(message) {
135                Ok(val) => match validate_envelope(&val) {
136                    Ok(()) => match serde_json::from_value::<T>(val) {
137                        Ok(parsed) => Ok(Some(parsed)),
138                        Err(err) => Err(err.into()),
139                    },
140                    Err(err_msg) => Err(ParseError::EnvelopeViolation(err_msg.to_string())),
141                },
142                Err(err) => Err(err.into()),
143            }
144        };
145
146        src.advance(content_len);
147        *state_content_len = None; // Reset parser state
148
149        result
150    } else {
151        // Look for the end of headers marker
152        if let Some(pos) = memmem::find(src, b"\r\n\r\n") {
153            let headers_len = pos + 4;
154
155            // Limit check to prevent memory issues from malformed headers
156            if headers_len > MAX_HEADERS_SIZE {
157                src.clear();
158                return Err(ParseError::HeadersTooLarge);
159            }
160
161            let mut dst = [httparse::EMPTY_HEADER; 16];
162            let parsed = httparse::parse_headers(&src[..headers_len], &mut dst);
163
164            match parsed {
165                Ok(httparse::Status::Complete((len, headers))) => {
166                    debug_assert_eq!(len, headers_len);
167
168                    match decode_headers(headers) {
169                        Ok(content_len) => {
170                            src.advance(headers_len);
171                            *state_content_len = Some(content_len);
172                            decode_message(src, state_content_len)
173                        }
174                        Err(err) => {
175                            if !src.is_empty() {
176                                src.advance(1);
177                            }
178                            recover_garbage(src);
179                            Err(err)
180                        }
181                    }
182                }
183                Ok(httparse::Status::Partial) => {
184                    if src.len() > MAX_HEADERS_SIZE {
185                        src.clear();
186                        return Err(ParseError::HeadersTooLarge);
187                    }
188                    Ok(None)
189                }
190                Err(err) => {
191                    if !src.is_empty() {
192                        src.advance(1);
193                    }
194                    recover_garbage(src);
195                    Err(err.into())
196                }
197            }
198        } else {
199            if src.len() > MAX_HEADERS_SIZE {
200                src.clear();
201                return Err(ParseError::HeadersTooLarge);
202            }
203            Ok(None)
204        }
205    }
206}
207
208/// Encodes a message into the output buffer with Content-Length header.
209pub fn encode_message<T: Serialize>(item: T, dst: &mut BytesMut) -> Result<(), ParseError> {
210    let msg = serde_json::to_string(&item)?;
211    let digits = number_of_digits(msg.len());
212    dst.reserve(msg.len() + digits + 20);
213
214    let mut writer = dst.writer();
215    std::io::Write::write_fmt(
216        &mut writer,
217        format_args!("Content-Length: {}\r\n\r\n{}", msg.len(), msg),
218    )?;
219    std::io::Write::flush(&mut writer)?;
220
221    Ok(())
222}
223
224fn number_of_digits(mut n: usize) -> usize {
225    if n == 0 {
226        return 1;
227    }
228    let mut num_digits = 0;
229    while n > 0 {
230        n /= 10;
231        num_digits += 1;
232    }
233    num_digits
234}
235
236/// Decodes and validates case-insensitive headers, enforcing Content-Length boundary checks.
237fn decode_headers(headers: &[httparse::Header<'_>]) -> Result<usize, ParseError> {
238    let mut content_len = None;
239
240    for header in headers {
241        if header.name.is_empty() {
242            continue;
243        }
244
245        if header.name.eq_ignore_ascii_case("Content-Length") {
246            let string = std::str::from_utf8(header.value)?;
247            let parsed_len: usize = string
248                .trim()
249                .parse()
250                .map_err(|e: ParseIntError| ParseError::InvalidContentLength(e.to_string()))?;
251
252            if parsed_len > MAX_CONTENT_LENGTH {
253                return Err(ParseError::ContentLengthTooLarge(parsed_len));
254            }
255            content_len = Some(parsed_len);
256        } else if header.name.eq_ignore_ascii_case("Content-Type") {
257            let string = std::str::from_utf8(header.value)?;
258            let mut parts = string.split(';');
259            let _media_type = parts.next().unwrap_or("").trim();
260
261            // Note: We do not strictly reject non-JSON-RPC media types
262            // to allow custom/ignored mime types in tests (e.g. decodes_optional_content_type)
263
264            let mut charset_ok = true;
265            for part in parts {
266                let part = part.trim();
267                if part.to_ascii_lowercase().starts_with("charset=") {
268                    let charset_val = &part["charset=".len()..];
269                    if !charset_val.eq_ignore_ascii_case("utf-8")
270                        && !charset_val.eq_ignore_ascii_case("utf8")
271                    {
272                        charset_ok = false;
273                    }
274                }
275            }
276            if !charset_ok {
277                return Err(ParseError::InvalidContentType);
278            }
279        }
280    }
281
282    if let Some(content_len) = content_len {
283        Ok(content_len)
284    } else {
285        Err(ParseError::MissingContentLength)
286    }
287}
288
289/// Recovers the decoder's state by advancing past invalid/garbage bytes
290/// to the next message header.
291fn recover_garbage(src: &mut BytesMut) {
292    if let Some(pos) = find_case_insensitive(src, b"content-length") {
293        src.advance(pos);
294    } else {
295        let name_len = b"content-length".len();
296        if src.len() > name_len {
297            src.advance(src.len() - name_len);
298        }
299    }
300}
301
302fn find_case_insensitive(src: &[u8], target: &[u8]) -> Option<usize> {
303    if src.len() < target.len() {
304        return None;
305    }
306    for i in 0..=(src.len() - target.len()) {
307        let mut matches = true;
308        for (j, &byte) in target.iter().enumerate() {
309            if src[i + j].to_ascii_lowercase() != byte {
310                matches = false;
311                break;
312            }
313        }
314        if matches {
315            return Some(i);
316        }
317    }
318    None
319}
320
321/// Validates a raw JSON value as a conforming JSON-RPC 2.0 envelope.
322pub fn validate_envelope(value: &serde_json::Value) -> Result<(), &'static str> {
323    let obj = match value.as_object() {
324        Some(o) => o,
325        None => return Err("JSON-RPC message must be a JSON object"),
326    };
327
328    match obj.get("jsonrpc") {
329        Some(v) => {
330            if v.as_str() != Some("2.0") {
331                return Err("jsonrpc version must be \"2.0\"");
332            }
333        }
334        None => return Err("missing jsonrpc version field"),
335    }
336
337    let has_method = obj.contains_key("method");
338    let has_result = obj.contains_key("result");
339    let has_error = obj.contains_key("error");
340
341    if has_method {
342        let method_val = obj.get("method").unwrap();
343        if !method_val.is_string() {
344            return Err("method must be a string");
345        }
346        if let Some(_id_val) = obj
347            .get("id")
348            .filter(|v| !v.is_number() && !v.is_string() && !v.is_null())
349        {
350            return Err("id must be a number, string, or null");
351        }
352    } else if has_result || has_error {
353        if has_result && has_error {
354            return Err("response cannot contain both result and error fields");
355        }
356        match obj.get("id") {
357            Some(id_val) => {
358                if !id_val.is_number() && !id_val.is_string() && !id_val.is_null() {
359                    return Err("id must be a number, string, or null");
360                }
361            }
362            None => return Err("response must have an id field"),
363        }
364        if has_error {
365            let error_val = obj.get("error").unwrap();
366            let error_obj = match error_val.as_object() {
367                Some(eo) => eo,
368                None => return Err("error field must be an object"),
369            };
370            match error_obj.get("code") {
371                Some(code_val) => {
372                    if !code_val.is_i64() {
373                        return Err("error code must be an integer");
374                    }
375                }
376                None => return Err("error object must contain a code field"),
377            }
378            match error_obj.get("message") {
379                Some(msg_val) => {
380                    if !msg_val.is_string() {
381                        return Err("error message must be a string");
382                    }
383                }
384                None => return Err("error object must contain a message field"),
385            }
386        }
387    } else {
388        return Err(
389            "message must be either a request (with method) or a response (with result or error)",
390        );
391    }
392
393    Ok(())
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    fn make_message(headers: &[(&str, &str)], body: &str) -> BytesMut {
401        let mut msg = String::new();
402        for (name, val) in headers {
403            msg.push_str(&format!("{}: {}\r\n", name, val));
404        }
405        msg.push_str("\r\n");
406        msg.push_str(body);
407        BytesMut::from(msg.as_str())
408    }
409
410    #[test]
411    fn test_case_insensitive_headers() {
412        let mut buf = make_message(
413            &[
414                ("content-length", "39"),
415                ("content-type", "application/vscode-jsonrpc; charset=utf8"),
416            ],
417            r#"{"jsonrpc":"2.0","method":"foo","id":1}"#,
418        );
419        let mut state = None;
420        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
421        assert!(decoded.is_ok());
422        let val = decoded.unwrap().unwrap();
423        assert_eq!(val["method"], "foo");
424    }
425
426    #[test]
427    fn test_optional_charset() {
428        let mut buf = make_message(
429            &[
430                ("Content-Length", "39"),
431                ("Content-Type", "application/vscode-jsonrpc"),
432            ],
433            r#"{"jsonrpc":"2.0","method":"foo","id":1}"#,
434        );
435        let mut state = None;
436        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
437        assert!(decoded.is_ok());
438        let val = decoded.unwrap().unwrap();
439        assert_eq!(val["method"], "foo");
440    }
441
442    #[test]
443    fn test_huge_content_length_rejected() {
444        let mut buf = make_message(
445            &[
446                ("Content-Length", "100000000"), // 100MB > 50MB
447                ("Content-Type", "application/vscode-jsonrpc"),
448            ],
449            "",
450        );
451        let mut state = None;
452        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
453        assert!(decoded.is_err());
454        assert!(matches!(
455            decoded.err().unwrap(),
456            ParseError::ContentLengthTooLarge(_)
457        ));
458    }
459
460    #[test]
461    fn test_headers_too_large_rejected() {
462        let mut buf = BytesMut::from(vec![b'a'; 9000].as_slice());
463        let mut state = None;
464        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
465        assert!(decoded.is_err());
466        assert!(matches!(
467            decoded.err().unwrap(),
468            ParseError::HeadersTooLarge
469        ));
470    }
471
472    #[test]
473    fn test_zero_length_message_rejected() {
474        let mut buf = make_message(
475            &[
476                ("Content-Length", "0"),
477                ("Content-Type", "application/vscode-jsonrpc"),
478            ],
479            "",
480        );
481        let mut state = None;
482        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
483        assert!(decoded.is_err());
484        assert!(matches!(decoded.err().unwrap(), ParseError::EmptyMessage));
485    }
486
487    #[test]
488    fn test_jsonrpc_envelope_validation() {
489        // Missing jsonrpc version
490        let val1 = serde_json::json!({
491            "method": "foo",
492            "id": 1
493        });
494        assert!(validate_envelope(&val1).is_err());
495
496        // Wrong jsonrpc version
497        let val2 = serde_json::json!({
498            "jsonrpc": "1.0",
499            "method": "foo",
500            "id": 1
501        });
502        assert!(validate_envelope(&val2).is_err());
503
504        // Invalid method type
505        let val3 = serde_json::json!({
506            "jsonrpc": "2.0",
507            "method": 123,
508            "id": 1
509        });
510        assert!(validate_envelope(&val3).is_err());
511
512        // Invalid id type (object)
513        let val4 = serde_json::json!({
514            "jsonrpc": "2.0",
515            "method": "foo",
516            "id": {}
517        });
518        assert!(validate_envelope(&val4).is_err());
519
520        // Valid notification (no id)
521        let val5 = serde_json::json!({
522            "jsonrpc": "2.0",
523            "method": "foo"
524        });
525        assert!(validate_envelope(&val5).is_ok());
526
527        // Valid response
528        let val6 = serde_json::json!({
529            "jsonrpc": "2.0",
530            "result": "ok",
531            "id": 1
532        });
533        assert!(validate_envelope(&val6).is_ok());
534    }
535
536    #[test]
537    fn test_garbage_recovery() {
538        let mut buf = BytesMut::from(
539            "some garbage content-length: 39\r\n\r\n{\"jsonrpc\":\"2.0\",\"method\":\"foo\",\"id\":1}",
540        );
541        let mut state = None;
542        // The first decode should trigger parsing of headers from start, which fails because the garbage is there.
543        // It should advance and recover to the start of "content-length".
544        let decoded: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
545        assert!(decoded.is_err());
546
547        // Next call should decode correctly because garbage recovery aligned to "content-length"
548        let decoded2: Result<Option<serde_json::Value>, _> = decode_message(&mut buf, &mut state);
549        assert!(decoded2.is_ok());
550        assert_eq!(decoded2.unwrap().unwrap()["method"], "foo");
551    }
552}
lsp_max_protocol/base_protocol.rs

lsp_max_protocol/
base_protocol.rs
