Allow disabling free memory poisoning in ASAN builds

Summary: There's a lot of overhead when adding ASAN to every cachelib memory allocate/free.  Allow disabling memory poisoning through the config, even in ASAN builds (this lets us keep existing ASAN environments).

Reviewed By: AlnisM

Differential Revision: D101399129

fbshipit-source-id: 1d36cfe302b6cd280ab0305ff2e3174f5d3d18d3

Author: rlyerly

cachelib/allocator/CacheAllocator.h

@@ -25,7 +25,6 @@
 #include <folly/logging/xlog.h>
 #include <folly/synchronization/SanitizeThread.h>
 #include <gtest/gtest.h>
-#include <sanitizer/asan_interface.h>
 
 #include <chrono>
 #include <functional>
@@ -1631,23 +1630,25 @@ class CacheAllocator : public CacheBase {
       return {};
     }
 #if FOLLY_SANITIZE_ADDRESS
-    // Copy the key out so we can use it without disabling instrumentation for
-    // the whole find path. Don't use std utilities (like std::string or
-    // std::memcpy) because they still have ASAN enabled. If it's a small key,
-    // store on the stack. Otherwise malloc on the heap.
     char stackBuf[KAllocation::kKeyMaxLenSmall];
     std::unique_ptr<char[]> heapBuf;
-    char* buffer;
-    if (key.size() <= KAllocation::kKeyMaxLenSmall) {
-      buffer = stackBuf;
-    } else {
-      heapBuf = std::unique_ptr<char[]>(new char[key.size()]);
-      buffer = heapBuf.get();
-    }
-    for (size_t i = 0; i < key.size(); i++) {
-      buffer[i] = key[i];
+    if (config_.isSlabAsanPoisoningEnabled()) {
+      // Copy the key out so we can use it without disabling instrumentation for
+      // the whole find path. Don't use std utilities (like std::string or
+      // std::memcpy) because they still have ASAN enabled. If it's a small key,
+      // store on the stack. Otherwise malloc on the heap.
+      char* buffer;
+      if (key.size() <= KAllocation::kKeyMaxLenSmall) {
+        buffer = stackBuf;
+      } else {
+        heapBuf = std::unique_ptr<char[]>(new char[key.size()]);
+        buffer = heapBuf.get();
+      }
+      for (size_t i = 0; i < key.size(); i++) {
+        buffer[i] = key[i];
+      }
+      key = Key(buffer, key.size());
     }
-    key = Key(buffer, key.size());
 #endif
     if (Item::getRequiredSize(key, /* size */ 0) > allocSize) {
       return {};
@@ -2182,7 +2183,7 @@ class CacheAllocator : public CacheBase {
                   config.reduceFragmentationInAllocationClass)
             : config.defaultAllocSizes,
         config.enableZeroedSlabAllocs, config.disableFullCoredump,
-        config.lockMemory};
+        config.lockMemory, config.isSlabAsanPoisoningEnabled()};
   }
 
   // starts one of the cache workers passing the current instance and the args
@@ -2766,8 +2767,8 @@ CacheAllocator<CacheTrait>::restoreMemoryAllocator() {
           ->attachShm(detail::kShmCacheName, config_.slabMemoryBaseAddr,
                       createShmCacheOpts())
           .addr,
-      config_.getCacheSize(),
-      config_.disableFullCoredump);
+      config_.getCacheSize(), config_.disableFullCoredump,
+      config_.isSlabAsanPoisoningEnabled());
 }
 
 template <typename CacheTrait>

cachelib/allocator/CacheAllocatorConfig.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <folly/CPortability.h>
 #include <folly/Optional.h>
 #include <folly/json/DynamicConverter.h>
 #include <folly/json/json.h>
@@ -370,6 +371,21 @@ class CacheAllocatorConfig {
     return skipPromoteChildrenWhenParentFailed;
   }
 
+  // Enable or disable ASAN slab poisoning. When enabled, freed slab memory is
+  // poisoned so that ASAN can detect use-after-free bugs. This is disabled by
+  // default to avoid overhead in ASAN builds. Enable this for testing or when
+  // actively debugging memory issues. This is a no-op in non-ASAN builds.
+  CacheAllocatorConfig& setSlabAsanPoisoning(bool enable);
+
+  // @return whether slab ASAN poisoning is enabled
+  bool isSlabAsanPoisoningEnabled() const noexcept {
+#if FOLLY_SANITIZE_ADDRESS
+    return enableSlabAsanPoisoning;
+#else
+    return false;
+#endif
+  }
+
   // Enable aggregating pool stats to a single stat
   //
   // When enabled, pool stats from all pools will be aggregated into a single
