[llvm] 6280c29 - [Bitstream] Add assert to ReadVBR and ReadVBR64

[Jan Korous via llvm-commits]

Author: Jan Korous
Date: 2022-02-15T17:12:35-08:00
New Revision: 6280c29a870173601a7befb3b45f92025727c41d

URL: https://github.com/llvm/llvm-project/commit/6280c29a870173601a7befb3b45f92025727c41d
DIFF: https://github.com/llvm/llvm-project/commit/6280c29a870173601a7befb3b45f92025727c41d.diff

LOG: [Bitstream] Add assert to ReadVBR and ReadVBR64

We want to prevent UB potentially caused by left-shifting by type bit-width.

Differential Revision: https://reviews.llvm.org/D119307

Added: 
    

Modified: 
    llvm/include/llvm/Bitstream/BitstreamReader.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Bitstream/BitstreamReader.h b/llvm/include/llvm/Bitstream/BitstreamReader.h
index afe327cffdba9..d2e4b28591a20 100644
--- a/llvm/include/llvm/Bitstream/BitstreamReader.h
+++ b/llvm/include/llvm/Bitstream/BitstreamReader.h
@@ -227,21 +227,25 @@ class SimpleBitstreamCursor {
     return R;
   }
 
-  Expected<uint32_t> ReadVBR(unsigned NumBits) {
+  Expected<uint32_t> ReadVBR(const unsigned NumBits) {
     Expected<unsigned> MaybeRead = Read(NumBits);
     if (!MaybeRead)
       return MaybeRead;
     uint32_t Piece = MaybeRead.get();
 
-    if ((Piece & (1U << (NumBits-1))) == 0)
+    assert(NumBits <= 32 && NumBits >= 1 && "Invalid NumBits value");
+    const uint32_t MaskBitOrder = (NumBits - 1);
+    const uint32_t Mask = 1UL << MaskBitOrder;
+
+    if ((Piece & Mask) == 0)
       return Piece;
 
     uint32_t Result = 0;
     unsigned NextBit = 0;
     while (true) {
-      Result |= (Piece & ((1U << (NumBits-1))-1)) << NextBit;
+      Result |= (Piece & (Mask - 1)) << NextBit;
 
-      if ((Piece & (1U << (NumBits-1))) == 0)
+      if ((Piece & Mask) == 0)
         return Result;
 
       NextBit += NumBits-1;
@@ -258,21 +262,24 @@ class SimpleBitstreamCursor {
 
   // Read a VBR that may have a value up to 64-bits in size. The chunk size of
   // the VBR must still be <= 32 bits though.
-  Expected<uint64_t> ReadVBR64(unsigned NumBits) {
+  Expected<uint64_t> ReadVBR64(const unsigned NumBits) {
     Expected<uint64_t> MaybeRead = Read(NumBits);
     if (!MaybeRead)
       return MaybeRead;
     uint32_t Piece = MaybeRead.get();
+    assert(NumBits <= 32 && NumBits >= 1 && "Invalid NumBits value");
+    const uint32_t MaskBitOrder = (NumBits - 1);
+    const uint32_t Mask = 1UL << MaskBitOrder;
 
-    if ((Piece & (1U << (NumBits-1))) == 0)
+    if ((Piece & Mask) == 0)
       return uint64_t(Piece);
 
     uint64_t Result = 0;
     unsigned NextBit = 0;
     while (true) {
-      Result |= uint64_t(Piece & ((1U << (NumBits-1))-1)) << NextBit;
+      Result |= uint64_t(Piece & (Mask - 1)) << NextBit;
 
-      if ((Piece & (1U << (NumBits-1))) == 0)
+      if ((Piece & Mask) == 0)
         return Result;
 
       NextBit += NumBits-1;