BitstreamWriter: Better assertion when a block's abbrev code size is too small

[jordan_rose]

Summary

Before: ((Val & ~(~0U >> (32-NumBits))) == 0 && "High bits set!")
After: (Abbrev <= (~0U >> (32-CurCodeSize)) && "Block code length is too small")

The same check, but earlier and with a better error message for this specific use case. (I was originally going to put this in EmitAbbrev instead, but realized that it's possible for some overly clever soul to have a block that just never uses high-abbrev records.)

[jordan_rose]

Reviving this old review. I'd like *someone* to tell me I didn't miss something.

[lebedev.ri]

See inline notes.
Also, the edge case with CurCodeSize == 32, if that is possible, should be added.

include/llvm/Bitcode/BitstreamWriter.h

303

What is the expected range of CurCodeSize values?
This has UB if CurCodeSize == 32.
It would be much safer to just use the variant with subtraction:

assert(Abbrev <= (~(~0U >> (32-CurCodeSize))) && "Block code length is too small");

unittests/Bitcode/BitstreamWriterTest.cpp

65

I'd expect to see three tests.
One last valid one, the invalid one, and the next after the first invalid one.
So 3U, 4U, 5U.

[jordan_rose]

include/llvm/Bitcode/BitstreamWriter.h

303

Good point. I don't think anyone's trying to use 32-bit record codes, but the spec doesn't seem to impose a limit. I'll switch it.

unittests/Bitcode/BitstreamWriterTest.cpp

65

This is just an arbitrary data value. What makes this too big is the CodeLen of 2 above. But I can certainly test overflowing a CodeLen of 3 as well.

[jordan_rose]

unittests/Bitcode/BitstreamWriterTest.cpp

65

Oh, I see what you mean. Values 0-3 have a special meaning in bitstream mode, so I can't actually do that test, but I'll add one for overflowing 3.

[lebedev.ri]

unittests/Bitcode/BitstreamWriterTest.cpp

65

Values 0-3 have a special meaning in bitstream mode

Right, so do that for CodeLen = 3, it will still be 3 values :)

[jordan_rose]

Took recommendation about the form of the check, pumped up the tests.

[lebedev.ri]

unittests/Bitcode/BitstreamWriterTest.cpp

104–211

This looks rather monstrous, with a lot of duplication.
I was suggesting to just do:

TEST(BitstreamWriterTest, trapOnSmallAbbrevUsingBlockInfo_CodeLength3_Ok) {
  SmallString<64> Buffer;
  BitstreamWriter W(Buffer);
  W.EnterBlockInfoBlock();
  auto Abbrev = std::make_shared<BitCodeAbbrev>();
  Abbrev->Add(BitCodeAbbrevOp(2U);
  (void)W.EmitBlockInfoAbbrev(bitc::FIRST_APPLICATION_BLOCKID,
                              std::move(Abbrev));
  W.ExitBlock();
  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/3);
  W.EmitRecordWithAbbrev(Abbrev, makeArrayRef(2U));
  W.ExitBlock();
}

TEST(BitstreamWriterTest, trapOnSmallAbbrevUsingBlockInfo_CodeLength3_LastOk) {
  SmallString<64> Buffer;
  BitstreamWriter W(Buffer);
  W.EnterBlockInfoBlock();
  auto Abbrev = std::make_shared<BitCodeAbbrev>();
  Abbrev->Add(BitCodeAbbrevOp(3U);
  (void)W.EmitBlockInfoAbbrev(bitc::FIRST_APPLICATION_BLOCKID,
                              std::move(Abbrev));
  W.ExitBlock();
  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/3);
  W.EmitRecordWithAbbrev(Abbrev, makeArrayRef(3U));
  W.ExitBlock();
}

TEST(BitstreamWriterTest, trapOnSmallAbbrevUsingBlockInfo_CodeLength3_NotOk) {
  SmallString<64> Buffer;
  BitstreamWriter W(Buffer);
  W.EnterBlockInfoBlock();
  auto Abbrev = std::make_shared<BitCodeAbbrev>();
  Abbrev->Add(BitCodeAbbrevOp(4U);
  (void)W.EmitBlockInfoAbbrev(bitc::FIRST_APPLICATION_BLOCKID,
                              std::move(Abbrev));
  W.ExitBlock();
  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/3);
  W.EmitRecordWithAbbrev(Abbrev, makeArrayRef(4U));
  W.ExitBlock();

TEST(BitstreamWriterTest, trapOnSmallAbbrevUsingBlockInfo_CodeLength3_NoOverFlow) {
  SmallString<64> Buffer;
  BitstreamWriter W(Buffer);
  W.EnterBlockInfoBlock();
  auto Abbrev = std::make_shared<BitCodeAbbrev>();
  Abbrev->Add(BitCodeAbbrevOp(4294967295U); // i32 -1 
  (void)W.EmitBlockInfoAbbrev(bitc::FIRST_APPLICATION_BLOCKID,
                              std::move(Abbrev));
  W.ExitBlock();
  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/32);
  W.EmitRecordWithAbbrev(Abbrev, makeArrayRef(4294967295U));
  W.ExitBlock();
}

[jordan_rose]

unittests/Bitcode/BitstreamWriterTest.cpp

104–211

The values in the abbreviation have nothing to do with the abbreviation codes themselves; in order to get to code 8 you have to actually register 5 abbreviations. (0-3 are reserved.) I could probably pull out some of this into helper functions but I didn't feel like it was going to be super important.

[lebedev.ri]

I would still like to see a test with CodeLen=32, and even a single call to EmitRecordWithAbbrevImpl()
(ubsan would then have caught the ub in the initial approach), but otherwise this seems reasonable to me.
I have no expertise in this code, so *please* wait for someone else to review this, too.

include/llvm/Bitcode/BitstreamWriter.h

303

I think you want to move this assert to before the first use of Abbrev (before line 301)

unittests/Bitcode/BitstreamWriterTest.cpp

104–211

Oh duh, i knew that, but kinda forgot :)

