Adjusting Transport.Read to return an llvm::Expected<std::optional<protocol::Message>> to try to refine the caller.

ashgti · ashgti · commit 7d9eeb178bba · 2025-03-11T13:08:14.000-07:00
An error reading or validating a message should now result in an `llvm::Error`.

EOF is now only acceptable at the start of a message, otherwise we return an error indiciating we encountered a partial message.
diff --git a/lldb/test/API/tools/lldb-dap/io/TestDAP_io.py b/lldb/test/API/tools/lldb-dap/io/TestDAP_io.py
@@ -2,6 +2,8 @@
 Test lldb-dap IO handling.
 """
 
+import sys
+
 from lldbsuite.test.decorators import *
 import lldbdap_testcase
 import dap_server
@@ -19,18 +21,18 @@ def cleanup():
             if process.poll() is None:
                 process.terminate()
                 process.wait()
-            stdout_data = process.stdout.read()
-            stderr_data = process.stderr.read()
-            print("========= STDOUT =========")
-            print(stdout_data)
-            print("========= END =========")
-            print("========= STDERR =========")
-            print(stderr_data)
-            print("========= END =========")
-            print("========= DEBUG ADAPTER PROTOCOL LOGS =========")
+            stdout_data = process.stdout.read().decode()
+            stderr_data = process.stderr.read().decode()
+            print("========= STDOUT =========", file=sys.stderr)
+            print(stdout_data, file=sys.stderr)
+            print("========= END =========", file=sys.stderr)
+            print("========= STDERR =========", file=sys.stderr)
+            print(stderr_data, file=sys.stderr)
+            print("========= END =========", file=sys.stderr)
+            print("========= DEBUG ADAPTER PROTOCOL LOGS =========", file=sys.stderr)
             with open(log_file_path, "r") as file:
-                print(file.read())
-            print("========= END =========")
+                print(file.read(), file=sys.stderr)
+            print("========= END =========", file=sys.stderr)
 
         # Execute the cleanup function during test case tear down.
         self.addTearDownHook(cleanup)
@@ -45,14 +47,23 @@ def test_eof_immediately(self):
         process.stdin.close()
         self.assertEqual(process.wait(timeout=5.0), 0)
 
+    def test_invalid_header(self):
+        """
+        lldb-dap handles invalid message headers.
+        """
+        process = self.launch()
+        process.stdin.write(b"not the corret message header")
+        process.stdin.close()
+        self.assertEqual(process.wait(timeout=5.0), 1)
+
     def test_partial_header(self):
         """
         lldb-dap handles parital message headers.
         """
         process = self.launch()
         process.stdin.write(b"Content-Length: ")
         process.stdin.close()
-        self.assertEqual(process.wait(timeout=5.0), 0)
+        self.assertEqual(process.wait(timeout=5.0), 1)
 
     def test_incorrect_content_length(self):
         """
@@ -61,13 +72,13 @@ def test_incorrect_content_length(self):
         process = self.launch()
         process.stdin.write(b"Content-Length: abc")
         process.stdin.close()
-        self.assertEqual(process.wait(timeout=5.0), 0)
+        self.assertEqual(process.wait(timeout=5.0), 1)
 
     def test_partial_content_length(self):
         """
         lldb-dap handles partial messages.
         """
         process = self.launch()
-        process.stdin.write(b"Content-Length: 10{")
+        process.stdin.write(b"Content-Length: 10\r\n\r\n{")
         process.stdin.close()
