Use ELF notes to set the desired memory tagging level.

Use a note in executables to specify
(none|sync|async) heap tagging level. To be extended with (heap x stack x
globals) in the future. A missing note disables all tagging.

Bug: b/135772972
Test: bionic-unit-tests (in a future change)

Change-Id: Iab145a922c7abe24cdce17323f9e0c1063cc1321
diff --git a/libc/bionic/heap_tagging.cpp b/libc/bionic/heap_tagging.cpp
index 2c5d4d8..72f8f1d 100644
--- a/libc/bionic/heap_tagging.cpp
+++ b/libc/bionic/heap_tagging.cpp
@@ -42,30 +42,29 @@
 
 void SetDefaultHeapTaggingLevel() {
 #if defined(__aarch64__)
-#ifdef ANDROID_EXPERIMENTAL_MTE
-  // First, try enabling MTE in asynchronous mode, with tag 0 excluded. This will fail if the kernel
-  // or hardware doesn't support MTE, and we will fall back to just enabling tagged pointers in
-  // syscall arguments.
-  if (prctl(PR_SET_TAGGED_ADDR_CTRL,
-            PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_ASYNC | (0xfffe << PR_MTE_TAG_SHIFT), 0, 0,
-            0) == 0) {
-    heap_tagging_level = M_HEAP_TAGGING_LEVEL_ASYNC;
-    return;
-  }
-#endif // ANDROID_EXPERIMENTAL_MTE
-
-  // Allow the kernel to accept tagged pointers in syscall arguments. This is a no-op (kernel
-  // returns -EINVAL) if the kernel doesn't understand the prctl.
-  if (prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0) == 0) {
 #if !__has_feature(hwaddress_sanitizer)
-    heap_tagging_level = M_HEAP_TAGGING_LEVEL_TBI;
-    __libc_globals.mutate([](libc_globals* globals) {
-      // Arrange for us to set pointer tags to POINTER_TAG, check tags on
-      // deallocation and untag when passing pointers to the allocator.
-      globals->heap_pointer_tag = (reinterpret_cast<uintptr_t>(POINTER_TAG) << TAG_SHIFT) |
-                                  (0xffull << CHECK_SHIFT) | (0xffull << UNTAG_SHIFT);
-    });
-#endif  // hwaddress_sanitizer
+  heap_tagging_level = __libc_shared_globals()->initial_heap_tagging_level;
+#endif
+  switch (heap_tagging_level) {
+    case M_HEAP_TAGGING_LEVEL_TBI:
+      __libc_globals.mutate([](libc_globals* globals) {
+        // Arrange for us to set pointer tags to POINTER_TAG, check tags on
+        // deallocation and untag when passing pointers to the allocator.
+        globals->heap_pointer_tag = (reinterpret_cast<uintptr_t>(POINTER_TAG) << TAG_SHIFT) |
+                                    (0xffull << CHECK_SHIFT) | (0xffull << UNTAG_SHIFT);
+      });
+      break;
+#if defined(ANDROID_EXPERIMENTAL_MTE) && defined(USE_SCUDO)
+    case M_HEAP_TAGGING_LEVEL_SYNC:
+      scudo_malloc_set_track_allocation_stacks(1);
+      break;
+
+    case M_HEAP_TAGGING_LEVEL_NONE:
+      scudo_malloc_disable_memory_tagging();
+      break;
+#endif  // ANDROID_EXPERIMENTAL_MTE
+    default:
+      break;
   }
 #endif  // aarch64
 }
diff --git a/libc/bionic/libc_init_common.cpp b/libc/bionic/libc_init_common.cpp
index 80adbbe..f2c3f1c 100644
--- a/libc/bionic/libc_init_common.cpp
+++ b/libc/bionic/libc_init_common.cpp
@@ -86,7 +86,7 @@
   _thread_arc4_lock();
 }
 
-static void __libc_init_malloc_fill_contents() {
+void __libc_init_scudo() {
 // TODO(b/158870657) make this unconditional when all devices support SCUDO.
 #if defined(USE_SCUDO)
 #if defined(SCUDO_PATTERN_FILL_CONTENTS)
@@ -95,6 +95,7 @@
   scudo_malloc_set_zero_contents(1);
 #endif
 #endif
+  SetDefaultHeapTaggingLevel();
 }
 
 __BIONIC_WEAK_FOR_NATIVE_BRIDGE
@@ -119,9 +120,6 @@
   __system_properties_init(); // Requires 'environ'.
   __libc_init_fdsan(); // Requires system properties (for debug.fdsan).
   __libc_init_fdtrack();
-
-  __libc_init_malloc_fill_contents();
-  SetDefaultHeapTaggingLevel();
 }
 
 void __libc_init_fork_handler() {
diff --git a/libc/bionic/libc_init_common.h b/libc/bionic/libc_init_common.h
index be7526f..a899089 100644
--- a/libc/bionic/libc_init_common.h
+++ b/libc/bionic/libc_init_common.h
@@ -28,6 +28,7 @@
 
 #pragma once
 
+#include <stdint.h>
 #include <sys/cdefs.h>
 
 typedef void init_func_t(int, char*[], char*[]);
@@ -57,6 +58,8 @@
 
 __LIBC_HIDDEN__ void __libc_init_common();
 
+__LIBC_HIDDEN__ void __libc_init_scudo();
+
 __LIBC_HIDDEN__ void __libc_init_AT_SECURE(char** envp);
 
 // The fork handler must be initialised after __libc_init_malloc, as
diff --git a/libc/bionic/libc_init_dynamic.cpp b/libc/bionic/libc_init_dynamic.cpp
index c9da02e..175fa3e 100644
--- a/libc/bionic/libc_init_dynamic.cpp
+++ b/libc/bionic/libc_init_dynamic.cpp
@@ -90,6 +90,7 @@
 
   __libc_init_globals();
   __libc_init_common();
+  __libc_init_scudo();
 
   // Hooks for various libraries to let them know that we're starting up.
   __libc_globals.mutate(__libc_init_malloc);
diff --git a/libc/bionic/libc_init_static.cpp b/libc/bionic/libc_init_static.cpp
index 4a73918..3705606 100644
--- a/libc/bionic/libc_init_static.cpp
+++ b/libc/bionic/libc_init_static.cpp
@@ -44,6 +44,8 @@
 #include "private/bionic_elf_tls.h"
 #include "private/bionic_globals.h"
 #include "platform/bionic/macros.h"
+#include "private/bionic_asm.h"
+#include "private/bionic_asm_note.h"
 #include "private/bionic_tls.h"
 #include "private/KernelArgumentBlock.h"
 
@@ -158,6 +160,94 @@
   layout.finish_layout();
 }
 
