scudo: Move the management of the UseMemoryTagging bit out of the Primary. NFCI.

[pcc]

Summary

The primary and secondary allocators will need to share this bit,
so move the management of the bit to the combined allocator and
make useMemoryTagging() a free function.

Depends on D93729

[eugenis]

LGTM

compiler-rt/lib/scudo/standalone/combined.h

void *allocate(UNUSED uptr Size) {
       Chunk::UnpackedHeader Header = {};
       Header.ClassId = QuarantineClassId & Chunk::ClassIdMask;
       Header.SizeOrUnusedBytes = sizeof(QuarantineBatch);
       Header.State = Chunk::State::Allocated;
       Chunk::storeHeader(Allocator.Cookie, Ptr, &Header);
 
       // Reset tag to 0 as this chunk may have been previously used for a tagged
       // user allocation.
-      if (UNLIKELY(Allocator.useMemoryTagging()))
+      if (UNLIKELY(useMemoryTagging<Params>(Allocator.Primary.Options.load())))
         storeTags(reinterpret_cast<uptr>(Ptr),
                   reinterpret_cast<uptr>(Ptr) + sizeof(QuarantineBatch));
 
       return Ptr;
     }
 
     void deallocate(void *Ptr) {
       const uptr QuarantineClassId = SizeClassMap::getClassIdBySize(
void initLinkerInitialized() {
     if (getFlags()->zero_contents)
       Primary.Options.setFillContentsMode(ZeroFill);
     else if (getFlags()->pattern_fill_contents)
       Primary.Options.setFillContentsMode(PatternOrZeroFill);
     if (getFlags()->dealloc_type_mismatch)
       Primary.Options.set(OptionBit::DeallocTypeMismatch);
     if (getFlags()->delete_size_mismatch)
       Primary.Options.set(OptionBit::DeleteSizeMismatch);
+    if (allocatorSupportsMemoryTagging<Params>() &&
+        systemSupportsMemoryTagging())
+      Primary.Options.set(OptionBit::UseMemoryTagging);
     Primary.Options.set(OptionBit::UseOddEvenTags);
 
     QuarantineMaxChunkSize =
         static_cast<u32>(getFlags()->quarantine_max_chunk_size);
 
     Stats.initLinkerInitialized();
     const s32 ReleaseToOsIntervalMs = getFlags()->release_to_os_interval_ms;
     Primary.initLinkerInitialized(ReleaseToOsIntervalMs);
#endif // GWP_ASAN_HOOKS
   //   the last two items).
   void commitBack(TSD<ThisT> *TSD) {
     Quarantine.drain(&TSD->QuarantineCache,
                      QuarantineCallback(*this, TSD->Cache));
     TSD->Cache.destroy(&Stats);
   }
 
   ALWAYS_INLINE void *untagPointerMaybe(void *Ptr) {
-    if (Primary.SupportsMemoryTagging)
+    if (allocatorSupportsMemoryTagging<Params>())
       return reinterpret_cast<void *>(
           untagPointer(reinterpret_cast<uptr>(Ptr)));
     return Ptr;
   }
 
   NOINLINE u32 collectStackTrace() {
 #ifdef HAVE_ANDROID_UNSAFE_FRAME_POINTER_CHASE
     // Discard collectStackTrace() frame and allocator function frame.
if (LIKELY(ClassId)) {
       // allocations. We assume that guard pages are enough to protect these
       // allocations.
       //
       // FIXME: When the kernel provides a way to set the background tag of a
       // mapping, we should be able to tag secondary allocations as well.
       //
       // When memory tagging is enabled, zeroing the contents is done as part of
       // setting the tag.
-      if (UNLIKELY(useMemoryTagging(Options))) {
+      if (UNLIKELY(useMemoryTagging<Params>(Options))) {
         uptr PrevUserPtr;
         Chunk::UnpackedHeader Header;
         const uptr BlockSize = PrimaryT::getSizeByClassId(ClassId);
         const uptr BlockEnd = BlockUptr + BlockSize;
         // If possible, try to reuse the UAF tag that was set by deallocate().
         // For simplicity, only reuse tags if we have the same start address as
         // the previous allocation. This handles the majority of cases since
         // most allocations will not be more aligned than the minimum alignment.
#endif // GWP_ASAN_HOOKS
     if (reinterpret_cast<uptr>(OldPtr) + NewSize <= BlockEnd) {
       if (NewSize > OldSize || (OldSize - NewSize) < getPageSizeCached()) {
         Chunk::UnpackedHeader NewHeader = OldHeader;
         NewHeader.SizeOrUnusedBytes =
             (ClassId ? NewSize
                      : BlockEnd - (reinterpret_cast<uptr>(OldPtr) + NewSize)) &
             Chunk::SizeOrUnusedBytesMask;
         Chunk::compareExchangeHeader(Cookie, OldPtr, &NewHeader, &OldHeader);
-        if (UNLIKELY(ClassId && useMemoryTagging(Options))) {
+        if (UNLIKELY(ClassId && useMemoryTagging<Params>(Options))) {
           resizeTaggedChunk(reinterpret_cast<uptr>(OldTaggedPtr) + OldSize,
                             reinterpret_cast<uptr>(OldTaggedPtr) + NewSize,
                             BlockEnd);
           storeAllocationStackMaybe(Options, OldPtr);
         }
         return OldTaggedPtr;
       }
     }
void iterateOverChunks(uptr Base, uptr Size, iterate_callback Callback,
     auto Lambda = [this, From, To, Callback, Arg](uptr Block) {
       if (Block < From || Block >= To)
         return;
       uptr Chunk;
       Chunk::UnpackedHeader Header;
       if (getChunkFromBlock(Block, &Chunk, &Header) &&
           Header.State == Chunk::State::Allocated) {
         uptr TaggedChunk = Chunk;
-        if (useMemoryTagging(Primary.Options.load()))
+        if (useMemoryTagging<Params>(Primary.Options.load()))
           TaggedChunk = loadTag(Chunk);
         Callback(TaggedChunk, getSize(reinterpret_cast<void *>(Chunk), &Header),
                  Arg);
       }
     };
     Primary.iterateOverBlocks(Lambda);
     Secondary.iterateOverBlocks(Lambda);
 #ifdef GWP_ASAN_HOOKS
#endif // GWP_ASAN_HOOKS
     if (!Ptr || !isAligned(reinterpret_cast<uptr>(Ptr), MinAlignment))
       return false;
     Ptr = untagPointerMaybe(const_cast<void *>(Ptr));
     Chunk::UnpackedHeader Header;
     return Chunk::isValid(Cookie, Ptr, &Header) &&
            Header.State == Chunk::State::Allocated;
   }
 
-  bool useMemoryTagging() const {
-    return useMemoryTagging(Primary.Options.load());
+  bool useMemoryTaggingTestOnly() const {
+    return useMemoryTagging<Params>(Primary.Options.load());
   }
-  static bool useMemoryTagging(Options Options) {
-    return PrimaryT::useMemoryTagging(Options);
+  void disableMemoryTagging() {
+    if (allocatorSupportsMemoryTagging<Params>())
+      Primary.Options.clear(OptionBit::UseMemoryTagging);
   }
 
-  void disableMemoryTagging() { Primary.disableMemoryTagging(); }
-
   void setTrackAllocationStacks(bool Track) {
     initThreadMaybe();
     if (Track)
       Primary.Options.set(OptionBit::TrackAllocationStacks);
     else
       Primary.Options.clear(OptionBit::TrackAllocationStacks);
   }
 
#endif // GWP_ASAN_HOOKS
   }
 
   static void getErrorInfo(struct scudo_error_info *ErrorInfo,
                            uintptr_t FaultAddr, const char *DepotPtr,
                            const char *RegionInfoPtr, const char *Memory,
                            const char *MemoryTags, uintptr_t MemoryAddr,
                            size_t MemorySize) {
     *ErrorInfo = {};
-    if (!PrimaryT::SupportsMemoryTagging ||
+    if (!allocatorSupportsMemoryTagging<Params>() ||
         MemoryAddr + MemorySize < MemoryAddr)
       return;
 
     uptr UntaggedFaultAddr = untagPointer(FaultAddr);
     u8 FaultAddrTag = extractTag(FaultAddr);
     BlockInfo Info =
         PrimaryT::findNearestBlock(RegionInfoPtr, UntaggedFaultAddr);
 
private:
   static const uptr MaxAlignmentLog = 24U; // 16 MB seems reasonable.
   static const uptr MinAlignment = 1UL << MinAlignmentLog;
   static const uptr MaxAlignment = 1UL << MaxAlignmentLog;
   static const uptr MaxAllowedMallocSize =
       FIRST_32_SECOND_64(1UL << 31, 1ULL << 40);
 
   static_assert(MinAlignment >= sizeof(Chunk::PackedHeader),
                 "Minimal alignment must at least cover a chunk header.");
-  static_assert(!PrimaryT::SupportsMemoryTagging ||
+  static_assert(!allocatorSupportsMemoryTagging<Params>() ||
                     MinAlignment >= archMemoryTagGranuleSize(),
                 "");
 
   static const u32 BlockMarker = 0x44554353U;
 
   // These are indexes into an "array" of 32-bit values that store information
   // inline with a chunk that is relevant to diagnosing memory tag faults, where
   // 0 corresponds to the address of the user memory. This means that negative
if (LIKELY(Header->ClassId))
       return SizeOrUnusedBytes;
     return SecondaryT::getBlockEnd(getBlockBegin(Ptr, Header)) -
            reinterpret_cast<uptr>(Ptr) - SizeOrUnusedBytes;
   }
 
   void quarantineOrDeallocateChunk(Options Options, void *Ptr,
                                    Chunk::UnpackedHeader *Header, uptr Size) {
     Chunk::UnpackedHeader NewHeader = *Header;
-    if (UNLIKELY(NewHeader.ClassId && useMemoryTagging(Options))) {
+    if (UNLIKELY(NewHeader.ClassId && useMemoryTagging<Params>(Options))) {
       u8 PrevTag = extractTag(loadTag(reinterpret_cast<uptr>(Ptr)));
       if (!TSDRegistry.getDisableMemInit()) {
         uptr TaggedBegin, TaggedEnd;
         const uptr OddEvenMask = computeOddEvenMaskForPointerMaybe(
             Options, reinterpret_cast<uptr>(getBlockBegin(Ptr, &NewHeader)),
             SizeClassMap::getSizeByClassId(NewHeader.ClassId));
         // Exclude the previous tag so that immediate use after free is detected
         // 100% of the time.

compiler-rt/lib/scudo/standalone/options.h

 //===-- options.h -----------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCUDO_OPTIONS_H_
 #define SCUDO_OPTIONS_H_
 
 #include "atomic_helpers.h"
 #include "common.h"
+#include "memtag.h"
 
 namespace scudo {
 
 enum class OptionBit {
   MayReturnNull,
   FillContents0of2,
   FillContents1of2,
   DeallocTypeMismatch,
struct Options {
   bool get(OptionBit Opt) const { return Val & (1U << static_cast<u32>(Opt)); }
 
   FillContentsMode getFillContentsMode() const {
     return static_cast<FillContentsMode>(
         (Val >> static_cast<u32>(OptionBit::FillContents0of2)) & 3);
   }
 };
 
+template <typename Config> bool useMemoryTagging(Options Options) {
+  return allocatorSupportsMemoryTagging<Config>() &&
+         Options.get(OptionBit::UseMemoryTagging);
+}
+
 struct AtomicOptions {
   atomic_u32 Val;
 
 public:
   Options load() const { return Options{atomic_load_relaxed(&Val)}; }
 
   void clear(OptionBit Opt) {
     atomic_fetch_and(&Val, ~(1U << static_cast<u32>(Opt)),

compiler-rt/lib/scudo/standalone/primary32.h

public:
   // The bytemap can only track UINT8_MAX - 1 classes.
   static_assert(SizeClassMap::LargestClassId <= (UINT8_MAX - 1), "");
   // Regions should be large enough to hold the largest Block.
   static_assert((1UL << Config::PrimaryRegionSizeLog) >= SizeClassMap::MaxSize,
                 "");
   typedef SizeClassAllocator32<Config> ThisT;
   typedef SizeClassAllocatorLocalCache<ThisT> CacheT;
   typedef typename CacheT::TransferBatch TransferBatch;
-  static const bool SupportsMemoryTagging = false;
 
   static uptr getSizeByClassId(uptr ClassId) {
     return (ClassId == SizeClassMap::BatchClassId)
                ? sizeof(TransferBatch)
                : SizeClassMap::getSizeByClassId(ClassId);
   }
 
   static bool canAllocate(uptr Size) { return Size <= SizeClassMap::MaxSize; }
for (uptr I = 0; I < NumClasses; I++) {
         continue;
       SizeClassInfo *Sci = getSizeClassInfo(I);
       ScopedLock L(Sci->Mutex);
       TotalReleasedBytes += releaseToOSMaybe(Sci, I, /*Force=*/true);
     }
     return TotalReleasedBytes;
   }
 
-  static bool useMemoryTagging(UNUSED Options Options) { return false; }
-  void disableMemoryTagging() {}
-
   const char *getRegionInfoArrayAddress() const { return nullptr; }
   static uptr getRegionInfoArraySize() { return 0; }
 
   static BlockInfo findNearestBlock(UNUSED const char *RegionInfoData,
                                     UNUSED uptr Ptr) {
     return {};
   }
 

compiler-rt/lib/scudo/standalone/primary64.h

 // released if the platform allows for it.
 
 template <typename Config> class SizeClassAllocator64 {
 public:
   typedef typename Config::SizeClassMap SizeClassMap;
   typedef SizeClassAllocator64<Config> ThisT;
   typedef SizeClassAllocatorLocalCache<ThisT> CacheT;
   typedef typename CacheT::TransferBatch TransferBatch;
-  static const bool SupportsMemoryTagging =
-      allocatorSupportsMemoryTagging<Config>();
 
   static uptr getSizeByClassId(uptr ClassId) {
     return (ClassId == SizeClassMap::BatchClassId)
                ? sizeof(TransferBatch)
                : SizeClassMap::getSizeByClassId(ClassId);
   }
 
   static bool canAllocate(uptr Size) { return Size <= SizeClassMap::MaxSize; }
for (uptr I = 0; I < NumClasses; I++) {
       RegionInfo *Region = getRegionInfo(I);
       // The actual start of a region is offseted by a random number of pages.
       Region->RegionBeg =
           getRegionBaseByClassId(I) + (getRandomModN(&Seed, 16) + 1) * PageSize;
       Region->RandState = getRandomU32(&Seed);
       Region->ReleaseInfo.LastReleaseAtNs = Time;
     }
     setOption(Option::ReleaseInterval, static_cast<sptr>(ReleaseToOsInterval));
-
-    if (SupportsMemoryTagging && systemSupportsMemoryTagging())
-      Options.set(OptionBit::UseMemoryTagging);
   }
   void init(s32 ReleaseToOsInterval) {
     memset(this, 0, sizeof(*this));
     initLinkerInitialized(ReleaseToOsInterval);
   }
 
   void unmapTestOnly() {
     unmap(reinterpret_cast<void *>(PrimaryBase), PrimarySize, UNMAP_ALL, &Data);
for (uptr I = 0; I < NumClasses; I++) {
         continue;
       RegionInfo *Region = getRegionInfo(I);
       ScopedLock L(Region->Mutex);
       TotalReleasedBytes += releaseToOSMaybe(Region, I, /*Force=*/true);
     }
     return TotalReleasedBytes;
   }
 
-  static bool useMemoryTagging(Options Options) {
-    return SupportsMemoryTagging && Options.get(OptionBit::UseMemoryTagging);
-  }
-  void disableMemoryTagging() { Options.clear(OptionBit::UseMemoryTagging); }
-
   const char *getRegionInfoArrayAddress() const {
     return reinterpret_cast<const char *>(RegionInfoArray);
   }
 
   static uptr getRegionInfoArraySize() { return sizeof(RegionInfoArray); }
 
   static BlockInfo findNearestBlock(const char *RegionInfoData, uptr Ptr) {
     const RegionInfo *RegionInfoArray =
if (UNLIKELY(TotalUserBytes > MappedUser)) {
         }
         return nullptr;
       }
       if (MappedUser == 0)
         Region->Data = Data;
       if (!map(reinterpret_cast<void *>(RegionBeg + MappedUser), UserMapSize,
                "scudo:primary",
                MAP_ALLOWNOMEM | MAP_RESIZABLE |
-                   (useMemoryTagging(Options.load()) ? MAP_MEMTAG : 0),
+                   (useMemoryTagging<Config>(Options.load()) ? MAP_MEMTAG : 0),
                &Region->Data))
         return nullptr;
       Region->MappedUser += UserMapSize;
       C->getStats().add(StatMapped, UserMapSize);
     }
 
     const u32 NumberOfBlocks = Min(
         MaxNumBatches * MaxCount,

compiler-rt/lib/scudo/standalone/tests/combined_test.cpp

const scudo::uptr NeededSize =
       scudo::roundUpTo(Size, MinAlignment) +
       ((Alignment > MinAlignment) ? Alignment : scudo::Chunk::getHeaderSize());
   return AllocatorT::PrimaryT::canAllocate(NeededSize);
 }
 
 template <class AllocatorT>
 bool isTaggedAllocation(AllocatorT *Allocator, scudo::uptr Size,
                         scudo::uptr Alignment) {
-  return Allocator->useMemoryTagging() &&
+  return Allocator->useMemoryTaggingTestOnly() &&
          scudo::systemDetectsMemoryTagFaultsTestOnly() &&
          isPrimaryAllocation<AllocatorT>(Size, Alignment);
 }
 
 template <class AllocatorT>
 void checkMemoryTaggingMaybe(AllocatorT *Allocator, void *P, scudo::uptr Size,
                              scudo::uptr Alignment) {
   if (!isTaggedAllocation(Allocator, Size, Alignment))
template <class Config> static void testAllocator() {
   for (scudo::uptr SizeLog = 0U; SizeLog <= 20U; SizeLog++) {
     for (scudo::uptr Delta = 0U; Delta <= 4U; Delta++) {
       const scudo::uptr Size = (1U << SizeLog) + Delta * 128U;
       void *P = Allocator->allocate(Size, Origin, 1U << MinAlignLog, false);
       EXPECT_NE(P, nullptr);
       for (scudo::uptr I = 0; I < Size; I++) {
         unsigned char V = (reinterpret_cast<unsigned char *>(P))[I];
         if (isPrimaryAllocation<AllocatorT>(Size, 1U << MinAlignLog) &&
-            !Allocator->useMemoryTagging())
+            !Allocator->useMemoryTaggingTestOnly())
           ASSERT_EQ(V, scudo::PatternFillByte);
         else
           ASSERT_TRUE(V == scudo::PatternFillByte || V == 0);
       }
       memset(P, 0xaa, Size);
       Allocator->deallocate(P, Origin, Size);
     }
   }
if (!UseQuarantine) {
     while (!V.empty()) {
       Allocator->deallocate(V.back(), Origin);
       V.pop_back();
     }
   }
 
   Allocator->releaseToOS();
 
-  if (Allocator->useMemoryTagging() &&
+  if (Allocator->useMemoryTaggingTestOnly() &&
       scudo::systemDetectsMemoryTagFaultsTestOnly()) {
     // Check that use-after-free is detected.
     for (scudo::uptr SizeLog = 0U; SizeLog <= 20U; SizeLog++) {
       const scudo::uptr Size = 1U << SizeLog;
       if (!isTaggedAllocation(Allocator.get(), Size, 1))
         continue;
       // UAF detection is probabilistic, so we repeat the test up to 256 times
       // if necessary. With 15 possible tags this means a 1 in 15^256 chance of
TEST(ScudoCombinedTest, FullRegion) {
   EXPECT_EQ(FailedAllocationsCount, 0U);
 }
 
 TEST(ScudoCombinedTest, OddEven) {
   using AllocatorT = TestAllocator<scudo::AndroidConfig>;
   using SizeClassMap = AllocatorT::PrimaryT::SizeClassMap;
   auto Allocator = std::unique_ptr<AllocatorT>(new AllocatorT());
 
-  if (!Allocator->useMemoryTagging())
+  if (!Allocator->useMemoryTaggingTestOnly())
     return;
 
   auto CheckOddEven = [](scudo::uptr P1, scudo::uptr P2) {
     scudo::uptr Tag1 = scudo::extractTag(scudo::loadTag(P1));
     scudo::uptr Tag2 = scudo::extractTag(scudo::loadTag(P2));
     EXPECT_NE(Tag1 % 2, Tag2 % 2);
   };