[lldb] Add a Stream variant that escapes null bytes

[JDevlieghere]

Summary

In D122025 Jim suggested that we should sanitize the input to the command return object to retain the guarantee that we don't have null bytes in the output. This patch implements that through a Stream wrapper that escapes null bytes.

[jingham]

Nice! Thanks for doing this.

[labath]

Although I can't say exactly what would be the problem with this, this seems like a strange thing to do, and I'm not aware of any other stream class which would handle nul (and _only_ nul) characters specially.

lldb/include/lldb/Interpreter/CommandReturnObject.h

public:
 
   bool GetSuppressImmediateOutput() const;
 
   void SetSuppressImmediateOutput(bool b);
 
 private:
   enum { eStreamStringIndex = 0, eImmediateStreamIndex = 1 };
 
-  StreamTee m_out_stream;
-  StreamTee m_err_stream;
+  NullEscaped<StreamTee> m_out_stream;
+  NullEscaped<StreamTee> m_err_stream;
 
   lldb::ReturnStatus m_status = lldb::eReturnStatusStarted;
 
   bool m_did_change_process_state = false;
   bool m_suppress_immediate_output = false;
 
   /// If true, then the input handle from the debugger will be hooked up.
   bool m_interactive = true;
 };
 
 } // namespace lldb_private
 
 #endif // LLDB_INTERPRETER_COMMANDRETURNOBJECT_H

lldb/include/lldb/Utility/Stream.h

public:
     RawOstreamForward(Stream &target, bool colors = false)
         : llvm::raw_ostream(/*unbuffered*/ true), m_target(target) {
       enable_colors(colors);
     }
   };
   RawOstreamForward m_forwarder;
 };
 
+template <typename T> class NullEscaped : public T {
+public:
+  using T::T;
+  ~NullEscaped() = default;
+
+  static_assert(std::is_base_of<Stream, T>::value, "T must be a Stream");
+
+protected:
+  size_t WriteImpl(const void *s, size_t length) override {
+    llvm::StringRef input(static_cast<const char *>(s), length);
+    llvm::SmallVector<llvm::StringRef, 1> parts;
+    input.split(parts, '\0');
+    size_t written = T::WriteImpl(parts[0].data(), parts[0].size());
+    for (llvm::StringRef part : llvm::drop_begin(parts)) {
+      written += T::WriteImpl("\\0", 2);
+      written += T::WriteImpl(part.data(), part.size());
+    }
+    return written;
+  }
+};
+
 /// Output an address value to this stream.
 ///
 /// Put an address \a addr out to the stream with optional \a prefix and \a
 /// suffix strings.
 ///
 /// \param[in] s
 ///     The output stream.
 ///

lldb/unittests/Utility/StreamTest.cpp

struct StreamTest : ::testing::Test {
   // We return here a std::string because that way gtest can print better
   // assertion messages.
   std::string TakeValue() {
     std::string result = s.GetString().str();
     s.Clear();
     return result;
   }
 };
+
+struct NullEscapedTest : ::testing::Test {
+  NullEscaped<StreamString> s;
+  std::string TakeValue() {
+    std::string result = s.GetString().str();
+    s.Clear();
+    return result;
   }
+};
+} // namespace
 
 namespace {
 // A StreamTest where we expect the Stream output to be binary.
 struct BinaryStreamTest : StreamTest {
-  void SetUp() override {
-    s.GetFlags().Set(Stream::eBinary);
-  }
+  void SetUp() override { s.GetFlags().Set(Stream::eBinary); }
 };
-}
+} // namespace
 
 TEST_F(StreamTest, AddressPrefix) {
   DumpAddress(s.AsRawOstream(), 0x1, 1, "foo");
   EXPECT_EQ("foo0x01", TakeValue());
 }
 
 TEST_F(StreamTest, AddressEmptyPrefix) {
   DumpAddress(s.AsRawOstream(), 0x1, 1, nullptr);

 }
 
 TEST_F(StreamTest, PutSLEB128) {
   auto bytes = s.PutSLEB128(0x1985LL);
   EXPECT_EQ(6U, s.GetWrittenBytes());
   EXPECT_EQ("0x6533", TakeValue());
   EXPECT_EQ(6U, bytes);
 }
+
+TEST_F(NullEscapedTest, NonNull) {
+  s.PutChar(' ');
+  EXPECT_EQ(1U, s.GetWrittenBytes());
+  EXPECT_EQ(" ", TakeValue());
+
+  s.PutCString(llvm::StringRef("\0", 1));
+  EXPECT_EQ(2U, s.GetWrittenBytes());
+  EXPECT_EQ("\\0", TakeValue());
+
+  s.PutCString(llvm::StringRef("foo\0bar\0baz", 11));
+  EXPECT_EQ(13U, s.GetWrittenBytes());
+  EXPECT_EQ("foo\\0bar\\0baz", TakeValue());
+
+  s.PutCString(llvm::StringRef("foo\0bar\0baz\0", 12));
+  EXPECT_EQ(15U, s.GetWrittenBytes());
+  EXPECT_EQ("foo\\0bar\\0baz\\0", TakeValue());
+}