[BOLT] Enable PLT analysis for aarch64

[yota9]

Summary

This patch enables PLT analysis for aarch64. It is used by the static
relocations in order to provide final symbol address of PLT entry for some
instructions like ADRP.

Vladislav Khmelevsky,
Advanced Software Technology Lab, Huawei

[yota9]

Remove debug output + clang format

[yota9]

Code-refactor. Hello @maksfb! I will extremely need this patch in the nearest future, I hope we could submit it soon :) Thank you!

[yota9]

nit

[yota9]

Fix x86 disassembly

[yota9]

Gentle ping to reviewers :)

[yota9]

minor comment fixes. And taking the opportunity gentle ping :)

[rafauler]

bolt/include/bolt/Core/MCPlusBuilder.h

1332

bolt/include/bolt/Rewrite/RewriteInstance.h

252

256

260

bolt/lib/Rewrite/RewriteInstance.cpp

1288

EntryOffset is not a loop-carried dependency, better to declare it inside the loop.

1288

1293

1294

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

731

746

747

I guess you meant

[rafauler]

nit: Title of the diff should also be changed to "[BOLT] Enable PLT analysis for aarch64"

[yota9]

Address comments

[yota9]

Hello @rafaelauler ! Thank you for your comments! Sorry for the long reply, something insane is going on I just don't have enough time and power these days :(

bolt/lib/Rewrite/RewriteInstance.cpp

1288

Agree, thanks

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

747

Yes, strange that I've skipped this, thanks!

[rafauler]

Sure, Vladislav! Take all the time you need, I hope things improve soon. Stay safe!

Thanks for working on this, changes look good.

[rafauler]

Before committing this, could you please clang-format line 1378? The linter is complaining. Please update this diff with the final version.

[yota9]

@rafaelauler Sure, thanks! Also thank you for you kind words :)

bolt/include/bolt/Core/MCPlusBuilder.h

analyzeIndirectBranch(MCInst &Instruction, InstructionIterator Begin,
                         InstructionIterator End, const unsigned PtrSize,
                         MCInst *&MemLocInstr, unsigned &BaseRegNum,
                         unsigned &IndexRegNum, int64_t &DispValue,
                         const MCExpr *&DispExpr, MCInst *&PCRelBaseOut) const {
     llvm_unreachable("not implemented");
     return IndirectBranchType::UNKNOWN;
   }
 
