[llvm-xray] Support for PIE

[phosek]

Summary

When the instrumented binary is linked as PIE, we need to apply the
relative relocations to sleds. This is handled by the dynamic linker
at runtime, but when processing the file we have to do it ourselves.

[dberris]

LGTM -- Thanks @phosek!

[MaskRay]

llvm/lib/XRay/InstrumentationMap.cpp

90 ↗

(On Diff #177648)

Section.relocations() applies on either SHT_REL or SHT_RELA. I think the code can be applied on just SHT_RELA.

if (Section.isBSS() || Section.isVirtual()) is not necessary.

[MaskRay]

llvm/lib/XRay/InstrumentationMap.cpp

97 ↗

(On Diff #177648)

The code only works on x86_64 and ppc64le. So the relative relocations can just be enumerated here.

[joerg]

llvm/lib/XRay/InstrumentationMap.cpp

97 ↗

(On Diff #177648)

Please explicitly check for the relocation types you care about. They are not the only relocations you can find in a PIE binary, even with static linking.

[dberris]

LGTM still, thanks @phosek!

[MaskRay]

llvm/lib/XRay/InstrumentationMap.cpp

105 ↗

(On Diff #178162)

It can be simplified with
if (auto AddendOrErr = object::ELFRelocationRef(Reloc).getAddend())

llvm/trunk/lib/XRay/InstrumentationMap.cpp

 //===- InstrumentationMap.cpp - XRay Instrumentation Map ------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Implementation of the InstrumentationMap type for XRay sleds.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/XRay/InstrumentationMap.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/DataExtractor.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>

 
 Optional<uint64_t> InstrumentationMap::getFunctionAddr(int32_t FuncId) const {
   auto I = FunctionAddresses.find(FuncId);
   if (I != FunctionAddresses.end())
     return I->second;
   return None;
 }
 
+using RelocMap = DenseMap<uint64_t, uint64_t>;
+
 static Error
 loadObj(StringRef Filename, object::OwningBinary<object::ObjectFile> &ObjFile,
           InstrumentationMap::SledContainer &Sleds,
           InstrumentationMap::FunctionAddressMap &FunctionAddresses,
           InstrumentationMap::FunctionAddressReverseMap &FunctionIds) {
   InstrumentationMap Map;
 
   // Find the section named "xray_instr_map".
if (I == Sections.end())
     return make_error<StringError>(
         "Failed to find XRay instrumentation map.",
         std::make_error_code(std::errc::executable_format_error));
 
   if (I->getContents(Contents))
     return errorCodeToError(
         std::make_error_code(std::errc::executable_format_error));
 
+  RelocMap Relocs;
+  if (ObjFile.getBinary()->isELF()) {
+    uint32_t RelrRelocationType = [](object::ObjectFile *ObjFile) {
+      if (const auto *ELFObj = dyn_cast<object::ELF32LEObjectFile>(ObjFile))
+        return ELFObj->getELFFile()->getRelrRelocationType();
+      else if (const auto *ELFObj = dyn_cast<object::ELF32BEObjectFile>(ObjFile))
+        return ELFObj->getELFFile()->getRelrRelocationType();
+      else if (const auto *ELFObj = dyn_cast<object::ELF64LEObjectFile>(ObjFile))
+        return ELFObj->getELFFile()->getRelrRelocationType();
+      else if (const auto *ELFObj = dyn_cast<object::ELF64BEObjectFile>(ObjFile))
+        return ELFObj->getELFFile()->getRelrRelocationType();
+      else
+        return static_cast<uint32_t>(0);
+    }(ObjFile.getBinary());
+
+    for (const object::SectionRef &Section : Sections) {
+      for (const object::RelocationRef &Reloc : Section.relocations()) {
+        if (Reloc.getType() != RelrRelocationType)
+          continue;
+        if (auto AddendOrErr = object::ELFRelocationRef(Reloc).getAddend())
+          Relocs.insert({Reloc.getOffset(), *AddendOrErr});
+      }
+    }
+  }
+
   // Copy the instrumentation map data into the Sleds data structure.
   auto C = Contents.bytes_begin();
   static constexpr size_t ELF64SledEntrySize = 32;
 
   if ((C - Contents.bytes_end()) % ELF64SledEntrySize != 0)
     return make_error<StringError>(
         Twine("Instrumentation map entries not evenly divisible by size of "
               "an XRay sled entry in ELF64."),
         std::make_error_code(std::errc::executable_format_error));
 
+  auto RelocateOrElse = [&](uint32_t Offset, uint64_t Address) {
+    if (!Address) {
+      uint64_t A = I->getAddress() + C - Contents.bytes_begin() + Offset;
+      RelocMap::const_iterator R = Relocs.find(A);
+      if (R != Relocs.end())
+        return R->second;
+    }
+    return Address;
+  };
+
   int32_t FuncId = 1;
   uint64_t CurFn = 0;
   for (; C != Contents.bytes_end(); C += ELF64SledEntrySize) {
     DataExtractor Extractor(
         StringRef(reinterpret_cast<const char *>(C), ELF64SledEntrySize), true,
         8);
     Sleds.push_back({});
     auto &Entry = Sleds.back();
     uint32_t OffsetPtr = 0;
-    Entry.Address = Extractor.getU64(&OffsetPtr);
-    Entry.Function = Extractor.getU64(&OffsetPtr);
+    Entry.Address = RelocateOrElse(OffsetPtr, Extractor.getU64(&OffsetPtr));
+    Entry.Function = RelocateOrElse(OffsetPtr, Extractor.getU64(&OffsetPtr));
     auto Kind = Extractor.getU8(&OffsetPtr);
     static constexpr SledEntry::FunctionKinds Kinds[] = {
         SledEntry::FunctionKinds::ENTRY, SledEntry::FunctionKinds::EXIT,
         SledEntry::FunctionKinds::TAIL,
         SledEntry::FunctionKinds::LOG_ARGS_ENTER,
         SledEntry::FunctionKinds::CUSTOM_EVENT};
     if (Kind >= sizeof(Kinds))
       return errorCodeToError(

llvm/trunk/test/tools/llvm-xray/X86/extract-instrmap-pie.ll

+; This test makes sure we can extract the instrumentation map from an
+; XRay-instrumented PIE file.
+;
+; RUN: llvm-xray extract %S/Inputs/elf64-pie.bin -s | FileCheck %s
+
+; CHECK:      ---
+; CHECK-NEXT: - { id: 1, address: 0x00000000000299C0, function: 0x00000000000299C0, kind: function-enter, always-instrument: true, function-name: {{.*foo.*}} }
+; CHECK-NEXT: - { id: 1, address: 0x00000000000299D0, function: 0x00000000000299C0, kind: function-exit, always-instrument: true, function-name: {{.*foo.*}} }
+; CHECK-NEXT: - { id: 2, address: 0x00000000000299E0, function: 0x00000000000299E0, kind: function-enter, always-instrument: true, function-name: {{.*bar.*}} }
+; CHECK-NEXT: - { id: 2, address: 0x00000000000299F6, function: 0x00000000000299E0, kind: function-exit, always-instrument: true, function-name: {{.*bar.*}} }
+; CHECK-NEXT: ...