Harden writer of munged bitcode for fuzzing

lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp

Index: lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
diff --git a/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp b/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
index d036c15c53848c3c674eaa18ebf7da6078a2340b..c2107db6dac6d4f96687e4cf3f3fff636a194a43 100644
--- a/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
+++ b/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
@@ -14,35 +14,54 @@
 
 #include "llvm/Bitcode/NaCl/NaClBitstreamWriter.h"
 #include "llvm/Bitcode/NaCl/NaClReaderWriter.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
 namespace {
 
-// For debugging. When true, shows each PNaCl record that is
-// emitted to the bitcode file.
-static bool DebugEmitRecord = false;
+// For debugging. When true, shows trace information of emitting
+// bitcode.
+static bool DebugEmit = false;
+
+// Description of current block scope
+struct BlockScope {
+  BlockScope(unsigned CurBlockID, size_t AbbrevIndexLimit)
+      : CurBlockID(CurBlockID), AbbrevIndexLimit(AbbrevIndexLimit) {}
+  unsigned CurBlockID;
+  // Defines the maximun value for abbreviation indices in block.

[jvoung (off chromium), 2015/05/13 00:58:16]

maximun -> maximum

[Karl, 2015/05/13 21:39:24]

Done.

+  size_t AbbrevIndexLimit;
+  void print(raw_ostream &Out) const {
+    Out << "BlockScope(ID=" << CurBlockID << ", AbbrevIndexLimit="
+        << AbbrevIndexLimit << ")";
+  }
+};
+
+inline raw_ostream &operator<<(raw_ostream &Out, const BlockScope &Scope) {
+  Scope.print(Out);
+  return Out;
+}
 
 // State of bitcode writing.
 struct WriteState {
-  // The block ID associated with the block being written.
-  int WriteBlockID = -1;
+  // The block ID associated with records not in any block.
+  static const unsigned UnknownWriteBlockID = UINT_MAX;
   // The SetBID for the blockinfo block.
-  int SetBID = -1;
-  // The stack of maximum abbreviation indices allowed by block enter record.
-  SmallVector<uint64_t, 3> AbbrevIndexLimitStack;
+  unsigned SetBID = UnknownWriteBlockID;
+  // The stack of scopes the writer is in.
+  SmallVector<BlockScope, 3> ScopeStack;
   // The set of write flags to use.
   const NaClMungedBitcode::WriteFlags &Flags;
   // The results of the attempted write.
   NaClMungedBitcode::WriteResults Results;
 
-  WriteState(const NaClMungedBitcode::WriteFlags &Flags) : Flags(Flags) {}
+  WriteState(const NaClMungedBitcode::WriteFlags &Flags) : Flags(Flags) {
+    BlockScope Scope(UnknownWriteBlockID, naclbitc::DEFAULT_MAX_ABBREV);
+    ScopeStack.push_back(Scope);
+  }
 
   // Returns stream to print error message to.
-  raw_ostream &Error() {
-    ++Results.NumErrors;
-    return Flags.getErrStream() << "Error (Block " << WriteBlockID << "): ";
-  }
+  raw_ostream &Error();
 
   // Returns stream to print error message to, assuming that
   // the error message can be repaired if Flags.TryToRecover is true.
@@ -52,6 +71,21 @@ struct WriteState {
     return Error();
   }
 
+  bool atOutermostScope() {
+    assert(!ScopeStack.empty());
+    return ScopeStack.size() == 1;
+  }
+
+  unsigned getCurWriteBlockID() const {
+    assert(!ScopeStack.empty());
+    return ScopeStack.back().CurBlockID;
+  }
+
+  unsigned getCurAbbrevIndexLimit() const {
+    assert(!ScopeStack.empty());
+    return ScopeStack.back().AbbrevIndexLimit;
+  }
+
   // Converts the abbreviation record to the corresponding abbreviation.
   // Returns nullptr if unable to build abbreviation.
   NaClBitCodeAbbrev *buildAbbrev(const NaClBitcodeAbbrevRecord &Record);
@@ -61,19 +95,92 @@ struct WriteState {
   bool emitRecord(NaClBitstreamWriter &Writer,
                   const NaClBitcodeAbbrevRecord &Record);
 
-  // Adds any missing end blocks to written bitcode.
-  void writeMissingEndBlocks(NaClBitstreamWriter &Writer) {
-    while (!AbbrevIndexLimitStack.empty()) {
-      Writer.ExitBlock();
-      AbbrevIndexLimitStack.pop_back();
-    }
+  // Enter the given block
+  bool enterBlock(NaClBitstreamWriter &Writer, unsigned WriteBlockID,
+                  unsigned NumBits, const NaClBitcodeAbbrevRecord &Record);
+
+  // Exit current block and return to previous block. Silently recovers
+  // if at outermost scope.
+  void exitBlock(NaClBitstreamWriter &Writer) {
+    Writer.ExitBlock();
+    ScopeStack.pop_back();

[jvoung (off chromium), 2015/05/13 00:58:17]

Hmm, I don't see how this prevents popping too much? ScopeStack is a plain
SmallVector?

[Karl, 2015/05/13 21:39:24]

Hmm, I'm confused too. Earlier versions had the check.

Looking more closely, I see that the callers guarantee that this won't happen.
However, a better solution is to put it here, and generate an error if there is
a mismatch.

+    if (DebugEmit)
+      printScopeStack(errs());
+  }
+
+  // Completes the write.
+  NaClMungedBitcode::WriteResults &finish(NaClBitstreamWriter &Writer,
+                                          bool RecoverSilently = false);
+
+  void printScopeStack(raw_ostream &Out) {
+    Out << "Scope Stack:\n";
+    for (auto Scope : ScopeStack)

[jvoung (off chromium), 2015/05/13 00:58:17]

auto &Scope ? Not sure if this is one of those cases where it might have made a
copy instead.

[Karl, 2015/05/13 21:39:23]

Done.

+      Out << "  " << Scope << "\n";
   }
 };
 
+raw_ostream &WriteState::Error() {
+  ++Results.NumErrors;
+  raw_ostream &ErrStrm = Flags.getErrStream();
+  unsigned WriteBlockID = getCurWriteBlockID();
+  ErrStrm << "Error (Block ";
+  if (WriteBlockID == UnknownWriteBlockID)
+    ErrStrm << "unknown";
+  else
+    ErrStrm << WriteBlockID;
+  return ErrStrm <<  "): ";
+}
+
+bool WriteState::enterBlock(NaClBitstreamWriter &Writer, unsigned WriteBlockID,
+                            unsigned NumBits,
+                            const NaClBitcodeAbbrevRecord &Record) {
+  uint64_t MaxAbbrev = (static_cast<uint64_t>(1) << NumBits) - 1;

[jvoung (off chromium), 2015/05/13 00:58:17]

maybe assert NumBits < 64 or some smaller limit -- unless it's checked somewhere
else

[Karl, 2015/05/13 21:39:23]

Good point. The code reader/writer assumes 32. Adding test.

+  BlockScope Scope(WriteBlockID, MaxAbbrev);
+  ScopeStack.push_back(Scope);
+  if (DebugEmit)
+    printScopeStack(errs());
+  if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID) {
+    unsigned DefaultMaxBits =
+        NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV);
+    if (NumBits != DefaultMaxBits) {
+      RecoverableError()
+          << "Numbits entry for abbreviations record not "
+          << DefaultMaxBits << " but found " << NumBits <<
+          ": " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
+    }
+    Writer.EnterBlockInfoBlock();
+  } else {
+    NaClBitcodeSelectorAbbrev CurCodeLen(MaxAbbrev);
+    Writer.EnterSubblock(WriteBlockID, CurCodeLen);
+  }
+  return true;
+}
+
+NaClMungedBitcode::WriteResults & WriteState::finish(

[jvoung (off chromium), 2015/05/13 00:58:17]

flush the & in the return type to the right

[Karl, 2015/05/13 21:39:24]

Done.

+    NaClBitstreamWriter &Writer, bool RecoverSilently) {
+  // Be sure blocks are balanced.
+  while (!atOutermostScope()) {
+    if (!RecoverSilently)
+      RecoverableError() << "Missing close block.\n";
+    exitBlock(Writer);
+  }
+
+  // Be sure that generated bitcode buffer is word aligned.
+  if (Writer.GetCurrentBitNo() % 4 * CHAR_BIT) {
+    if (!RecoverSilently)
+      RecoverableError() << "Written bitstream not word aligned\n";
+    // Force a repair so that the bitstream writer doesn't crash.
+    Writer.FlushToWord();
+  }
+  return Results;
+}
+
 bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
                             const NaClBitcodeAbbrevRecord &Record) {
   size_t NumValues = Record.Values.size();
-  if (DebugEmitRecord) {
+  if (DebugEmit) {
     errs() << "Emit " << Record.Abbrev << ": <" << Record.Code;
     for (size_t i = 0; i < NumValues; ++i) {
       errs() << ", " << Record.Values[i];
@@ -83,8 +190,9 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
 
   switch (Record.Code) {
   case naclbitc::BLK_CODE_ENTER: {
-    unsigned NumBits = NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV);
-    WriteBlockID = -1;
+    unsigned MinNumBits = NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV);
+    unsigned WriteBlockID = UnknownWriteBlockID;
+    unsigned NumBits = MinNumBits;
     if (Record.Abbrev != naclbitc::ENTER_SUBBLOCK) {
       RecoverableError()
           << "Uses illegal abbreviation index in enter block record: "
@@ -92,43 +200,49 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       if (!Flags.getTryToRecover())
         return false;
     }
-    if (NumValues == 2) {
-      WriteBlockID = Record.Values[0];
-      NumBits = Record.Values[1];
-      if (NumBits > 32 || NumBits < 2) {
-        RecoverableError()
-            << "Bit size " << NumBits << " for record should be "
-            << (NumBits > 32 ? "<= 32" : ">= 2") << ": " << Record << "\n";
-        if (!Flags.getTryToRecover())
-          return false;
-        NumBits = 32;
-      }
-    } else {
-      Error() << "Values for enter record should be of size 2, but found "
-              << NumValues << ": " << Record << "\n";
-      return false;
+    if (NumValues != 2) {
+      RecoverableError()
+          << "Values for enter record should be of size 2, but found "
+          << NumValues << ": " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
     }
-    uint64_t MaxAbbrev = (static_cast<uint64_t>(1) << NumBits) - 1;
-    AbbrevIndexLimitStack.push_back(MaxAbbrev);
-    if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID) {
-      unsigned DefaultMaxBits =
-          NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV);
-      if (NumBits != DefaultMaxBits) {
-        RecoverableError()
-            << "Numbits entry for abbreviations record not "
-            << DefaultMaxBits << " but found " << NumBits <<
-            ": " << Record << "\n";
-        if (!Flags.getTryToRecover())
-          return false;
-      }
-      Writer.EnterBlockInfoBlock();
-    } else {
-      NaClBitcodeSelectorAbbrev CurCodeLen(MaxAbbrev);
-      Writer.EnterSubblock(WriteBlockID, CurCodeLen);
+    if (NumValues == 0)
+      return enterBlock(Writer, WriteBlockID, NumBits, Record);
+    WriteBlockID = Record.Values[0];
+    if (Record.Values[0] > UINT_MAX) {
+      RecoverableError() << "Block id must be less than << " << UINT_MAX
+                         << ": " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
+      WriteBlockID = UnknownWriteBlockID;
+      enterBlock(Writer, WriteBlockID, NumBits, Record);

[jvoung (off chromium), 2015/05/13 00:58:17]

return enterBlock(...);

To avoid falling through and doing another enterBlock ()

[jvoung (off chromium), 2015/05/13 00:58:17]

At this point, it's still possible for NumValues == 2, so technically you could
have set "NumBits = Record.Values[1];"

[Karl, 2015/05/13 21:39:24]

Removing call to writeBlock(). Rather, let the code fall through to get NumBits

     }
-    break;
+    if (NumValues < 2)
+      return enterBlock(Writer, WriteBlockID, NumBits, Record);
+    NumBits = Record.Values[1];
+    if (NumBits > 32 || MinNumBits < 2) {

[jvoung (off chromium), 2015/05/13 00:58:17]

"MinNumBits < 2" MinNumBits is pretty much a constant, so I'm not sure that
makes sense to check...

Should that be "NumBits < MinNumBits" ?

[Karl, 2015/05/13 21:39:24]

Good catch. Using MinNumBits.

+      raw_ostream &ErrStrm = RecoverableError();
+      ErrStrm << "Bit size " << NumBits << " for record should be ";
+      if (NumBits > 32)
+        ErrStrm << "<= 32";
+      else
+        ErrStrm << ">= " << MinNumBits;
+      ErrStrm << ": " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
+      NumBits = 32;
+    }
+    return enterBlock(Writer, WriteBlockID, NumBits, Record);
   }
   case naclbitc::BLK_CODE_EXIT:
+    if (atOutermostScope()) {
+      RecoverableError()
+          << "Extraneous exit block: " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
+      break;
+    }
     if (Record.Abbrev != naclbitc::END_BLOCK) {
       RecoverableError()
           << "Uses illegal abbreviation index in exit block record: "
@@ -142,9 +256,7 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       if (!Flags.getTryToRecover())
         return false;
     }
-    if (!AbbrevIndexLimitStack.empty())
-      AbbrevIndexLimitStack.pop_back();
-    Writer.ExitBlock();
+    exitBlock(Writer);

[jvoung (off chromium), 2015/05/13 22:24:05]

btw: I was about to suggest adding LLVM_ATTRIBUTE_UNUSED_RESULT to the new
methods with bool results, to get a warning if the return value is mistakenly
ignored to help catch some of these + the previous case where you had
enterBlock() without the return and just a fallthrough.

Not sure number of false positives vs the benefit, though, so up to you.

[Karl, 2015/05/14 17:08:11]

Added and fixed code (3 cases). In all three, they shouldn't happen, but forced
a non-recoverable error should they fail for some reason. This makes the code
safer, in case code gets modified and the conditions are no longer assured.

     break;
   case naclbitc::BLK_CODE_DEFINE_ABBREV: {
     if (Record.Abbrev != naclbitc::DEFINE_ABBREV) {
@@ -154,10 +266,17 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       if (!Flags.getTryToRecover())
         return false;
     }
+    if (getCurWriteBlockID() != naclbitc::BLOCKINFO_BLOCK_ID
+        && Writer.getMaxCurAbbrevIndex() >=  getCurAbbrevIndexLimit()) {

[jvoung (off chromium), 2015/05/13 00:58:17]

">=  "
">= " 

[Karl, 2015/05/13 21:39:23]

Done.

[jvoung (off chromium), 2015/05/13 22:24:06]

Ping =) >=

[Karl, 2015/05/14 17:08:11]

Done.

+      RecoverableError() << "Exceeds abbreviation index limit of "
+                         << getCurAbbrevIndexLimit() << ": " << Record << "\n";
+      // Recover by not writing.
+      return Flags.getTryToRecover();
+    }
     NaClBitCodeAbbrev *Abbrev = buildAbbrev(Record);
     if (Abbrev == NULL)
       return Flags.getTryToRecover();
-    if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID) {
+    if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID) {
       Writer.EmitBlockInfoAbbrev(SetBID, Abbrev);
     } else {
       Writer.EmitAbbrev(Abbrev);
@@ -170,24 +289,27 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       Writer.Emit(Value, 8);
     break;
   default:
-    if (AbbrevIndexLimitStack.empty()) {
-      Error() << "Can't write record outside record block: " << Record << "\n";
-      return false;
-    }
     bool UsesDefaultAbbrev = Record.Abbrev == naclbitc::UNABBREV_RECORD;
+    if (atOutermostScope()) {
+      RecoverableError() << "Record outside block: " << Record << "\n";
+      if (!Flags.getTryToRecover())
+        return false;
+      // Create a dummy block to hold record.
+      enterBlock(Writer, UnknownWriteBlockID, naclbitc::DEFAULT_MAX_ABBREV,
+                 Record);
+      UsesDefaultAbbrev = true;
+    }
     if (!UsesDefaultAbbrev
         && !Writer.isUserRecordAbbreviation(Record.Abbrev)) {
       // Illegal abbreviation index found.
       if (Flags.getWriteBadAbbrevIndex()) {
         Error() << "Uses illegal abbreviation index: " << Record << "\n";
         // Generate bad abbreviation index so that the bitcode reader
-        // can be tested.
+        // can be tested, and then quit.
         Results.WroteBadAbbrevIndex = true;
         Writer.EmitCode(Record.Abbrev);
-        // Note: We need to close blocks or the bitcode Writer will terminate
-        // due to assertions.
-        writeMissingEndBlocks(Writer);
-        return false;
+        finish(Writer, /* RecoverSilently = */ true);
+        return true;
       }
       RecoverableError() << "Uses illegal abbreviation index: "
                          << Record << "\n";
@@ -195,7 +317,7 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
         return false;
       UsesDefaultAbbrev = true;
     }
-    if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID
+    if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID
         && Record.Code == naclbitc::BLOCKINFO_CODE_SETBID) {
       // Note: SetBID records are handled by Writer->EmitBlockInfoAbbrev,
       // based on the SetBID value. Don't bother to generate SetBID record here.
@@ -322,12 +444,5 @@ NaClMungedBitcode::WriteResults NaClMungedBitcode::writeMaybeRepair(
     if (!State.emitRecord(Writer, Record))
       break;
   }
-  if (!State.AbbrevIndexLimitStack.empty()) {
-    State.RecoverableError()
-        << "Bitcode missing " << State.AbbrevIndexLimitStack.size()
-        << " close blocks.\n";
-    if (Flags.getTryToRecover())
-      State.writeMissingEndBlocks(Writer);
-  }
-  return State.Results;
+  return State.finish(Writer);
 }