diff --git a/crates/bamts-cli/src/lsp.rs b/crates/bamts-cli/src/lsp.rs index 0904e241..03654f7c 100644 --- a/crates/bamts-cli/src/lsp.rs +++ b/crates/bamts-cli/src/lsp.rs @@ -5,7 +5,7 @@ use std::{ collections::{BTreeMap, BTreeSet}, - io::{self, BufRead, Write}, + io::{self, BufRead, Read, Write}, path::{Path, PathBuf}, sync::Arc, }; @@ -20,6 +20,7 @@ use bamts_compiler::source::{SourceText, TextRange, Utf16Pos}; use serde_json::{Value, json}; const MAX_MESSAGE_BYTES: usize = 16 * 1024 * 1024; +const MAX_HEADER_BYTES: usize = 8 * 1024; const REQUEST_CANCELLED: i32 = -32800; /// How the stdio loop finished. @@ -778,19 +779,46 @@ fn percent_decode(input: &str) -> Result { fn read_message(input: &mut impl BufRead) -> io::Result>> { let mut content_length = None; + let mut header_bytes = 0; loop { let mut header = String::new(); - if input.read_line(&mut header)? == 0 { - return Ok(None); + // Bound the read itself so a peer cannot grow an unterminated line. + let remaining = (MAX_HEADER_BYTES - header_bytes + 1) as u64; + let read = input.take(remaining).read_line(&mut header)?; + if read == 0 { + return if header_bytes == 0 { + Ok(None) + } else { + Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "stream ended inside an LSP header", + )) + }; + } + header_bytes += read; + if header_bytes > MAX_HEADER_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "LSP header exceeds 8 KiB", + )); } let header = header.trim_end_matches(['\r', '\n']); if header.is_empty() { break; } let Some((name, value)) = header.split_once(':') else { - continue; + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "LSP header line has no separator", + )); }; if name.eq_ignore_ascii_case("Content-Length") { + if content_length.is_some() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "duplicate Content-Length", + )); + } content_length = Some( value .trim() @@ -2433,6 +2461,61 @@ mod tests { assert_eq!(error.kind(), io::ErrorKind::InvalidData); } + #[test] + fn framing_bounds_header_reads_before_a_newline() { + for prefix in ["", "Content-Length: 2\r\n", "X: y\r\n"] { + let bytes = format!("{prefix}X-Padding: {}", "x".repeat(16 * 1024)); + let mut input = Cursor::new(bytes); + let error = read_message(&mut input).expect_err("header budget"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(input.position(), 8193, "reject at the first excess byte"); + } + } + + #[test] + fn framing_rejects_duplicate_lengths_and_malformed_headers() { + for bytes in [ + "Content-Length: 2\r\nContent-Length: 2\r\n\r\n{}", + "Content-Length: 9\r\ncontent-length: 2\r\n\r\n{}", + "Not a header\r\nContent-Length: 2\r\n\r\n{}", + ] { + let error = read_message(&mut Cursor::new(bytes)).expect_err(bytes); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + } + } + + #[test] + fn framing_distinguishes_eof_from_a_truncated_header() { + assert_eq!( + read_message(&mut Cursor::new(b"")).expect("clean EOF"), + None + ); + for bytes in ["Content-Length: 2", "Content-Length: 2\r\n"] { + let error = read_message(&mut Cursor::new(bytes)).expect_err("truncated header"); + assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); + } + } + + #[test] + fn framing_preserves_exact_budget_and_fragmented_frame_boundaries() { + let mut bytes = b"Content-Length: 2\r\nX-Padding: ".to_vec(); + bytes.resize(8192 - 4, b'x'); + bytes.extend_from_slice(b"\r\n\r\n{}"); + bytes.extend_from_slice(b"Content-Length: 2\r\n\r\n[]"); + for capacity in [1, 2, 3, 7, 13, 512, 8192] { + let mut input = io::BufReader::with_capacity(capacity, Cursor::new(&bytes)); + assert_eq!( + read_message(&mut input).expect("first frame"), + Some(b"{}".to_vec()) + ); + assert_eq!( + read_message(&mut input).expect("second frame"), + Some(b"[]".to_vec()) + ); + assert_eq!(read_message(&mut input).expect("clean EOF"), None); + } + } + /// A `BufRead` that serves framed messages in segments, running each /// segment's action just before its bytes become readable. This lets a /// test mutate the filesystem strictly between two processed messages. diff --git a/crates/bamts-cli/tests/cli.rs b/crates/bamts-cli/tests/cli.rs index dff54312..13902f08 100644 --- a/crates/bamts-cli/tests/cli.rs +++ b/crates/bamts-cli/tests/cli.rs @@ -250,6 +250,110 @@ process.stdout.write(process.env.BAMTS_AOT_ENTRYPOINT === undefined ? "hidden\n" assert_eq!(aot.stdout, jit.stdout); } +#[cfg(unix)] +#[test] +fn lsp_rejects_malformed_headers_without_dispatching_the_next_frame() { + for mut input in [ + vec![b'x'; 8 * 1024 + 1], + b"Content-Length: 2\r\ncontent-length: 2\r\n\r\n{}".to_vec(), + b"Not a header\r\nContent-Length: 2\r\n\r\n{}".to_vec(), + ] { + let directory = ScratchDirectory::new(); + input.extend(framed( + br#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"capabilities":{}}}"#, + )); + let (mut peer, child_stdin) = UnixStream::pair().expect("socket pair for child stdin"); + let child = directory + .command() + .arg("--lsp") + .current_dir(&directory.path) + .stdin(Stdio::from(OwnedFd::from(child_stdin))) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("LSP child starts"); + peer.write_all(&input).expect("malformed input is written"); + let output = wait_for_output(child, "bamts --lsp malformed header"); + drop(peer); + assert_eq!(output.status.code(), Some(1)); + assert!( + output.stdout.is_empty(), + "ambiguous framing must not dispatch JSON" + ); + assert!(!stderr(&output).contains("panicked")); + } +} + +#[cfg(unix)] +#[test] +fn lsp_fragmented_typescript_edits_publish_errors_then_recover() { + for chunk_size in [1, 3, 17, 4096] { + let directory = ScratchDirectory::new(); + directory.write("main.ts", "const value: number = 7;\n"); + let uri = format!("file://{}", directory.path.join("main.ts").display()); + let requests = [ + serde_json::json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"capabilities":{}}}), + serde_json::json!({"jsonrpc":"2.0","method":"initialized","params":{}}), + serde_json::json!({"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{ + "uri":uri,"languageId":"typescript","version":1,"text":"const value: number = \"bad\";\n" + }}}), + serde_json::json!({"jsonrpc":"2.0","method":"textDocument/didChange","params":{ + "textDocument":{"uri":uri,"version":2},"contentChanges":[{"text":"const value: number = 7;\n"}] + }}), + serde_json::json!({"jsonrpc":"2.0","id":2,"method":"shutdown","params":null}), + serde_json::json!({"jsonrpc":"2.0","method":"exit"}), + ]; + let input: Vec = requests + .iter() + .flat_map(|request| framed(&serde_json::to_vec(request).expect("request serializes"))) + .collect(); + let mut child = directory + .command() + .arg("--lsp") + .current_dir(&directory.path) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("LSP child starts"); + let mut stdin = child.stdin.take().expect("piped stdin"); + for chunk in input.chunks(chunk_size) { + stdin.write_all(chunk).expect("fragment is written"); + } + drop(stdin); + let output = wait_for_output(child, "bamts --lsp fragmented edits"); + assert_success(&output, "bamts --lsp fragmented edits"); + let responses = decode_frames(&output.stdout); + let diagnostics: Vec<_> = responses + .iter() + .filter(|response| response["method"] == "textDocument/publishDiagnostics") + .map(|response| { + response["params"]["diagnostics"] + .as_array() + .expect("diagnostics") + }) + .collect(); + assert_eq!(diagnostics.len(), 2, "{responses:?}"); + assert!( + diagnostics[0] + .iter() + .any(|diagnostic| diagnostic["code"] == "BAMTS-C004"), + "the invalid TypeScript must be checked: {responses:?}" + ); + assert!( + diagnostics[1] + .iter() + .all(|diagnostic| diagnostic["severity"] != 1), + "the valid edit must clear stale errors: {responses:?}" + ); + assert!( + responses + .iter() + .any(|response| response["id"] == 2 && response.get("result").is_some()) + ); + } +} + #[test] fn aot_and_jit_execute_non_decimal_bigint_literals() { let project = ScratchDirectory::new();