[lldb] Refactoring JSONTransport into an abstract RPC Message Handler…

… and transport layer. (llvm#153121)

This abstracts the base Transport handler to have a MessageHandler
component and allows us to generalize both JSON-RPC 2.0 for MCP (or an
LSP) and DAP format.

This should allow us to create clearly defined clients and servers for
protocols, both for testing and for RPC between the lldb instances and
an lldb-mcp multiplexer.

This basic model is inspiried by the clangd/Transport.h file and the
mlir/lsp-server-support/Transport.h that are both used for LSP servers
within the llvm project.

Additionally, this helps with testing by subclassing `Transport` to
allow us to simplify sending/receiving messages without needing to use a
toJSON/fromJSON and a pair of pipes, see `TestTransport` in
DAP/TestBase.h.

(cherry picked from commit 538bd83)

lldb/include/lldb/Host/JSONTransport.h

@@ -13,29 +13,25 @@
 #ifndef LLDB_HOST_JSONTRANSPORT_H
 #define LLDB_HOST_JSONTRANSPORT_H
 
+#include "lldb/Host/MainLoop.h"
 #include "lldb/Host/MainLoopBase.h"
 #include "lldb/Utility/IOObject.h"
 #include "lldb/Utility/Status.h"
 #include "lldb/lldb-forward.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_ostream.h"
 #include <string>
 #include <system_error>
+#include <variant>
 #include <vector>
 
 namespace lldb_private {
 
-class TransportEOFError : public llvm::ErrorInfo<TransportEOFError> {
-public:
-  static char ID;
-
-  TransportEOFError() = default;
-  void log(llvm::raw_ostream &OS) const override;
-  std::error_code convertToErrorCode() const override;
-};
-
 class TransportUnhandledContentsError
     : public llvm::ErrorInfo<TransportUnhandledContentsError> {
 public:
@@ -54,112 +50,214 @@ class TransportUnhandledContentsError
   std::string m_unhandled_contents;
 };
 
-class TransportInvalidError : public llvm::ErrorInfo<TransportInvalidError> {
+/// A transport is responsible for maintaining the connection to a client
+/// application, and reading/writing structured messages to it.
+///
+/// Transports have limited thread safety requirements:
+///  - Messages will not be sent concurrently.
+///  - Messages MAY be sent while Run() is reading, or its callback is active.
+template <typename Req, typename Resp, typename Evt> class Transport {
 public:
-  static char ID;
-
-  TransportInvalidError() = default;
+  using Message = std::variant<Req, Resp, Evt>;
+
+  virtual ~Transport() = default;
+
+  /// Sends an event, a message that does not require a response.
+  virtual llvm::Error Send(const Evt &) = 0;
+  /// Sends a request, a message that expects a response.
+  virtual llvm::Error Send(const Req &) = 0;
+  /// Sends a response to a specific request.
+  virtual llvm::Error Send(const Resp &) = 0;
+
+  /// Implemented to handle incoming messages. (See Run() below).
+  class MessageHandler {
+  public:
+    virtual ~MessageHandler() = default;
+    /// Called when an event is received.
+    virtual void Received(const Evt &) = 0;
+    /// Called when a request is received.
+    virtual void Received(const Req &) = 0;
+    /// Called when a response is received.
+    virtual void Received(const Resp &) = 0;
+
+    /// Called when an error occurs while reading from the transport.
+    ///
+    /// NOTE: This does *NOT* indicate that a specific request failed, but that
+    /// there was an error in the underlying transport.
+    virtual void OnError(llvm::Error) = 0;
+
+    /// Called on EOF or client disconnect.
+    virtual void OnClosed() = 0;
+  };
+
+  using MessageHandlerSP = std::shared_ptr<MessageHandler>;
+
+  /// RegisterMessageHandler registers the Transport with the given MainLoop and
+  /// handles any incoming messages using the given MessageHandler.
+  ///
+  /// If an unexpected error occurs, the MainLoop will be terminated and a log
+  /// message will include additional information about the termination reason.
+  virtual llvm::Expected<MainLoop::ReadHandleUP>
+  RegisterMessageHandler(MainLoop &loop, MessageHandler &handler) = 0;
 
-  void log(llvm::raw_ostream &OS) const override;
-  std::error_code convertToErrorCode() const override;
+protected:
+  template <typename... Ts> inline auto Logv(const char *Fmt, Ts &&...Vals) {
+    Log(llvm::formatv(Fmt, std::forward<Ts>(Vals)...).str());
+  }
+  virtual void Log(llvm::StringRef message) = 0;
 };
 
-/// A transport class that uses JSON for communication.
-class JSONTransport {
+/// A JSONTransport will encode and decode messages using JSON.
+template <typename Req, typename Resp, typename Evt>
+class JSONTransport : public Transport<Req, Resp, Evt> {
 public:
-  using ReadHandleUP = MainLoopBase::ReadHandleUP;
-  template <typename T>
-  using Callback = std::function<void(MainLoopBase &, const llvm::Expected<T>)>;
-
-  JSONTransport(lldb::IOObjectSP input, lldb::IOObjectSP output);
-  virtual ~JSONTransport() = default;
-
-  /// Transport is not copyable.
-  /// @{
-  JSONTransport(const JSONTransport &rhs) = delete;
-  void operator=(const JSONTransport &rhs) = delete;
-  /// @}
-
-  /// Writes a message to the output stream.
-  template <typename T> llvm::Error Write(const T &t) {
-    const std::string message = llvm::formatv("{0}", toJSON(t)).str();
-    return WriteImpl(message);
+  using Transport<Req, Resp, Evt>::Transport;
+  using MessageHandler = typename Transport<Req, Resp, Evt>::MessageHandler;
+
+  JSONTransport(lldb::IOObjectSP in, lldb::IOObjectSP out)
+      : m_in(in), m_out(out) {}
+
+  llvm::Error Send(const Evt &evt) override { return Write(evt); }
+  llvm::Error Send(const Req &req) override { return Write(req); }
+  llvm::Error Send(const Resp &resp) override { return Write(resp); }
+
+  llvm::Expected<MainLoop::ReadHandleUP>
+  RegisterMessageHandler(MainLoop &loop, MessageHandler &handler) override {
+    Status status;
+    MainLoop::ReadHandleUP read_handle = loop.RegisterReadObject(
+        m_in,
+        std::bind(&JSONTransport::OnRead, this, std::placeholders::_1,
+                  std::ref(handler)),
+        status);
+    if (status.Fail()) {
+      return status.takeError();
+    }
+    return read_handle;
   }
 
-  /// Registers the transport with the MainLoop.
-  template <typename T>
-  llvm::Expected<ReadHandleUP> RegisterReadObject(MainLoopBase &loop,
-                                                  Callback<T> read_cb) {
-    Status error;
-    ReadHandleUP handle = loop.RegisterReadObject(
-        m_input,
-        [read_cb, this](MainLoopBase &loop) {
-          char buf[kReadBufferSize];
-          size_t num_bytes = sizeof(buf);
-          if (llvm::Error error = m_input->Read(buf, num_bytes).takeError()) {
-            read_cb(loop, std::move(error));
-            return;
-          }
-          if (num_bytes)
-            m_buffer.append(std::string(buf, num_bytes));
-
-          // If the buffer has contents, try parsing any pending messages.
-          if (!m_buffer.empty()) {
-            llvm::Expected<std::vector<std::string>> messages = Parse();
-            if (llvm::Error error = messages.takeError()) {
-              read_cb(loop, std::move(error));
-              return;
-            }
-
-            for (const auto &message : *messages)
-              if constexpr (std::is_same<T, std::string>::value)
-                read_cb(loop, message);
-              else
-                read_cb(loop, llvm::json::parse<T>(message));
-          }
-
-          // On EOF, notify the callback after the remaining messages were
-          // handled.
-          if (num_bytes == 0) {
-            if (m_buffer.empty())
-              read_cb(loop, llvm::make_error<TransportEOFError>());
-            else
-              read_cb(loop, llvm::make_error<TransportUnhandledContentsError>(
-                                std::string(m_buffer)));
-          }
-        },
-        error);
-    if (error.Fail())
-      return error.takeError();
-    return handle;
-  }
+  /// Public for testing purposes, otherwise this should be an implementation
+  /// detail.
+  static constexpr size_t kReadBufferSize = 1024;
 
 protected:
-  template <typename... Ts> inline auto Logv(const char *Fmt, Ts &&...Vals) {
-    Log(llvm::formatv(Fmt, std::forward<Ts>(Vals)...).str());
+  virtual llvm::Expected<std::vector<std::string>> Parse() = 0;
+  virtual std::string Encode(const llvm::json::Value &message) = 0;
+  llvm::Error Write(const llvm::json::Value &message) {
+    this->Logv("<-- {0}", message);
+    std::string output = Encode(message);
+    size_t bytes_written = output.size();
+    return m_out->Write(output.data(), bytes_written).takeError();
   }
-  virtual void Log(llvm::StringRef message);
 
-  virtual llvm::Error WriteImpl(const std::string &message) = 0;
-  virtual llvm::Expected<std::vector<std::string>> Parse() = 0;
+  llvm::SmallString<kReadBufferSize> m_buffer;
 
-  static constexpr size_t kReadBufferSize = 1024;
+private:
+  void OnRead(MainLoopBase &loop, MessageHandler &handler) {
+    char buf[kReadBufferSize];
+    size_t num_bytes = sizeof(buf);
+    if (Status status = m_in->Read(buf, num_bytes); status.Fail()) {
+      handler.OnError(status.takeError());
+      return;
+    }
+
+    if (num_bytes)
+      m_buffer.append(llvm::StringRef(buf, num_bytes));
+
+    // If the buffer has contents, try parsing any pending messages.
+    if (!m_buffer.empty()) {
+      llvm::Expected<std::vector<std::string>> raw_messages = Parse();
+      if (llvm::Error error = raw_messages.takeError()) {
+        handler.OnError(std::move(error));
+        return;
+      }
+
+      for (const std::string &raw_message : *raw_messages) {
+        llvm::Expected<typename Transport<Req, Resp, Evt>::Message> message =
+            llvm::json::parse<typename Transport<Req, Resp, Evt>::Message>(
+                raw_message);
+        if (!message) {
+          handler.OnError(message.takeError());
+          return;
+        }
+
+        std::visit([&handler](auto &&msg) { handler.Received(msg); }, *message);
+      }
+    }
+
+    // Check if we reached EOF.
+    if (num_bytes == 0) {
+      // EOF reached, but there may still be unhandled contents in the buffer.
+      if (!m_buffer.empty())
+        handler.OnError(llvm::make_error<TransportUnhandledContentsError>(
+            std::string(m_buffer.str())));
+      handler.OnClosed();
+    }
+  }
 
-  lldb::IOObjectSP m_input;
-  lldb::IOObjectSP m_output;
-  llvm::SmallString<kReadBufferSize> m_buffer;
+  lldb::IOObjectSP m_in;
+  lldb::IOObjectSP m_out;
 };
 
 /// A transport class for JSON with a HTTP header.
-class HTTPDelimitedJSONTransport : public JSONTransport {
+template <typename Req, typename Resp, typename Evt>
+class HTTPDelimitedJSONTransport : public JSONTransport<Req, Resp, Evt> {
 public:
-  HTTPDelimitedJSONTransport(lldb::IOObjectSP input, lldb::IOObjectSP output)
-      : JSONTransport(input, output) {}
-  virtual ~HTTPDelimitedJSONTransport() = default;
+  using JSONTransport<Req, Resp, Evt>::JSONTransport;
 
 protected:
-  llvm::Error WriteImpl(const std::string &message) override;
-  llvm::Expected<std::vector<std::string>> Parse() override;
+  /// Encodes messages based on
+  /// https://microsoft.github.io/debug-adapter-protocol/overview#base-protocol
+  std::string Encode(const llvm::json::Value &message) override {
+    std::string output;
+    std::string raw_message = llvm::formatv("{0}", message).str();
+    llvm::raw_string_ostream OS(output);
+    OS << kHeaderContentLength << kHeaderFieldSeparator << ' '
+       << std::to_string(raw_message.size()) << kEndOfHeader << raw_message;
+    return output;
+  }
+
+  /// Parses messages based on
+  /// https://microsoft.github.io/debug-adapter-protocol/overview#base-protocol
+  llvm::Expected<std::vector<std::string>> Parse() override {
+    std::vector<std::string> messages;
+    llvm::StringRef buffer = this->m_buffer;
+    while (buffer.contains(kEndOfHeader)) {
+      auto [headers, rest] = buffer.split(kEndOfHeader);
+      size_t content_length = 0;
+      // HTTP Headers are formatted like `<field-name> ':' [<field-value>]`.
+      for (const llvm::StringRef &header :
+           llvm::split(headers, kHeaderSeparator)) {
+        auto [key, value] = header.split(kHeaderFieldSeparator);
+        // 'Content-Length' is the only meaningful key at the moment. Others are
+        // ignored.
+        if (!key.equals_insensitive(kHeaderContentLength))
+          continue;
+
+        value = value.trim();
+        if (!llvm::to_integer(value, content_length, 10)) {
+          // Clear the buffer to avoid re-parsing this malformed message.
+          this->m_buffer.clear();
+          return llvm::createStringError(std::errc::invalid_argument,
+                                         "invalid content length: %s",
+                                         value.str().c_str());
+        }
+      }
+
+      // Check if we have enough data.
+      if (content_length > rest.size())
+        break;
+
+      llvm::StringRef body = rest.take_front(content_length);
+      buffer = rest.drop_front(content_length);
+      messages.emplace_back(body.str());
+      this->Logv("--> {0}", body);
+    }
+
+    // Store the remainder of the buffer for the next read callback.
+    this->m_buffer = buffer.str();
+
+    return std::move(messages);
+  }
 
   static constexpr llvm::StringLiteral kHeaderContentLength = "Content-Length";
   static constexpr llvm::StringLiteral kHeaderFieldSeparator = ":";
@@ -168,15 +266,31 @@ class HTTPDelimitedJSONTransport : public JSONTransport {
 };
 
 /// A transport class for JSON RPC.
-class JSONRPCTransport : public JSONTransport {
+template <typename Req, typename Resp, typename Evt>
+class JSONRPCTransport : public JSONTransport<Req, Resp, Evt> {
 public:
-  JSONRPCTransport(lldb::IOObjectSP input, lldb::IOObjectSP output)
-      : JSONTransport(input, output) {}
-  virtual ~JSONRPCTransport() = default;
+  using JSONTransport<Req, Resp, Evt>::JSONTransport;
 
 protected:
-  llvm::Error WriteImpl(const std::string &message) override;
-  llvm::Expected<std::vector<std::string>> Parse() override;
+  std::string Encode(const llvm::json::Value &message) override {
+    return llvm::formatv("{0}{1}", message, kMessageSeparator).str();
+  }
+
+  llvm::Expected<std::vector<std::string>> Parse() override {
+    std::vector<std::string> messages;
+    llvm::StringRef buf = this->m_buffer;
+    while (buf.contains(kMessageSeparator)) {
+      auto [raw_json, rest] = buf.split(kMessageSeparator);
+      buf = rest;
+      messages.emplace_back(raw_json.str());
+      this->Logv("--> {0}", raw_json);
+    }
+
+    // Store the remainder of the buffer for the next read callback.
+    this->m_buffer = buf.str();
+
+    return messages;
+  }
 
   static constexpr llvm::StringLiteral kMessageSeparator = "\n";
 };