[jordan_rose]

Thanks, Roman. Too bad this is such a stable piece of LLVM, or else I'd have more people qualified to review!

include/llvm/Bitcode/BitstreamWriter.h

303

It shouldn't matter in this case. The first use (and the above assertion) is for indexing into a write-time data structure; the second is for the actual representation in the bitstream.

[lebedev.ri]

include/llvm/Bitcode/BitstreamWriter.h

303

Point being, if you don't move it, you won't be able to add a test for 32-bit CurCodeSize overflow.

[jordan_rose]

include/llvm/Bitcode/BitstreamWriter.h

303

I have to admit I'm not really interested in that. It's worth writing code that won't silently invoke undefined behavior, but in practice no one is actually going to have billions of abbreviations, and it's not interesting to me which assertion fires in that case.

include/llvm/Bitcode/BitstreamWriter.h

private:
   /// the code.
   template <typename uintty>
   void EmitRecordWithAbbrevImpl(unsigned Abbrev, ArrayRef<uintty> Vals,
                                 StringRef Blob, Optional<unsigned> Code) {
     const char *BlobData = Blob.data();
     unsigned BlobLen = (unsigned) Blob.size();
     unsigned AbbrevNo = Abbrev-bitc::FIRST_APPLICATION_ABBREV;
     assert(AbbrevNo < CurAbbrevs.size() && "Invalid abbrev #!");
+    assert(Abbrev < (1U << CurCodeSize) && "Block code length is too small");
     const BitCodeAbbrev *Abbv = CurAbbrevs[AbbrevNo].get();
 
     EmitCode(Abbrev);
 
     unsigned i = 0, e = static_cast<unsigned>(Abbv->getNumOperandInfos());
     if (Code) {
       assert(e && "Expected non-empty abbreviation");
       const BitCodeAbbrevOp &Op = Abbv->getOperandInfo(i++);

unittests/Bitcode/BitstreamWriterTest.cpp

 
 TEST(BitstreamWriterTest, emitBlob4ByteAligned) {
   SmallString<64> Buffer;
   BitstreamWriter W(Buffer);
   W.emitBlob("str0", /* ShouldEmitSize */ false);
   EXPECT_EQ(StringRef("str0"), Buffer);
 }
 
+#ifndef NDEBUG
+TEST(BitstreamWriterTest, trapOnSmallAbbrev) {
+  SmallString<64> Buffer;
+  BitstreamWriter W(Buffer);
+  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/2);
+  auto Abbrev = std::make_shared<BitCodeAbbrev>();
+  Abbrev->Add(BitCodeAbbrevOp(42U));
+  unsigned AbbrevCode = W.EmitAbbrev(std::move(Abbrev));
+  EXPECT_DEATH({ W.EmitRecordWithAbbrev(AbbrevCode, makeArrayRef(42U)); },
+               "Block code length is too small");
+  W.ExitBlock();
+}
+
+TEST(BitstreamWriterTest, trapOnSmallAbbrevUsingBlockInfo) {
+  SmallString<64> Buffer;
+  BitstreamWriter W(Buffer);
+  W.EnterBlockInfoBlock();
+  auto Abbrev = std::make_shared<BitCodeAbbrev>();
+  Abbrev->Add(BitCodeAbbrevOp(42U));
+  unsigned AbbrevCode = W.EmitBlockInfoAbbrev(bitc::FIRST_APPLICATION_BLOCKID,
+                                              std::move(Abbrev));
+  W.ExitBlock();
+  W.EnterSubblock(bitc::FIRST_APPLICATION_BLOCKID, /*CodeLen*/2);
+  EXPECT_DEATH({ W.EmitRecordWithAbbrev(AbbrevCode, makeArrayRef(42U)); },
+               "Block code length is too small");
+  W.ExitBlock();
+}
+#endif
+
 } // end namespace