[Support] Fix UB in BumpPtrAllocator when first allocation is zero.

[sammccall]

Summary

BumpPtrAllocator::Allocate() is marked attribute((returns_nonnull)) when the
compiler supports it, which makes it UB to return null.

When there have been no allocations yet, the current slab is [nullptr, nullptr).
A zero-sized allocation fits in this range, and so Allocate(0, 1) returns null.

There's no explicit docs whether Allocate(0) is valid. I think we have to assume
that it is:

• the implementation tries to support it (e.g. >= tests instead of >)
• malloc(0) is allowed
• requiring each callsite to do a check is bug-prone
• I found real LLVM code that makes zero-sized allocations

[sammccall]

Here's an example of an innocuous-looking Allocate() that can be zero-sized: https://github.com/llvm/llvm-project/blob/16dcbb53dc7968a3752661aac731172ebe0faf64/clang-tools-extra/pseudo/lib/Forest.cpp#L130

Found this by running a fuzzer on clang-pseudo, in an environment that instruments returns_nonnull functions (probably it's UBSan)

[hokein]

llvm/include/llvm/Support/Allocator.h

142–145

mention this specific behavior of Allocate(0) in the doc: Allocate(0) returns a non-NULL zero-sized pointer (Dereferencing this pointer is UB).

llvm/include/llvm/Support/Allocator.h

void Reset() {
     __asan_poison_memory_region(*Slabs.begin(), computeSlabSize(0));
     DeallocateSlabs(std::next(Slabs.begin()), Slabs.end());
     Slabs.erase(std::next(Slabs.begin()), Slabs.end());
   }
 
   /// Allocate space at the specified alignment.
   // This method is *not* marked noalias, because
   // SpecificBumpPtrAllocator::DestroyAll() loops over all allocations, and
   // that loop is not based on the Allocate() return value.
+  //
+  // Allocate(0, N) is valid, it returns a non-null pointer (which should not
+  // be dereferenced).
   LLVM_ATTRIBUTE_RETURNS_NONNULL void *Allocate(size_t Size, Align Alignment) {
     // Keep track of how many bytes we've allocated.
     BytesAllocated += Size;
 
     size_t Adjustment = offsetToAlignedAddr(CurPtr, Alignment);
     assert(Adjustment + Size >= Size && "Adjustment + Size must not overflow");
 
     size_t SizeToAllocate = Size;
 #if LLVM_ADDRESS_SANITIZER_BUILD
     // Add trailing bytes as a "red zone" under ASan.
     SizeToAllocate += RedZoneSize;
 #endif
 
     // Check if we have enough space.
-    if (Adjustment + SizeToAllocate <= size_t(End - CurPtr)) {
+    if (Adjustment + SizeToAllocate <= size_t(End - CurPtr)
+        // We can't return nullptr even for a zero-sized allocation!
+        && CurPtr != nullptr) {
       char *AlignedPtr = CurPtr + Adjustment;
       CurPtr = AlignedPtr + SizeToAllocate;
       // Update the allocation point of this memory block in MemorySanitizer.
       // Without this, MemorySanitizer messages for values originated from here
       // will point to the allocation of the entire slab.
       __msan_allocated_memory(AlignedPtr, Size);
       // Similarly, tell ASan about this space.
       __asan_unpoison_memory_region(AlignedPtr, Size);

llvm/unittests/Support/AllocatorTest.cpp

TEST(AllocatorTest, TestAlignment) {
   a = (uintptr_t)Alloc.Allocate(1, 32);
   EXPECT_EQ(0U, a & 31);
   a = (uintptr_t)Alloc.Allocate(1, 64);
   EXPECT_EQ(0U, a & 63);
   a = (uintptr_t)Alloc.Allocate(1, 128);
   EXPECT_EQ(0U, a & 127);
 }
 
+// Test zero-sized allocations.
+// In general we don't need to allocate memory for these.
+// However Allocate never returns null, so if the first allocation is zero-sized
+// we end up creating a slab for it.
+TEST(AllocatorTest, TestZero) {
+  BumpPtrAllocator Alloc;
+  EXPECT_EQ(0u, Alloc.GetNumSlabs());
+  EXPECT_EQ(0u, Alloc.getBytesAllocated());
+
+  void *Empty = Alloc.Allocate(0, 1);
+  EXPECT_NE(Empty, nullptr) << "Allocate is __attribute__((returns_nonnull))";
+  EXPECT_EQ(1u, Alloc.GetNumSlabs()) << "Allocated a slab to point to";
+  EXPECT_EQ(0u, Alloc.getBytesAllocated());
+
+  void *Large = Alloc.Allocate(4096, 1);
+  EXPECT_EQ(1u, Alloc.GetNumSlabs());
+  EXPECT_EQ(4096u, Alloc.getBytesAllocated());
+  EXPECT_EQ(Empty, Large);
+
+  void *Empty2 = Alloc.Allocate(0, 1);
+  EXPECT_NE(Empty2, nullptr);
+  EXPECT_EQ(1u, Alloc.GetNumSlabs());
+  EXPECT_EQ(4096u, Alloc.getBytesAllocated());
+}
+
 // Test allocating just over the slab size.  This tests a bug where before the
 // allocator incorrectly calculated the buffer end pointer.
 TEST(AllocatorTest, TestOverflow) {
   BumpPtrAllocator Alloc;
 
   // Fill the slab right up until the end pointer.
   Alloc.Allocate(4096, 1);
   EXPECT_EQ(1U, Alloc.GetNumSlabs());