[scudo] Split the code paths in quarantineOrDeallocateChunk()

[Chia-hungDuan]

Summary

Move block untagging logic to another function and sink untagged pointer
to quarantineOrDeallocateChunk().

[cferris]

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

1175–1176

it might be better if you move this into the else down below. Since the first if sets this unconditionally, I think it would be easier to follow to have it set in both the if and else.

1178

I don't think this initializing is necessary since BlockBegin is set in both the if and else statement below.

1210–1216

The Quarantine code path looks like it's kind of odd.

I think it would make this code a lot more readable if the qurantine case was moved to the location of the BypassQuarantine bool init.

It would look like this:

if (!BypassQuarantine) {
  NewHeader.State = Chunk::State::Quarantined;
  bool UnlockRequired;
  auto *TSD = TSDRegistry.getTSDAndLock(&UnlockRequired);
  Quarantine.put(&TSD->getQuarantineCache(),
                   QuarantineCallback(*this, TSD->getCache()), Ptr, Size);
  if (UnlockRequired)
      TSD->unlock();
  return;
}

I think you might need to call unTagBlock if memory tagging is enabled in the quarantine path.

Right now this new change does some extra work when in the quarantine path, but if you split it out then that won't happen. And adding the unTagBlock call might also do some extra work though so I'm not sure what's the best way to handle this.

It might be better to do that in a follow-on though. I'll leave that up to you, or since it's my suggestion, I can also do this.

[Chia-hungDuan]

Address review comments

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

1210–1216

I agree with the code motion of quarantine block. As we discussed, we may also want to config quarantine by compilation flag. Thus I'm planning to do the code motion with the config change in a later CL.

I've fixed the part that we did more work in quarantine path. Now it should be the same logic as before.

[cferris]

LGTM.

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

#endif // GWP_ASAN_HOOKS
     }
 
     const uptr Size = getSize(Ptr, &Header);
     if (DeleteSize && Options.get(OptionBit::DeleteSizeMismatch)) {
       if (UNLIKELY(DeleteSize != Size))
         reportDeleteSizeMismatch(Ptr, DeleteSize, Size);
     }
 