@@ -718,6 +734,13 @@ class CacheAllocatorConfig {
   // If true, aggregate pool stats into a single stat before exporting
   bool aggregatePoolStats{false};
 
+#if FOLLY_SANITIZE_ADDRESS
+  // If true, slab memory is ASAN-poisoned on free and unpoisoned on allocation.
+  // Disabled by default to avoid overhead in ASAN builds. Enable for testing
+  // or when actively debugging memory issues.
+  bool enableSlabAsanPoisoning{false};
+#endif
+
   friend CacheT;
 
  private:
@@ -1214,6 +1237,15 @@ CacheAllocatorConfig<T>& CacheAllocatorConfig<T>::setNumShards(size_t shards) {
   return *this;
 }
 
+template <typename T>
+CacheAllocatorConfig<T>& CacheAllocatorConfig<T>::setSlabAsanPoisoning(
+    [[maybe_unused]] bool enable) {
+#if FOLLY_SANITIZE_ADDRESS
+  enableSlabAsanPoisoning = enable;
+#endif
+  return *this;
+}
+
 template <typename T>
 CacheAllocatorConfig<T>& CacheAllocatorConfig<T>::enableAggregatePoolStats() {
   aggregatePoolStats = true;

cachelib/allocator/memory/AllocationClass.cpp

@@ -18,7 +18,6 @@
 
 #include <folly/Try.h>
 #include <folly/logging/xlog.h>
-#include <sanitizer/asan_interface.h>
 
 #include "cachelib/allocator/memory/SlabAllocator.h"
 #include "cachelib/common/Exceptions.h"
@@ -52,7 +51,8 @@ AllocationClass::AllocationClass(ClassId classId,
           slabAlloc_.createPtrCompressor<FreeAlloc, CompressedPtr5B>()
 #if FOLLY_SANITIZE_ADDRESS
               ,
-          allocationSize_
+          // Only enable ASAN poisoning in the free list if configured
+          slabAlloc_.isAsanPoisoningEnabled() ? allocationSize_ : 0
 #endif
       } {
   checkState();
@@ -110,7 +110,8 @@ AllocationClass::AllocationClass(
           slabAlloc_.createPtrCompressor<FreeAlloc, CompressedPtr5B>()
 #if FOLLY_SANITIZE_ADDRESS
               ,
-          allocationSize_
+          // Only enable ASAN poisoning in the free list if configured
+          slabAlloc_.isAsanPoisoningEnabled() ? allocationSize_ : 0
 #endif
           ),
       canAllocate_(*object.canAllocate()) {
@@ -133,22 +134,22 @@ AllocationClass::AllocationClass(
 
   // Free slabs are entirely unused
   for (auto* slab : freeSlabs_) {
-    ASAN_POISON_MEMORY_REGION(slab->memoryAtOffset(0), Slab::kSize);
+    slabAlloc_.asanPoisonMemoryRegion(slab->memoryAtOffset(0), Slab::kSize);
   }
 
   // Poison the internal fragmentation at the tail of fully allocated slabs
   const uint32_t usableSize = getAllocsPerSlab() * allocationSize_;
   if (usableSize < Slab::kSize) {
     for (auto* slab : allocatedSlabs_) {
-      ASAN_POISON_MEMORY_REGION(slab->memoryAtOffset(usableSize),
-                                Slab::kSize - usableSize);
+      slabAlloc_.asanPoisonMemoryRegion(slab->memoryAtOffset(usableSize),
+                                        Slab::kSize - usableSize);
     }
   }
 
   // Poison the uncarved region of currSlab_
   if (currSlab_ != nullptr && currOffset_ < Slab::kSize) {
-    ASAN_POISON_MEMORY_REGION(currSlab_->memoryAtOffset(currOffset_),
-                              Slab::kSize - currOffset_);
+    slabAlloc_.asanPoisonMemoryRegion(currSlab_->memoryAtOffset(currOffset_),
+                                      Slab::kSize - currOffset_);
   }
 }
 
@@ -177,7 +178,7 @@ void* AllocationClass::allocateFromCurrentSlabLocked() noexcept {
   XDCHECK(canAllocateFromCurrentSlabLocked());
   void* ret = currSlab_->memoryAtOffset(currOffset_);
   currOffset_ += allocationSize_;
-  ASAN_UNPOISON_MEMORY_REGION(ret, allocationSize_);
+  slabAlloc_.asanUnpoisonMemoryRegion(ret, allocationSize_);
   return ret;
 }
 
@@ -208,7 +209,7 @@ void* AllocationClass::allocateLocked() {
     FreeAlloc* ret = freedAllocations_.getHead();
     XDCHECK(ret != nullptr);
     freedAllocations_.pop();
-    ASAN_UNPOISON_MEMORY_REGION(ret, allocationSize_);
+    slabAlloc_.asanUnpoisonMemoryRegion(ret, allocationSize_);
     return reinterpret_cast<void*>(ret);
   }
 
@@ -687,7 +688,7 @@ void AllocationClass::free(void* memory) {
         throw std::invalid_argument(
             folly::sformat("Allocation {} is already marked as free", memory));
       }
-      ASAN_POISON_MEMORY_REGION(memory, allocationSize_);
+      slabAlloc_.asanPoisonMemoryRegion(memory, allocationSize_);
       allocState[idx] = true;
       return;
     }

cachelib/allocator/memory/MemoryAllocator.cpp

@@ -65,15 +65,17 @@ MemoryAllocator::MemoryAllocator(Config config,
     : config_(std::move(config)),
       slabAllocator_(memoryStart,
                      memSize,
-                     {config_.disableFullCoredump, config_.lockMemory}),
+                     {config_.disableFullCoredump, config_.lockMemory,
+                      config_.enableAsanPoisoning}),
       memoryPoolManager_(slabAllocator_) {
   checkConfig(config_);
 }
 
 MemoryAllocator::MemoryAllocator(Config config, size_t memSize)
     : config_(std::move(config)),
       slabAllocator_(memSize,
-                     {config_.disableFullCoredump, config_.lockMemory}),
+                     {config_.disableFullCoredump, config_.lockMemory,
+                      config_.enableAsanPoisoning}),
       memoryPoolManager_(slabAllocator_) {
   checkConfig(config_);
 }
@@ -82,16 +84,19 @@ MemoryAllocator::MemoryAllocator(
     const serialization::MemoryAllocatorObject& object,
     void* memoryStart,
     size_t memSize,
-    bool disableCoredump)
+    bool disableCoredump,
+    bool enableAsanPoisoning)
     : config_(std::set<uint32_t>{object.allocSizes()->begin(),
                                  object.allocSizes()->end()},
               *object.enableZeroedSlabAllocs(),
               disableCoredump,
-              *object.lockMemory()),
+              *object.lockMemory(),
+              enableAsanPoisoning),
       slabAllocator_(*object.slabAllocator(),
                      memoryStart,
                      memSize,
-                     {config_.disableFullCoredump, config_.lockMemory}),
+                     {config_.disableFullCoredump, config_.lockMemory,
+                      config_.enableAsanPoisoning}),
       memoryPoolManager_(*object.memoryPoolManager(), slabAllocator_) {
   checkConfig(config_);
 }

