Make "disable memory mitigations" and "set heap tagging level" more available.

These were only available internally via android_mallopt(), but they're
likely to be needed by more code in future, so move them into mallopt().

This change leaves the android_mallopt() options for now, but I plan on
coming back to remove them after I've switched the handful of callers
over to mallopt() instead.

Bug: http://b/135772972
Test: treehugger
Change-Id: Ia154614069a7623c6aca85975a91e6a156f04759
diff --git a/libc/bionic/heap_tagging.cpp b/libc/bionic/heap_tagging.cpp
index 49b02da..9805fbd 100644
--- a/libc/bionic/heap_tagging.cpp
+++ b/libc/bionic/heap_tagging.cpp
@@ -96,12 +96,7 @@
 }
 
 // Requires `g_heap_tagging_lock` to be held.
-bool SetHeapTaggingLevel(void* arg, size_t arg_size) {
-  if (arg_size != sizeof(HeapTaggingLevel)) {
-    return false;
-  }
-
-  auto tag_level = *reinterpret_cast<HeapTaggingLevel*>(arg);
+bool SetHeapTaggingLevel(HeapTaggingLevel tag_level) {
   if (tag_level == heap_tagging_level) {
     return true;
   }
diff --git a/libc/bionic/heap_tagging.h b/libc/bionic/heap_tagging.h
index db45fc1..110b6ed 100644
--- a/libc/bionic/heap_tagging.h
+++ b/libc/bionic/heap_tagging.h
@@ -42,5 +42,5 @@
 
 // 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);
+bool SetHeapTaggingLevel(HeapTaggingLevel level);
 HeapTaggingLevel GetHeapTaggingLevel();
diff --git a/libc/bionic/malloc_common.cpp b/libc/bionic/malloc_common.cpp
index 0ee12a7..955cf18 100644
--- a/libc/bionic/malloc_common.cpp
+++ b/libc/bionic/malloc_common.cpp
@@ -102,6 +102,15 @@
 }
 
 extern "C" int mallopt(int param, int value) {
+  // Some are handled by libc directly rather than by the allocator.
+  if (param == M_BIONIC_SET_HEAP_TAGGING_LEVEL) {
+    ScopedPthreadMutexLocker locker(&g_heap_tagging_lock);
+    return SetHeapTaggingLevel(static_cast<HeapTaggingLevel>(value));
+  }
+  if (param == M_BIONIC_DISABLE_MEMORY_MITIGATIONS) {
+    return DisableMemoryMitigations(value);
+  }
+  // The rest we pass on...
   auto dispatch_table = GetDispatchTable();
   if (__predict_false(dispatch_table != nullptr)) {
     return dispatch_table->mallopt(param, value);
@@ -316,9 +325,8 @@
   if (opcode == M_SET_ALLOCATION_LIMIT_BYTES) {
     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_SET_HEAP_TAGGING_LEVEL && arg_size == sizeof(HeapTaggingLevel)) {
+    return mallopt(M_BIONIC_SET_HEAP_TAGGING_LEVEL, *reinterpret_cast<HeapTaggingLevel*>(arg));
   }
   if (opcode == M_INITIALIZE_GWP_ASAN) {
     if (arg == nullptr || arg_size != sizeof(bool)) {
@@ -329,8 +337,8 @@
       return MaybeInitGwpAsan(globals, *reinterpret_cast<bool*>(arg));
     });
   }
-  if (opcode == M_DISABLE_MEMORY_MITIGATIONS) {
-    return DisableMemoryMitigations(arg, arg_size);
+  if (opcode == M_DISABLE_MEMORY_MITIGATIONS && arg_size == sizeof(int)) {
+    return mallopt(M_BIONIC_DISABLE_MEMORY_MITIGATIONS, reinterpret_cast<intptr_t>(arg));
   }
   errno = ENOTSUP;
   return false;
diff --git a/libc/bionic/malloc_common_dynamic.cpp b/libc/bionic/malloc_common_dynamic.cpp
index 2d6a1bb..93d4d53 100644
--- a/libc/bionic/malloc_common_dynamic.cpp
+++ b/libc/bionic/malloc_common_dynamic.cpp
@@ -523,9 +523,8 @@
     }
     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_SET_HEAP_TAGGING_LEVEL && arg_size == sizeof(HeapTaggingLevel)) {
+    return mallopt(M_BIONIC_SET_HEAP_TAGGING_LEVEL, *reinterpret_cast<HeapTaggingLevel*>(arg));
   }
   if (opcode == M_INITIALIZE_GWP_ASAN) {
     if (arg == nullptr || arg_size != sizeof(bool)) {
@@ -536,8 +535,8 @@
       return MaybeInitGwpAsan(globals, *reinterpret_cast<bool*>(arg));
     });
   }
