[llvm] 3c86642 - [Bitstream] Reject implausibly large reservations

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-02-07T12:16:12+01:00
New Revision: 3c86642edd28f1ce970882edaba8dce468ec7401

URL: https://github.com/llvm/llvm-project/commit/3c86642edd28f1ce970882edaba8dce468ec7401
DIFF: https://github.com/llvm/llvm-project/commit/3c86642edd28f1ce970882edaba8dce468ec7401.diff

LOG: [Bitstream] Reject implausibly large reservations

If we're trying to reserve more memory than bits in the stream,
reject this early to avoid OOM.

Added: 
    llvm/test/Bitcode/Inputs/size-not-plausible.bc

Modified: 
    llvm/include/llvm/Bitstream/BitstreamReader.h
    llvm/lib/Bitstream/Reader/BitstreamReader.cpp
    llvm/test/Bitcode/invalid.test

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Bitstream/BitstreamReader.h b/llvm/include/llvm/Bitstream/BitstreamReader.h
index 91955c27364a..10670648f2fc 100644
--- a/llvm/include/llvm/Bitstream/BitstreamReader.h
+++ b/llvm/include/llvm/Bitstream/BitstreamReader.h
@@ -299,6 +299,13 @@ class SimpleBitstreamCursor {
 
   /// Skip to the end of the file.
   void skipToEnd() { NextChar = BitcodeBytes.size(); }
+
+  /// Check whether a reservation of Size elements is plausible.
+  bool isSizePlausible(size_t Size) const {
+    // Don't allow reserving more elements than the number of bits, assuming
+    // at least one bit is needed to encode an element.
+    return Size < BitcodeBytes.size() * 8;
+  }
 };
 
 /// When advancing through a bitstream cursor, each advance can discover a few

diff  --git a/llvm/lib/Bitstream/Reader/BitstreamReader.cpp b/llvm/lib/Bitstream/Reader/BitstreamReader.cpp
index ffeb506154f9..f9247909dc3e 100644
--- a/llvm/lib/Bitstream/Reader/BitstreamReader.cpp
+++ b/llvm/lib/Bitstream/Reader/BitstreamReader.cpp
@@ -222,6 +222,8 @@ Expected<unsigned> BitstreamCursor::readRecord(unsigned AbbrevID,
     if (!MaybeNumElts)
       return MaybeNumElts.takeError();
     uint32_t NumElts = MaybeNumElts.get();
+    if (!isSizePlausible(NumElts))
+      return error("Size is not plausible");
     Vals.reserve(Vals.size() + NumElts);
 
     for (unsigned i = 0; i != NumElts; ++i)
@@ -275,6 +277,8 @@ Expected<unsigned> BitstreamCursor::readRecord(unsigned AbbrevID,
       if (!MaybeNumElts)
         return MaybeNumElts.takeError();
       uint32_t NumElts = MaybeNumElts.get();
+      if (!isSizePlausible(NumElts))
+        return error("Size is not plausible");
       Vals.reserve(Vals.size() + NumElts);
 
       // Get the element encoding.

diff  --git a/llvm/test/Bitcode/Inputs/size-not-plausible.bc b/llvm/test/Bitcode/Inputs/size-not-plausible.bc
new file mode 100644
index 000000000000..bbc424c7ab5f
Binary files /dev/null and b/llvm/test/Bitcode/Inputs/size-not-plausible.bc 
diff er

diff  --git a/llvm/test/Bitcode/invalid.test b/llvm/test/Bitcode/invalid.test
index 67546bf36bfb..ef229de32018 100644
--- a/llvm/test/Bitcode/invalid.test
+++ b/llvm/test/Bitcode/invalid.test
@@ -251,3 +251,8 @@ RUN: not llvm-dis -disable-output %p/Inputs/invalid-abbrev-number.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-ABBREV-NUMBER %s
 
 INVALID-ABBREV-NUMBER: Invalid abbrev number
+
+RUN: not llvm-dis -disable-output %p/Inputs/size-not-plausible.bc 2>&1 | \
+RUN:   FileCheck --check-prefix=SIZE-NOT-PLAUSIBLE %s
+
+SIZE-NOT-PLAUSIBLE: Size is not plausible