[BOLT] Enable PLT analysis for aarch64

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

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

5 files changed, 248 insertions, 58 deletions

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 93970e02a4d1..c5376ca1d3fc 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1327,6 +1327,16 @@ public:
     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,
diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index 3b9df2c76ae2..63df47c7bf55 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -246,6 +246,19 @@ private:
   /// 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);        \
@@ -473,7 +486,7 @@ private:
   /// 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.
@@ -485,10 +498,7 @@ private:
   };
 
   /// 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.
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 7bc5711c970d..d25f1205afdb 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -1174,8 +1174,7 @@ void RewriteInstance::discoverFileObjects() {
   processDynamicRelocations();
 
   // Process PLT section.
-  if (BC->TheTriple->getArch() == Triple::x86_64)
-    disassemblePLT();
+  disassemblePLT();
 
   // See if we missed any functions marked by FDE.
   for (const auto &FDEI : CFIRdWrt->getFDEs()) {
@@ -1245,69 +1244,149 @@ void RewriteInstance::discoverFileObjects() {
   }
 }
 
-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();
-
-    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);
-        }
+void RewriteInstance::createPLTBinaryFunction(uint64_t TargetAddress,
+                                              uint64_t EntryAddress,
+                                              uint64_t EntrySize) {
+  if (!TargetAddress)
+    return;
 
-        // Check if the entry size needs adjustment.
-        if (EntryOffset == 0 && BC->MIB->isTerminateBranch(Instruction) &&
-            EntrySize == 8)
-          EntrySize = 16;
+  const Relocation *Rel = BC->getDynamicRelocationAt(TargetAddress);
+  if (!Rel || !Rel->Symbol)
+    return;
 
-        if (BC->MIB->isIndirectBranch(Instruction))
-          break;
+  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();
+  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;
+    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;
+    MCInst Instruction;
+    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;
       }
 
-      if (InstrOffset + InstrSize > EntryOffset + EntrySize)
-        continue;
+      const uint64_t EntryAddress = SectionAddress + EntryOffset;
+      const uint64_t TargetAddress = BC->MIB->analyzePLTEntry(
+          Instruction, Instructions.begin(), Instructions.end(), EntryAddress);
 
-      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);
-      }
+      createPLTBinaryFunction(TargetAddress, EntryAddress, EntrySize);
+      break;
+    }
 
-      const Relocation *Rel = BC->getDynamicRelocationAt(TargetAddress);
-      if (!Rel || !Rel->Symbol)
-        continue;
+    // Branch instruction
+    InstrOffset += InstrSize;
 
-      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);
+    // 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;
+    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);
     }
   };
 
+  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);
+      // Check if the entry size needs adjustment.
+      if (EntryOffset == 0 && BC->MIB->isTerminateBranch(Instruction) &&
+          EntrySize == 8)
+        EntrySize = 16;
+
+      if (BC->MIB->isIndirectBranch(Instruction))
+        break;
+
+      InstrOffset += InstrSize;
+    }
+
+    if (InstrOffset + InstrSize > EntryOffset + EntrySize)
+      continue;
+
+    uint64_t TargetAddress;
+    if (!BC->MIB->evaluateMemOperandTarget(Instruction, TargetAddress,
+                                           SectionAddress + InstrOffset,
+                                           InstrSize)) {
+      errs() << "BOLT-ERROR: error evaluating PLT instruction at offset 0x"
+             << Twine::utohexstr(SectionAddress + InstrOffset) << '\n';
+      exit(1);
+    }
+
+    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)
@@ -1786,6 +1865,11 @@ bool RewriteInstance::analyzeRelocation(
     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
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 018bf09adafc..41c19c113f0c 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -709,6 +709,67 @@ public:
     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:
diff --git a/bolt/test/runtime/AArch64/plt.c b/bolt/test/runtime/AArch64/plt.c
new file mode 100644
index 000000000000..296599fff6d7
--- /dev/null
+++ b/bolt/test/runtime/AArch64/plt.c
@@ -0,0 +1,25 @@
+// 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;
+}