+#ifdef __aarch64__
+static bool __read_memtag_note(const ElfW(Nhdr)* note, const char* name, const char* desc,
+                               unsigned* result) {
+  if (note->n_namesz != 8 || strncmp(name, "Android", 8) != 0) {
+    return false;
+  }
+  if (note->n_type != NT_TYPE_MEMTAG) {
+    return false;
+  }
+  if (note->n_descsz != 4) {
+    async_safe_fatal("unrecognized android.memtag note: n_descsz = %d, expected 4", note->n_descsz);
+  }
+  *result = *reinterpret_cast<const ElfW(Word)*>(desc);
+  return true;
+}
+
+static unsigned __get_memtag_note(const ElfW(Phdr)* phdr_start, size_t phdr_ct,
+                                  const ElfW(Addr) load_bias) {
+  for (size_t i = 0; i < phdr_ct; ++i) {
+    const ElfW(Phdr)* phdr = &phdr_start[i];
+    if (phdr->p_type != PT_NOTE) {
+      continue;
+    }
+    ElfW(Addr) p = load_bias + phdr->p_vaddr;
+    ElfW(Addr) note_end = load_bias + phdr->p_vaddr + phdr->p_memsz;
+    while (p + sizeof(ElfW(Nhdr)) <= note_end) {
+      const ElfW(Nhdr)* note = reinterpret_cast<const ElfW(Nhdr)*>(p);
+      p += sizeof(ElfW(Nhdr));
+      const char* name = reinterpret_cast<const char*>(p);
+      p += align_up(note->n_namesz, 4);
+      const char* desc = reinterpret_cast<const char*>(p);
+      p += align_up(note->n_descsz, 4);
+      if (p > note_end) {
+        break;
+      }
+      unsigned ret;
+      if (__read_memtag_note(note, name, desc, &ret)) {
+        return ret;
+      }
+    }
+  }
+  return 0;
+}
+
+// Figure out the desired memory tagging mode (sync/async, heap/globals/stack) for this executable.
+// This function is called from the linker before the main executable is relocated.
+__attribute__((no_sanitize("hwaddress", "memtag"))) void __libc_init_mte(const void* phdr_start,
+                                                                         size_t phdr_ct,
+                                                                         uintptr_t load_bias) {
+  unsigned v =
+      __get_memtag_note(reinterpret_cast<const ElfW(Phdr)*>(phdr_start), phdr_ct, load_bias);
+
+  if (v & ~(NT_MEMTAG_LEVEL_MASK | NT_MEMTAG_HEAP)) {
+    async_safe_fatal("unrecognized android.memtag note: desc = %d", v);
+  }
+
+  if (v & NT_MEMTAG_HEAP) {
+    unsigned memtag_level = v & NT_MEMTAG_LEVEL_MASK;
+    unsigned long arg = PR_TAGGED_ADDR_ENABLE | (0xfffe << PR_MTE_TAG_SHIFT);
+    HeapTaggingLevel level;
+    switch (memtag_level) {
+      case NT_MEMTAG_LEVEL_ASYNC:
+        arg |= PR_MTE_TCF_ASYNC;
+        level = M_HEAP_TAGGING_LEVEL_ASYNC;
+        break;
+      case NT_MEMTAG_LEVEL_DEFAULT:
+      case NT_MEMTAG_LEVEL_SYNC:
+        arg |= PR_MTE_TCF_SYNC;
+        level = M_HEAP_TAGGING_LEVEL_SYNC;
+        break;
+      default:
+        async_safe_fatal("unrecognized android.memtag note: level = %d", memtag_level);
+    }
+
+    if (prctl(PR_SET_TAGGED_ADDR_CTRL, arg, 0, 0, 0) == 0) {
+      __libc_shared_globals()->initial_heap_tagging_level = level;
+      return;
+    }
+  }
+
+  if (prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0) == 0) {
+    __libc_shared_globals()->initial_heap_tagging_level = M_HEAP_TAGGING_LEVEL_TBI;
+  }
+}
+#else   // __aarch64__
+void __libc_init_mte(const void*, size_t, uintptr_t) {}
+#endif  // __aarch64__
+
 __noreturn static void __real_libc_init(void *raw_args,
                                         void (*onexit)(void) __unused,
                                         int (*slingshot)(int, char**, char**),
@@ -175,6 +265,9 @@
   layout_static_tls(args);
   __libc_init_main_thread_final();
   __libc_init_common();
+  __libc_init_mte(reinterpret_cast<ElfW(Phdr)*>(getauxval(AT_PHDR)), getauxval(AT_PHNUM),
+                  /*load_bias = */ 0);
+  __libc_init_scudo();
   __libc_init_fork_handler();
 
   call_ifunc_resolvers();