[BOLT] Support instrumentation hook via DT_INIT_ARRAY

This commit follows implementation of instrumentation hook via
DT_FINI_ARRAY (#67348) and extends it for instrumentation
initialization hooking.

Initialization has has differences compared to finalization:
- Executables always use ELF entry point address. Update code
checks it and updates init_array entry if ELF is shared library
(have no interp entry) and have no DT_INIT entry. Also this
commit introduces "instrumentation-no-use-entry-point" option
to avoid that avoids usage of ELF entry point and force
usage of DT_INIT or DT_INIT_ARRAY, e.g. in case of libc
instrumentation.
- Shared library init_array entries relocations usually has
R_AARCH64_ABS64 type on AArch64 binaries. We check relocation
type and adjust methods for reading init_array relocations in
discovery and update methods.

Author: Vasily Leonenko

bolt/include/bolt/Core/BinaryContext.h

@@ -807,6 +807,15 @@ class BinaryContext {
   /// the execution of the binary is completed.
   std::optional<uint64_t> FiniFunctionAddress;
 
+  /// DT_INIT.
+  std::optional<uint64_t> InitAddress;
+
+  /// DT_INIT_ARRAY. Only used when DT_INIT is not set.
+  std::optional<uint64_t> InitArrayAddress;
+
+  /// DT_INIT_ARRAYSZ. Only used when DT_INIT is not set.
+  std::optional<uint64_t> InitArraySize;
+
   /// DT_FINI.
   std::optional<uint64_t> FiniAddress;
 

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -93,11 +93,20 @@ class RewriteInstance {
   /// section allocations if found.
   void discoverBOLTReserved();
 
+  /// Check whether we should use DT_INIT or DT_INIT_ARRAY for instrumentation.
+  /// DT_INIT is preferred; DT_INIT_ARRAY is only used when no DT_INIT entry was
+  /// found.
+  Error discoverRtInitAddress();
+
   /// Check whether we should use DT_FINI or DT_FINI_ARRAY for instrumentation.
   /// DT_FINI is preferred; DT_FINI_ARRAY is only used when no DT_FINI entry was
   /// found.
   Error discoverRtFiniAddress();
 
+  /// If DT_INIT_ARRAY is used for instrumentation, update the relocation of its
+  /// first entry to point to the instrumentation library's init address.
+  void updateRtInitReloc();
+
   /// If DT_FINI_ARRAY is used for instrumentation, update the relocation of its
   /// first entry to point to the instrumentation library's fini address.
   void updateRtFiniReloc();

bolt/lib/Passes/Instrumentation.cpp

@@ -34,6 +34,12 @@ cl::opt<std::string> InstrumentationFilename(
              "/tmp/prof.fdata)"),
     cl::init("/tmp/prof.fdata"), cl::Optional, cl::cat(BoltInstrCategory));
 
+cl::opt<bool> InstrumentNoUseEntryPoint(
+    "instrumentation-no-use-entry-point",
+    cl::desc("Do not use ELF entry point for instrumentation "
+             "initialization (e.g. for libc) (default: false)"),
+    cl::init(false), cl::Optional, cl::cat(BoltInstrCategory));
+
 cl::opt<std::string> InstrumentationBinpath(
     "instrumentation-binpath",
     cl::desc("path to instrumented binary in case if /proc/self/map_files "

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -80,6 +80,7 @@ namespace opts {
 extern cl::list<std::string> HotTextMoveSections;
 extern cl::opt<bool> Hugify;
 extern cl::opt<bool> Instrument;
+extern cl::opt<bool> InstrumentNoUseEntryPoint;
 extern cl::opt<bool> KeepNops;
 extern cl::opt<bool> Lite;
 extern cl::list<std::string> ReorderData;
@@ -737,9 +738,12 @@ Error RewriteInstance::run() {
   adjustCommandLineOptions();
   discoverFileObjects();
 
-  if (opts::Instrument && !BC->IsStaticExecutable)
+  if (opts::Instrument && !BC->IsStaticExecutable) {
+    if (Error E = discoverRtInitAddress())
+      return E;
     if (Error E = discoverRtFiniAddress())
       return E;
+  }
 
   preprocessProfileData();
 
@@ -781,8 +785,10 @@ Error RewriteInstance::run() {
 
   updateMetadata();
 
-  if (opts::Instrument && !BC->IsStaticExecutable)
+  if (opts::Instrument && !BC->IsStaticExecutable) {
+    updateRtInitReloc();
     updateRtFiniReloc();
+  }
 
   if (opts::OutputFilename == "/dev/null") {
     BC->outs() << "BOLT-INFO: skipping writing final binary to disk\n";
@@ -1407,6 +1413,55 @@ void RewriteInstance::discoverBOLTReserved() {
   NextAvailableAddress = BC->BOLTReserved.start();
 }
 
+Error RewriteInstance::discoverRtInitAddress() {
+  if (BC->HasInterpHeader && !opts::InstrumentNoUseEntryPoint)
+    return Error::success();
+
+  // Use DT_INIT if it's available.
+  if (BC->InitAddress) {
+    BC->StartFunctionAddress = BC->InitAddress;
+    return Error::success();
+  }
+
+  if (!BC->InitArrayAddress || !BC->InitArraySize) {
+    return createStringError(std::errc::not_supported,
+                             "Instrumentation of shared library needs either "
+                             "DT_INIT or DT_INIT_ARRAY");
+  }
+
+  if (*BC->InitArraySize < BC->AsmInfo->getCodePointerSize()) {
+    return createStringError(std::errc::not_supported,
+                             "Need at least 1 DT_FINI_ARRAY slot");
+  }
+
+  ErrorOr<BinarySection &> InitArraySection =
+      BC->getSectionForAddress(*BC->InitArrayAddress);
+  if (auto EC = InitArraySection.getError())
+    return errorCodeToError(EC);
+
+  if (const Relocation *Reloc = InitArraySection->getDynamicRelocationAt(0)) {
+    if (Reloc->isRelative()) {
+      BC->StartFunctionAddress = Reloc->Addend;
+    } else {
+      MCSymbol *Sym = Reloc->Symbol;
+      assert(Sym && "Failed to locate symbol for 0 entry of .init_array");
+      const BinaryFunction *BF = BC->getFunctionForSymbol(Sym);
+      assert(BF &&
+             "Failed to locate binary function for 0 entry of .init_array");
+      BC->StartFunctionAddress = BF->getAddress() + Reloc->Addend;
+    }
+    return Error::success();
+  }
+
+  if (const Relocation *Reloc = InitArraySection->getRelocationAt(0)) {
+    BC->StartFunctionAddress = Reloc->Value;
+    return Error::success();
+  }
+
+  return createStringError(std::errc::not_supported,
+                           "No relocation for first DT_INIT_ARRAY slot");
+}
+
 Error RewriteInstance::discoverRtFiniAddress() {
   // Use DT_FINI if it's available.
   if (BC->FiniAddress) {
@@ -1444,6 +1499,58 @@ Error RewriteInstance::discoverRtFiniAddress() {
                            "No relocation for first DT_FINI_ARRAY slot");
 }
 
+void RewriteInstance::updateRtInitReloc() {
+  if (BC->HasInterpHeader && !opts::InstrumentNoUseEntryPoint)
+    return;
+
+  // Updating DT_INIT is handled by patchELFDynamic.
+  if (BC->InitAddress)
+    return;
+
+  const RuntimeLibrary *RT = BC->getRuntimeLibrary();
+  if (!RT || !RT->getRuntimeStartAddress())
+    return;
+
+  if (!BC->InitArrayAddress)
+    return;
+
+  assert(BC->InitArrayAddress && BC->InitArraySize &&
+         "inconsistent .init_array state");
+
+  ErrorOr<BinarySection &> InitArraySection =
+      BC->getSectionForAddress(*BC->InitArrayAddress);
+  assert(InitArraySection && ".init_array removed");
+
+  if (std::optional<Relocation> Reloc =
+          InitArraySection->takeDynamicRelocationAt(0)) {
+    if (Reloc->isRelative()) {
+      assert(Reloc->Addend == BC->StartFunctionAddress &&
+             "inconsistent .init_array dynamic relocation");
+      Reloc->Addend = RT->getRuntimeStartAddress();
+      InitArraySection->addDynamicRelocation(*Reloc);
+    } else {
+      MCSymbol *Sym = Reloc->Symbol;
+      assert(Sym && "Failed to locate symbol for 0 entry of .init_array");
+      const BinaryFunction *BF = BC->getFunctionForSymbol(Sym);
+      assert(BF &&
+             "Failed to locate binary function for 0 entry of .init_array");
+      assert(BF->getAddress() + Reloc->Addend == BC->StartFunctionAddress &&
+             "inconsistent .init_array dynamic relocation");
+      InitArraySection->addDynamicRelocation(Relocation{
+          /*Offset*/ 0, /*Symbol*/ nullptr, /*Type*/ Relocation::getAbs64(),
+          /*Addend*/ RT->getRuntimeStartAddress(), /*Value*/ 0});
+    }
+  }
+  // Update the static relocation by adding a pending relocation which will get
+  // patched when flushPendingRelocations is called in rewriteFile. Note that
+  // flushPendingRelocations will calculate the value to patch as
+  // "Symbol + Addend". Since we don't have a symbol, just set the addend to the
+  // desired value.
+  InitArraySection->addPendingRelocation(Relocation{
+      /*Offset*/ 0, /*Symbol*/ nullptr, /*Type*/ Relocation::getAbs64(),
+      /*Addend*/ RT->getRuntimeStartAddress(), /*Value*/ 0});
+}
+
 void RewriteInstance::updateRtFiniReloc() {
   // Updating DT_FINI is handled by patchELFDynamic.
   if (BC->FiniAddress)
@@ -4838,7 +4945,8 @@ void RewriteInstance::patchELFSectionHeaderTable(ELFObjectFile<ELFT> *File) {
   ELFEhdrTy NewEhdr = Obj.getHeader();
 
   if (BC->HasRelocations) {
-    if (RuntimeLibrary *RtLibrary = BC->getRuntimeLibrary())
+    RuntimeLibrary *RtLibrary = BC->getRuntimeLibrary();
+    if (RtLibrary && !opts::InstrumentNoUseEntryPoint)
       NewEhdr.e_entry = RtLibrary->getRuntimeStartAddress();
     else
       NewEhdr.e_entry = getNewFunctionAddress(NewEhdr.e_entry);
@@ -5684,7 +5792,8 @@ void RewriteInstance::patchELFDynamic(ELFObjectFile<ELFT> *File) {
         if (uint64_t Addr = RtLibrary->getRuntimeFiniAddress())
           NewDE.d_un.d_ptr = Addr;
       }
-      if (RtLibrary && Dyn.getTag() == ELF::DT_INIT && !BC->HasInterpHeader) {
+      if (RtLibrary && Dyn.getTag() == ELF::DT_INIT &&
+          (!BC->HasInterpHeader || opts::InstrumentNoUseEntryPoint)) {
         if (auto Addr = RtLibrary->getRuntimeStartAddress()) {
           LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Set DT_INIT to 0x"
                             << Twine::utohexstr(Addr) << '\n');
@@ -5760,6 +5869,13 @@ Error RewriteInstance::readELFDynamic(ELFObjectFile<ELFT> *File) {
         LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Set start function address\n");
         BC->StartFunctionAddress = Dyn.getPtr();
       }
+      BC->InitAddress = Dyn.getPtr();
+      break;
+    case ELF::DT_INIT_ARRAY:
+      BC->InitArrayAddress = Dyn.getPtr();
+      break;
+    case ELF::DT_INIT_ARRAYSZ:
+      BC->InitArraySize = Dyn.getPtr();
       break;
     case ELF::DT_FINI:
       BC->FiniAddress = Dyn.getPtr();