From 57a3254668730a3ca6c49bdd20e98d40ed255606 Mon Sep 17 00:00:00 2001
From: Tyler Wilding <xTVaser@users.noreply.github.com>
Date: Sat, 1 Apr 2023 18:40:21 -0500
Subject: [PATCH] lsp: fix route handling json error (#2448)

---
 lsp/handlers/lsp_router.cpp |  6 +++---
 lsp/main.cpp                | 38 ++++++++++++++++++++-----------------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/lsp/handlers/lsp_router.cpp b/lsp/handlers/lsp_router.cpp
index 48de534ef0..33bb7c1a19 100644
--- a/lsp/handlers/lsp_router.cpp
+++ b/lsp/handlers/lsp_router.cpp
@@ -64,7 +64,7 @@ std::optional<std::vector<std::string>> LSPRouter::route_message(
     const MessageBuffer& message_buffer,
     AppState& appstate) {
   const json& body = message_buffer.body();
-  auto& method = body["method"];
+  const auto method = body.at("method").get<std::string>();
 
   // If the workspace has not yet been initialized but the client sends a
   // message that doesn't have method "initialize" then we'll return an error
@@ -76,7 +76,7 @@ std::optional<std::vector<std::string>> LSPRouter::route_message(
   }
 
   // Exit early if we can't handle the route
-  if (m_routes.count(method) == 0) {
+  if (m_routes.find(method) == m_routes.end()) {
     lg::warn("Method not supported '{}'", method);
     auto error = {make_response(
         error_resp(ErrorCodes::MethodNotFound, fmt::format("Method '{}' not supported", method)))};
@@ -84,7 +84,7 @@ std::optional<std::vector<std::string>> LSPRouter::route_message(
   }
 
   try {
-    auto route = m_routes[method];
+    auto& route = m_routes.at(method);
     std::vector<json> resp_bodies;
     // Handle the request/notificiation
     switch (route.m_route_type) {
diff --git a/lsp/main.cpp b/lsp/main.cpp
index fee908efe0..057721feb6 100644
--- a/lsp/main.cpp
+++ b/lsp/main.cpp
@@ -83,28 +83,32 @@ int main(int argc, char** argv) {
   _setmode(_fileno(stdin), _O_BINARY);
 #endif
 
-  char c;
-  MessageBuffer message_buffer;
-  while (std::cin.get(c)) {
-    message_buffer.handle_char(c);
+  try {
+    char c;
+    MessageBuffer message_buffer;
+    while (std::cin.get(c)) {
+      message_buffer.handle_char(c);
 
-    if (message_buffer.message_completed()) {
-      json body = message_buffer.body();
-      auto method_name = body["method"].get<std::string>();
-      lg::info(">>> Received message of method '{}'", method_name);
-      auto responses = lsp_router.route_message(message_buffer, appstate);
-      if (responses) {
-        for (const auto& response : responses.value()) {
-          std::cout << response.c_str() << std::flush;
-          if (appstate.verbose) {
-            lg::debug("<<< Sending message: {}", response);
-          } else {
-            lg::info("<<< Sending message of method '{}'", method_name);
+      if (message_buffer.message_completed()) {
+        json body = message_buffer.body();
+        auto method_name = body["method"].get<std::string>();
+        lg::info(">>> Received message of method '{}'", method_name);
+        auto responses = lsp_router.route_message(message_buffer, appstate);
+        if (responses) {
+          for (const auto& response : responses.value()) {
+            std::cout << response.c_str() << std::flush;
+            if (appstate.verbose) {
+              lg::debug("<<< Sending message: {}", response);
+            } else {
+              lg::info("<<< Sending message of method '{}'", method_name);
+            }
           }
         }
+        message_buffer.clear();
       }
-      message_buffer.clear();
     }
+  } catch (std::exception& e) {
+    lg::error("Unexpected LSP Exception occured - {}", e.what());
   }
 
   return 0;