cachelib/allocator/memory/MemoryAllocator.h

@@ -88,11 +88,13 @@ class MemoryAllocator {
     Config(std::set<uint32_t> sizes,
            bool zeroOnRelease,
            bool disableCoredump,
-           bool _lockMemory)
+           bool _lockMemory,
+           bool _enableAsanPoisoning)
         : allocSizes(std::move(sizes)),
           enableZeroedSlabAllocs(zeroOnRelease),
           disableFullCoredump(disableCoredump),
-          lockMemory(_lockMemory) {}
+          lockMemory(_lockMemory),
+          enableAsanPoisoning(_enableAsanPoisoning) {}
 
     // Hint to determine the allocation class sizes
     std::set<uint32_t> allocSizes;
@@ -110,6 +112,10 @@ class MemoryAllocator {
     // allocator is not shared, user needs to ensure there are appropriate
     // rlimits setup to lock the memory.
     bool lockMemory{false};
+
+    // When true, slab memory is ASAN-poisoned on free and unpoisoned on
+    // allocation so ASAN can detect use-after-free bugs.
+    bool enableAsanPoisoning{false};
   };
 
   // Creates a memory allocator out of the caller allocated memory region. The
@@ -147,10 +153,12 @@ class MemoryAllocator {
   // @param memSize         the size of the memory region that was originally
   //                        used to create this memory allocator
   // @param disableCoredump exclude mapped region from core dumps
+  // @param enableAsanPoisoning When true, slab memory is ASAN-poisoned
   MemoryAllocator(const serialization::MemoryAllocatorObject& object,
                   void* memoryStart,
                   size_t memSize,
-                  bool disableCoredump);
+                  bool disableCoredump,
+                  bool enableAsanPoisoning);
 
   MemoryAllocator(const MemoryAllocator&) = delete;
   MemoryAllocator& operator=(const MemoryAllocator&) = delete;

cachelib/allocator/memory/MemoryPool.cpp

@@ -16,8 +16,6 @@
 
 #include "cachelib/allocator/memory/MemoryPool.h"
 
