Fix abbreviation handling in munged bitcode writer.

lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp

 //===- NaClBitcodeMungeWriter.cpp - Write munged bitcode --------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Implements method NaClMungedBitcode.write(), which writes out a munged
 // list of bitcode records using a bitstream writer.
 
 #include "llvm/Bitcode/NaCl/NaClBitcodeMungeUtils.h"
 
 #include "llvm/Bitcode/NaCl/NaClBitstreamWriter.h"
 #include "llvm/Bitcode/NaCl/NaClReaderWriter.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <set>
+
 using namespace llvm;
 
 namespace {
 
 // For debugging. When true, shows trace information of emitting
 // 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 << ")";
   }
 };
 
 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 records not in any block.
   static const unsigned UnknownWriteBlockID = UINT_MAX;
   // The SetBID for the blockinfo block.
   unsigned SetBID = UnknownWriteBlockID;
   // The stack of scopes the writer is in.
   SmallVector<BlockScope, 3> ScopeStack;

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

btw, I don't remember if this 3 was "optimized" for the most common case, but in
some recent refactoring you added an sentinel extra layer, so the 3 may no
longer be "optimized".

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

Good point. It was optimized as you suggested. Increasing to 4.

   // The set of write flags to use.
   const NaClMungedBitcode::WriteFlags &Flags;
   // The results of the attempted write.
   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),
         BlockMinBits(NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV)) {
     BlockScope Scope(UnknownWriteBlockID, naclbitc::DEFAULT_MAX_ABBREV);
     ScopeStack.push_back(Scope);
   }
 
   // Returns stream to print error message to.
   raw_ostream &Error();
@@ skipping 14 matching lines @@
   unsigned getCurWriteBlockID() const {
     assert(!ScopeStack.empty());
     return ScopeStack.back().CurBlockID;
   }
 
   unsigned getCurAbbrevIndexLimit() const {
     assert(!ScopeStack.empty());
     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);
 
   // 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
   enterBlock(NaClBitstreamWriter &Writer, uint64_t WriteBlockID,
              uint64_t NumBits, const NaClBitcodeAbbrevRecord &Record);
 
   // Exit current block and return to previous block. Silently recovers
   // if at outermost scope.
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
   exitBlock(NaClBitstreamWriter &Writer,
             const NaClBitcodeAbbrevRecord *Record = nullptr) {
     if (atOutermostScope())
       return false;
     Writer.ExitBlock();
     ScopeStack.pop_back();
     if (DebugEmit)
       printScopeStack(errs());
     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);
 
   void printScopeStack(raw_ostream &Out) {
     Out << "Scope Stack:\n";
     for (auto &Scope : ScopeStack)
       Out << "  " << Scope << "\n";
   }
 };
@@ skipping 47 matching lines @@
         return false;
     }
     Writer.EnterBlockInfoBlock();
   } else {
     NaClBitcodeSelectorAbbrev CurCodeLen(MaxAbbrev);
     Writer.EnterSubblock(WriteBlockID, CurCodeLen);
   }
   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.
   while (!atOutermostScope()) {
     if (!RecoverSilently)
       RecoverableError() << "Missing close block.\n";
     if (!exitBlock(Writer)) {
       Error() << "Failed to add missing close block at end of file.\n";
       break;
     }
@@ skipping 70 matching lines @@
               << Record << "\n";
       return false;
     }
     break;
   }
   case naclbitc::BLK_CODE_DEFINE_ABBREV: {
     if (Record.Abbrev != naclbitc::DEFINE_ABBREV) {
       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 {
       Writer.EmitAbbrev(Abbrev);
     }
     break;
   }
   case naclbitc::BLK_CODE_HEADER:
     // Note: There is no abbreviation index here. Ignore.
     for (uint64_t Value : Record.Values)
       Writer.Emit(Value, 8);
     break;
   default:
     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.
       if (!enterBlock(Writer, UnknownWriteBlockID,
                       naclbitc::DEFAULT_MAX_ABBREV, Record)) {
         Error() << "Failed to recover from record outside block\n";
         return false;
       }
       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
         // can be tested, and then quit.
         Results.WroteBadAbbrevIndex = true;
         Writer.EmitCode(Record.Abbrev);
         bool RecoverSilently = true;
         finish(Writer, RecoverSilently);
         return false;
       }
       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) {
       // Note: SetBID records are handled by Writer->EmitBlockInfoAbbrev,
       // based on the SetBID value. Don't bother to generate SetBID record here.
       // Rather just set SetBID and let call to Writer->EmitBlockInfoAbbrev
       // generate the SetBID record.
       if (NumValues != 1) {
         Error() << "SetBID record expects 1 value but found "
                 << NumValues << ": " << Record << "\n";
         return false;
       }
       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;
 }
 
 static NaClBitCodeAbbrev *deleteAbbrev(NaClBitCodeAbbrev *Abbrev) {
   Abbrev->dropRef();
   return nullptr;
 }
 
 NaClBitCodeAbbrev *WriteState::buildAbbrev(
@@ skipping 10 matching lines @@
   size_t NumAbbreviations = Record.Values[Index++];
   if (NumAbbreviations == 0) {
     RecoverableError() << "Abbreviation must contain at least one operator: "
                        << Record << "\n";
     return deleteAbbrev(Abbrev);
   }
   for (size_t Count = 0; Count < NumAbbreviations; ++Count) {
     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++]) {
     case 1:
       if (Index >= NumValues) {
         RecoverableError() << "Malformed literal abbreviation: "
                            << Record << "\n";
         return deleteAbbrev(Abbrev);
       }
       Abbrev->Add(NaClBitCodeAbbrevOp(Record.Values[Index++]));
@@ skipping 36 matching lines @@
         Abbrev->Add(NaClBitCodeAbbrevOp(NaClBitCodeAbbrevOp::Char6));
         break;
       default:
         RecoverableError() << "Unknown abbreviation kind " << Kind
                            << ": " << Record << "\n";
         return deleteAbbrev(Abbrev);
       }
       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);
     }
   }
   return Abbrev;
 }
 
 } // end of anonymous namespace.
 
 NaClMungedBitcode::WriteResults NaClMungedBitcode::writeMaybeRepair(
     SmallVectorImpl<char> &Buffer, bool AddHeader,
     const WriteFlags &Flags) const {
   NaClBitstreamWriter Writer(Buffer);
   WriteState State(Flags);
   if (AddHeader) {
     NaClWriteHeader(Writer, true);
   }
   for (const NaClBitcodeAbbrevRecord &Record : *this) {
     if (!State.emitRecord(Writer, Record))
       break;
   }
   bool RecoverSilently =
       State.Results.NumErrors > 0 && !Flags.getTryToRecover();
   return State.finish(Writer, RecoverSilently);
 }