Fix abbreviation handling in munged bitcode writer.

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 b536fe3c848a4d81241cb4b7f6e39c49313451de..6ad7ac6536ba755d570ee4a04438c4e9f457c3ec 100644
--- a/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
+++ b/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
@@ -16,6 +16,8 @@
 #include "llvm/Bitcode/NaCl/NaClReaderWriter.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <set>
+
 using namespace llvm;
 
 namespace {
@@ -24,16 +26,26 @@ namespace {
 // bitcode.
 static bool DebugEmit = false;
 
+// Maximum bits that can be printed out using method
+// NaClBitstreamWriter::Emit().

[jvoung (off chromium), 2015/05/20 22:40:09]

Seems like this should be part of NaClBitstreamWriter and used there too.

[Karl, 2015/05/21 18:22:07]

I agree. I originally put it there, but I didn't like how it read. Will put it
back in.

+static const unsigned MaxEmitFixedBits = 32;
+
 // Description of current block scope
 struct BlockScope {
   BlockScope(unsigned CurBlockID, size_t AbbrevIndexLimit)
-      : CurBlockID(CurBlockID), AbbrevIndexLimit(AbbrevIndexLimit) {}
+      : CurBlockID(CurBlockID), AbbrevIndexLimit(AbbrevIndexLimit),
+        OmittedAbbreviations(false) {}
   unsigned CurBlockID;
   // Defines the maximum value for abbreviation indices in block.
   size_t AbbrevIndexLimit;
+  // Defines if an abbreviation definition was omitted (i.e. not
+  // written) from this block. Used to turn off writing further
+  // abbreviation definitions for this block.
+  bool OmittedAbbreviations;
   void print(raw_ostream &Out) const {
     Out << "BlockScope(ID=" << CurBlockID << ", AbbrevIndexLimit="
-        << AbbrevIndexLimit << ")";
+        << AbbrevIndexLimit << "OmittedAbbreviations = "

[jvoung (off chromium), 2015/05/20 22:40:10]

= spacing consistency

previous ones just did X=Y

[Karl, 2015/05/21 18:22:08]

Done.

+        << OmittedAbbreviations << ")";
   }
 };
 
@@ -56,6 +68,9 @@ struct WriteState {
   NaClMungedBitcode::WriteResults Results;
   // The minimum number of bits allowed to be specified in a block.
   const unsigned BlockMinBits;
+  // The set of block IDs for which abbreviation definitions have been
+  // omitted in the blockinfo block.
+  std::set<unsigned> BlocksWithOmittedAbbrevs;

[jvoung (off chromium), 2015/05/20 22:40:10]

Would DenseSet be better?

I suppose there are only so many block types, so this doesn't grow much.

[Karl, 2015/05/21 18:22:07]

I did not expect the set to get much larger than 8, and hence, I think a set is
a better choice in this case.

 
   WriteState(const NaClMungedBitcode::WriteFlags &Flags)
       : Flags(Flags),
@@ -90,6 +105,24 @@ struct WriteState {
     return ScopeStack.back().AbbrevIndexLimit;
   }
 
+  // Returns whether any abbreviation definitions were not written to
+  // the bitcode buffer.
+  bool curBlockHasOmittedAbbreviations() const {
+    assert(!ScopeStack.empty());
+    return ScopeStack.back().OmittedAbbreviations
+        || BlocksWithOmittedAbbrevs.count(getCurWriteBlockID());
+  }
+
+  // Marks that an abbreviation definition is being omitted (i.e. not
+  // written) for the current block.
+  void markCurrentBlockWithOmittedAbbreviations() {
+    assert(!ScopeStack.empty());
+    ScopeStack.back().OmittedAbbreviations = true;
+    if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID)
+      BlocksWithOmittedAbbrevs.insert(SetBID);
+      return;

[jvoung (off chromium), 2015/05/20 22:40:10]

The return doesn't seem to do much so may be better to remove.

It's indented even though it's not part of the then-statement (no curly), so
that's a bit confusing =)

[Karl, 2015/05/21 18:22:08]

Hmm. I guess this is a leftover of an earlier version. Removing the return.

+  }
+
   // Converts the abbreviation record to the corresponding abbreviation.
   // Returns nullptr if unable to build abbreviation.
   NaClBitCodeAbbrev *buildAbbrev(const NaClBitcodeAbbrevRecord &Record);
@@ -97,7 +130,8 @@ struct WriteState {
   // Emits the given record to the bitcode file. Returns true if
   // successful.
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
-  emitRecord(NaClBitstreamWriter &Writer, const NaClBitcodeAbbrevRecord &Record);
+  emitRecord(NaClBitstreamWriter &Writer,
+             const NaClBitcodeAbbrevRecord &Record);
 
   // Enter the given block
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
@@ -118,6 +152,22 @@ struct WriteState {
     return true;
   }
 
+  void WriteRecord(NaClBitstreamWriter &Writer,
+                   const NaClBitcodeAbbrevRecord &Record,
+                   bool UsesDefaultAbbrev) {
+    if (UsesDefaultAbbrev)
+      Writer.EmitRecord(Record.Code, Record.Values);
+    else
+      Writer.EmitRecord(Record.Code, Record.Values, Record.Abbrev);
+
+  }
+
+  // Returns true if the abbreviation index defines an anbbreviation

[jvoung (off chromium), 2015/05/20 22:40:09]

anbbreviation -> ?

[Karl, 2015/05/21 18:22:07]

Fixed.

+  // that can be applied to the record.
+  bool LLVM_ATTRIBUTE_UNUSED_RESULT
+  canApplyAbbreviation(NaClBitstreamWriter &Writer,
+                       const NaClBitcodeAbbrevRecord &Record);
+
   // Completes the write.
   NaClMungedBitcode::WriteResults &finish(NaClBitstreamWriter &Writer,
                                           bool RecoverSilently);
@@ -185,6 +235,58 @@ bool WriteState::enterBlock(NaClBitstreamWriter &Writer, uint64_t WriteBlockID,
   return true;
 }
 
+bool WriteState::canApplyAbbreviation(
+    NaClBitstreamWriter &Writer, const NaClBitcodeAbbrevRecord &Record) {

[Karl, 2015/05/20 20:32:04]

Added following two lines to make sure that the block defined enough bits to
apply abbreviation.

[jvoung (off chromium), 2015/05/20 22:40:09]

Does the Writer.getAbbreviation(...); if (Abbrev == nullptr) cover this, or does
this first check help?

[Karl, 2015/05/21 18:22:07]

This check helps. I didn't see the case until the fuzzer generated it. Because
global abbreviations are generated without knowledge of the number of bits
allowed for an individual block, and the abbreviations are copied in blindly,
this can lead to an error.

The bitstream writer also doesn't check that there is a mismatch between the
number of defined abbreviations, and the number of bits defined for the block.
Hence I added this test.

I decided to move this logic back into the bitstream writer, since it should
really check for this complex condition and fail. Otherwise, subtle bugs could
be added to pnacl bitcode files that won't manifest themselves until the file is
read.

+  if (Record.Abbrev > getCurAbbrevIndexLimit())
+    return false;
+
+  const NaClBitCodeAbbrev *Abbrev = Writer.getAbbreviation(Record.Abbrev);
+  if (Abbrev == nullptr)
+    return false;
+
+  // Merge record code into values and then match abbreviation.
+  NaClBitcodeValues Values(Record);
+  size_t ValueIndex = 0;
+  size_t ValuesSize = Values.size();
+  size_t AbbrevIndex = 0;
+  size_t AbbrevSize = Abbrev->getNumOperandInfos();
+  bool FoundArray = false;
+  while (ValueIndex < ValuesSize && AbbrevIndex < AbbrevSize) {
+    const NaClBitCodeAbbrevOp *Op = &Abbrev->getOperandInfo(AbbrevIndex++);
+    uint64_t Value = Values[ValueIndex++];
+    if (Op->getEncoding() == NaClBitCodeAbbrevOp::Array) {
+      if (AbbrevIndex + 1 != AbbrevSize)
+        return false;
+      Op = &Abbrev->getOperandInfo(AbbrevIndex);
+      --AbbrevIndex; // i.e. don't advance to next abbreviation op.
+      FoundArray = true;
+    }
+    switch (Op->getEncoding()) {
+    case NaClBitCodeAbbrevOp::Literal:
+      if (Value != Op->getValue())
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::Fixed:
+      if (Value >= (static_cast<uint64_t>(1) < MaxEmitFixedBits)

[jvoung (off chromium), 2015/05/20 22:40:10]

I assume this should actually be a shift and not (1 < MaxEmitFixedBits)?

[Karl, 2015/05/21 18:22:07]

Yes. Fixing.

+          || NaClBitsNeededForValue(Value) > Op->getValue())
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::VBR:
+      if (Op->getValue() < 2)
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::Array:
+      // This should not happen

[jvoung (off chromium), 2015/05/20 22:40:10]

llvm_unreachable(...);

[Karl, 2015/05/21 18:22:08]

Done.

+      return false;
+    case NaClBitCodeAbbrevOp::Char6:
+      if (!Op->isChar6(Value))
+        return false;
+      continue;
+    }
+  }
+  return ValueIndex == ValuesSize && (FoundArray || AbbrevIndex == AbbrevSize);
+}
+
 NaClMungedBitcode::WriteResults &WriteState::finish(
     NaClBitstreamWriter &Writer, bool RecoverSilently) {
   // Be sure blocks are balanced.
@@ -275,19 +377,29 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       RecoverableError()
           << "Uses illegal abbreviation index in define abbreviation record: "
           << Record << "\n";
+      markCurrentBlockWithOmittedAbbreviations();
       if (!Flags.getTryToRecover())

[jvoung (off chromium), 2015/05/20 22:40:10]

Might be simpler to just return now like the other cases, since it will just
fall through and hit the "if (curBlockHasOmmitted...) { ... return ...; }".

Then you can have the "if (curBlock...)" check as the first thing, which seems
more straightforward.

[Karl, 2015/05/21 18:22:07]

I agree that the check for omitted abbreviations should be the first case, and
moved it up.

For the case of the wrong abbreviation index, I did the wrong thing. The simple
recovery is to just ignore the abbreviation index and generate it.

Further, we can remove the limit on bitsize check, since method
Writer.getAbbreviation() now checks this when it is used.

         return false;
     }
+    if (curBlockHasOmittedAbbreviations()) {
+      // If reached, a previous abbreviation for the block was omitted. Can't
+      // generate more abbreviations without having to fix abbreviation indices.
+      RecoverableError() << "Ignoring abbreviation: " << Record << "\n";
+      return Flags.getTryToRecover();
+    }
     if (getCurWriteBlockID() != naclbitc::BLOCKINFO_BLOCK_ID
         && Writer.getMaxCurAbbrevIndex() >= getCurAbbrevIndexLimit()) {
       RecoverableError() << "Exceeds abbreviation index limit of "
                          << getCurAbbrevIndexLimit() << ": " << Record << "\n";
       // Recover by not writing.
+      markCurrentBlockWithOmittedAbbreviations();
       return Flags.getTryToRecover();
     }
     NaClBitCodeAbbrev *Abbrev = buildAbbrev(Record);
-    if (Abbrev == NULL)
+    if (Abbrev == nullptr) {
+      markCurrentBlockWithOmittedAbbreviations();
       return Flags.getTryToRecover();
+    }
     if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID) {
       Writer.EmitBlockInfoAbbrev(SetBID, Abbrev);
     } else {
@@ -314,9 +426,15 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       }
       UsesDefaultAbbrev = true;
     }
-    if (!UsesDefaultAbbrev
-        && !Writer.isUserRecordAbbreviation(Record.Abbrev)) {
-      // Illegal abbreviation index found.
+    if (!UsesDefaultAbbrev && !canApplyAbbreviation(Writer, Record)) {
+      // Can't apply abbreviation, figure out why.
+      if (Writer.getAbbreviation(Record.Abbrev) != nullptr) {
+        RecoverableError() << "Abbreviation doesn't apply to record: "

[jvoung (off chromium), 2015/05/20 22:40:09]

So here the abbreviation exists, but doesn't apply.

[Karl, 2015/05/21 18:22:07]

Yes.

+                           << Record << "\n";
+        UsesDefaultAbbrev = true;
+        WriteRecord(Writer, Record, UsesDefaultAbbrev);

[jvoung (off chromium), 2015/05/20 22:40:09]

This and the one below might be the first case where WriteRecord() happens
before checking Flags.getTryToRecover(). So, is that the intention to allow
writing some additional recovery-based-amount when the flag isn't set?

[Karl, 2015/05/21 18:22:08]

Done.

+        return Flags.getTryToRecover();
+      }
       if (Flags.getWriteBadAbbrevIndex()) {

[jvoung (off chromium), 2015/05/20 22:40:10]

Maybe add a comment here. This CL adds the "figure out why" comment, but this
case doesn't apply to the "figure out why". I.e., it's not clear how much of the
following code is for the "figure out why".

Alternatively, make the "figure out why" more precise and have it only apply to
the first if-block.

[Karl, 2015/05/21 18:22:08]

Removed the first comment, and added a comment to this case to explain.

         Error() << "Uses illegal abbreviation index: " << Record << "\n";
         // Generate bad abbreviation index so that the bitcode reader
@@ -329,9 +447,9 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       }
       RecoverableError() << "Uses illegal abbreviation index: "
                          << Record << "\n";
-      if (!Flags.getTryToRecover())
-        return false;
       UsesDefaultAbbrev = true;
+      WriteRecord(Writer, Record, UsesDefaultAbbrev);
+      return Flags.getTryToRecover();
     }
     if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID
         && Record.Code == naclbitc::BLOCKINFO_CODE_SETBID) {
@@ -347,10 +465,8 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       SetBID = Record.Values[0];
       return true;
     }
-    if (UsesDefaultAbbrev)
-      Writer.EmitRecord(Record.Code, Record.Values);
-    else
-      Writer.EmitRecord(Record.Code, Record.Values, Record.Abbrev);
+    WriteRecord(Writer, Record, UsesDefaultAbbrev);
+    break;
   }
   return true;
 }
@@ -381,7 +497,7 @@ NaClBitCodeAbbrev *WriteState::buildAbbrev(
     if (Index >= NumValues) {
       RecoverableError()
           << "Malformed abbreviation found. Expects " << NumAbbreviations
-          << " operands but ound " << Count << ": " << Record << "\n";
+          << " operands but found " << Count << ": " << Record << "\n";
       return deleteAbbrev(Abbrev);
     }
     switch (Record.Values[Index++]) {
@@ -438,8 +554,8 @@ NaClBitCodeAbbrev *WriteState::buildAbbrev(
       break;
     }
     default:
-      RecoverableError() << "Error: Bad literal flag " << Record.Values[Index]
-                         << ": " << Record << "\n";
+      RecoverableError() << "Error: Bad abbreviation operand encoding "
+                         << Record.Values[Index-1] << ": " << Record << "\n";
       return deleteAbbrev(Abbrev);
     }
   }