🐛 fix(core): reject invalid LSP framing

lsp-server/src/core/server.cppm

@@ -49,7 +49,20 @@ export namespace lsp::core
             scheduler::async_executor::TaskHandle handle;
         };
 
-        std::optional<std::string> ReadMessage();
+        enum class ReadStatus
+        {
+            kMessage,
+            kEndOfStream,
+            kFramingError,
+        };
+
+        struct ReadResult
+        {
+            ReadStatus status;
+            std::string message;
+        };
+
+        ReadResult ReadMessage();
         void HandleMessage(const std::string& raw_message);
         void SendMessage(const std::string& message);
 
@@ -182,13 +195,19 @@ namespace lsp::core
 
             try
             {
-                auto message = ReadMessage();
+                auto read_result = ReadMessage();
-                if (!message)
+                if (read_result.status == ReadStatus::kEndOfStream)
                 {
                     spdlog::info("End of input stream, exiting main loop");
                     break;
                 }
-                HandleMessage(*message);
+                if (read_result.status == ReadStatus::kFramingError)
+                {
+                    spdlog::error("Invalid LSP message framing");
+                    exit_code_ = 1;
+                    break;
+                }
+                HandleMessage(read_result.message);
             }
             catch (const std::exception& error)
             {
@@ -209,40 +228,79 @@ namespace lsp::core
         return exit_code_;
     }
 
-    std::optional<std::string> LspServer::ReadMessage()
+    LspServer::ReadResult LspServer::ReadMessage()
     {
+        constexpr std::size_t kMaxMessageSize = 16U * 1024U * 1024U;
+        constexpr std::string_view kContentLengthPrefix = "Content-Length: ";
+        constexpr std::string_view kContentType =
+            "Content-Type: application/vscode-jsonrpc; charset=utf-8";
+
         std::string line;
         std::size_t content_length = 0;
+        bool read_header = false;
+        bool has_content_length = false;
+        bool has_content_type = false;
 
-        while (std::getline(input_, line))
+        while (true)
         {
-            if (!line.empty() && line.back() == '\r')
+            if (!std::getline(input_, line))
+            {
+                if (!read_header && input_.eof())
+                    return { ReadStatus::kEndOfStream, {} };
+                return { ReadStatus::kFramingError, {} };
+            }
+
+            read_header = true;
+            if (line.empty() || line.back() != '\r')
+                return { ReadStatus::kFramingError, {} };
             line.pop_back();
 
             if (line.empty())
                 break;
 
-            if (line.rfind("Content-Length:", 0) != 0)
+            if (line.starts_with(kContentLengthPrefix))
-                continue;
+            {
+                if (has_content_length)
+                    return { ReadStatus::kFramingError, {} };
 
-            std::string length = line.substr(std::string_view("Content-Length:").size());
+                const std::string_view value(line.data() + kContentLengthPrefix.size(),
-            const auto start = length.find_first_not_of(' ');
+                                             line.size() - kContentLengthPrefix.size());
-            if (start == std::string::npos)
+                if (value.empty())
-                return std::nullopt;
+                    return { ReadStatus::kFramingError, {} };
 
-            content_length = std::stoul(length.substr(start));
+                const auto [end, error] =
+                    std::from_chars(value.data(), value.data() + value.size(), content_length);
+                if (error != std::errc{} || end != value.data() + value.size() ||
+                    content_length == 0 || content_length > kMaxMessageSize)
+                {
+                    return { ReadStatus::kFramingError, {} };
                 }
 
-        if (content_length == 0)
+                has_content_length = true;
-            return std::nullopt;
+                continue;
+            }
+
+            if (line == kContentType)
+            {
+                if (has_content_type)
+                    return { ReadStatus::kFramingError, {} };
+                has_content_type = true;
+                continue;
+            }
+
+            return { ReadStatus::kFramingError, {} };
+        }
+
+        if (!has_content_length)
+            return { ReadStatus::kFramingError, {} };
 
         std::string body(content_length, '\0');
         input_.read(body.data(), static_cast<std::streamsize>(content_length));
         if (input_.gcount() != static_cast<std::streamsize>(content_length))
-            return std::nullopt;
+            return { ReadStatus::kFramingError, {} };
 
         spdlog::trace("Received message: {}", body);
-        return body;
+        return { ReadStatus::kMessage, std::move(body) };
     }
 
     void LspServer::HandleMessage(const std::string& raw_message)

lsp-server/test/test_core_server.py

@@ -18,6 +18,15 @@ def raw_frame(body: bytes) -> bytes:
     return f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body
 
 
+def content_type_frame(message: dict) -> bytes:
+    body = json.dumps(message, separators=(",", ":")).encode("utf-8")
+    header = (
+        f"Content-Length: {len(body)}\r\n"
+        "Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n\r\n"
+    ).encode("ascii")
+    return header + body
+
+
 def read_messages(data: bytes) -> list[dict]:
     messages = []
     offset = 0
@@ -294,6 +303,43 @@ def assert_cancellation_and_failures(server: Path) -> None:
         client.kill()
 
 
+def assert_framing(server: Path) -> None:
+    invalid_payloads = [
+        b"Content-Length: 2junk\r\n\r\n{}",
+        b"Content-Length: 2\r\nContent-Length: 2\r\n\r\n{}",
+        b"Content-Length: 0\r\n\r\n",
+        b"Content-Length: 16777217\r\n\r\n",
+        b"Content-Length: 10\r\n\r\n{}",
+        b"Content-Length: 2\n\n{}",
+        b"X-Length: 2\r\n\r\n{}",
+        b"Content-Length: 2\r\nContent-Type: application/json\r\n\r\n{}",
+    ]
+    for index, payload in enumerate(invalid_payloads):
+        result = run_raw(server, payload)
+        if result.returncode != 1:
+            raise RuntimeError(f"invalid framing case {index} should return code 1")
+        if result.stdout:
+            raise RuntimeError(f"invalid framing case {index} should not process its body")
+
+    valid_payload = b"".join([
+        content_type_frame({
+            "jsonrpc": "2.0",
+            "id": 31,
+            "method": "initialize",
+            "params": {},
+        }),
+        frame({"jsonrpc": "2.0", "id": 32, "method": "shutdown"}),
+        frame({"jsonrpc": "2.0", "method": "exit"}),
+    ])
+    valid = run_raw(server, valid_payload)
+    if valid.returncode != 0:
+        raise RuntimeError("canonical Content-Type should be accepted")
+    if response_by_id(valid.stdout, 31).get("error"):
+        raise RuntimeError("initialize with canonical Content-Type should succeed")
+    if response_by_id(valid.stdout, 32).get("result", "missing") is not None:
+        raise RuntimeError("continuous framing should reach shutdown")
+
+
 def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--server", type=Path, required=True)
@@ -302,6 +348,7 @@ def main() -> int:
     assert_lifecycle(args.server)
     assert_json_rpc_errors(args.server)
     assert_cancellation_and_failures(args.server)
+    assert_framing(args.server)
     return 0