+  /// Analyze branch \p Instruction in PLT section and try to determine
+  /// associated got entry address.
+  virtual uint64_t analyzePLTEntry(MCInst &Instruction,
+                                   InstructionIterator Begin,
+                                   InstructionIterator End,
+                                   uint64_t BeginPC) const {
+    llvm_unreachable("not implemented");
+    return 0;
+  }
+
   virtual bool analyzeVirtualMethodCall(InstructionIterator Begin,
                                         InstructionIterator End,
                                         std::vector<MCInst *> &MethodFetchInsns,
                                         unsigned &VtableRegNum,
                                         unsigned &BaseRegNum,
                                         uint64_t &MethodOffset) const {
     llvm_unreachable("not implemented");
     return false;

bolt/include/bolt/Rewrite/RewriteInstance.h

private:
   void rewriteNoteSections();
 
   /// Write .eh_frame_hdr.
   void writeEHFrameHeader();
 
   /// Disassemble and create function entries for PLT.
   void disassemblePLT();
 
+  /// Auxiliary function to create .plt BinaryFunction on \p EntryAddres
+  /// with the \p EntrySize size. \p TargetAddress is the .got entry
+  /// associated address.
+  void createPLTBinaryFunction(uint64_t TargetAddress, uint64_t EntryAddress,
+                               uint64_t EntrySize);
+
+  /// Disassemble aarch64-specific .plt \p Section auxiliary function
+  void disassemblePLTSectionAArch64(BinarySection &Section);
+
+  /// Disassemble X86-specific .plt \p Section auxiliary function. \p EntrySize
+  /// is the expected .plt \p Section entry function size.
+  void disassemblePLTSectionX86(BinarySection &Section, uint64_t EntrySize);
+
   /// ELF-specific part. TODO: refactor into new class.
 #define ELF_FUNCTION(FUNC)                                                     \
   template <typename ELFT> void FUNC(object::ELFObjectFile<ELFT> *Obj);        \
   void FUNC() {                                                                \
     if (auto *ELF32LE = dyn_cast<object::ELF32LEObjectFile>(InputFile))        \
       return FUNC(ELF32LE);                                                    \
     if (auto *ELF64LE = dyn_cast<object::ELF64LEObjectFile>(InputFile))        \
       return FUNC(ELF64LE);                                                    \
private:
 
   /// Number of local symbols in newly written symbol table.
   uint64_t NumLocalSymbols{0};
 
   /// Information on special Procedure Linkage Table sections. There are
   /// multiple variants generated by different linkers.
   struct PLTSectionInfo {
     const char *Name;
-    uint64_t EntrySize;
+    uint64_t EntrySize{0};
   };
 
   /// Different types of X86-64 PLT sections.
   const PLTSectionInfo X86_64_PLTSections[4] = {
       { ".plt", 16 },
       { ".plt.got", 8 },
       { ".plt.sec", 8 },
       { nullptr, 0 }
   };
 
   /// AArch64 PLT sections.
-  const PLTSectionInfo AArch64_PLTSections[2] = {
-      { ".plt", 16 },
-      { nullptr, 0 }
-  };
+  const PLTSectionInfo AArch64_PLTSections[2] = {{".plt"}, {nullptr}};
 
   /// Return PLT information for a section with \p SectionName or nullptr
   /// if the section is not PLT.
   const PLTSectionInfo *getPLTSectionInfo(StringRef SectionName) {
     const PLTSectionInfo *PLTSI = nullptr;
     switch (BC->TheTriple->getArch()) {
     default:
       break;

bolt/lib/Rewrite/RewriteInstance.cpp

void RewriteInstance::discoverFileObjects() {
   }
 
   // Read dynamic relocation first as their presence affects the way we process
   // static relocations. E.g. we will ignore a static relocation at an address
   // that is a subject to dynamic relocation processing.
   processDynamicRelocations();
 
   // Process PLT section.
-  if (BC->TheTriple->getArch() == Triple::x86_64)
-    disassemblePLT();
 
   // See if we missed any functions marked by FDE.
   for (const auto &FDEI : CFIRdWrt->getFDEs()) {
     const uint64_t Address = FDEI.first;
     const dwarf::FDE *FDE = FDEI.second;
     const BinaryFunction *BF = BC->getBinaryFunctionAtAddress(Address);
     if (BF)
       continue;
if (opts::LinuxKernelMode) {
     // Read all special linux kernel sections and their relocations
     processLKSections();
   } else {
     // Read all relocations now that we have binary functions mapped.
     processRelocations();
   }
 }
 
-void RewriteInstance::disassemblePLT() {
-  auto analyzeOnePLTSection = [&](BinarySection &Section, uint64_t EntrySize) {
-    const uint64_t PLTAddress = Section.getAddress();
-    StringRef PLTContents = Section.getContents();
-    ArrayRef<uint8_t> PLTData(
-        reinterpret_cast<const uint8_t *>(PLTContents.data()),
-        Section.getSize());
-    const unsigned PtrSize = BC->AsmInfo->getCodePointerSize();
+void RewriteInstance::createPLTBinaryFunction(uint64_t TargetAddress,
+                                              uint64_t EntryAddress,
+                                              uint64_t EntrySize) {
+  if (!TargetAddress)
+    return;
+
+  const Relocation *Rel = BC->getDynamicRelocationAt(TargetAddress);
+  if (!Rel || !Rel->Symbol)
+    return;
+
+  const unsigned PtrSize = BC->AsmInfo->getCodePointerSize();
+  ErrorOr<BinarySection &> Section = BC->getSectionForAddress(EntryAddress);
+  assert(Section && "cannot get section for address");
+  BinaryFunction *BF = BC->createBinaryFunction(
+      Rel->Symbol->getName().str() + "@PLT", *Section, EntryAddress, 0,
+      EntrySize, Section->getAlignment());
+  MCSymbol *TargetSymbol = BC->registerNameAtAddress(
+      Rel->Symbol->getName().str() + "@GOT", TargetAddress, PtrSize, PtrSize);
+  BF->setPLTSymbol(TargetSymbol);
+}
+
+void RewriteInstance::disassemblePLTSectionAArch64(BinarySection &Section) {
+  const uint64_t SectionAddress = Section.getAddress();
+  const uint64_t SectionSize = Section.getSize();
+      reinterpret_cast<const uint8_t *>(PLTContents.data()), SectionSize);
 
-    for (uint64_t EntryOffset = 0; EntryOffset + EntrySize <= Section.getSize();
-         EntryOffset += EntrySize) {
-      uint64_t InstrOffset = EntryOffset;
-      uint64_t InstrSize;
-      MCInst Instruction;
-      while (InstrOffset < EntryOffset + EntrySize) {
-        uint64_t InstrAddr = PLTAddress + InstrOffset;
-        if (!BC->DisAsm->getInstruction(Instruction, InstrSize,
-                                        PLTData.slice(InstrOffset), InstrAddr,
-                                        nulls())) {
-          errs() << "BOLT-ERROR: unable to disassemble instruction in PLT "
-                    "section "
-                 << Section.getName() << " at offset 0x"
-                 << Twine::utohexstr(InstrOffset) << '\n';
-          exit(1);
-        }
+  auto disassembleInstruction = [&](uint64_t InstrOffset, MCInst &Instruction,
+                                    uint64_t &InstrSize) {
+    const uint64_t InstrAddr = SectionAddress + InstrOffset;
+    if (!BC->DisAsm->getInstruction(Instruction, InstrSize,
+                                    PLTData.slice(InstrOffset), InstrAddr,
+                                    nulls())) {
+      errs() << "BOLT-ERROR: unable to disassemble instruction in PLT section "
+             << Section.getName() << " at offset 0x"
+             << Twine::utohexstr(InstrOffset) << '\n';
+      exit(1);
+    }
+  };
+
+  uint64_t InstrOffset = 0;
+  // Locate new plt entry
+  while (InstrOffset < SectionSize) {
+    InstructionListType Instructions;
+    uint64_t EntryOffset = InstrOffset;
+    uint64_t EntrySize = 0;
+    uint64_t InstrSize;
+    // Loop through entry instructions
+    while (InstrOffset < SectionSize) {
+      disassembleInstruction(InstrOffset, Instruction, InstrSize);
+      EntrySize += InstrSize;
+      if (!BC->MIB->isIndirectBranch(Instruction)) {
+        Instructions.emplace_back(Instruction);
+        InstrOffset += InstrSize;
+        continue;
+      }
+
+      const uint64_t EntryAddress = SectionAddress + EntryOffset;
+      const uint64_t TargetAddress = BC->MIB->analyzePLTEntry(
+          Instruction, Instructions.begin(), Instructions.end(), EntryAddress);
+
+      createPLTBinaryFunction(TargetAddress, EntryAddress, EntrySize);
+      break;
+    }
+
+    // Branch instruction
+    InstrOffset += InstrSize;
+
+    // Skip nops if any
+    while (InstrOffset < SectionSize) {
+      disassembleInstruction(InstrOffset, Instruction, InstrSize);
+      if (!BC->MIB->isNoop(Instruction))
+        break;
+
+      InstrOffset += InstrSize;
+    }
+  }
+}
+
+void RewriteInstance::disassemblePLTSectionX86(BinarySection &Section,
+                                               uint64_t EntrySize) {
+  const uint64_t SectionAddress = Section.getAddress();
+  const uint64_t SectionSize = Section.getSize();
+  StringRef PLTContents = Section.getContents();
+  ArrayRef<uint8_t> PLTData(
+      reinterpret_cast<const uint8_t *>(PLTContents.data()), SectionSize);
+
+  auto disassembleInstruction = [&](uint64_t InstrOffset, MCInst &Instruction,
+                                    uint64_t &InstrSize) {
+    const uint64_t InstrAddr = SectionAddress + InstrOffset;
+      errs() << "BOLT-ERROR: unable to disassemble instruction in PLT section "
+  };
 
-        // Check if the entry size needs adjustment.
-        if (EntryOffset == 0 && BC->MIB->isTerminateBranch(Instruction) &&
-            EntrySize == 8)
-          EntrySize = 16;
+  for (uint64_t EntryOffset = 0; EntryOffset + EntrySize <= SectionSize;
+       EntryOffset += EntrySize) {
+    MCInst Instruction;
+    uint64_t InstrSize, InstrOffset = EntryOffset;
+    while (InstrOffset < EntryOffset + EntrySize) {
+      disassembleInstruction(InstrOffset, Instruction, InstrSize);
 
-        if (BC->MIB->isIndirectBranch(Instruction))
-          break;
 
-        InstrOffset += InstrSize;
-      }
 
-      if (InstrOffset + InstrSize > EntryOffset + EntrySize)
-        continue;
 
-      uint64_t TargetAddress;
-      if (!BC->MIB->evaluateMemOperandTarget(Instruction, TargetAddress,
-                                             PLTAddress + InstrOffset,
-                                             InstrSize)) {
-        errs() << "BOLT-ERROR: error evaluating PLT instruction at offset 0x"
-               << Twine::utohexstr(PLTAddress + InstrOffset) << '\n';
-        exit(1);
-      }
+                                           SectionAddress + InstrOffset,
+             << Twine::utohexstr(SectionAddress + InstrOffset) << '\n';
 
-      const Relocation *Rel = BC->getDynamicRelocationAt(TargetAddress);
-      if (!Rel || !Rel->Symbol)
-        continue;
-
-      BinaryFunction *BF = BC->createBinaryFunction(
-          Rel->Symbol->getName().str() + "@PLT", Section,
-          PLTAddress + EntryOffset, 0, EntrySize, Section.getAlignment());
-      MCSymbol *TargetSymbol =
-          BC->registerNameAtAddress(Rel->Symbol->getName().str() + "@GOT",
-                                    TargetAddress, PtrSize, PtrSize);
-      BF->setPLTSymbol(TargetSymbol);
-    }
+    createPLTBinaryFunction(TargetAddress, SectionAddress + EntryOffset,
+                            EntrySize);
+}
+
+void RewriteInstance::disassemblePLT() {
+  auto analyzeOnePLTSection = [&](BinarySection &Section, uint64_t EntrySize) {
+    if (BC->isAArch64())
+      return disassemblePLTSectionAArch64(Section);
+    return disassemblePLTSectionX86(Section, EntrySize);
   };
 
   for (BinarySection &Section : BC->allocatableSections()) {
     const PLTSectionInfo *PLTSI = getPLTSectionInfo(Section.getName());
     if (!PLTSI)
       continue;
 
     analyzeOnePLTSection(Section, PLTSI->EntrySize);
if (SymbolIter == InputFile->symbol_end()) {
     IsSectionRelocation = false;
   } else {
     const SymbolRef &Symbol = *SymbolIter;
     SymbolName = std::string(cantFail(Symbol.getName()));
     SymbolAddress = cantFail(Symbol.getAddress());
     SkipVerification = (cantFail(Symbol.getType()) == SymbolRef::ST_Other);
     // Section symbols are marked as ST_Debug.
     IsSectionRelocation = (cantFail(Symbol.getType()) == SymbolRef::ST_Debug);
+    // Check for PLT entry registered with symbol name
+    if (!SymbolAddress && IsAArch64) {
+      BinaryData *BD = BC->getBinaryDataByName(SymbolName + "@PLT");
+      SymbolAddress = BD ? BD->getAddress() : 0;
+    }
   }
   // For PIE or dynamic libs, the linker may choose not to put the relocation
   // result at the address if it is a X86_64_64 one because it will emit a
   // dynamic relocation (X86_RELATIVE) for the dynamic linker and loader to
   // resolve it at run time. The static relocation result goes as the addend
   // of the dynamic relocation in this case. We can't verify these cases.
   // FIXME: perhaps we can try to find if it really emitted a corresponding
   // RELATIVE relocation at this offset with the correct value as the addend.

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

IndirectBranchType analyzeIndirectBranch(
 
     MemLocInstrOut = MemLocInstr;
     DispValueOut = DispValue;
     DispExprOut = DispExpr;
     PCRelBaseOut = PCRelBase;
     return IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE;
   }
 
+  ///  Matches PLT entry pattern and returns the associated GOT entry address.
+  ///  Typical PLT entry looks like the following:
+  ///
+  ///    adrp    x16, 230000
+  ///    ldr     x17, [x16, #3040]
+  ///    add     x16, x16, #0xbe0
+  ///    br      x17
+  ///
+  uint64_t analyzePLTEntry(MCInst &Instruction, InstructionIterator Begin,
+                           InstructionIterator End,
+                           uint64_t BeginPC) const override {
+    // Check branch instruction
+    MCInst *Branch = &Instruction;
+    assert(Branch->getOpcode() == AArch64::BR && "Unexpected opcode");
+
+    DenseMap<const MCInst *, SmallVector<llvm::MCInst *, 4>> UDChain =
+        computeLocalUDChain(Branch, Begin, End);
+
+    // Match ldr instruction
+    SmallVector<MCInst *, 4> &BranchUses = UDChain[Branch];
+    if (BranchUses.size() < 1 || BranchUses[0] == nullptr)
+      return 0;
+
+    // Check ldr instruction
+    const MCInst *Ldr = BranchUses[0];
+    if (Ldr->getOpcode() != AArch64::LDRXui)
+      return 0;
+
+    // Get ldr value
+    const unsigned ScaleLdr = 8; // LDRX operates on 8 bytes segments
+    assert(Ldr->getOperand(2).isImm() && "Unexpected ldr operand");
+    const uint64_t Offset = Ldr->getOperand(2).getImm() * ScaleLdr;
+
+    // Match adrp instruction
+    SmallVector<MCInst *, 4> &LdrUses = UDChain[Ldr];
+    if (LdrUses.size() < 2 || LdrUses[1] == nullptr)
+      return 0;
+
+    // Check adrp instruction
+    MCInst *Adrp = LdrUses[1];
+    if (Adrp->getOpcode() != AArch64::ADRP)
+      return 0;
+
+    // Get adrp instruction PC
+    const unsigned InstSize = 4;
+    uint64_t AdrpPC = BeginPC;
+    for (InstructionIterator It = Begin; It != End; ++It) {
+      if (&(*It) == Adrp)
+        break;
+      AdrpPC += InstSize;
+    }
+
+    // Get adrp value
+    uint64_t Base;
+    assert(Adrp->getOperand(1).isImm() && "Unexpected adrp operand");
+    bool Ret = evaluateMemOperandTarget(*Adrp, Base, AdrpPC, InstSize);
+    assert(Ret && "Failed to evaluate adrp");
+
+    return Base + Offset;
+  }
+
   unsigned getInvertedBranchOpcode(unsigned Opcode) const {
     switch (Opcode) {
     default:
       llvm_unreachable("Failed to invert branch opcode");
       return Opcode;
     case AArch64::TBZW:     return AArch64::TBNZW;
     case AArch64::TBZX:     return AArch64::TBNZX;
     case AArch64::TBNZW:    return AArch64::TBZW;

bolt/test/runtime/AArch64/plt.c

+// This test checks that the pointers to PLT are properly updated.
+
+// RUN: %clang %cflags %s -fuse-ld=lld \
+// RUN:    -o %t.exe -Wl,-q
+// RUN: llvm-bolt %t.exe -o %t.bolt.exe -use-old-text=0 -lite=0
+// RUN: %t.bolt.exe
+
+#include <string.h>
+
+void *(*memcpy_p)(void *dest, const void *src, size_t n);
+void *(*memset_p)(void *dest, int c, size_t n);
+
+int main() {
+  int a = 0xdeadbeef, b = 0;
+
+  memcpy_p = memcpy;
+  memcpy_p(&b, &a, sizeof(b));
+  if (b != 0xdeadbeef)
+    return 1;
+
+  memset_p = memset;
+  memset_p(&a, 0, sizeof(a));
+  if (a != 0)
+    return 1;
+}