[scudo][standalone] Implement TSD registry disabling

[cryptoad]

Summary

In order to implement malloc_{enable|disable} we were just disabling
(or really locking) the Primary and the Secondary. That meant that
allocations could still be serviced from the TSD as long as the cache
wouldn't have to be filled from the Primary.

This wasn't working out for Android tests, so this change implements
registry disabling (eg: locking) so that getTSDAndLock doesn't
return a TSD if the allocator is disabled. This also means that the
Primary doesn't have to be disabled in this situation.

For the Shared Registry, we loop through all the TSDs and lock them.
For the Exclusive Registry, we add a Disabled boolean to the Registry
that forces getTSDAndLock to use the Fallback TSD instead of the
thread local one. Disabling the Registry is then done by locking the
Fallback TSD and setting the boolean in question (I don't think this
needed an atomic variable but I might be wrong).

I clang-formatted the whole thing as usual hence the couple of extra
whiteline changes in this CL.

[cryptoad]

Extra whiteline.

[eugenis]

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

51–52

is not this a data race?

[cryptoad]

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

51–52

I think it's OK for allocations to succeed until disable() returns, which I assume is the racey part you are referring to?

[eugenis]

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

51–52

If the read of Disabled is not an acquire-read, the thread may not observe the change for a very long time. If you want any guarantee there, both the reading and writing threads need to do some kind of synchronize-y thing.

This sounds like a hard problem if we don't want to sacrifice performance of the fast path.

[cryptoad]

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

51–52

Right. I'll update this to use an atomic and run some quick tests.

[cryptoad]

Use an atomic_u8 for Disabled in the Exclusive Registry.
This should ensure that changes to the variable are visible to the
threads in a somewhat timely manner.

A run of rpc2-benchmark with this new change doesn't show noticeable
performance degradation.

EDIT: doesn't, not does :(

[eugenis]

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

83

I'm also concerned about the shared tsd code path. Is locking the tsd enough to ensure that no allocations will happen in the future? It looks like threads hold the tsd lock for as short a time as possible, and do the rest of the work such as updating the chuck state later without synchronization. This work may not be finished by the time malloc_disable returns. Could malloc_iterate be confused if it sees the allocator in that state?

[cryptoad]

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

83

You are correct, a thread could still be updating a block by the time disable returns.

Hopefully the window is small enough for the amount of potential confusion to be low. Given that this is mostly used in the context of libmemunreachable, which is mostly for debugging purposes, I think a best-effort approach as opposed to a catch-all that might have stronger performance implication is acceptable.

To be fair there might be something better that I haven't thought about as well.

[cryptoad]

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

83

PS: I can add a TODO stating so.

[eugenis]

LGTM

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

83

I think this can be fixed properly with one or several strategically placed atomic stores in the block update code such that malloc_iterate can always recognize the state the block is in. We've made changes like this in asan & lsan some time ago, mostly in response to bad reports of racy use-after-frees.

This stuff is really hard to reason about. I think this change is OK as a first step, but please leave a TODO.

[cryptoad]

Add a comment for disable() with a TODO to improve the behavior.

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

 #include "local_cache.h"
 #include "quarantine.h"
 #include "report.h"
 #include "secondary.h"
 #include "string_utils.h"
 #include "tsd.h"
 
 #ifdef GWP_ASAN_HOOKS
-# include "gwp_asan/guarded_pool_allocator.h"
+#include "gwp_asan/guarded_pool_allocator.h"
 // GWP-ASan is declared here in order to avoid indirect call overhead. It's also
 // instantiated outside of the Allocator class, as the allocator is only
 // zero-initialised. GWP-ASan requires constant initialisation, and the Scudo
 // allocator doesn't have a constexpr constructor (see discussion here:
 // https://reviews.llvm.org/D69265#inline-624315).
 static gwp_asan::GuardedPoolAllocator GuardedAlloc;
 #endif // GWP_ASAN_HOOKS
 
#endif // GWP_ASAN_HOOKS
     if (NewPtr) {
       const uptr OldSize = getSize(OldPtr, &OldHeader);
       memcpy(NewPtr, OldPtr, Min(NewSize, OldSize));
       quarantineOrDeallocateChunk(OldPtr, &OldHeader, OldSize);
     }
     return NewPtr;
   }
 
-  // TODO(kostyak): while this locks the Primary & Secondary, it still allows
-  //                pointers to be fetched from the TSD. We ultimately want to
-  //                lock the registry as well. For now, it's good enough.
+  // 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() {
     initThreadMaybe();
