[lld] Synthesize metadata for MTE globals

As per the ABI at
https://github.com/ARM-software/abi-aa/blob/main/memtagabielf64/memtagabielf64.rst,
this patch interprets the SHT_AARCH64_MEMTAG_GLOBALS_STATIC section,
which contains R_NONE relocations to tagged globals, and emits a
SHT_AARCH64_MEMTAG_GLOBALS_DYNAMIC section, with the correct
DT_AARCH64_MEMTAG_GLOBALS and DT_AARCH64_MEMTAG_GLOBALSSZ dynamic
entries. This section describes, in a uleb-encoded stream, global memory
ranges that should be tagged with MTE.

We are also out of bits to spare in the LLD Symbol class. As a result,
I've reused the 'needsTocRestore' bit, which is a PPC64 only feature.
Now, it's also used for 'isTagged' on AArch64.

An entry in SHT_AARCH64_MEMTAG_GLOBALS_STATIC is practically a guarantee
from an objfile that all references to the linked symbol are through the
GOT, and meet correct alignment requirements. As a result, we go through
all symbols and make sure that, for all symbols $SYM, all object files
that reference $SYM also have a SHT_AARCH64_MEMTAG_GLOBALS_STATIC entry
for $SYM. If this isn't the case, we demote the symbol to being
untagged. Symbols that are imported from other DSOs should always be
fine, as they're GOT-referenced (and thus the GOT entry either has the
correct tag or not, depending on whether it's tagged in the defining DSO
or not).

Additionally hand-tested by building {libc, libm, libc++, libm, and libnetd}
on Android with some experimental MTE globals support in the
linker/libc.

Reviewed By: MaskRay, peter.smith

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

Author: hctim

lld/ELF/Arch/AArch64.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "InputFiles.h"
 #include "OutputSections.h"
 #include "Symbols.h"
 #include "SyntheticSections.h"
