Harden writer of munged bitcode for fuzzing

lib/Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp

 //===--- Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp - Bitcode Munger -----===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Bitcode writer/munger implementation for testing.
@@ skipping 11 matching lines @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MemoryBuffer.h"
 
 #include <memory>
 
 using namespace llvm;
 
 // For debugging. When true, shows each test being run.
 static bool TraceTestRuns = false;
 
-void NaClBitcodeMunger::setupTest(
+bool NaClBitcodeMunger::setupTest(
     const char *TestName, const uint64_t Munges[], size_t MungesSize,
     bool AddHeader) {
   assert(DumpStream == nullptr && "Test run with DumpStream already defined");
   assert(MungedInput.get() == nullptr
          && "Test run with MungedInput already defined");
   FoundErrors = false;
   DumpResults.clear(); // Throw away any previous results.
   std::string DumpBuffer;
   DumpStream = new raw_string_ostream(DumpResults);
   MungedInputBuffer.clear();
 
   if (TraceTestRuns) {
     errs() << "*** Run test: " << TestName << "\n";
   }
 
   MungedBitcode.munge(Munges, MungesSize, RecordTerminator);
   WriteFlags.setErrStream(getDumpStream());
   NaClMungedBitcode::WriteResults Results =
       MungedBitcode.writeMaybeRepair(MungedInputBuffer, AddHeader, WriteFlags);
   if (Results.NumErrors != 0
       && !(WriteFlags.getTryToRecover()
            && Results.NumRepairs == Results.NumErrors)
       && !(WriteFlags.getWriteBadAbbrevIndex()
-           && Results.WroteBadAbbrevIndex && Results.NumErrors == 1))
-    report_fatal_error("Unable to generate bitcode file due to write errors");
+           && Results.WroteBadAbbrevIndex && Results.NumErrors == 1)) {
+    Error() << "Unable to generate bitcode file due to write errors\n";
+    return false;
+  }
 
   // Add null terminator, so that we meet the requirements of the
   // MemoryBuffer API.
   MungedInputBuffer.push_back('\0');
 
   MungedInput = MemoryBuffer::getMemBuffer(
       StringRef(MungedInputBuffer.data(), MungedInputBuffer.size()-1),
       TestName);
+  return true;
 }
 
-void NaClBitcodeMunger::cleanupTest() {
+bool NaClBitcodeMunger::cleanupTest() {
   RunAsDeathTest = false;
   WriteFlags.reset();
   MungedBitcode.removeEdits();
   MungedInput.reset();
   assert(DumpStream && "Dump stream removed before cleanup!");
   DumpStream->flush();
   delete DumpStream;
   DumpStream = nullptr;
+  return !FoundErrors;
 }
 
 // Return the next line of input (including eoln), starting from
 // Pos. Then increment Pos past the end of that line.
 static std::string getLine(const std::string &Input, size_t &Pos) {
   std::string Line;
   if (Pos >= Input.size()) {
     Pos = std::string::npos;
     return Line;
   }
@@ skipping 16 matching lines @@
     std::string Line = getLine(DumpResults, LastLinePos);
     if (LastLinePos == std::string::npos) break;
     size_t Pos = Line.find(Substring);
     if (Pos != std::string::npos && (!MustBePrefix || Pos == 0)) {
       Messages.append(Line);
     }
   }
   return Messages;
 }
 
+bool NaClWriteMunger::runTest(const char* Name, const uint64_t Munges[],
+                              size_t MungesSize) {
+  bool AddHeader = true;
+  if (!setupTest(Name, Munges, MungesSize, AddHeader))
+    return cleanupTest();
+  MemoryBufferRef InputRef(MungedInput->getMemBufferRef());
+  NaClMungedBitcode WrittenBitcode(MemoryBuffer::getMemBuffer(InputRef));
+  WrittenBitcode.print(WriteFlags.getErrStream());
+  return cleanupTest();
+}
+
 bool NaClObjDumpMunger::runTestWithFlags(
     const char *Name, const uint64_t Munges[], size_t MungesSize,
     bool AddHeader, bool NoRecords, bool NoAssembly) {
-  setupTest(Name, Munges, MungesSize, AddHeader);
+  if (!setupTest(Name, Munges, MungesSize, AddHeader))
+    return cleanupTest();
 
   if (NaClObjDump(MungedInput.get()->getMemBufferRef(),
                   getDumpStream(), NoRecords, NoAssembly))
     FoundErrors = true;
-  cleanupTest();
-  return !FoundErrors;
+  return cleanupTest();
 }
 
 bool NaClParseBitcodeMunger::runTest(
     const char *Name, const uint64_t Munges[], size_t MungesSize,
     bool VerboseErrors) {
   bool AddHeader = true;
-  setupTest(Name, Munges, MungesSize, AddHeader);
+  if (!setupTest(Name, Munges, MungesSize, AddHeader))
+    return cleanupTest();
+
   LLVMContext &Context = getGlobalContext();
+
   raw_ostream *VerboseStrm = VerboseErrors ? &getDumpStream() : nullptr;
   ErrorOr<Module *> ModuleOrError =
       NaClParseBitcodeFile(MungedInput->getMemBufferRef(), Context,
                            VerboseStrm);
   if (ModuleOrError) {
     if (VerboseErrors)
       getDumpStream() << "Successful parse!\n";
     delete ModuleOrError.get();
   } else {
     Error() << ModuleOrError.getError().message() << "\n";
   }
-  cleanupTest();
-  return !FoundErrors;
+  return cleanupTest();
 }
 
 bool NaClCompressMunger::runTest(const char* Name, const uint64_t Munges[],
                                  size_t MungesSize) {
   bool AddHeader = true;
-  setupTest(Name, Munges, MungesSize, AddHeader);
+  if (!setupTest(Name, Munges, MungesSize, AddHeader))
+    return cleanupTest();
+
   NaClBitcodeCompressor Compressor;
-  bool Result = Compressor.compress(MungedInput.get(), getDumpStream());
-  cleanupTest();
-  return Result;
+  if (!Compressor.compress(MungedInput.get(), getDumpStream()))
+    Error() << "Unable to compress\n";
+  return cleanupTest();
 }