-        self.assertEqual(process.wait(timeout=5.0), 0)
+        self.assertEqual(process.wait(timeout=5.0), 1)
diff --git a/lldb/tools/lldb-dap/DAP.cpp b/lldb/tools/lldb-dap/DAP.cpp
@@ -769,11 +769,15 @@ llvm::Error DAP::Loop() {
     StopEventHandlers();
   });
   while (!disconnecting) {
-    std::optional<protocol::Message> next = transport.Read();
+    llvm::Expected<std::optional<protocol::Message>> next = transport.Read();
     if (!next)
+      return next.takeError();
+
+    // nullopt on EOF
+    if (!*next)
       break;
 
-    if (!HandleObject(*next)) {
+    if (!HandleObject(**next)) {
       return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                      "unhandled packet");
     }
diff --git a/lldb/tools/lldb-dap/Transport.cpp b/lldb/tools/lldb-dap/Transport.cpp
@@ -25,16 +25,19 @@ using namespace lldb_private;
 using namespace lldb_dap;
 using namespace lldb_dap::protocol;
 
-static Expected<std::string> ReadFull(IOObject *descriptor, size_t length) {
+/// ReadFull attempts to read the specified number of bytes. If EOF is
+/// encountered, '' is returned.
+static Expected<std::string> ReadFull(IOObject &descriptor, size_t length) {
   std::string data;
   data.resize(length);
-  auto status = descriptor->Read(data.data(), length);
+  auto status = descriptor.Read(data.data(), length);
   if (status.Fail())
     return status.takeError();
+  // Return the actual number of bytes read.
   return data.substr(0, length);
 }
 
-static Expected<std::string> ReadUntil(IOObject *descriptor,
+static Expected<std::string> ReadUntil(IOObject &descriptor,
                                        StringRef delimiter) {
   std::string buffer;
   buffer.reserve(delimiter.size() + 1);
@@ -43,9 +46,9 @@ static Expected<std::string> ReadUntil(IOObject *descriptor,
         ReadFull(descriptor, buffer.empty() ? delimiter.size() : 1);
     if (auto Err = next.takeError())
       return std::move(Err);
-    // '' is returned on EOF.
+    // Return "" if EOF is encountered.
     if (next->empty())
-      return buffer;
+      return "";
     buffer += *next;
   }
   return buffer.substr(0, buffer.size() - delimiter.size());
@@ -65,65 +68,44 @@ Transport::Transport(StringRef client_name, std::ofstream *log,
     : m_client_name(client_name), m_log(log), m_input(std::move(input)),
       m_output(std::move(output)) {}
 
-std::optional<Message> Transport::Read() {
-  if (!m_input || !m_input->IsValid()) {
-    DAP_LOG(m_log, "({0}) input is closed", m_client_name);
-    return std::nullopt;
-  }
+Expected<std::optional<Message>> Transport::Read() {
+  if (!m_input || !m_input->IsValid())
+    return createStringError("transport output is closed");
+
   IOObject *input = m_input.get();
   Expected<std::string> message_header =
-      ReadFull(input, kHeaderContentLength.size());
-  if (!message_header) {
-    DAP_LOG_ERROR(m_log, message_header.takeError(), "({1}) read failed: {0}",
-                  m_client_name);
-    return std::nullopt;
-  }
-
+      ReadFull(*input, kHeaderContentLength.size());
+  if (!message_header)
+    return message_header.takeError();
   // '' returned on EOF.
   if (message_header->empty())
     return std::nullopt;
-  if (*message_header != kHeaderContentLength) {
-    DAP_LOG(m_log, "({0}) read failed: expected '{1}' and got '{2}'",
-            m_client_name, kHeaderContentLength, *message_header);
-    return std::nullopt;
-  }
+  if (*message_header != kHeaderContentLength)
+    return createStringError(formatv("expected '{0}' and got '{1}'",
+                                     kHeaderContentLength, *message_header)
+                                 .str());
 
-  Expected<std::string> raw_length = ReadUntil(input, kHeaderSeparator);
-  if (!raw_length) {
-    DAP_LOG_ERROR(m_log, raw_length.takeError(), "({1}) read failed: {0}",
-                  m_client_name);
-    return std::nullopt;
-  }
+  Expected<std::string> raw_length = ReadUntil(*input, kHeaderSeparator);
+  if (!raw_length)
+    return raw_length.takeError();
+  if (raw_length->empty())
+    return createStringError("unexpected EOF parsing DAP header");
 
   size_t length;
-  if (!to_integer(*raw_length, length)) {
-    DAP_LOG(m_log, "({0}) read failed: invalid content length {1}",
-            m_client_name, *raw_length);
-    return std::nullopt;
-  }
+  if (!to_integer(*raw_length, length))
+    return createStringError(
+        formatv("invalid content length {0}", *raw_length).str());
 
-  Expected<std::string> raw_json = ReadFull(input, length);
-  if (!raw_json) {
-    DAP_LOG_ERROR(m_log, raw_json.takeError(), "({1}) read failed: {0}",
-                  m_client_name);
-    return std::nullopt;
-  }
-  if (raw_json->length() != length) {
-    DAP_LOG(m_log, "({0}) read failed: expected {1} bytes and got {2} bytes",
-            m_client_name, length, raw_json->length());
-    return std::nullopt;
-  }
+  Expected<std::string> raw_json = ReadFull(*input, length);
+  if (!raw_json)
+    return raw_json.takeError();
+  // If we got less than the expected number of bytes then we hit EOF.
+  if (raw_json->length() != length)
+    return createStringError("unexpected EOF parse DAP message body");
 
   DAP_LOG(m_log, "<-- ({0}) {1}", m_client_name, *raw_json);
 
-  llvm::Expected<Message> message = json::parse<Message>(*raw_json);
-  if (!message) {
-    DAP_LOG_ERROR(m_log, message.takeError(), "({1}) read failed: {0}",
-                  m_client_name);
-    return std::nullopt;
-  }
-
-  return std::move(*message);
+  return json::parse<Message>(*raw_json);
 }
 
 Error Transport::Write(const Message &message) {
diff --git a/lldb/tools/lldb-dap/Transport.h b/lldb/tools/lldb-dap/Transport.h
@@ -41,7 +41,7 @@ class Transport {
   llvm::Error Write(const protocol::Message &M);
 
   /// Reads the next Debug Adater Protocol message from the input stream.
-  std::optional<protocol::Message> Read();
+  llvm::Expected<std::optional<protocol::Message>> Read();
 
 private:
   llvm::StringRef m_client_name;
diff --git a/lldb/tools/lldb-dap/lldb-dap.cpp b/lldb/tools/lldb-dap/lldb-dap.cpp
@@ -345,7 +345,8 @@ serveConnection(const Socket::SocketProtocol &protocol, const std::string &name,
 
       if (auto Err = dap.Loop()) {
         llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(),
-                                    "DAP session error: ");
+                                    "DAP session (" + client_name +
+                                        ") error: ");
       }
 
       DAP_LOG(log, "({0}) client disconnected", client_name);
@@ -587,7 +588,7 @@ int main(int argc, char *argv[]) {
 
   if (auto Err = dap.Loop()) {
     llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(),
-                                "DAP session error: ");
+                                "DAP session (" + client_name + ") error: ");
     return EXIT_FAILURE;
   }
   return EXIT_SUCCESS;
