[Tagged Pointers] Allow probing the current TP level w/ locking.

aosp/1484976 introduced a breaking change where
DisableMemoryMitigations() now indiscriminately turns tagged pointers
off. When android_mallopt(M_DISABLE_MEMORY_MITIGATIONS) is called, the
correct behaviour is:
 - In SYNC/ASYNC MTE mode -> disable all tagged pointers.
 - If all tagged pointers are already disabled -> nop.
 - If we're in TBI mode -> nop (keep the TBI mode as-is).

In order to do that we have to allow probing of the current heap tagging
mode. In order to prevent TOCTOU between GetHeapTaggingLevel() and
SetHeapTaggingLevel(), we expose a global mutex that should be held when
calling these functions.

Bug: 174263432
Test: atest CtsTaggingHostTestCases on Flame
Change-Id: Ia96f7269d542c9041270458806aee36766d2fbbb
diff --git a/libc/bionic/heap_tagging.cpp b/libc/bionic/heap_tagging.cpp
index 7601ddd..2c5d4d8 100644
--- a/libc/bionic/heap_tagging.cpp
+++ b/libc/bionic/heap_tagging.cpp
@@ -30,16 +30,14 @@
 #include "malloc_common.h"
 #include "malloc_tagged_pointers.h"
 
+#include <bionic/pthread_internal.h>
 #include <platform/bionic/malloc.h>
 #include <platform/bionic/mte_kernel.h>
 
-#include <bionic/pthread_internal.h>
-
-#include "private/ScopedPthreadMutexLocker.h"
-
 extern "C" void scudo_malloc_disable_memory_tagging();
 extern "C" void scudo_malloc_set_track_allocation_stacks(int);
 
+// Protected by `g_heap_tagging_lock`.
 static HeapTaggingLevel heap_tagging_level = M_HEAP_TAGGING_LEVEL_NONE;
 
 void SetDefaultHeapTaggingLevel() {
@@ -94,10 +92,15 @@
 }
 #endif
 
-bool SetHeapTaggingLevel(void* arg, size_t arg_size) {
-  static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
-  ScopedPthreadMutexLocker locker(&mutex);
+pthread_mutex_t g_heap_tagging_lock = PTHREAD_MUTEX_INITIALIZER;
 
+// Requires `g_heap_tagging_lock` to be held.
+HeapTaggingLevel GetHeapTaggingLevel() {
+  return heap_tagging_level;
+}
+
+// Requires `g_heap_tagging_lock` to be held.
+bool SetHeapTaggingLevel(void* arg, size_t arg_size) {
   if (arg_size != sizeof(HeapTaggingLevel)) {
     return false;
   }
diff --git a/libc/bionic/heap_tagging.h b/libc/bionic/heap_tagging.h
index 2aaf608..db45fc1 100644
--- a/libc/bionic/heap_tagging.h
+++ b/libc/bionic/heap_tagging.h
@@ -28,7 +28,19 @@
 
 #pragma once
 
+#include <bionic/pthread_internal.h>
+#include <platform/bionic/malloc.h>
 #include <stddef.h>
 
+// Expected to be called in a single-threaded context during libc init, so no
+// synchronization required.
 void SetDefaultHeapTaggingLevel();
+
+// Lock for the heap tagging level. You may find ScopedPthreadMutexLocker
+// useful for RAII on this lock.
+extern pthread_mutex_t g_heap_tagging_lock;
+
+// These functions can be called in a multithreaded context, and thus should
+// only be called when holding the `g_heap_tagging_lock`.
 bool SetHeapTaggingLevel(void* arg, size_t arg_size);
+HeapTaggingLevel GetHeapTaggingLevel();
diff --git a/libc/bionic/malloc_common.cpp b/libc/bionic/malloc_common.cpp
index 6b7006d..0ee12a7 100644
--- a/libc/bionic/malloc_common.cpp
+++ b/libc/bionic/malloc_common.cpp
@@ -38,8 +38,9 @@
 #include <stdint.h>
 #include <stdio.h>
 
-#include <private/bionic_config.h>
 #include <platform/bionic/malloc.h>
+#include <private/ScopedPthreadMutexLocker.h>
+#include <private/bionic_config.h>
 
 #include "gwp_asan_wrappers.h"
 #include "heap_tagging.h"
@@ -316,6 +317,7 @@
     return LimitEnable(arg, arg_size);
   }
   if (opcode == M_SET_HEAP_TAGGING_LEVEL) {
+    ScopedPthreadMutexLocker locker(&g_heap_tagging_lock);
     return SetHeapTaggingLevel(arg, arg_size);
   }
   if (opcode == M_INITIALIZE_GWP_ASAN) {
diff --git a/libc/bionic/malloc_common_dynamic.cpp b/libc/bionic/malloc_common_dynamic.cpp
index eeeaff9..2d6a1bb 100644
--- a/libc/bionic/malloc_common_dynamic.cpp
+++ b/libc/bionic/malloc_common_dynamic.cpp
@@ -58,6 +58,7 @@
 #include <android/dlext.h>
 
 #include <platform/bionic/malloc.h>
+#include <private/ScopedPthreadMutexLocker.h>
 #include <private/bionic_config.h>
 #include <private/bionic_defs.h>
 #include <private/bionic_malloc_dispatch.h>
@@ -523,6 +524,7 @@
     return FreeMallocLeakInfo(reinterpret_cast<android_mallopt_leak_info_t*>(arg));
   }
   if (opcode == M_SET_HEAP_TAGGING_LEVEL) {
+    ScopedPthreadMutexLocker locker(&g_heap_tagging_lock);
     return SetHeapTaggingLevel(arg, arg_size);
   }
   if (opcode == M_INITIALIZE_GWP_ASAN) {
diff --git a/libc/bionic/memory_mitigation_state.cpp b/libc/bionic/memory_mitigation_state.cpp
index abb1e8d..4761d88 100644
--- a/libc/bionic/memory_mitigation_state.cpp
+++ b/libc/bionic/memory_mitigation_state.cpp
@@ -39,8 +39,10 @@
 #include <bionic/malloc.h>
 #include <bionic/mte.h>
 
+#include <private/ScopedPthreadMutexLocker.h>
+#include <private/ScopedRWLock.h>
+
 #include "heap_tagging.h"
-#include "private/ScopedRWLock.h"
 #include "pthread_internal.h"
 
 extern "C" void scudo_malloc_set_zero_contents(int zero_contents);
@@ -54,8 +56,13 @@
   scudo_malloc_set_zero_contents(0);
 #endif
 
-  HeapTaggingLevel level = M_HEAP_TAGGING_LEVEL_NONE;
-  SetHeapTaggingLevel(reinterpret_cast<void*>(&level), sizeof(level));
+  ScopedPthreadMutexLocker locker(&g_heap_tagging_lock);
+
+  HeapTaggingLevel current_level = GetHeapTaggingLevel();
+  if (current_level != M_HEAP_TAGGING_LEVEL_NONE && current_level != M_HEAP_TAGGING_LEVEL_TBI) {
+    HeapTaggingLevel level = M_HEAP_TAGGING_LEVEL_NONE;
+    SetHeapTaggingLevel(reinterpret_cast<void*>(&level), sizeof(level));
+  }
 
   return true;
 }