-  if (opcode == M_DISABLE_MEMORY_MITIGATIONS) {
-    return DisableMemoryMitigations(arg, arg_size);
+  if (opcode == M_DISABLE_MEMORY_MITIGATIONS && arg_size == sizeof(int)) {
+    return mallopt(M_BIONIC_DISABLE_MEMORY_MITIGATIONS, reinterpret_cast<intptr_t>(arg));
   }
   // Try heapprofd's mallopt, as it handles options not covered here.
   return HeapprofdMallopt(opcode, arg, arg_size);
diff --git a/libc/bionic/memory_mitigation_state.cpp b/libc/bionic/memory_mitigation_state.cpp
index 4761d88..b262f86 100644
--- a/libc/bionic/memory_mitigation_state.cpp
+++ b/libc/bionic/memory_mitigation_state.cpp
@@ -47,8 +47,8 @@
 
 extern "C" void scudo_malloc_set_zero_contents(int zero_contents);
 
-bool DisableMemoryMitigations(void* arg, size_t arg_size) {
-  if (arg || arg_size) {
+bool DisableMemoryMitigations(int arg) {
+  if (arg != 0) {
     return false;
   }
 
@@ -61,7 +61,7 @@
   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));
+    SetHeapTaggingLevel(level);
   }
 
   return true;
diff --git a/libc/bionic/memory_mitigation_state.h b/libc/bionic/memory_mitigation_state.h
index ffa1912..047b6ad 100644
--- a/libc/bionic/memory_mitigation_state.h
+++ b/libc/bionic/memory_mitigation_state.h
@@ -30,4 +30,4 @@
 
 #include <stddef.h>
 
-bool DisableMemoryMitigations(void* arg, size_t arg_size);
+bool DisableMemoryMitigations(int arg);
diff --git a/libc/include/malloc.h b/libc/include/malloc.h
index a237254..598a50b 100644
--- a/libc/include/malloc.h
+++ b/libc/include/malloc.h
@@ -204,6 +204,60 @@
 #define M_TSDS_COUNT_MAX (-202)
 
 /**
+ * mallopt() option to disable heap initialization across the whole process.
+ * If the hardware supports memory tagging, this also disables memory tagging.
+ * May be called at any time, including when multiple threads are running. The
+ * value is unused but must be set to 0.
+ *
+ * Note that these memory mitigations are only implemented in scudo and
+ * therefore this will have no effect when using another allocator (such as
+ * jemalloc on Android Go devices).
+ *
+ * Available since API level 31.
+ */
+#define M_BIONIC_DISABLE_MEMORY_MITIGATIONS (-203)
+
+/**
+ * mallopt() option to change the heap tagging state. May be called at any
+ * time, including when multiple threads are running.
+ * The value must be one of the M_HEAP_TAGGING_LEVEL_ constants.
+ *
+ * Available since API level 31.
+ */
+#define M_BIONIC_SET_HEAP_TAGGING_LEVEL (-204)
+
+/**
+ * Constants for use with the M_BIONIC_SET_HEAP_TAGGING_LEVEL mallopt() option.
+ */
+enum HeapTaggingLevel {
+  /**
+   * Disable heap tagging and memory tag checks (if supported).
+   * Heap tagging may not be re-enabled after being disabled.
+   */
+  M_HEAP_TAGGING_LEVEL_NONE = 0,
+#define M_HEAP_TAGGING_LEVEL_NONE M_HEAP_TAGGING_LEVEL_NONE
+  /**
+   * Address-only tagging. Heap pointers have a non-zero tag in the
+   * most significant ("top") byte which is checked in free(). Memory
+   * accesses ignore the tag using arm64's Top Byte Ignore (TBI) feature.
+   */
+  M_HEAP_TAGGING_LEVEL_TBI = 1,
+#define M_HEAP_TAGGING_LEVEL_TBI M_HEAP_TAGGING_LEVEL_TBI
+  /**
+   * Enable heap tagging and asynchronous memory tag checks (if supported).
+   * Disable stack trace collection.
+   */
+  M_HEAP_TAGGING_LEVEL_ASYNC = 2,
+#define M_HEAP_TAGGING_LEVEL_ASYNC M_HEAP_TAGGING_LEVEL_ASYNC
+  /**
+   * Enable heap tagging and synchronous memory tag checks (if supported).
+   * Enable stack trace collection.
+   */
+  M_HEAP_TAGGING_LEVEL_SYNC = 3,
+#define M_HEAP_TAGGING_LEVEL_SYNC M_HEAP_TAGGING_LEVEL_SYNC
+};
+
+/**
  * [mallopt(3)](http://man7.org/linux/man-pages/man3/mallopt.3.html) modifies
  * heap behavior. Values of `__option` are the `M_` constants from this header.
  *
diff --git a/libc/platform/bionic/malloc.h b/libc/platform/bionic/malloc.h
index 56badf0..313cbdf 100644
--- a/libc/platform/bionic/malloc.h
+++ b/libc/platform/bionic/malloc.h
@@ -28,6 +28,7 @@
 
 #pragma once
 
+#include <malloc.h>
 #include <stdbool.h>
 #include <stdint.h>
 
@@ -114,21 +115,6 @@
 #define M_DISABLE_MEMORY_MITIGATIONS M_DISABLE_MEMORY_MITIGATIONS
 };
 
-enum HeapTaggingLevel {
-  // Disable heap tagging and memory tag checks if supported. Heap tagging may not be re-enabled
-  // after being disabled.
-  M_HEAP_TAGGING_LEVEL_NONE = 0,
-  // Address-only tagging. Heap pointers have a non-zero tag in the most significant byte which is
-  // checked in free(). Memory accesses ignore the tag.
-  M_HEAP_TAGGING_LEVEL_TBI = 1,
-  // Enable heap tagging and asynchronous memory tag checks if supported. Disable stack trace
-  // collection.
-  M_HEAP_TAGGING_LEVEL_ASYNC = 2,
-  // Enable heap tagging and synchronous memory tag checks if supported. Enable stack trace
-  // collection.
-  M_HEAP_TAGGING_LEVEL_SYNC = 3,
-};
-
 // Manipulates bionic-specific handling of memory allocation APIs such as
 // malloc. Only for use by the Android platform itself.
 //