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/Android.bp b/libc/Android.bp
index d739ce6..9dcd180 100644
--- a/libc/Android.bp
+++ b/libc/Android.bp
@@ -2145,6 +2145,28 @@
     defaults: ["crt_defaults"],
 }
 
+cc_library_static {
+    name: "note_memtag_heap_async",
+    arch: {
+      arm64: {
+        srcs: ["arch-arm64/bionic/note_memtag_heap_async.S"],
+      }
+    },
+
+    defaults: ["crt_defaults"],
+}
+
+cc_library_static {
+    name: "note_memtag_heap_sync",
+    arch: {
+      arm64: {
+        srcs: ["arch-arm64/bionic/note_memtag_heap_sync.S"],
+      }
+    },
+
+    defaults: ["crt_defaults"],
+}
+
 // ========================================================
 // NDK headers.
 // ========================================================
diff --git a/libc/NOTICE b/libc/NOTICE
index f7d73d8..7041e5c 100644
--- a/libc/NOTICE
+++ b/libc/NOTICE
@@ -1083,6 +1083,34 @@
 
 -------------------------------------------------------------------
 
+Copyright (C) 2021 The Android Open Source Project
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+ * Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+ * Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in
+   the documentation and/or other materials provided with the
+   distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGE.
+
+-------------------------------------------------------------------
+
 Copyright (c) 1980, 1983, 1988, 1993
    The Regents of the University of California.  All rights reserved.
 
diff --git a/libc/arch-arm64/bionic/note_memtag_heap_async.S b/libc/arch-arm64/bionic/note_memtag_heap_async.S
new file mode 100644
index 0000000..208115c
--- /dev/null
+++ b/libc/arch-arm64/bionic/note_memtag_heap_async.S
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in
+ *    the documentation and/or other materials provided with the
+ *    distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+ * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+ * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+ * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include <private/bionic_asm.h>
+#include <private/bionic_asm_note.h>
+
+__bionic_asm_custom_note_gnu_section()
+
+  .section ".note.android.memtag", "a", %note
+  .p2align 2
+  .long 1f - 0f         // int32_t namesz
+  .long 3f - 2f         // int32_t descsz
+  .long NT_TYPE_MEMTAG  // int32_t type
+0:
+  .asciz "Android"      // char name[]
+1:
+  .p2align 2
+2:
+  .long (NT_MEMTAG_LEVEL_ASYNC | NT_MEMTAG_HEAP) // value
+3:
+  .p2align 2
diff --git a/libc/arch-arm64/bionic/note_memtag_heap_sync.S b/libc/arch-arm64/bionic/note_memtag_heap_sync.S
new file mode 100644
index 0000000..d71ad02
--- /dev/null
+++ b/libc/arch-arm64/bionic/note_memtag_heap_sync.S
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in
+ *    the documentation and/or other materials provided with the
+ *    distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+ * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+ * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+ * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include <private/bionic_asm.h>
+#include <private/bionic_asm_note.h>
+
+__bionic_asm_custom_note_gnu_section()
+
+  .section ".note.android.memtag", "a", %note
+  .p2align 2
+  .long 1f - 0f         // int32_t namesz
+  .long 3f - 2f         // int32_t descsz
+  .long NT_TYPE_MEMTAG  // int32_t type
+0:
+  .asciz "Android"      // char name[]
+1:
+  .p2align 2
+2:
+  .long (NT_MEMTAG_LEVEL_SYNC | NT_MEMTAG_HEAP) // value
+3:
+  .p2align 2
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();
diff --git a/libc/private/bionic_asm_arm64.h b/libc/private/bionic_asm_arm64.h
index c11732a..e73f25d 100644
--- a/libc/private/bionic_asm_arm64.h
+++ b/libc/private/bionic_asm_arm64.h
@@ -69,4 +69,10 @@
     .long (__bionic_asm_aarch64_feature_pac | \
            __bionic_asm_aarch64_feature_bti); \
     .long 0; \
-    .popsection; \
+    .popsection;
+
+#define NT_MEMTAG_LEVEL_MASK 3
+#define NT_MEMTAG_LEVEL_DEFAULT 0
+#define NT_MEMTAG_LEVEL_ASYNC 1
+#define NT_MEMTAG_LEVEL_SYNC 2
+#define NT_MEMTAG_HEAP 4
diff --git a/libc/private/bionic_asm_note.h b/libc/private/bionic_asm_note.h
index c335c5a..80b8f01 100644
--- a/libc/private/bionic_asm_note.h
+++ b/libc/private/bionic_asm_note.h
@@ -30,3 +30,4 @@
 
 #define NT_TYPE_IDENT 1
 #define NT_TYPE_KUSER 3
+#define NT_TYPE_MEMTAG 4
diff --git a/libc/private/bionic_globals.h b/libc/private/bionic_globals.h
index 5c9b726..8132261 100644
--- a/libc/private/bionic_globals.h
+++ b/libc/private/bionic_globals.h
@@ -41,6 +41,8 @@
 #include "private/bionic_vdso.h"
 #include "private/WriteProtected.h"
 
+#include <platform/bionic/malloc.h>
+
 struct libc_globals {
   vdso_entry vdso[VDSO_END];
   long setjmp_cookie;
@@ -106,6 +108,8 @@
 
   const char* scudo_stack_depot = nullptr;
   const char* scudo_region_info = nullptr;
+
+  HeapTaggingLevel initial_heap_tagging_level = M_HEAP_TAGGING_LEVEL_NONE;
 };
 
 __LIBC_HIDDEN__ libc_shared_globals* __libc_shared_globals();