[BitcodeReader] Use tighter upper bound to validate forward references.

[fhahn]

Summary

At the moment, bitcode files with invalid forward reference can easily
cause the bitcode reader to run out of memory, by creating a forward
reference with a very high index.

We can use the size of the bitcode file as an upper bound, because a
valid bitcode file can never contain more records. This should be
sufficient to fail early in most cases. The only exception is large
files with invalid forward references close to the file size.

There are a couple of clusterfuzz runs that fail with out-of-memory
because of very high forward references and they should be fixed by this
patch.

A concrete example for this is D64507, which causes out-of-memory on
systems with low memory, like the hexagon upstream bots.

I am not entirely sure about the way we get the size of the stream in
this patch. Maybe it would be better to get NumWords, when we enter the
module block?

[fhahn]

llvm/test/Bitcode/pr18704.ll

3

The invalid forward reference in this test is covered by the upper bound check now. Is there a tool that can take the output of llvm-bcanalyzer like in the test case and create a bitcode file?

[jfb]

llvm/include/llvm/Bitstream/BitstreamReader.h

298

I'd rather call this "sizeInBytes" or something. I get confused by bits and bytes :)

llvm/include/llvm/Bitstream/BitstreamReader.h

if (sizeof(word_t) > 4 &&
       CurWord >>= BitsInCurWord-32;
       BitsInCurWord = 32;
       return;
     }
 
     BitsInCurWord = 0;
   }
 
