Add red zones to BumpPtrAllocator under ASan

[jordan_rose]

Summary

To help catch buffer overruns, this patch changes BumpPtrAllocator to insert an extra unused byte between allocations when building with ASan. SpecificBumpPtrAllocator opts out of this behavior, since it needs to destroy its items later by walking the allocated memory.

I wrote this up while trying to catch a bug. It didn't actually catch anything new, but I didn't want it to go to waste. Do people think it's generally useful?

[pete]

Very cool. Justin and I looked at ASan on BumpPtrAllocator a while ago and I think this would have been useful then too. I think the reason you didn't find issues is that he fixed many of them already.

I think we could support SpecificBumpPtrAllocator but it would require walking the objects by 'sizeof(T) + RedZoneSize'. RedZoneSize itself would also have to be set to make sure we continue to have appropriate alignment. Not something you need to do now, but could be interesting in future.

So LGTM. Thanks!

[jordan_rose]

Thanks, Pete! Will land soon, then.

I think we could support SpecificBumpPtrAllocator but it would require walking the objects by 'sizeof(T) + RedZoneSize'. RedZoneSize itself would also have to be set to make sure we continue to have appropriate alignment. Not something you need to do now, but could be interesting in future.

FWIW, this doesn't work as is because SpecificBumpPtrAllocator lets you allocate *arrays* of objects as well as single objects, and the arrays obviously can't have red zones between elements.

[pete]

In D30723#695798, @jordan_rose wrote:

Thanks, Pete! Will land soon, then.

I think we could support SpecificBumpPtrAllocator but it would require walking the objects by 'sizeof(T) + RedZoneSize'. RedZoneSize itself would also have to be set to make sure we continue to have appropriate alignment. Not something you need to do now, but could be interesting in future.

FWIW, this doesn't work as is because SpecificBumpPtrAllocator lets you allocate *arrays* of objects as well as single objects, and the arrays obviously can't have red zones between elements.

Good point. Never thought of that. Wouldn't be surprised if there's a whole bunch of users who never use arrays and could benefit, but still something to think about more.

[jordan_rose]

Committed in rL297310.

include/llvm/Support/Allocator.h

BumpPtrAllocatorImpl(T &&Allocator)
       : CurPtr(nullptr), End(nullptr), BytesAllocated(0),
         Allocator(std::forward<T &&>(Allocator)) {}
 
   // Manually implement a move constructor as we must clear the old allocator's
   // slabs as a matter of correctness.
   BumpPtrAllocatorImpl(BumpPtrAllocatorImpl &&Old)
       : CurPtr(Old.CurPtr), End(Old.End), Slabs(std::move(Old.Slabs)),
         CustomSizedSlabs(std::move(Old.CustomSizedSlabs)),