@@ -377,6 +378,20 @@ void AArch64::relocate(uint8_t *loc, const Relocation &rel,
     write32(loc, val);
     break;
   case R_AARCH64_ABS64:
+    // AArch64 relocations to tagged symbols have extended semantics, as
+    // described here:
+    // https://github.com/ARM-software/abi-aa/blob/main/memtagabielf64/memtagabielf64.rst#841extended-semantics-of-r_aarch64_relative.
+    // tl;dr: encode the symbol's special addend in the place, which is an
+    // offset to the point where the logical tag is derived from. Quick hack, if
+    // the addend is within the symbol's bounds, no need to encode the tag
+    // derivation offset.
+    if (rel.sym && rel.sym->isTagged() &&
+        (rel.addend < 0 ||
+         rel.addend >= static_cast<int64_t>(rel.sym->getSize())))
+      write64(loc, -rel.addend);
+    else
+      write64(loc, val);
+    break;
   case R_AARCH64_PREL64:
     write64(loc, val);
     break;
@@ -745,6 +760,12 @@ bool AArch64Relaxer::tryRelaxAdrpLdr(const Relocation &adrpRel,
   return true;
 }
 
+// Tagged symbols have upper address bits that are added by the dynamic loader,
+// and thus need the full 64-bit GOT entry. Do not relax such symbols.
+static bool needsGotForMemtag(const Relocation &rel) {
+  return rel.sym->isTagged() && needsGot(rel.expr);
+}
+
 void AArch64::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
   uint64_t secAddr = sec.getOutputSection()->addr;
   if (auto *s = dyn_cast<InputSection>(&sec))
@@ -756,6 +777,12 @@ void AArch64::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
     const uint64_t val =
         sec.getRelocTargetVA(sec.file, rel.type, rel.addend,
                              secAddr + rel.offset, *rel.sym, rel.expr);
+
+    if (needsGotForMemtag(rel)) {
+      relocate(loc, rel, val);
+      continue;
+    }
+
     switch (rel.expr) {
     case R_AARCH64_GOT_PAGE_PC:
       if (i + 1 < size &&
@@ -950,3 +977,107 @@ static TargetInfo *getTargetInfo() {
 }
 
 TargetInfo *elf::getAArch64TargetInfo() { return getTargetInfo(); }
+
+template <class ELFT>
+static void
+addTaggedSymbolReferences(InputSectionBase &sec,
+                          DenseMap<Symbol *, unsigned> &referenceCount) {
+  assert(sec.type == SHT_AARCH64_MEMTAG_GLOBALS_STATIC);
+
+  const RelsOrRelas<ELFT> rels = sec.relsOrRelas<ELFT>();
+  if (rels.areRelocsRel())
+    error("non-RELA relocations are not allowed with memtag globals");
+
+  for (const typename ELFT::Rela &rel : rels.relas) {
+    Symbol &sym = sec.getFile<ELFT>()->getRelocTargetSym(rel);
+    // Linker-synthesized symbols such as __executable_start may be referenced
+    // as tagged in input objfiles, and we don't want them to be tagged. A
+    // cheap way to exclude them is the type check, but their type is
+    // STT_NOTYPE. In addition, this save us from checking untaggable symbols,
+    // like functions or TLS symbols.
+    if (sym.type != STT_OBJECT)
+      continue;
+    // STB_LOCAL symbols can't be referenced from outside the object file, and
+    // thus don't need to be checked for references from other object files.
+    if (sym.binding == STB_LOCAL) {
+      sym.setIsTagged(true);
+      continue;
+    }
+    ++referenceCount[&sym];
+  }
+  sec.markDead();
+}
+
+// A tagged symbol must be denoted as being tagged by all references and the
+// chosen definition. For simplicity, here, it must also be denoted as tagged
+// for all definitions. Otherwise:
+//
+//  1. A tagged definition can be used by an untagged declaration, in which case
+//     the untagged access may be PC-relative, causing a tag mismatch at
+//     runtime.
+//  2. An untagged definition can be used by a tagged declaration, where the
+//     compiler has taken advantage of the increased alignment of the tagged
+//     declaration, but the alignment at runtime is wrong, causing a fault.
+//
+// Ideally, this isn't a problem, as any TU that imports or exports tagged
+// symbols should also be built with tagging. But, to handle these cases, we
+// demote the symbol to be untagged.
+void lld::elf::createTaggedSymbols(const SmallVector<ELFFileBase *, 0> &files) {
+  assert(config->emachine == EM_AARCH64 &&
+         config->androidMemtagMode != ELF::NT_MEMTAG_LEVEL_NONE);
+
+  // First, collect all symbols that are marked as tagged, and count how many
+  // times they're marked as tagged.
+  DenseMap<Symbol *, unsigned> taggedSymbolReferenceCount;
+  for (InputFile* file : files) {
+    if (file->kind() != InputFile::ObjKind)
+      continue;
+    for (InputSectionBase *section : file->getSections()) {
+      if (!section || section->type != SHT_AARCH64_MEMTAG_GLOBALS_STATIC ||
+          section == &InputSection::discarded)
+        continue;
+      invokeELFT(addTaggedSymbolReferences, *section,
+                 taggedSymbolReferenceCount);
+    }
+  }
+
+  // Now, go through all the symbols. If the number of declarations +
+  // definitions to a symbol exceeds the amount of times they're marked as
+  // tagged, it means we have an objfile that uses the untagged variant of the
+  // symbol.
+  for (InputFile *file : files) {
+    if (file->kind() != InputFile::BinaryKind &&
+        file->kind() != InputFile::ObjKind)
+      continue;
+
+    for (Symbol *symbol : file->getSymbols()) {
+      // See `addTaggedSymbolReferences` for more details.
+      if (symbol->type != STT_OBJECT ||
+          symbol->binding == STB_LOCAL)
+        continue;
+      auto it = taggedSymbolReferenceCount.find(symbol);
+      if (it == taggedSymbolReferenceCount.end()) continue;
+      unsigned &remainingAllowedTaggedRefs = it->second;
+      if (remainingAllowedTaggedRefs == 0) {
+        taggedSymbolReferenceCount.erase(it);
+        continue;
+      }
+      --remainingAllowedTaggedRefs;
+    }
+  }
+
+  // `addTaggedSymbolReferences` has already checked that we have RELA
+  // relocations, the only other way to get written addends is with
+  // --apply-dynamic-relocs.
+  if (!taggedSymbolReferenceCount.empty() && config->writeAddends)
+    error("--apply-dynamic-relocs cannot be used with MTE globals");
+
+  // Now, `taggedSymbolReferenceCount` should only contain symbols that are
+  // defined as tagged exactly the same amount as it's referenced, meaning all
+  // uses are tagged.
+  for (auto &[symbol, remainingTaggedRefs] : taggedSymbolReferenceCount) {
+    assert(remainingTaggedRefs == 0 &&
+            "Symbol is defined as tagged more times than it's used");
+    symbol->setIsTagged(true);
+  }
+}

lld/ELF/Arch/PPC64.cpp

@@ -1556,7 +1556,7 @@ void PPC64::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
         break;
 
       // Patch a nop (0x60000000) to a ld.
-      if (rel.sym->needsTocRestore) {
+      if (rel.sym->needsTocRestore()) {
         // gcc/gfortran 5.4, 6.3 and earlier versions do not add nop for
         // recursive calls even if the function is preemptible. This is not
         // wrong in the common case where the function is not preempted at

lld/ELF/Driver.cpp

@@ -786,13 +786,6 @@ static int getMemtagMode(opt::InputArgList &args) {
     return ELF::NT_MEMTAG_LEVEL_NONE;
   }
 
-  if (!config->androidMemtagHeap && !config->androidMemtagStack) {
-    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")
     return ELF::NT_MEMTAG_LEVEL_SYNC;
   if (memtagModeArg == "async")
@@ -2940,6 +2933,12 @@ void LinkerDriver::link(opt::InputArgList &args) {
   // partition.
   copySectionsIntoPartitions();
 
+  if (config->emachine == EM_AARCH64 &&
+      config->androidMemtagMode != ELF::NT_MEMTAG_LEVEL_NONE) {
+    llvm::TimeTraceScope timeScope("Process memory tagged symbols");
+    createTaggedSymbols(ctx.objectFiles);
+  }
+
   // Create synthesized sections such as .got and .plt. This is called before
   // processSectionCommands() so that they can be placed by SECTIONS commands.
   invokeELFT(createSyntheticSections,);

lld/ELF/Relocations.cpp

@@ -53,6 +53,7 @@
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/BinaryFormat/ELF.h"
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/Support/Endian.h"
 #include <algorithm>
@@ -199,10 +200,7 @@ static bool needsPlt(RelExpr expr) {
                R_PPC32_PLTREL, R_PPC64_CALL_PLT>(expr);
 }
 
-// Returns true if Expr refers a GOT entry. Note that this function
-// returns false for TLS variables even though they need GOT, because
-// TLS variables uses GOT differently than the regular variables.
-static bool needsGot(RelExpr expr) {
+bool lld::elf::needsGot(RelExpr expr) {
   return oneof<R_GOT, R_GOT_OFF, R_MIPS_GOT_LOCAL_PAGE, R_MIPS_GOT_OFF,
                R_MIPS_GOT_OFF32, R_AARCH64_GOT_PAGE_PC, R_GOT_PC, R_GOTPLT,
                R_AARCH64_GOT_PAGE, R_LOONGARCH_GOT, R_LOONGARCH_GOT_PAGE_PC>(
@@ -859,6 +857,23 @@ static void addRelativeReloc(InputSectionBase &isec, uint64_t offsetInSec,
                              RelType type) {
   Partition &part = isec.getPartition();
 
+  if (sym.isTagged()) {
+    std::lock_guard<std::mutex> lock(relocMutex);
+    part.relaDyn->addRelativeReloc(target->relativeRel, isec, offsetInSec, sym,
+                                   addend, type, expr);
+    // With MTE globals, we always want to derive the address tag by `ldg`-ing
+    // the symbol. When we have a RELATIVE relocation though, we no longer have
+    // a reference to the symbol. Because of this, when we have an addend that
+    // puts the result of the RELATIVE relocation out-of-bounds of the symbol
+    // (e.g. the addend is outside of [0, sym.getSize()]), the AArch64 MemtagABI
+    // says we should store the offset to the start of the symbol in the target
+    // field. This is described in further detail in:
+    // https://github.com/ARM-software/abi-aa/blob/main/memtagabielf64/memtagabielf64.rst#841extended-semantics-of-r_aarch64_relative
+    if (addend < 0 || static_cast<uint64_t>(addend) >= sym.getSize())
+      isec.relocations.push_back({expr, type, offsetInSec, addend, &sym});
+    return;
+  }
+
   // Add a relative relocation. If relrDyn section is enabled, and the
   // relocation offset is guaranteed to be even, add the relocation to
   // the relrDyn section, otherwise add it to the relaDyn section.
@@ -1645,6 +1660,10 @@ void elf::postScanRelocations() {
     auto flags = sym.flags.load(std::memory_order_relaxed);
     if (handleNonPreemptibleIfunc(sym, flags))
       return;
+
+    if (sym.isTagged() && sym.isDefined())
+      mainPart->memtagDescriptors->addSymbol(sym);
+
     if (!sym.needsDynReloc())
       return;
     sym.allocateAux();

lld/ELF/Relocations.h

@@ -220,6 +220,11 @@ ArrayRef<RelTy> sortRels(ArrayRef<RelTy> rels, SmallVector<RelTy, 0> &storage) {
   }
   return rels;
 }
+
+// Returns true if Expr refers a GOT entry. Note that this function returns
+// false for TLS variables even though they need GOT, because TLS variables uses
+// GOT differently than the regular variables.
+bool needsGot(RelExpr expr);
 } // namespace lld::elf
 
 #endif

lld/ELF/Symbols.h

@@ -254,8 +254,8 @@ class Symbol {
   Symbol(Kind k, InputFile *file, StringRef name, uint8_t binding,
          uint8_t stOther, uint8_t type)
       : file(file), nameData(name.data()), nameSize(name.size()), type(type),
-        binding(binding), stOther(stOther), symbolKind(k),
-        exportDynamic(false) {}
+        binding(binding), stOther(stOther), symbolKind(k), exportDynamic(false),
+        archSpecificBit(false) {}
 
   void overwrite(Symbol &sym, Kind k) const {
     if (sym.traced)
@@ -279,9 +279,18 @@ class Symbol {
   // True if defined relative to a section discarded by ICF.
   uint8_t folded : 1;
 
-  // True if a call to this symbol needs to be followed by a restore of the
-  // PPC64 toc pointer.
-  uint8_t needsTocRestore : 1;
+  // Allow reuse of a bit between architecture-exclusive symbol flags.
+  // - needsTocRestore(): On PPC64, true if a call to this symbol needs to be
+  //   followed by a restore of the toc pointer.
+  // - isTagged(): On AArch64, true if the symbol needs special relocation and
+  //   metadata semantics because it's tagged, under the AArch64 MemtagABI.
+  uint8_t archSpecificBit : 1;
+  bool needsTocRestore() const { return archSpecificBit; }
+  bool isTagged() const { return archSpecificBit; }
+  void setNeedsTocRestore(bool v) { archSpecificBit = v; }
+  void setIsTagged(bool v) {
+    archSpecificBit = v;
+  }
 
   // True if this symbol is defined by a symbol assignment or wrapped by --wrap.
   //

lld/ELF/SyntheticSections.cpp

@@ -1453,6 +1453,10 @@ DynamicSection<ELFT>::computeContents() {
       addInt(DT_AARCH64_MEMTAG_MODE, config->androidMemtagMode == NT_MEMTAG_LEVEL_ASYNC);
       addInt(DT_AARCH64_MEMTAG_HEAP, config->androidMemtagHeap);
       addInt(DT_AARCH64_MEMTAG_STACK, config->androidMemtagStack);
+      if (mainPart->memtagDescriptors->isNeeded()) {
+        addInSec(DT_AARCH64_MEMTAG_GLOBALS, *mainPart->memtagDescriptors);
+        addInt(DT_AARCH64_MEMTAG_GLOBALSSZ, mainPart->memtagDescriptors->getSize());
+      }
     }
   }
 
@@ -3900,6 +3904,76 @@ size_t PackageMetadataNote::getSize() const {
          alignTo(config->packageMetadata.size() + 1, 4);
 }
 
+// Helper function, return the size of the ULEB128 for 'v', optionally writing
+// it to `*(buf + offset)` if `buf` is non-null.
+static size_t computeOrWriteULEB128(uint64_t v, uint8_t *buf, size_t offset) {
+  if (buf)
+    return encodeULEB128(v, buf + offset);
+  return getULEB128Size(v);
+}
+
+// https://github.com/ARM-software/abi-aa/blob/main/memtagabielf64/memtagabielf64.rst#83encoding-of-sht_aarch64_memtag_globals_dynamic
+constexpr uint64_t kMemtagStepSizeBits = 3;
+constexpr uint64_t kMemtagGranuleSize = 16;
+static size_t createMemtagDescriptors(const SmallVector<const Symbol *, 0> &symbols,
+                                      uint8_t *buf = nullptr) {
+  size_t sectionSize = 0;
+  uint64_t lastGlobalEnd = 0;
+
+  for (const Symbol *sym : symbols) {
+    if (!includeInSymtab(*sym))
+      continue;
+    const uint64_t addr = sym->getVA();
+    const uint64_t size = sym->getSize();
+
+    if (addr <= kMemtagGranuleSize && buf != nullptr)
+      errorOrWarn("address of the tagged symbol \"" + sym->getName() +
+                  "\" falls in the ELF header. This is indicative of a "
+                  "compiler/linker bug");
+    if (addr % kMemtagGranuleSize != 0)
+      errorOrWarn("address of the tagged symbol \"" + sym->getName() +
+                  "\" at 0x" + Twine::utohexstr(addr) +
+                  "\" is not granule (16-byte) aligned");
+    if (size == 0)
+      errorOrWarn("size of the tagged symbol \"" + sym->getName() +
+                  "\" is not allowed to be zero");
+    if (size % kMemtagGranuleSize != 0)
+      errorOrWarn("size of the tagged symbol \"" + sym->getName() +
+                  "\" (size 0x" + Twine::utohexstr(size) +
+                  ") is not granule (16-byte) aligned");
+
+    const uint64_t sizeToEncode = size / kMemtagGranuleSize;
+    const uint64_t stepToEncode = ((addr - lastGlobalEnd) / kMemtagGranuleSize)
+                                  << kMemtagStepSizeBits;
+    if (sizeToEncode < (1 << kMemtagStepSizeBits)) {
+      sectionSize += computeOrWriteULEB128(stepToEncode | sizeToEncode, buf, sectionSize);
+    } else {
+      sectionSize += computeOrWriteULEB128(stepToEncode, buf, sectionSize);
+      sectionSize += computeOrWriteULEB128(sizeToEncode - 1, buf, sectionSize);
+    }
+    lastGlobalEnd = addr + size;
+  }
+
+  return sectionSize;
+}
+
+bool MemtagDescriptors::updateAllocSize() {
+  size_t oldSize = getSize();
+  std::stable_sort(symbols.begin(), symbols.end(),
+                   [](const Symbol *s1, const Symbol *s2) {
+                     return s1->getVA() < s2->getVA();
+                   });
+  return oldSize != getSize();
+}
+
+void MemtagDescriptors::writeTo(uint8_t *buf) {
+  createMemtagDescriptors(symbols, buf);
+}
+
+size_t MemtagDescriptors::getSize() const {
+  return createMemtagDescriptors(symbols);
+}
+
 InStruct elf::in;
 
 std::vector<Partition> elf::partitions;