[Bitstream] Fix UB in left-shift in ReadVBR

[jkorous]

Summary

This patch aims to address: https://reviews.llvm.org/D119182

This commit from today improves the situation but doesn't fix the UB.
https://github.com/llvm/llvm-project/commit/67348c8acfc205785996f2aea21b442f4b76f2c2

I am trying to prevent the UB in a mechanical fashion. Possibly there's a better solution and the error messages can very likely be worded better. I am not familiar with the code - open to suggestions!

When running the test against ubsanitized build (-fsanitize=integer) it still triggers couple issues but not the one that lead to the test getting disabled in https://reviews.llvm.org/D119182

> bin/llvm-lit -v ../repo.git/llvm/test/Bitcode/invalid-no-ubsan.test
/tmp/llvm-project/include/llvm/Support/DJB.h:23:12: runtime error: left shift of 4042302735 by 5 places cannot be represented in type 'uint32_t' (aka 'unsigned int')
/tmp/llvm-project/include/llvm/Support/DJB.h:23:18: runtime error: unsigned integer overflow: 504668640 + 4042302735 cannot be represented in type 'unsigned int'
/tmp/llvm-project/include/llvm/Option/ArgList.h:135:42: runtime error: negation of 1 cannot be represented in type 'unsigned int'
/tmp/llvm-project/include/llvm/ADT/DenseMap.h:807:86: runtime error: unsigned integer overflow: 0 - 1 cannot be represented in type 'unsigned int'
/tmp/llvm-project/lib/Option/ArgList.cpp:68:18: runtime error: negation of 1 cannot be represented in type 'unsigned int'
/tmp/llvm-project/lib/Support/StringRef.cpp:299:45: runtime error: unsigned integer overflow: 0 - 1 cannot be represented in type 'unsigned long'
/tmp/llvm-project/include/llvm/ADT/StringRef.h:864:74: runtime error: unsigned integer overflow: 18446744073709551615 + 1 cannot be represented in type 'unsigned long'

[browneee]

Thanks for following up on D119182!

llvm/include/llvm/Bitstream/BitstreamReader.h

230

nit: Consider improving the error message for this test by adding some context (i.e, it is the size which has invalid numbits) to the error, at the caller - here.

237

nit: Provide the bad value of NumBits in the error message?

246–247

Why not simpler condition?

if (NextBit >= 32) { ... }

And move later in the loop for a little less overhead (the first iteration will always be NextBit==0).

llvm/test/Bitcode/invalid-no-ubsan.test

7 ↗

(On Diff #407021)

With these tests fixed, there is no reason for a separate file (separate file is so UNSUPPORTED: ubsan only applies to these, and not the rest).

These tests should be moved back to invalid.test.

[nikic]

Can you please rebase this on main and clarify which UB this addresses that is not covered by either https://github.com/llvm/llvm-project/commit/67348c8acfc205785996f2aea21b442f4b76f2c2 or https://github.com/llvm/llvm-project/commit/bf17cb294af422512b522575f3f74ab7964bf911? Please note that -fsanitize=integer includes a number of checks that are not actually undefined behavior. Some of the warnings quoted in the summary definitely cannot be fixed (e.g. DJB overflow is expected).

I believe a correct NumBits value is supposed to be a caller invariant (though this should probably be at least an assertion), which previously was not enforced in some cases due to an incorrect MaxChunkSize value.

[jkorous]

rebased + improved error messages

[jkorous]

clang-format

[jkorous]

llvm/include/llvm/Bitstream/BitstreamReader.h

246–247

You're absolutely right.

[jkorous]

@nikic Sorry, my statement was incorrect. What I meant is that the method doesn't prevent UB via assert or runtime check of NumBits value.
I'm happy to go with assert if that'd be a better fit.

[browneee]

LGTM. Improves error message, and the new local variables are also an improvement.

[nikic]

In D119307#3316281, @jkorous wrote:

@nikic Sorry, my statement was incorrect. What I meant is that the method doesn't prevent UB via assert or runtime check of NumBits value.
I'm happy to go with assert if that'd be a better fit.

Yes, I think an assert would be a better fit here semantically.

[jkorous]

switched to assert
added the assert also to ReadVBR64

[jkorous]

I noticed ReadVBR64 method below does the same arithmetics and keeping the implementation in sync seems like a no-brainer.

Ultimately the error message propagation is unrelated to the assert but it does seem like a minor improvement - I'll just land it as a separate commit.

I plan to land this tomorrow unless there's some additional feedback.

Thank you both for the opportunity to learn a bit about this part of llvm-project! :)

llvm/include/llvm/Bitstream/BitstreamReader.h

Expected<word_t> Read(unsigned NumBits) {
 
     BitsInCurWord -= BitsLeft;
 
     R |= R2 << (NumBits - BitsLeft);
 
     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;
       if (NextBit >= 32)
         return createStringError(std::errc::illegal_byte_sequence,
                                  "Unterminated VBR");
 
       MaybeRead = Read(NumBits);
       if (!MaybeRead)
         return MaybeRead;
       Piece = MaybeRead.get();
     }
   }
 
   // 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;
       if (NextBit >= 64)
         return createStringError(std::errc::illegal_byte_sequence,
                                  "Unterminated VBR");
 
       MaybeRead = Read(NumBits);