-        BytesAllocated(Old.BytesAllocated),
+        BytesAllocated(Old.BytesAllocated), RedZoneSize(Old.RedZoneSize),
         Allocator(std::move(Old.Allocator)) {
     Old.CurPtr = Old.End = nullptr;
     Old.BytesAllocated = 0;
     Old.Slabs.clear();
     Old.CustomSizedSlabs.clear();
   }
 
   ~BumpPtrAllocatorImpl() {
     DeallocateSlabs(Slabs.begin(), Slabs.end());
     DeallocateCustomSizedSlabs();
   }
 
   BumpPtrAllocatorImpl &operator=(BumpPtrAllocatorImpl &&RHS) {
     DeallocateSlabs(Slabs.begin(), Slabs.end());
     DeallocateCustomSizedSlabs();
 
     CurPtr = RHS.CurPtr;
     End = RHS.End;
     BytesAllocated = RHS.BytesAllocated;
+    RedZoneSize = RHS.RedZoneSize;
     Slabs = std::move(RHS.Slabs);
     CustomSizedSlabs = std::move(RHS.CustomSizedSlabs);
     Allocator = std::move(RHS.Allocator);
 
     RHS.CurPtr = RHS.End = nullptr;
     RHS.BytesAllocated = 0;
     RHS.Slabs.clear();
     RHS.CustomSizedSlabs.clear();
Allocate(size_t Size, size_t Alignment) {
     assert(Alignment > 0 && "0-byte alignnment is not allowed. Use 1 instead.");
 
     // Keep track of how many bytes we've allocated.
     BytesAllocated += Size;
 
     size_t Adjustment = alignmentAdjustment(CurPtr, Alignment);
     assert(Adjustment + Size >= Size && "Adjustment + Size must not overflow");
 
+    size_t SizeToAllocate = Size;
+#if LLVM_ADDRESS_SANITIZER_BUILD
+    // Add a trailing byte as a "red zone" under ASan.
+    SizeToAllocate += RedZoneSize;
+#endif
+
     // Check if we have enough space.
-    if (Adjustment + Size <= size_t(End - CurPtr)) {
+    if (Adjustment + SizeToAllocate <= size_t(End - CurPtr)) {
       char *AlignedPtr = CurPtr + Adjustment;
-      CurPtr = AlignedPtr + Size;
+      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);
       return AlignedPtr;
     }
 
     // If Size is really big, allocate a separate slab for it.
-    size_t PaddedSize = Size + Alignment - 1;
+    size_t PaddedSize = SizeToAllocate + Alignment - 1;
     if (PaddedSize > SizeThreshold) {
       void *NewSlab = Allocator.Allocate(PaddedSize, 0);
       // We own the new slab and don't want anyone reading anyting other than
       // pieces returned from this method.  So poison the whole slab.
       __asan_poison_memory_region(NewSlab, PaddedSize);
       CustomSizedSlabs.push_back(std::make_pair(NewSlab, PaddedSize));
 
       uintptr_t AlignedAddr = alignAddr(NewSlab, Alignment);
       assert(AlignedAddr + Size <= (uintptr_t)NewSlab + PaddedSize);
       char *AlignedPtr = (char*)AlignedAddr;
       __msan_allocated_memory(AlignedPtr, Size);
       __asan_unpoison_memory_region(AlignedPtr, Size);
       return AlignedPtr;
     }
 
     // Otherwise, start a new slab and try again.
     StartNewSlab();
     uintptr_t AlignedAddr = alignAddr(CurPtr, Alignment);
-    assert(AlignedAddr + Size <= (uintptr_t)End &&
+    assert(AlignedAddr + SizeToAllocate <= (uintptr_t)End &&
            "Unable to allocate memory!");
     char *AlignedPtr = (char*)AlignedAddr;
-    CurPtr = AlignedPtr + Size;
+    CurPtr = AlignedPtr + SizeToAllocate;
     __msan_allocated_memory(AlignedPtr, Size);
     __asan_unpoison_memory_region(AlignedPtr, Size);
     return AlignedPtr;
   }
 
   // Pull in base class overloads.
   using AllocatorBase<BumpPtrAllocatorImpl>::Allocate;
 
for (auto I = Slabs.begin(), E = Slabs.end(); I != E; ++I)
       TotalMemory += computeSlabSize(std::distance(Slabs.begin(), I));
     for (auto &PtrAndSize : CustomSizedSlabs)
       TotalMemory += PtrAndSize.second;
     return TotalMemory;
   }
 
   size_t getBytesAllocated() const { return BytesAllocated; }
 
+  void setRedZoneSize(size_t NewSize) {
+    RedZoneSize = NewSize;
+  }
+
   void PrintStats() const {
     detail::printBumpPtrAllocatorStats(Slabs.size(), BytesAllocated,
                                        getTotalMemory());
   }
 
 private:
   /// \brief The current pointer into the current slab.
   ///
private:
   /// \brief Custom-sized slabs allocated for too-large allocation requests.
   SmallVector<std::pair<void *, size_t>, 0> CustomSizedSlabs;
 
   /// \brief How many bytes we've allocated.
   ///
   /// Used so that we can compute how much space was wasted.
   size_t BytesAllocated;
 
+  /// \brief The number of bytes to put between allocations when running under
+  /// a sanitizer.
+  size_t RedZoneSize = 1;
+
   /// \brief The allocator instance we use to get slabs of memory.
   AllocatorT Allocator;
 
   static size_t computeSlabSize(unsigned SlabIdx) {
     // Scale the actual allocated slab size based on the number of slabs
     // allocated. Every 128 slabs allocated, we double the allocated size to
     // reduce allocation frequency, but saturate at multiplying the slab size by
     // 2^30.

 /// allocated.
 ///
 /// This allows calling the destructor in DestroyAll() and when the allocator is
 /// destroyed.
 template <typename T> class SpecificBumpPtrAllocator {
   BumpPtrAllocator Allocator;
 
 public:
-  SpecificBumpPtrAllocator() = default;
+  SpecificBumpPtrAllocator() {
+    // Because SpecificBumpPtrAllocator walks the memory to call destructors,
+    // it can't have red zones between allocations.
+    Allocator.setRedZoneSize(0);
+  }
   SpecificBumpPtrAllocator(SpecificBumpPtrAllocator &&Old)
       : Allocator(std::move(Old.Allocator)) {}
   ~SpecificBumpPtrAllocator() { DestroyAll(); }
 
   SpecificBumpPtrAllocator &operator=(SpecificBumpPtrAllocator &&RHS) {
     Allocator = std::move(RHS.Allocator);
     return *this;
   }