[ELF][MTE] Add --android-memtag-* options to synthesize ELF notes

This ELF note is aarch64 and Android-specific. It specifies to the
dynamic loader that specific work should be scheduled to enable MTE
protection of stack and heap regions.

Current synthesis of the ".note.android.memtag" ELF note is done in the
Android build system. We'd like to move that to the compiler. This patch
adds the --memtag-stack, --memtag-heap, and --memtag-mode={async, sync,
none} flags to the linker, which synthesises the note for us.

Future changes will add -fsanitize=memtag* flags to clang which will
pass these through to lld.

Depends on D119381.

Differential Revision: https://reviews.llvm.org/D119384

lld/ELF/Config.h

@@ -347,6 +347,18 @@ struct Configuration {
 
   // 4 for ELF32, 8 for ELF64.
   int wordsize;
+
+  // Mode of MTE to write to the ELF note. Should be one of NT_MEMTAG_ASYNC (for
+  // async), NT_MEMTAG_SYNC (for sync), or NT_MEMTAG_LEVEL_NONE (for none). If
+  // async or sync is enabled, write the ELF note specifying the default MTE
+  // mode.
+  int androidMemtagMode;
+  // Signal to the dynamic loader to enable heap MTE.
+  bool androidMemtagHeap;
+  // Signal to the dynamic loader that this binary expects stack MTE. Generally,
+  // this means to map the primary and thread stacks as PROT_MTE. Note: This is
+  // not supported on Android 11 & 12.
+  bool androidMemtagStack;
 };
 
 // The only instance of Configuration struct.

lld/ELF/Driver.cpp

@@ -705,6 +705,28 @@ static StringRef getDynamicLinker(opt::InputArgList &args) {
   return arg->getValue();
 }
 
