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

[Sam McCall via llvm-commits]

Author: Sam McCall
Date: 2022-05-06T08:57:27+02:00
New Revision: ba0d50ad7ec664bfc79f691bfd358d6a0d815505

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

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

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

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

Added: 
    

Modified: 
    llvm/include/llvm/Support/Allocator.h
    llvm/unittests/Support/AllocatorTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Support/Allocator.h b/llvm/include/llvm/Support/Allocator.h
index ec5ed06b7fa4b..5ca0c9decac30 100644
--- a/llvm/include/llvm/Support/Allocator.h
+++ b/llvm/include/llvm/Support/Allocator.h
@@ -140,6 +140,9 @@ class BumpPtrAllocatorImpl
   // 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;
@@ -154,7 +157,9 @@ class BumpPtrAllocatorImpl
 #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.

diff  --git a/llvm/unittests/Support/AllocatorTest.cpp b/llvm/unittests/Support/AllocatorTest.cpp
index c41f597c404dd..665d1b7e38f03 100644
--- a/llvm/unittests/Support/AllocatorTest.cpp
+++ b/llvm/unittests/Support/AllocatorTest.cpp
@@ -99,6 +99,31 @@ TEST(AllocatorTest, TestAlignment) {
   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) {