-    Primary.disable();
+    TSDRegistry.disable();
     Secondary.disable();
   }
 
   void enable() {
     initThreadMaybe();
     Secondary.enable();
-    Primary.enable();
+    TSDRegistry.enable();
   }
 
   // The function returns the amount of bytes required to store the statistics,
   // which might be larger than the amount of bytes provided. Note that the
   // statistics buffer is not necessarily constant between calls to this
   // function. This can be called with a null buffer or zero size for buffer
   // sizing purposes.
   uptr getStats(char *Buffer, uptr Size) {

compiler-rt/lib/scudo/standalone/flags.cpp

 #define SCUDO_FLAG(Type, Name, DefaultValue, Description)                      \
   Parser->registerFlag(#Name, Description, FlagType::FT_##Type,                \
                        reinterpret_cast<void *>(&F->Name));
 #include "flags.inc"
 #undef SCUDO_FLAG
 
 #ifdef GWP_ASAN_HOOKS
 #define GWP_ASAN_OPTION(Type, Name, DefaultValue, Description)                 \
-  Parser->registerFlag("GWP_ASAN_"#Name, Description, FlagType::FT_##Type,     \
+  Parser->registerFlag("GWP_ASAN_" #Name, Description, FlagType::FT_##Type,    \
                        reinterpret_cast<void *>(&F->GWP_ASAN_##Name));
 #include "gwp_asan/options.inc"
 #undef GWP_ASAN_OPTION
 #endif // GWP_ASAN_HOOKS
 }
 
 static const char *getCompileDefinitionScudoDefaultOptions() {
 #ifdef SCUDO_DEFAULT_OPTIONS

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

TEST(ScudoWrappersCTest, MallocIterateBoundary) {
   malloc_iterate(Block - PageSize, PageSize, callback, nullptr);
   malloc_iterate(Block, PageSize, callback, nullptr);
   malloc_enable();
   EXPECT_EQ(Count, 1U);
 
   free(P);
 }
 
+// We expect heap operations within a disable/enable scope to deadlock.
+TEST(ScudoWrappersCTest, MallocDisableDeadlock) {
+  EXPECT_DEATH(
+      {
+        void *P = malloc(Size);
+        EXPECT_NE(P, nullptr);
+        free(P);
+        malloc_disable();
+        alarm(1);
+        P = malloc(Size);
+        malloc_enable();
+      },
+      "");
+}
+
 #if !SCUDO_FUCHSIA
 TEST(ScudoWrappersCTest, MallocInfo) {
   char Buffer[64];
   FILE *F = fmemopen(Buffer, sizeof(Buffer), "w+");
   EXPECT_NE(F, nullptr);
   errno = 0;
   EXPECT_EQ(malloc_info(0, F), 0);
   EXPECT_EQ(errno, 0);
   fclose(F);
   EXPECT_EQ(strncmp(Buffer, "<malloc version=\"scudo-", 23), 0);
 }
 #endif

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

template <class Allocator> struct TSDRegistryExT {
 
   ALWAYS_INLINE void initThreadMaybe(Allocator *Instance, bool MinimalInit) {
     if (LIKELY(State != ThreadState::NotInitialized))
       return;
     initThread(Instance, MinimalInit);
   }
 
   ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
-    if (LIKELY(State == ThreadState::Initialized)) {
+    if (LIKELY(State == ThreadState::Initialized &&
+               !atomic_load(&Disabled, memory_order_acquire))) {
       *UnlockRequired = false;
       return &ThreadTSD;
     }
     DCHECK(FallbackTSD);
     FallbackTSD->lock();
     *UnlockRequired = true;
     return FallbackTSD;
   }
 
+  // To disable the exclusive TSD registry, we effectively lock the fallback TSD
+  // and force all threads to attempt to use it instead of their local one.
+  void disable() {
+    FallbackTSD->lock();
+    atomic_store(&Disabled, 1U, memory_order_release);
+  }
+
+  void enable() {
+    atomic_store(&Disabled, 0U, memory_order_release);
+    FallbackTSD->unlock();
+  }
+
 private:
   void initOnceMaybe(Allocator *Instance) {
     ScopedLock L(Mutex);
     if (LIKELY(Initialized))
       return;
     initLinkerInitialized(Instance); // Sets Initialized.
   }
 
   // Using minimal initialization allows for global initialization while keeping
   // the thread specific structure untouched. The fallback structure will be
   // used instead.
   NOINLINE void initThread(Allocator *Instance, bool MinimalInit) {
     initOnceMaybe(Instance);
     if (UNLIKELY(MinimalInit))
       return;
     CHECK_EQ(
         pthread_setspecific(PThreadKey, reinterpret_cast<void *>(Instance)), 0);
     ThreadTSD.initLinkerInitialized(Instance);
     State = ThreadState::Initialized;
   }
 
   pthread_key_t PThreadKey;
   bool Initialized;
+  atomic_u8 Disabled;
   TSD<Allocator> *FallbackTSD;
   HybridMutex Mutex;
   static THREADLOCAL ThreadState State;
   static THREADLOCAL TSD<Allocator> ThreadTSD;
 
   friend void teardownThread<Allocator>(void *Ptr);
 };
 

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

ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
     *UnlockRequired = true;
     // Try to lock the currently associated context.
     if (TSD->tryLock())
       return TSD;
     // If that fails, go down the slow path.
     return getTSDAndLockSlow(TSD);
   }
 
+  void disable() {
+    for (u32 I = 0; I < NumberOfTSDs; I++)
+      TSDs[I].lock();
+  }
+
+  void enable() {
+    for (u32 I = 0; I < NumberOfTSDs; I++)
+      TSDs[I].unlock();
+  }
+
 private:
   ALWAYS_INLINE void setCurrentTSD(TSD<Allocator> *CurrentTSD) {
 #if _BIONIC
     *getAndroidTlsPtr() = reinterpret_cast<uptr>(CurrentTSD);
 #elif SCUDO_LINUX
     ThreadTSD = CurrentTSD;
 #else
     CHECK_EQ(