-#include <sanitizer/asan_interface.h>
-
 #include <algorithm>
 #include <numeric>
 #include <utility>
@@ -86,7 +84,7 @@ MemoryPool::MemoryPool(const serialization::MemoryPoolObject& object,
 
   for (auto freeSlabIdx : *object.freeSlabIdxs()) {
     auto* slab = slabAllocator_.getSlabForIdx(freeSlabIdx);
-    ASAN_POISON_MEMORY_REGION(slab->memoryAtOffset(0), Slab::kSize);
+    slabAllocator_.asanPoisonMemoryRegion(slab->memoryAtOffset(0), Slab::kSize);
     freeSlabs_.push_back(slab);
   }
   checkState();
@@ -406,7 +404,11 @@ void MemoryPool::releaseSlab(SlabReleaseMode mode,
                              ClassId receiverClassId) {
   if (zeroOnRelease) {
 #if FOLLY_SANITIZE_ADDRESS
-    memsetNoAsan(slab->memoryAtOffset(0), 0, Slab::kSize);
+    if (slabAllocator_.isAsanPoisoningEnabled()) {
+      memsetNoAsan(slab->memoryAtOffset(0), 0, Slab::kSize);
+    } else {
+      memset(slab->memoryAtOffset(0), 0, Slab::kSize);
+    }
 #else
     memset(slab->memoryAtOffset(0), 0, Slab::kSize);
 #endif

cachelib/allocator/memory/SlabAllocator.cpp

@@ -20,7 +20,6 @@
 #include <folly/Random.h>
 #include <folly/logging/xlog.h>
 #include <folly/synchronization/SanitizeThread.h>
-#include <sanitizer/asan_interface.h>
 #include <sys/mman.h>
 
 #include <chrono>
@@ -118,7 +117,12 @@ SlabAllocator::SlabAllocator(void* memoryStart,
       memorySize_(roundDownToSlabSize(memorySize)),
       slabMemoryStart_(computeSlabMemoryStart(memoryStart_, memorySize_)),
       nextSlabAllocation_(slabMemoryStart_),
-      ownsMemory_(ownsMemory) {
+      ownsMemory_(ownsMemory)
+#if FOLLY_SANITIZE_ADDRESS
+      ,
+      asanPoisoningEnabled_(config.enableAsanPoisoning)
+#endif
+{
   checkState();
 
   static_assert(!(sizeof(Slab) & (sizeof(Slab) - 1)),
@@ -132,7 +136,7 @@ SlabAllocator::SlabAllocator(void* memoryStart,
     memoryLocker_ = std::thread{[this]() { lockMemoryAsync(); }};
   }
 
-  ASAN_POISON_MEMORY_REGION(
+  asanPoisonMemoryRegion(
       slabMemoryStart_,
       reinterpret_cast<const uint8_t*>(getSlabMemoryEnd()) -
           reinterpret_cast<const uint8_t*>(slabMemoryStart_));
@@ -153,7 +157,12 @@ SlabAllocator::SlabAllocator(const serialization::SlabAllocatorObject& object,
       slabMemoryStart_(computeSlabMemoryStart(memoryStart_, memorySize_)),
       nextSlabAllocation_(getSlabForIdx(*object.nextSlabIdx())),
       canAllocate_(*object.canAllocate()),
-      ownsMemory_(false) {
+      ownsMemory_(false)
+#if FOLLY_SANITIZE_ADDRESS
+      ,
+      asanPoisoningEnabled_(config.enableAsanPoisoning)
+#endif
+{
   if (Slab::kSize != *object.slabSize()) {
     throw std::invalid_argument(folly::sformat(
         "current slab size {} does not match the previous one {}",
@@ -217,17 +226,17 @@ SlabAllocator::SlabAllocator(const serialization::SlabAllocatorObject& object,
   // Poison slabs that are not allocated. Tests depend on this running after
   // checkState().
 
-  ASAN_POISON_MEMORY_REGION(
+  asanPoisonMemoryRegion(
       nextSlabAllocation_,
       reinterpret_cast<const uint8_t*>(getSlabMemoryEnd()) -
           reinterpret_cast<const uint8_t*>(nextSlabAllocation_));
 
   for (const auto* slab : freeSlabs_) {
-    ASAN_POISON_MEMORY_REGION(slab, Slab::kSize);
+    asanPoisonMemoryRegion(slab, Slab::kSize);
   }
 
   for (const auto* slab : advisedSlabs_) {
-    ASAN_POISON_MEMORY_REGION(slab, Slab::kSize);
+    asanPoisonMemoryRegion(slab, Slab::kSize);
   }
 }