Fix abbreviation handling in munged bitcode writer.

The bitstream writer assumes several properties, and will fail
(either through assert or core dump) if these properties aren't
true. Fixing munged bitcode writer to handle the following
assumptions:

1) Don't write out an abbreviation if the abbreviation can't be
applied to the record.

2) Don't write out an abbreviation definition if the record doesn't
define a valid abbreviation (previous code was checking this, but
not correctly for the abbreviation operand).

3) Don't use abbreviations whose abbreviation index is larger than
couldn't be written out (case 2). Such indices will be broken because
the bitstream reader doesn't correctly match.

BUG=https://code.google.com/p/nativeclient/issues/detail?id=4169
[email protected]

Review URL: https://codereview.chromium.org/1146203003

Author: KarlSchimpf

include/llvm/Bitcode/NaCl/NaClBitstreamWriter.h

@@ -44,11 +44,14 @@ class NaClBitstreamWriter {
   std::vector<NaClBitCodeAbbrev*> CurAbbrevs;
 
   struct Block {
-    NaClBitcodeSelectorAbbrev PrevCodeSize;
-    unsigned StartSizeWord;
+    const NaClBitcodeSelectorAbbrev PrevCodeSize;
+    const unsigned StartSizeWord;
     std::vector<NaClBitCodeAbbrev*> PrevAbbrevs;
-    Block(const NaClBitcodeSelectorAbbrev& PCS, unsigned SSW)
-        : PrevCodeSize(PCS), StartSizeWord(SSW) {}
+    const unsigned AbbreviationIndexLimit;
+    Block(const NaClBitcodeSelectorAbbrev& PCS, unsigned SSW,
+          unsigned AbbreviationIndexLimit)
+        : PrevCodeSize(PCS), StartSizeWord(SSW),
+          AbbreviationIndexLimit(AbbreviationIndexLimit) {}
   };
 
   /// BlockScope - This tracks the current blocks that we have entered.
@@ -158,11 +161,15 @@ class NaClBitstreamWriter {
   // Basic Primitives for emitting bits to the stream.
   //===--------------------------------------------------------------------===//
 
+  // Max Number of bits that can be written using Emit.
+  static const unsigned MaxEmitNumBits = 32;
+
   void Emit(uint32_t Val, unsigned NumBits) {
-    assert(NumBits && NumBits <= 32 && "Invalid value size!");
-    assert((Val & ~(~0U >> (32-NumBits))) == 0 && "High bits set!");
+    assert(NumBits && NumBits <= MaxEmitNumBits && "Invalid value size!");
+    assert((Val &
+            ~(~0U >> (MaxEmitNumBits-NumBits))) == 0 && "High bits set!");
     CurValue |= Val << CurBit;
-    if (CurBit + NumBits < 32) {
+    if (CurBit + NumBits < MaxEmitNumBits) {
       CurBit += NumBits;
       return;
     }
@@ -171,19 +178,19 @@ class NaClBitstreamWriter {
     WriteWord(CurValue);
 
     if (CurBit)
-      CurValue = Val >> (32-CurBit);
+      CurValue = Val >> (MaxEmitNumBits-CurBit);
     else
       CurValue = 0;
-    CurBit = (CurBit+NumBits) & 31;
+    CurBit = (CurBit+NumBits) & (MaxEmitNumBits-1);
   }
 
   void Emit64(uint64_t Val, unsigned NumBits) {
-    if (NumBits <= 32)
-      Emit((uint32_t)Val, NumBits);
-    else {
-      Emit((uint32_t)Val, 32);
-      Emit((uint32_t)(Val >> 32), NumBits-32);
+    while (NumBits > MaxEmitNumBits) {
+      Emit((uint32_t)Val, MaxEmitNumBits);
+      Val >>= MaxEmitNumBits;
+      NumBits -= MaxEmitNumBits;
     }
+    Emit((uint32_t)Val, NumBits);
   }
 
   void flushToByte() {
@@ -287,7 +294,8 @@ class NaClBitstreamWriter {
 
     // Push the outer block's abbrev set onto the stack, start out with an
     // empty abbrev set.
-    BlockScope.push_back(Block(OldCodeSize, BlockSizeWordIndex));
+    BlockScope.push_back(Block(OldCodeSize, BlockSizeWordIndex,
+                               1 << CodeLen.NumBits));
     BlockScope.back().PrevAbbrevs.swap(CurAbbrevs);
 
     // If there is a blockinfo for this BlockID, add all the predefined abbrevs
@@ -389,9 +397,8 @@ class NaClBitstreamWriter {
   template<typename uintty>
   void EmitRecordWithAbbrevImpl(unsigned Abbrev,
                                 const AbbrevValues<uintty> &Vals) {
-    unsigned AbbrevNo = Abbrev-naclbitc::FIRST_APPLICATION_ABBREV;
-    assert(AbbrevNo < CurAbbrevs.size() && "Invalid abbrev #!");
-    NaClBitCodeAbbrev *Abbv = CurAbbrevs[AbbrevNo];
+    const NaClBitCodeAbbrev *Abbv = getAbbreviation(Abbrev);
+    assert(Abbv && "Abbreviation index is invalid");
 
     EmitCode(Abbrev);
 
@@ -421,11 +428,17 @@ class NaClBitstreamWriter {
 
 public:
 
-  /// Returns true if the given abbreviation index corresponds to a user-defined
-  /// abbreviation.
-  bool isUserRecordAbbreviation(unsigned Abbrev) const {
-    return Abbrev >= naclbitc::FIRST_APPLICATION_ABBREV
-        && Abbrev < (CurAbbrevs.size() + naclbitc::FIRST_APPLICATION_ABBREV);
+  /// Returns a pointer to the abbreviation currently associated with
+  /// the abbreviation index. Returns nullptr if no such abbreviation.
+  const NaClBitCodeAbbrev *getAbbreviation(unsigned Index) const {
+    if (Index < naclbitc::FIRST_APPLICATION_ABBREV)
+      return nullptr;
+    if (Index >= BlockScope.back().AbbreviationIndexLimit)
+      return nullptr;
+    unsigned AbbrevNo = Index - naclbitc::FIRST_APPLICATION_ABBREV;
+    if (AbbrevNo >= CurAbbrevs.size())
+      return nullptr;
+    return CurAbbrevs[AbbrevNo];
   }
 
   /// EmitRecord - Emit the specified record to the stream, using an abbrev if

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 {
@@ -27,13 +29,19 @@ static bool DebugEmit = false;
 // 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="
+        << OmittedAbbreviations << ")";
   }
 };
 
@@ -49,13 +57,16 @@ struct WriteState {
   // The SetBID for the blockinfo block.
   unsigned SetBID = UnknownWriteBlockID;
   // The stack of scopes the writer is in.
-  SmallVector<BlockScope, 3> ScopeStack;
+  SmallVector<BlockScope, 4> ScopeStack;
   // 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;
 
   WriteState(const NaClMungedBitcode::WriteFlags &Flags)
       : Flags(Flags),
@@ -90,14 +101,32 @@ 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);
+  }
+
   // 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
@@ -118,6 +147,21 @@ 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 abbreviation
+  // 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 +229,55 @@ bool WriteState::enterBlock(NaClBitstreamWriter &Writer, uint64_t WriteBlockID,
   return true;
 }
 
+bool WriteState::canApplyAbbreviation(
+    NaClBitstreamWriter &Writer, const NaClBitcodeAbbrevRecord &Record) {
+  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)
+                    << NaClBitstreamWriter::MaxEmitNumBits)
+          || NaClBitsNeededForValue(Value) > Op->getValue())
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::VBR:
+      if (Op->getValue() < 2)
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::Array:
+      llvm_unreachable("Array(Array) abbreviation is not legal!");
+    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.
@@ -271,23 +364,24 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
     break;
   }
   case naclbitc::BLK_CODE_DEFINE_ABBREV: {
+    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 (Record.Abbrev != naclbitc::DEFINE_ABBREV) {
       RecoverableError()
           << "Uses illegal abbreviation index in define abbreviation record: "
           << Record << "\n";
       if (!Flags.getTryToRecover())
         return false;
     }
-    if (getCurWriteBlockID() != naclbitc::BLOCKINFO_BLOCK_ID
-        && Writer.getMaxCurAbbrevIndex() >= getCurAbbrevIndexLimit()) {
-      RecoverableError() << "Exceeds abbreviation index limit of "
-                         << getCurAbbrevIndexLimit() << ": " << Record << "\n";
-      // Recover by not writing.
-      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,10 +408,21 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       }
       UsesDefaultAbbrev = true;
     }
-    if (!UsesDefaultAbbrev
-        && !Writer.isUserRecordAbbreviation(Record.Abbrev)) {
-      // Illegal abbreviation index found.
+    if (!UsesDefaultAbbrev && !canApplyAbbreviation(Writer, Record)) {
+      if (Writer.getAbbreviation(Record.Abbrev) != nullptr) {
+        RecoverableError() << "Abbreviation doesn't apply to record: "
+                           << Record << "\n";
+        UsesDefaultAbbrev = true;
+        if (!Flags.getTryToRecover())
+          return false;
+        WriteRecord(Writer, Record, UsesDefaultAbbrev);
+        return true;
+      }
       if (Flags.getWriteBadAbbrevIndex()) {
+        // The abbreviation is not understood by the bitcode writer,
+        // and the flag value implies that we should still write it
+        // out so that unit tests for this error condition can be
+        // applied.
         Error() << "Uses illegal abbreviation index: " << Record << "\n";
         // Generate bad abbreviation index so that the bitcode reader
         // can be tested, and then quit.
@@ -329,9 +434,11 @@ bool WriteState::emitRecord(NaClBitstreamWriter &Writer,
       }
       RecoverableError() << "Uses illegal abbreviation index: "
                          << Record << "\n";
+      UsesDefaultAbbrev = true;
       if (!Flags.getTryToRecover())
         return false;
-      UsesDefaultAbbrev = true;
+      WriteRecord(Writer, Record, UsesDefaultAbbrev);
+      return true;
     }
     if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID
         && Record.Code == naclbitc::BLOCKINFO_CODE_SETBID) {
@@ -347,10 +454,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 +486,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 +543,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);
     }
   }