From 6e28700ab1d876a9b01647782ce3c0ed4d8e0bb4 Mon Sep 17 00:00:00 2001
From: John Harrison <harjohn@google.com>
Date: Tue, 4 Mar 2025 19:09:28 +0100
Subject: [PATCH] [lldb-dap] Improving EOF handling on stream input and adding
 new unit tests (#129581)

This should improve the handling of EOF on stdin and adding some new
unit tests to malformed requests.
---
 lldb/test/API/tools/lldb-dap/io/TestDAP_io.py | 72 +++++++++++++++++++
 lldb/tools/lldb-dap/DAP.cpp                   |  6 +-
 lldb/tools/lldb-dap/IOStream.cpp              | 12 +++-
 3 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 lldb/test/API/tools/lldb-dap/io/TestDAP_io.py

diff --git a/lldb/test/API/tools/lldb-dap/io/TestDAP_io.py b/lldb/test/API/tools/lldb-dap/io/TestDAP_io.py
new file mode 100644
index 0000000000000..a39bd17ceb3b3
--- /dev/null
+++ b/lldb/test/API/tools/lldb-dap/io/TestDAP_io.py
@@ -0,0 +1,72 @@
+"""
+Test lldb-dap IO handling.
+"""
+
+from lldbsuite.test.decorators import *
+import lldbdap_testcase
+import dap_server
+
+
+class TestDAP_io(lldbdap_testcase.DAPTestCaseBase):
+    def launch(self):
+        log_file_path = self.getBuildArtifact("dap.txt")
+        process, _ = dap_server.DebugAdapterServer.launch(
+            executable=self.lldbDAPExec, log_file=log_file_path
+        )
+
+        def cleanup():
+            # If the process is still alive, terminate it.
+            if process.poll() is None:
+                process.terminate()
+                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 =========")
+                with open(log_file_path, "r") as file:
+                    print(file.read())
+                print("========= END =========")
+
+        # Execute the cleanup function during test case tear down.
+        self.addTearDownHook(cleanup)
+
+        return process
+
+    def test_eof_immediately(self):
+        """
+        lldb-dap handles EOF without any other input.
+        """
+        process = self.launch()
+        process.stdin.close()
+        self.assertEqual(process.wait(timeout=5.0), 0)
+
+    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)
+
+    def test_incorrect_content_length(self):
+        """
+        lldb-dap handles malformed content length headers.
+        """
+        process = self.launch()
+        process.stdin.write(b"Content-Length: abc")
+        process.stdin.close()
+        self.assertEqual(process.wait(timeout=5.0), 0)
+
+    def test_partial_content_length(self):
+        """
+        lldb-dap handles partial messages.
+        """
+        process = self.launch()
+        process.stdin.write(b"Content-Length: 10{")
+        process.stdin.close()
+        self.assertEqual(process.wait(timeout=5.0), 0)
diff --git a/lldb/tools/lldb-dap/DAP.cpp b/lldb/tools/lldb-dap/DAP.cpp
index 53c514b790f38..3dc9d6f5ca0a4 100644
--- a/lldb/tools/lldb-dap/DAP.cpp
+++ b/lldb/tools/lldb-dap/DAP.cpp
@@ -856,7 +856,11 @@ lldb::SBError DAP::Disconnect(bool terminateDebuggee) {
 }
 
 llvm::Error DAP::Loop() {
-  auto cleanup = llvm::make_scope_exit([this]() { StopEventHandlers(); });
+  auto cleanup = llvm::make_scope_exit([this]() {
+    if (output.descriptor)
+      output.descriptor->Close();
+    StopEventHandlers();
+  });
   while (!disconnecting) {
     llvm::json::Object object;
     lldb_dap::PacketStatus status = GetNextObject(object);
diff --git a/lldb/tools/lldb-dap/IOStream.cpp b/lldb/tools/lldb-dap/IOStream.cpp
index c6f1bfaf3b799..ee22a297ec248 100644
--- a/lldb/tools/lldb-dap/IOStream.cpp
+++ b/lldb/tools/lldb-dap/IOStream.cpp
@@ -35,16 +35,23 @@ bool InputStream::read_full(std::ofstream *log, size_t length,
   if (status.Fail())
     return false;
 
-  text += data;
+  text += data.substr(0, length);
   return true;
 }
 
 bool InputStream::read_line(std::ofstream *log, std::string &line) {
   line.clear();
   while (true) {
-    if (!read_full(log, 1, line))
+    std::string next;
+    if (!read_full(log, 1, next))
       return false;
 
+    // If EOF is encoutnered, '' is returned, break out of this loop.
+    if (next.empty())
+      return false;
+
+    line += next;
+
     if (llvm::StringRef(line).ends_with("\r\n"))
       break;
   }
@@ -60,6 +67,7 @@ bool InputStream::read_expected(std::ofstream *log, llvm::StringRef expected) {
     if (log)
       *log << "Warning: Expected '" << expected.str() << "', got '" << result
            << "\n";
+    return false;
   }
   return true;
 }