[lldb] Refactoring JSONTransport into an abstract RPC Message Handler and transport layer.

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";
 };