🐛 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();
-                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)