fix(lsp): bound headers and reject ambiguous framing

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<String, String> {
 
 fn read_message(input: &mut impl BufRead) -> io::Result<Option<Vec<u8>>> {
     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.

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<u8> = 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();