Harden writer of munged bitcode for fuzzing

unittests/Bitcode/NaClMungedIoTest.cpp

 //===- llvm/unittest/Bitcode/NaClMungedIoTest.cpp -------------------------===//
 //     Tests munging NaCl bitcode records.
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
@@ skipping 20 matching lines @@
   MungedBitcode.print(StrBuf);
   return StrBuf.str();
 }
 
 typedef SmallVector<char, 1024> TextBuffer;
 
 // Writes out a sequence of munged bitcode records, and writes them into
 // the text buffer. Returns a corresponding memory buffer containing
 // the munged bitcode records.
 std::unique_ptr<MemoryBuffer> writeMungedBitcode(
-    NaClMungedBitcode &Bitcode, TextBuffer &Buffer) {
-  Bitcode.write(Buffer, /* AddHeader = */ true);
+    NaClMungedBitcode &Bitcode, TextBuffer &Buffer,
+    NaClMungedBitcode::WriteFlags &Flags) {
+  Bitcode.write(Buffer, /* AddHeader = */ true, Flags);
   StringRef Input(Buffer.data(), Buffer.size());
   return MemoryBuffer::getMemBuffer(Input, "Test", false);
 }
 
+std::unique_ptr<MemoryBuffer> writeMungedBitcode(
+    NaClMungedBitcode &Bitcode, TextBuffer &Buffer) {
+  NaClMungedBitcode::WriteFlags Flags;
+  return writeMungedBitcode(Bitcode, Buffer, Flags);
+}
+
 // Write out the bitcode, parse it back, and return the resulting
 // munged bitcode.
 std::string parseWrittenMungedBitcode(NaClMungedBitcode &OutBitcode) {
   TextBuffer Buffer;
   NaClMungedBitcode InBitcode(writeMungedBitcode(OutBitcode, Buffer));
   return stringify(InBitcode);
 }
 
 // Sample toy bitcode records.
 const uint64_t Records[] = {
@@ skipping 96 matching lines @@
   for (size_t i = 2, e = Bitcode.getBaseRecords().size(); i < e; ++i)
     Bitcode.remove(i);
 
   // The expected output when stringifying this input.
   EXPECT_EQ(
       "       1: [65535, 8, 2]\n"
       "         3: [1, 1]\n",
       stringify(Bitcode));
 
   // Show that we can't write the bitcode correctly.
-  TextBuffer Buffer;
-  EXPECT_DEATH(writeMungedBitcode(Bitcode, Buffer),
-               ".*Unflushed data remaining.*");
+  TextBuffer WriteBuffer;
+  std::string LogBuffer;
+  raw_string_ostream StrBuf(LogBuffer);
+  NaClMungedBitcode::WriteFlags Flags;
+  Flags.setErrStream(StrBuf);
+  writeMungedBitcode(Bitcode, WriteBuffer, Flags);
+  EXPECT_EQ(
+      "Error (Block 8): Missing close block.\n",
+      StrBuf.str());
 }
 
 } // end of anonymous namespace