diff --git a/crates/lsp/src/lsp.rs b/crates/lsp/src/lsp.rs index e4b95bfb21..c4661a0793 100644 --- a/crates/lsp/src/lsp.rs +++ b/crates/lsp/src/lsp.rs @@ -31,6 +31,7 @@ use std::{ use std::{path::Path, process::Stdio}; use util::{ResultExt, TryFutureExt}; +const HEADER_DELIMITER: &'static [u8; 4] = b"\r\n\r\n"; const JSON_RPC_VERSION: &str = "2.0"; const CONTENT_LEN_HEADER: &str = "Content-Length: "; const LSP_REQUEST_TIMEOUT: Duration = Duration::from_secs(60 * 2); @@ -323,47 +324,17 @@ impl LanguageServer { loop { buffer.clear(); - if stdout.read_until(b'\n', &mut buffer).await? == 0 { - break; - }; + read_headers(&mut stdout, &mut buffer).await?; - if stdout.read_until(b'\n', &mut buffer).await? == 0 { - break; - }; + let headers = std::str::from_utf8(&buffer)?; - let header = std::str::from_utf8(&buffer)?; - let mut segments = header.lines(); - - let message_len: usize = segments - .next() - .with_context(|| { - format!("unable to find the first line of the LSP message header `{header}`") - })? - .strip_prefix(CONTENT_LEN_HEADER) - .with_context(|| format!("invalid LSP message header `{header}`"))? - .parse() - .with_context(|| { - format!("failed to parse Content-Length of LSP message header: `{header}`") - })?; - - if let Some(second_segment) = segments.next() { - match second_segment { - "" => (), // Header end - header_field => { - if header_field.starts_with("Content-Type:") { - stdout.read_until(b'\n', &mut buffer).await?; - } else { - anyhow::bail!( - "inside `{header}`, expected a Content-Type header field or a header ending CRLF, got `{second_segment:?}`" - ) - } - } - } - } else { - anyhow::bail!( - "unable to find the second line of the LSP message header `{header}`" - ); - } + let message_len = headers + .split("\n") + .find(|line| line.starts_with(CONTENT_LEN_HEADER)) + .and_then(|line| line.strip_prefix(CONTENT_LEN_HEADER)) + .ok_or_else(|| anyhow!("invalid LSP message header {headers:?}"))? + .trim_end() + .parse()?; buffer.resize(message_len, 0); stdout.read_exact(&mut buffer).await?; @@ -412,8 +383,6 @@ impl LanguageServer { // Don't starve the main thread when receiving lots of messages at once. smol::future::yield_now().await; } - - Ok(()) } async fn handle_stderr( @@ -1254,6 +1223,26 @@ impl FakeLanguageServer { } } +pub(self) async fn read_headers( + reader: &mut BufReader, + buffer: &mut Vec, +) -> Result<()> +where + Stdout: AsyncRead + Unpin + Send + 'static, +{ + loop { + if buffer.len() >= HEADER_DELIMITER.len() + && buffer[(buffer.len() - HEADER_DELIMITER.len())..] == HEADER_DELIMITER[..] + { + return Ok(()); + } + + if reader.read_until(b'\n', buffer).await? == 0 { + return Err(anyhow!("cannot read LSP message headers")); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -1327,4 +1316,28 @@ mod tests { drop(server); fake.receive_notification::().await; } + + #[gpui::test] + async fn test_read_headers() { + let mut buf = Vec::new(); + let mut reader = smol::io::BufReader::new(b"Content-Length: 123\r\n\r\n" as &[u8]); + read_headers(&mut reader, &mut buf).await.unwrap(); + assert_eq!(buf, b"Content-Length: 123\r\n\r\n"); + + let mut buf = Vec::new(); + let mut reader = smol::io::BufReader::new(b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n{\"somecontent\":123}" as &[u8]); + read_headers(&mut reader, &mut buf).await.unwrap(); + assert_eq!( + buf, + b"Content-Type: application/vscode-jsonrpc\r\nContent-Length: 1235\r\n\r\n" + ); + + let mut buf = Vec::new(); + let mut reader = smol::io::BufReader::new(b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n{\"somecontent\":true}" as &[u8]); + read_headers(&mut reader, &mut buf).await.unwrap(); + assert_eq!( + buf, + b"Content-Length: 1235\r\nContent-Type: application/vscode-jsonrpc\r\n\r\n" + ); + } }