diff --git a/lsp-server/src/core/server.cppm b/lsp-server/src/core/server.cppm index ab2fbcc..79ee8b6 100644 --- a/lsp-server/src/core/server.cppm +++ b/lsp-server/src/core/server.cppm @@ -49,7 +49,20 @@ export namespace lsp::core scheduler::async_executor::TaskHandle handle; }; - std::optional 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(); - if (!message) + auto read_result = ReadMessage(); + 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 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') - line.pop_back(); + 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)) + { + if (has_content_length) + return { ReadStatus::kFramingError, {} }; + + const std::string_view value(line.data() + kContentLengthPrefix.size(), + line.size() - kContentLengthPrefix.size()); + if (value.empty()) + return { ReadStatus::kFramingError, {} }; + + 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, {} }; + } + + has_content_length = true; continue; + } - std::string length = line.substr(std::string_view("Content-Length:").size()); - const auto start = length.find_first_not_of(' '); - if (start == std::string::npos) - return std::nullopt; + if (line == kContentType) + { + if (has_content_type) + return { ReadStatus::kFramingError, {} }; + has_content_type = true; + continue; + } - content_length = std::stoul(length.substr(start)); + return { ReadStatus::kFramingError, {} }; } - if (content_length == 0) - return std::nullopt; + if (!has_content_length) + return { ReadStatus::kFramingError, {} }; std::string body(content_length, '\0'); input_.read(body.data(), static_cast(content_length)); if (input_.gcount() != static_cast(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) diff --git a/lsp-server/test/test_core_server.py b/lsp-server/test/test_core_server.py index 7bcb4a5..875f637 100644 --- a/lsp-server/test/test_core_server.py +++ b/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