Revert "[MC] Move ELFWriter::createMemtagRelocs to AArch64ELFStreamer::finishImpl"

This reverts commit 9d63506.

There is a heap-use-after-free.

Author: MaskRay

llvm/lib/MC/ELFObjectWriter.cpp

@@ -193,6 +193,8 @@ struct ELFWriter {
   MCSectionELF *createRelocationSection(MCContext &Ctx,
                                         const MCSectionELF &Sec);
 
+  void createMemtagRelocs(MCAssembler &Asm);
+
   void writeSectionHeader(const MCAsmLayout &Layout,
                           const SectionIndexMapTy &SectionIndexMap,
                           const SectionOffsetsTy &SectionOffsets);
@@ -614,6 +616,23 @@ bool ELFWriter::isInSymtab(const MCAsmLayout &Layout, const MCSymbolELF &Symbol,
   return true;
 }
 
+void ELFWriter::createMemtagRelocs(MCAssembler &Asm) {
+  MCSectionELF *MemtagRelocs = nullptr;
+  for (const MCSymbol &Sym : Asm.symbols()) {
+    const auto &SymE = cast<MCSymbolELF>(Sym);
+    if (!SymE.isMemtag())
+      continue;
+    if (MemtagRelocs == nullptr) {
+      MemtagRelocs = OWriter.TargetObjectWriter->getMemtagRelocsSection(Asm.getContext());
+      if (MemtagRelocs == nullptr)
+        report_fatal_error("Tagged globals are not available on this architecture.");
+      Asm.registerSection(*MemtagRelocs);
+    }
+    ELFRelocationEntry Rec(0, &SymE, ELF::R_AARCH64_NONE, 0, nullptr, 0);
+    OWriter.Relocations[MemtagRelocs].push_back(Rec);
+  }
+}
+
 void ELFWriter::computeSymbolTable(
     MCAssembler &Asm, const MCAsmLayout &Layout,
     const SectionIndexMapTy &SectionIndexMap, const RevGroupMapTy &RevGroupMap,
@@ -1075,6 +1094,8 @@ uint64_t ELFWriter::writeObject(MCAssembler &Asm, const MCAsmLayout &Layout) {
       Ctx.getELFSection(".strtab", ELF::SHT_STRTAB, 0);
   StringTableIndex = addToSectionTable(StrtabSection);
 
+  createMemtagRelocs(Asm);
+
   RevGroupMapTy RevGroupMap;
   SectionIndexMapTy SectionIndexMap;
 

llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp

@@ -28,7 +28,7 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCObjectWriter.h"
-#include "llvm/MC/MCSectionELF.h"
+#include "llvm/MC/MCSection.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCSymbolELF.h"
@@ -183,7 +183,7 @@ class AArch64ELFStreamer : public MCELFStreamer {
                       std::move(Emitter)),
         MappingSymbolCounter(0), LastEMS(EMS_None) {}
 
-  void changeSection(MCSection *Section, uint32_t Subsection = 0) override {
+  void changeSection(MCSection *Section, uint32_t Subsection) override {
     // We have to keep track of the mapping symbol state of any sections we
     // use. Each one should start off as EMS_None, which is provided as the
     // default constructor by DenseMap::lookup.
@@ -248,9 +248,6 @@ class AArch64ELFStreamer : public MCELFStreamer {
     emitDataMappingSymbol();
     MCObjectStreamer::emitFill(NumBytes, FillValue, Loc);
   }
-
-  void finishImpl() override;
-
 private:
   enum ElfMappingSymbol {
     EMS_None,
@@ -287,27 +284,6 @@ class AArch64ELFStreamer : public MCELFStreamer {
   ElfMappingSymbol LastEMS;
 };
 
-void AArch64ELFStreamer::finishImpl() {
-  MCContext &Ctx = getContext();
-  auto &Asm = getAssembler();
-  MCSectionELF *MemtagSec = nullptr;
-  const auto *Zero = MCConstantExpr::create(0, Ctx);
-  for (const MCSymbol &Symbol : Asm.symbols()) {
-    const auto &Sym = cast<MCSymbolELF>(Symbol);
-    if (!Sym.isMemtag())
-      continue;
-    if (!MemtagSec) {
-      MemtagSec = Ctx.getELFSection(".memtag.globals.static",
-                                    ELF::SHT_AARCH64_MEMTAG_GLOBALS_STATIC, 0);
-      switchSection(MemtagSec);
-    }
-    auto *SRE = MCSymbolRefExpr::create(&Sym, MCSymbolRefExpr::VK_None, Ctx);
-    (void)MCObjectStreamer::emitRelocDirective(
-        *Zero, "BFD_RELOC_NONE", SRE, SMLoc(), *Ctx.getSubtargetInfo());
-  }
-  MCELFStreamer::finishImpl();
-}
-
 } // end anonymous namespace
 
 AArch64ELFStreamer &AArch64TargetELFStreamer::getStreamer() {