🐛 fix(core): reject invalid LSP framing
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user