🐛 fix(core): reject invalid LSP framing

This commit is contained in:
csh
2026-07-14 09:03:47 +08:00
parent 101b69e84f
commit 71b519793a
2 changed files with 123 additions and 18 deletions
+76 -18
View File
@@ -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();
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<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')
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<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)
+47
View File
@@ -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