-    quarantineOrDeallocateChunk(Options, TaggedPtr, &Header, Size);
+    quarantineOrDeallocateChunk(Options, TaggedPtr, Ptr, &Header, Size);
   }
 
   void *reallocate(void *OldPtr, uptr NewSize, uptr Alignment = MinAlignment) {
     initThreadMaybe();
 
     const Options Options = Primary.Options.load();
     if (UNLIKELY(NewSize >= MaxAllowedMallocSize)) {
       if (Options.get(OptionBit::MayReturnNull))
#endif // GWP_ASAN_HOOKS
 
     // Otherwise we allocate a new one, and deallocate the old one. Some
     // allocators will allocate an even larger chunk (by a fixed factor) to
     // allow for potential further in-place realloc. The gains of such a trick
     // are currently unclear.
     void *NewPtr = allocate(NewSize, Chunk::Origin::Malloc, Alignment);
     if (LIKELY(NewPtr)) {
       memcpy(NewPtr, OldTaggedPtr, Min(NewSize, OldSize));
-      quarantineOrDeallocateChunk(Options, OldTaggedPtr, &OldHeader, OldSize);
+      quarantineOrDeallocateChunk(Options, OldTaggedPtr, OldPtr, &OldHeader,
+                                  OldSize);
     }
     return NewPtr;
   }
 
   // TODO(kostyak): disable() is currently best-effort. There are some small
   //                windows of time when an allocation could still succeed after
   //                this function finishes. We will revisit that later.
   void disable() NO_THREAD_SAFETY_ANALYSIS {
initChunkWithMemoryTagging(const uptr ClassId, const Chunk::Origin Origin,
     Header.OriginOrWasZeroed = Origin & Chunk::OriginMask;
     Header.SizeOrUnusedBytes = SizeOrUnusedBytes & Chunk::SizeOrUnusedBytesMask;
     Chunk::storeHeader(Cookie, Ptr, &Header);
 
     return TaggedPtr;
   }
 
   void quarantineOrDeallocateChunk(const Options &Options, void *TaggedPtr,
+                                   void *HeaderTaggedPtr,
                                    Chunk::UnpackedHeader *Header,
                                    uptr Size) NO_THREAD_SAFETY_ANALYSIS {
-    void *Ptr = getHeaderTaggedPointer(TaggedPtr);
+    void *Ptr = HeaderTaggedPtr;
     Chunk::UnpackedHeader NewHeader = *Header;
     // If the quarantine is disabled, the actual size of a chunk is 0 or larger
     // than the maximum allowed, we return a chunk directly to the backend.
     // This purposefully underflows for Size == 0.
     const bool BypassQuarantine = !Quarantine.getCacheSize() ||
                                   ((Size - 1) >= QuarantineMaxChunkSize) ||
                                   !NewHeader.ClassId;
     if (BypassQuarantine)
       NewHeader.State = Chunk::State::Available;
     else
       NewHeader.State = Chunk::State::Quarantined;
-    NewHeader.OriginOrWasZeroed = useMemoryTagging<Config>(Options) &&
-                                  NewHeader.ClassId &&
-                                  !TSDRegistry.getDisableMemInit();
     Chunk::compareExchangeHeader(Cookie, Ptr, &NewHeader, Header);
 
+    void *BlockBegin;
     if (UNLIKELY(useMemoryTagging<Config>(Options))) {
-      u8 PrevTag = extractTag(reinterpret_cast<uptr>(TaggedPtr));
-      storeDeallocationStackMaybe(Options, Ptr, PrevTag, Size);
-      if (NewHeader.ClassId) {
-        if (!TSDRegistry.getDisableMemInit()) {
-          uptr TaggedBegin, TaggedEnd;
-          const uptr OddEvenMask = computeOddEvenMaskForPointerMaybe(
-              Options, reinterpret_cast<uptr>(getBlockBegin(Ptr, &NewHeader)),
-              NewHeader.ClassId);
-          // Exclude the previous tag so that immediate use after free is
-          // detected 100% of the time.
-          setRandomTag(Ptr, Size, OddEvenMask | (1UL << PrevTag), &TaggedBegin,
-                       &TaggedEnd);
-        }
-      }
+      NewHeader.OriginOrWasZeroed =
+          NewHeader.ClassId && !TSDRegistry.getDisableMemInit();
+      BlockBegin = unTagBlock(Options, TaggedPtr, HeaderTaggedPtr, &NewHeader,
+                              Size, BypassQuarantine);
+    } else {
+      NewHeader.OriginOrWasZeroed = 0U;
+      if (BypassQuarantine && allocatorSupportsMemoryTagging<Config>())
+        Ptr = untagPointer(Ptr);
+      BlockBegin = getBlockBegin(Ptr, &NewHeader);
     }
+
     if (BypassQuarantine) {
-      if (allocatorSupportsMemoryTagging<Config>())
-        Ptr = untagPointer(Ptr);
-      void *BlockBegin = getBlockBegin(Ptr, &NewHeader);
       const uptr ClassId = NewHeader.ClassId;
       if (LIKELY(ClassId)) {
         bool UnlockRequired;
         auto *TSD = TSDRegistry.getTSDAndLock(&UnlockRequired);
         const bool CacheDrained =
             TSD->getCache().deallocate(ClassId, BlockBegin);
         if (UnlockRequired)
           TSD->unlock();
         // When we have drained some blocks back to the Primary from TSD, that
         // implies that we may have the chance to release some pages as well.
         // Note that in order not to block other thread's accessing the TSD,
         // release the TSD first then try the page release.
         if (CacheDrained)
           Primary.tryReleaseToOS(ClassId, ReleaseToOS::Normal);
       } else {
-        if (UNLIKELY(useMemoryTagging<Config>(Options)))
-          storeTags(reinterpret_cast<uptr>(BlockBegin),
-                    reinterpret_cast<uptr>(Ptr));
         Secondary.deallocate(Options, BlockBegin);
       }
     } else {
       bool UnlockRequired;
       auto *TSD = TSDRegistry.getTSDAndLock(&UnlockRequired);
       Quarantine.put(&TSD->getQuarantineCache(),
                      QuarantineCallback(*this, TSD->getCache()), Ptr, Size);
       if (UnlockRequired)
         TSD->unlock();
     }
   }
 
+  NOINLINE void *unTagBlock(const Options &Options, void *TaggedPtr,
+                            void *HeaderTaggedPtr,
+                            Chunk::UnpackedHeader *Header, const uptr Size,
+                            bool BypassQuarantine) {
+    DCHECK(useMemoryTagging<Config>(Options));
+    void *Ptr = HeaderTaggedPtr;
+
+    const u8 PrevTag = extractTag(reinterpret_cast<uptr>(TaggedPtr));
+    storeDeallocationStackMaybe(Options, Ptr, PrevTag, Size);
+    if (Header->ClassId) {
+      if (!TSDRegistry.getDisableMemInit()) {
+        uptr TaggedBegin, TaggedEnd;
+        const uptr OddEvenMask = computeOddEvenMaskForPointerMaybe(
+            Options, reinterpret_cast<uptr>(getBlockBegin(Ptr, Header)),
+            Header->ClassId);
+        // Exclude the previous tag so that immediate use after free is
+        // detected 100% of the time.
+        setRandomTag(Ptr, Size, OddEvenMask | (1UL << PrevTag), &TaggedBegin,
+                     &TaggedEnd);
+      }
+    }
+
+    Ptr = untagPointer(Ptr);
+    void *BlockBegin = getBlockBegin(Ptr, Header);
+    if (BypassQuarantine && !Header->ClassId) {
+      storeTags(reinterpret_cast<uptr>(BlockBegin),
+                reinterpret_cast<uptr>(Ptr));
+    }
+
+    return BlockBegin;
+  }
+
   bool getChunkFromBlock(uptr Block, uptr *Chunk,
                          Chunk::UnpackedHeader *Header) {
     *Chunk =
         Block + getChunkOffsetFromBlock(reinterpret_cast<const char *>(Block));
     return Chunk::isValid(Cookie, reinterpret_cast<void *>(*Chunk), Header);
   }
 
   static uptr getChunkOffsetFromBlock(const char *Block) {