+static int getMemtagMode(opt::InputArgList &args) {
+  StringRef memtagModeArg = args.getLastArgValue(OPT_android_memtag_mode);
+  if (!config->androidMemtagHeap && !config->androidMemtagStack) {
+    if (!memtagModeArg.empty())
+      error("when using --android-memtag-mode, at least one of "
+            "--android-memtag-heap or "
+            "--android-memtag-stack is required");
+    return ELF::NT_MEMTAG_LEVEL_NONE;
+  }
+
+  if (memtagModeArg == "sync" || memtagModeArg.empty())
+    return ELF::NT_MEMTAG_LEVEL_SYNC;
+  if (memtagModeArg == "async")
+    return ELF::NT_MEMTAG_LEVEL_ASYNC;
+  if (memtagModeArg == "none")
+    return ELF::NT_MEMTAG_LEVEL_NONE;
+
+  error("unknown --android-memtag-mode value: \"" + memtagModeArg +
+        "\", should be one of {async, sync, none}");
+  return ELF::NT_MEMTAG_LEVEL_NONE;
+}
+
 static ICFLevel getICF(opt::InputArgList &args) {
   auto *arg = args.getLastArg(OPT_icf_none, OPT_icf_safe, OPT_icf_all);
   if (!arg || arg->getOption().getID() == OPT_icf_none)
@@ -1008,6 +1030,11 @@ static void readConfigs(opt::InputArgList &args) {
       args.hasFlag(OPT_allow_multiple_definition,
                    OPT_no_allow_multiple_definition, false) ||
       hasZOption(args, "muldefs");
+  config->androidMemtagHeap =
+      args.hasFlag(OPT_android_memtag_heap, OPT_no_android_memtag_heap, false);
+  config->androidMemtagStack = args.hasFlag(OPT_android_memtag_stack,
+                                            OPT_no_android_memtag_stack, false);
+  config->androidMemtagMode = getMemtagMode(args);
   config->auxiliaryList = args::getStrings(args, OPT_auxiliary);
   if (opt::Arg *arg =
           args.getLastArg(OPT_Bno_symbolic, OPT_Bsymbolic_non_weak_functions,

lld/ELF/Options.td

@@ -718,3 +718,15 @@ defm check_dynamic_relocations: BB<"check-dynamic-relocations",
   Flags<[HelpHidden]>;
 
 defm load_pass_plugins: EEq<"load-pass-plugin", "Load passes from plugin library">;
+
+// Hidden options, used by clang's -fsanitize=memtag-* options to emit an ELF
+// note to designate what kinds of memory (stack/heap) should be protected using
+// ARM's MTE on armv8.5+. A binary's desire for stack MTE can't be obtained
+// implicitly, so we have a specific bit in the note to signal to the loader to
+// remap the stack as PROT_MTE.
+defm android_memtag_stack: BB<"android-memtag-stack",
+    "Instruct the dynamic loader to prepare for MTE stack instrumentation", "">;
+defm android_memtag_heap: BB<"android-memtag-heap",
+    "Instruct the dynamic loader to enable MTE protection for the heap", "">;
+defm android_memtag_mode: EEq<"android-memtag-mode",
+    "Instruct the dynamic loader to start under MTE mode {async, sync, none}">;

lld/ELF/SyntheticSections.cpp

@@ -32,6 +32,7 @@
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/BinaryFormat/Dwarf.h"
+#include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFDebugPubTable.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/LEB128.h"
@@ -3847,6 +3848,35 @@ void InStruct::reset() {
   symTabShndx.reset();
 }
 
+constexpr char kMemtagAndroidNoteName[] = "Android";
+void MemtagAndroidNote::writeTo(uint8_t *buf) {
+  assert(sizeof(kMemtagAndroidNoteName) == 8); // ABI check for Android 11 & 12.
+  assert((config->androidMemtagStack || config->androidMemtagHeap) &&
+         "Should only be synthesizing a note if heap || stack is enabled.");
+
+  write32(buf, sizeof(kMemtagAndroidNoteName));
+  write32(buf + 4, sizeof(uint32_t));
+  write32(buf + 8, ELF::NT_ANDROID_TYPE_MEMTAG);
+  memcpy(buf + 12, kMemtagAndroidNoteName, sizeof(kMemtagAndroidNoteName));
+  buf += 12 + sizeof(kMemtagAndroidNoteName);
+
+  uint32_t value = 0;
+  value |= config->androidMemtagMode;
+  if (config->androidMemtagHeap)
+    value |= ELF::NT_MEMTAG_HEAP;
+  // Note, MTE stack is an ABI break. Attempting to run an MTE stack-enabled
+  // binary on Android 11 or 12 will result in a checkfail in the loader.
+  if (config->androidMemtagStack)
+    value |= ELF::NT_MEMTAG_STACK;
+  write32(buf, value); // note value
+}
+
+size_t MemtagAndroidNote::getSize() const {
+  return sizeof(llvm::ELF::Elf64_Nhdr) +
+         /*namesz=*/sizeof(kMemtagAndroidNoteName) +
+         /*descsz=*/sizeof(uint32_t);
+}
+
 InStruct elf::in;
 
 std::vector<Partition> elf::partitions;

lld/ELF/SyntheticSections.h

@@ -1187,6 +1187,18 @@ public:
   void writeTo(uint8_t *buf) override;
 };
 
+// See the following link for the Android-specific loader code that operates on
+// this section:
+// https://cs.android.com/android/platform/superproject/+/master:bionic/libc/bionic/libc_init_static.cpp;drc=9425b16978f9c5aa8f2c50c873db470819480d1d;l=192
+class MemtagAndroidNote : public SyntheticSection {
+public:
+  MemtagAndroidNote()
+      : SyntheticSection(llvm::ELF::SHF_ALLOC, llvm::ELF::SHT_NOTE,
+                         /*alignment=*/4, ".note.android.memtag") {}
+  void writeTo(uint8_t *buf) override;
+  size_t getSize() const override;
+};
+
 InputSection *createInterpSection();
 MergeInputSection *createCommentSection();
 template <class ELFT> void splitSections();
@@ -1217,6 +1229,7 @@ struct Partition {
   std::unique_ptr<EhFrameSection> ehFrame;
   std::unique_ptr<GnuHashTableSection> gnuHashTab;
   std::unique_ptr<HashTableSection> hashTab;
+  std::unique_ptr<MemtagAndroidNote> memtagAndroidNote;
   std::unique_ptr<RelocationBaseSection> relaDyn;
   std::unique_ptr<RelrBaseSection> relrDyn;
   std::unique_ptr<VersionDefinitionSection> verDef;

lld/ELF/Writer.cpp

@@ -362,6 +362,13 @@ template <class ELFT> void elf::createSyntheticSections() {
     part.dynSymTab =
         std::make_unique<SymbolTableSection<ELFT>>(*part.dynStrTab);
     part.dynamic = std::make_unique<DynamicSection<ELFT>>();
+
+    if (config->emachine == EM_AARCH64 &&
+        config->androidMemtagMode != ELF::NT_MEMTAG_LEVEL_NONE) {
+      part.memtagAndroidNote = std::make_unique<MemtagAndroidNote>();
+      add(*part.memtagAndroidNote);
+    }
+
     if (config->androidPackDynRelocs)
       part.relaDyn =
           std::make_unique<AndroidPackedRelocationSection<ELFT>>(relaDynName);

lld/test/ELF/aarch64-memtag-android-abi.s

@@ -0,0 +1,62 @@
+# REQUIRES: aarch64
+
+## Old versions of Android (Android 11 & 12) have very strict parsing logic on
+## the layout of the ELF note. This test ensures that backwards compatibility is
+## maintained, i.e. new versions of the linker will still produce binaries that
+## can be run on these versions of Android.
+
+# RUN: llvm-mc --filetype=obj -triple=aarch64-none-linux-android %s -o %t.o
+# RUN: ld.lld --android-memtag-mode=async --android-memtag-heap %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,NOSTACK,ASYNC
+
+# RUN: ld.lld --android-memtag-mode=sync --android-memtag-heap %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,NOSTACK,SYNC
+
+# RUN: ld.lld --android-memtag-mode=async --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,NOHEAP,STACK,ASYNC
+
+# RUN: ld.lld --android-memtag-mode=sync --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,NOHEAP,STACK,SYNC
+
+# RUN: ld.lld --android-memtag-mode=async --android-memtag-heap \
+# RUN:    --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,STACK,ASYNC
+
+# RUN: ld.lld --android-memtag-mode=sync --android-memtag-heap \
+# RUN:    --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,STACK,SYNC
+
+# RUN: ld.lld --android-memtag-heap %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,NOSTACK,SYNC
+
+# RUN: ld.lld --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,NOHEAP,STACK,SYNC
+
+# RUN: ld.lld --android-memtag-heap --android-memtag-stack %t.o -o %t
+# RUN: llvm-readelf -n %t | FileCheck %s --check-prefixes=NOTE,HEAP,STACK,SYNC
+
+# NOTE: .note.android.memtag
+# NOTE-NEXT: Owner
+# NOTE-NEXT: Android 0x00000004 NT_ANDROID_TYPE_MEMTAG (Android memory tagging
+# NOTE-SAME: information)
+# ASYNC-NEXT: Tagging Mode: ASYNC
+# SYNC-NEXT: Tagging Mode: SYNC
+# HEAP-NEXT: Heap: Enabled
+# NOHEAP-NEXT: Heap: Disabled
+## As of Android 12, stack MTE is unimplemented. However, we pre-emptively emit
+## a bit that signifies to the dynamic loader to map the primary and thread
+## stacks as PROT_MTE, in preparation for the bionic support.
+# STACK-NEXT:   Stack: Enabled
+# NOSTACK-NEXT: Stack: Disabled
+
+# RUN: not ld.lld --android-memtag-mode=asymm --android-memtag-heap 2>&1 | \
+# RUN:    FileCheck %s --check-prefix=BAD-MODE
+# BAD-MODE: unknown --android-memtag-mode value: "asymm", should be one of {async, sync, none}
+
+# RUN: not ld.lld --android-memtag-mode=async 2>&1 | \
+# RUN:    FileCheck %s --check-prefix=MISSING-STACK-OR-HEAP
+# MISSING-STACK-OR-HEAP: when using --android-memtag-mode, at least one of --android-memtag-heap or --android-memtag-stack is required
+
+.globl _start
+_start:
+  ret