+  /// Return the size of stream.
+  size_t size() const { return BitcodeBytes.size(); }
+
   /// Skip to the end of the file.
   void skipToEnd() { NextChar = BitcodeBytes.size(); }
 };
 
 /// When advancing through a bitstream cursor, each advance can discover a few
 /// different kinds of entries:
 struct BitstreamEntry {
   enum {
public:
   BitstreamCursor() = default;
   explicit BitstreamCursor(ArrayRef<uint8_t> BitcodeBytes)
       : SimpleBitstreamCursor(BitcodeBytes) {}
   explicit BitstreamCursor(StringRef BitcodeBytes)
       : SimpleBitstreamCursor(BitcodeBytes) {}
   explicit BitstreamCursor(MemoryBufferRef BitcodeBytes)
       : SimpleBitstreamCursor(BitcodeBytes) {}
 
-  using SimpleBitstreamCursor::canSkipToPos;
   using SimpleBitstreamCursor::AtEndOfStream;
+  using SimpleBitstreamCursor::canSkipToPos;
+  using SimpleBitstreamCursor::fillCurWord;
   using SimpleBitstreamCursor::getBitcodeBytes;
   using SimpleBitstreamCursor::GetCurrentBitNo;
   using SimpleBitstreamCursor::getCurrentByteNo;
   using SimpleBitstreamCursor::getPointerToByte;
   using SimpleBitstreamCursor::JumpToBit;
-  using SimpleBitstreamCursor::fillCurWord;
   using SimpleBitstreamCursor::Read;
   using SimpleBitstreamCursor::ReadVBR;
   using SimpleBitstreamCursor::ReadVBR64;
+  using SimpleBitstreamCursor::size;
 
   /// Return the number of bits used to encode an abbrev #.
   unsigned getAbbrevIDWidth() const { return CurCodeSize; }
 
   /// Flags that modify the behavior of advance().
   enum {
     /// If this flag is used, the advance() method does not automatically pop
     /// the block scope when the end of a block is reached.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

std::error_code llvm::errorToErrorCodeAndEmitErrors(LLVMContext &Ctx,
   }
   return std::error_code();
 }
 
 BitcodeReader::BitcodeReader(BitstreamCursor Stream, StringRef Strtab,
                              StringRef ProducerIdentification,
                              LLVMContext &Context)
     : BitcodeReaderBase(std::move(Stream), Strtab), Context(Context),
-      ValueList(Context) {
+      ValueList(Context, Stream.size()) {
   this->ProducerIdentification = ProducerIdentification;
 }
 
 Error BitcodeReader::materializeForwardReferencedFunctions() {
   if (WillMaterializeAllForwardRefs)
     return Error::success();
 
   // Prevent recursion.

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

struct {
     SmallDenseMap<MDString *, TempMDTuple, 1> Unknown;
     SmallDenseMap<MDString *, DICompositeType *, 1> Final;
     SmallDenseMap<MDString *, DICompositeType *, 1> FwdDecls;
     SmallVector<std::pair<TrackingMDRef, TempMDTuple>, 1> Arrays;
   } OldTypeRefs;
 
   LLVMContext &Context;
 
+  /// Maximum number of valid references. Forward references exceeding the
+  /// maximum must be invalid.
+  unsigned RefsUpperBound;
+
 public:
-  BitcodeReaderMetadataList(LLVMContext &C) : Context(C) {}
+  BitcodeReaderMetadataList(LLVMContext &C, size_t RefsUpperBound)
+      : Context(C),
+        RefsUpperBound(std::min((size_t)std::numeric_limits<unsigned>::max(),
+                                RefsUpperBound)) {}
 
   // vector compatibility methods
   unsigned size() const { return MetadataPtrs.size(); }
   void resize(unsigned N) { MetadataPtrs.resize(N); }
   void push_back(Metadata *MD) { MetadataPtrs.emplace_back(MD); }
   void clear() { MetadataPtrs.clear(); }
   Metadata *back() const { return MetadataPtrs.back(); }
   void pop_back() { MetadataPtrs.pop_back(); }
void BitcodeReaderMetadataList::assignValue(Metadata *MD, unsigned Idx) {
 
   // If there was a forward reference to this value, replace it.
   TempMDTuple PrevMD(cast<MDTuple>(OldMD.get()));
   PrevMD->replaceAllUsesWith(MD);
   ForwardReference.erase(Idx);
 }
 
 Metadata *BitcodeReaderMetadataList::getMetadataFwdRef(unsigned Idx) {
+  // Bail out for a clearly invalid value.
+  if (Idx >= RefsUpperBound)
+    return nullptr;
+
   if (Idx >= size())
     resize(Idx + 1);
 
   if (Metadata *MD = MetadataPtrs[Idx])
     return MD;
 
   // Track forward refs to be resolved later.
   ForwardReference.insert(Idx);
void upgradeDebugInfo() {
     upgradeCUVariables();
   }
 
 public:
   MetadataLoaderImpl(BitstreamCursor &Stream, Module &TheModule,
                      BitcodeReaderValueList &ValueList,
                      std::function<Type *(unsigned)> getTypeByID,
                      bool IsImporting)
-      : MetadataList(TheModule.getContext()), ValueList(ValueList),
-        Stream(Stream), Context(TheModule.getContext()), TheModule(TheModule),
-        getTypeByID(std::move(getTypeByID)), IsImporting(IsImporting) {}
+      : MetadataList(TheModule.getContext(), Stream.size()),
+        ValueList(ValueList), Stream(Stream), Context(TheModule.getContext()),
+        TheModule(TheModule), getTypeByID(std::move(getTypeByID)),
+        IsImporting(IsImporting) {}
 
   Error parseMetadata(bool ModuleLevel);
 
   bool hasFwdRefs() const { return MetadataList.hasFwdRefs(); }
 
   Metadata *getMetadataFwdRefOrLoad(unsigned ID) {
     if (ID < MDStringRef.size())
       return lazyLoadOneMDString(ID);

llvm/lib/Bitcode/Reader/ValueList.h

class BitcodeReaderValueList {
   /// ResolveConstantForwardRefs for more information about this.
   ///
   /// The key of this vector is the placeholder constant, the value is the slot
   /// number that holds the resolved value.
   using ResolveConstantsTy = std::vector<std::pair<Constant *, unsigned>>;
   ResolveConstantsTy ResolveConstants;
   LLVMContext &Context;
 
+  /// Maximum number of valid references. Forward references exceeding the
+  /// maximum must be invalid.
+  unsigned RefsUpperBound;
+
 public:
-  BitcodeReaderValueList(LLVMContext &C) : Context(C) {}
+  BitcodeReaderValueList(LLVMContext &C, size_t RefsUpperBound)
+      : Context(C),
+        RefsUpperBound(std::min((size_t)std::numeric_limits<unsigned>::max(),
+                                RefsUpperBound)) {}
 
   ~BitcodeReaderValueList() {
     assert(ResolveConstants.empty() && "Constants not resolved?");
   }
 
   // vector compatibility methods
   unsigned size() const { return ValuePtrs.size(); }
   void resize(unsigned N) {

llvm/lib/Bitcode/Reader/ValueList.cpp

if (Constant *PHC = dyn_cast<Constant>(&*OldV)) {
     // If there was a forward reference to this value, replace it.
     Value *PrevVal = OldV;
     OldV->replaceAllUsesWith(V);
     PrevVal->deleteValue();
   }
 }
 
 Constant *BitcodeReaderValueList::getConstantFwdRef(unsigned Idx, Type *Ty) {
+  // Bail out for a clearly invalid value.
+  if (Idx >= RefsUpperBound)
+    return nullptr;
+
   if (Idx >= size())
     resize(Idx + 1);
 
   if (Value *V = ValuePtrs[Idx]) {
     if (Ty != V->getType())
       report_fatal_error("Type mismatch in constant table!");
     return cast<Constant>(V);
   }
 
   // Create and return a placeholder, which will later be RAUW'd.
   Constant *C = new ConstantPlaceHolder(Ty, Context);
   ValuePtrs[Idx] = C;
   return C;
 }
 
 Value *BitcodeReaderValueList::getValueFwdRef(unsigned Idx, Type *Ty,
                                               Type **FullTy) {
-  // Bail out for a clearly invalid value. This would make us call resize(0)
-  if (Idx == std::numeric_limits<unsigned>::max())
+  // Bail out for a clearly invalid value.
+  if (Idx >= RefsUpperBound)
     return nullptr;
 
   if (Idx >= size())
     resize(Idx + 1);
 
   if (Value *V = ValuePtrs[Idx]) {
     // If the types don't match, it's invalid.
     if (Ty && Ty != V->getType())

llvm/test/Bitcode/pr18704.ll

 ; RUN:  not llvm-dis < %s.bc 2>&1 | FileCheck %s
 
-; CHECK: llvm-dis{{(\.EXE|\.exe)?}}: error: Never resolved value found in function
+; CHECK: llvm-dis{{(\.EXE|\.exe)?}}: error: Invalid record
 
 ; pr18704.ll.bc has an instruction referring to invalid type.
 ; The test checks that LLVM reports the error and doesn't access freed memory
 ; in doing so.
 
 ;<MODULE_BLOCK NumWords=217 BlockCodeSize=3>
 ;  <VERSION op0=1/>
 ;  <BLOCKINFO_BLOCK/>