[Bolt][Instrumentation] Add support for DT_INIT_ARRAY

Previously Bolt relied on ELF 'e_entry' field or DT_INIT to determine
the entry point of an ELF file for the instrumentation.
This PR aims to handle that case if an ELF file only contains
DT_INIT_ARRAY/DT_FINI_ARRAY sections.

Bolt is hooking its runtime function based on e_entry address
if the input is an ELF executable. When the input is a shared object,
Bolt takes address of DT_INIT if that exists. If it doesn't,
Bolt will use DT_INIT_ARRAY for hooking its runtime functions.

This PR follows the implementation of DT_FINI_ARRAY.

Author: kaadam

bolt/include/bolt/Core/BinaryContext.h

@@ -800,6 +800,15 @@ class BinaryContext {
   /// the execution of the binary is completed.
   std::optional<uint64_t> FiniFunctionAddress;
 
+  /// DT_INIT. Used when DT_INIT is available.
+  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/Rewrite/RewriteInstance.cpp

@@ -708,9 +708,13 @@ Error RewriteInstance::run() {
   adjustCommandLineOptions();
   discoverFileObjects();
 
-  if (opts::Instrument && !BC->IsStaticExecutable)
+  if (opts::Instrument && !BC->IsStaticExecutable) {
+    if (!BC->HasInterpHeader)
+      if (Error E = discoverRtInitAddress())
+        return E;
     if (Error E = discoverRtFiniAddress())
       return E;
+  }
 
   preprocessProfileData();
 
@@ -752,8 +756,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";
@@ -1381,6 +1387,46 @@ void RewriteInstance::discoverBOLTReserved() {
   NextAvailableAddress = BC->BOLTReserved.start();
 }
 
+Error RewriteInstance::discoverRtInitAddress() {
+  // Use init address if it is available.
+  if (BC->InitAddress) {
+    BC->StartFunctionAddress = BC->InitAddress;
+    return Error::success();
+  }
+
+  if (BC->InitArrayAddress || BC->InitArraySize) {
+    if (*BC->InitArraySize < BC->AsmInfo->getCodePointerSize()) {
+      return createStringError(std::errc::not_supported,
+                               "Need at least 1 DT_INIT_ARRAY slot");
+    }
+
+    ErrorOr<BinarySection &> InitArraySection =
+        BC->getSectionForAddress(*BC->InitArrayAddress);
+    if (auto EC = InitArraySection.getError())
+      return errorCodeToError(EC);
+
+    if (const Relocation *Reloc = InitArraySection->getDynamicRelocationAt(0)) {
+      BC->StartFunctionAddress = 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");
+  }
+
+  if (BC->StartFunctionAddress && BC->StartFunctionAddress.value() != 0)
+    return Error::success();
+
+  return createStringError(
+      std::errc::not_supported,
+      "Instrumentation needs any of ELF e_entry, DT_INIT or DT_INIT_ARRAY");
+}
+
 Error RewriteInstance::discoverRtFiniAddress() {
   // Use DT_FINI if it's available.
   if (BC->FiniAddress) {
@@ -1452,6 +1498,40 @@ void RewriteInstance::updateRtFiniReloc() {
       /*Addend*/ RT->getRuntimeFiniAddress(), /*Value*/ 0});
 }
 
+void RewriteInstance::updateRtInitReloc() {
+  // Updating DT_INIT is handled by patchELFDynamic.
+  if (BC->InitAddress || !BC->InitArrayAddress)
+    return;
+
+  const RuntimeLibrary *RT = BC->getRuntimeLibrary();
+  if (!RT || !RT->getRuntimeStartAddress())
+    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)) {
+    assert(Reloc->Addend == BC->StartFunctionAddress &&
+           "inconsistent .init_array dynamic relocation");
+    Reloc->Addend = RT->getRuntimeStartAddress();
+    InitArraySection->addDynamicRelocation(*Reloc);
+  }
+
+  // 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::registerFragments() {
   if (!BC->HasSplitFunctions ||
       opts::HeatmapMode == opts::HeatmapModeKind::HM_Exclusive)
@@ -5705,8 +5785,18 @@ Error RewriteInstance::readELFDynamic(ELFObjectFile<ELFT> *File) {
     switch (Dyn.d_tag) {
     case ELF::DT_INIT:
       if (!BC->HasInterpHeader) {
-        LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Set start function address\n");
-        BC->StartFunctionAddress = Dyn.getPtr();
+        LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Set init address\n");
+        BC->InitAddress = Dyn.getPtr();
+      }
+      break;
+    case ELF::DT_INIT_ARRAY:
+      if (!BC->HasInterpHeader) {
+        BC->InitArrayAddress = Dyn.getPtr();
+      }
+      break;
+    case ELF::DT_INIT_ARRAYSZ:
+      if (!BC->HasInterpHeader) {
+        BC->InitArraySize = Dyn.getPtr();
       }
       break;
     case ELF::DT_FINI:

bolt/lib/RuntimeLibs/InstrumentationRuntimeLibrary.cpp

@@ -51,12 +51,6 @@ void InstrumentationRuntimeLibrary::adjustCommandLineOptions(
     opts::JumpTables = JTS_MOVE;
     outs() << "BOLT-INFO: forcing -jump-tables=move for instrumentation\n";
   }
-  if (!BC.StartFunctionAddress) {
-    errs() << "BOLT-ERROR: instrumentation runtime libraries require a known "
-              "entry point of "
-              "the input binary\n";
-    exit(1);
-  }
 
   if (BC.IsStaticExecutable && !opts::InstrumentationSleepTime) {
     errs() << "BOLT-ERROR: instrumentation of static binary currently does not "
@@ -78,6 +72,13 @@ void InstrumentationRuntimeLibrary::adjustCommandLineOptions(
 
 void InstrumentationRuntimeLibrary::emitBinary(BinaryContext &BC,
                                                MCStreamer &Streamer) {
+  /*  if (!BC.StartFunctionAddress) {
+      errs() << "BOLT-ERROR: instrumentation runtime libraries require a known "
+                "entry point of "
+                "the input binary\n";
+      exit(1);
+    }*/
+
   MCSection *Section = BC.isELF()
                            ? static_cast<MCSection *>(BC.Ctx->getELFSection(
                                  ".bolt.instr.counters", ELF::SHT_PROGBITS,

bolt/test/AArch64/hook-init.s

@@ -0,0 +1,96 @@
+## Test the different ways of handling entry point for instrumentation.
+## Bolt is hooking its runtime function via Elf entry, DT_INIT or DT_INIT_ARRAYS.
+## Bolt uses Elf e_entry address for ELF executable, and DT_INIT address
+## for ELF shared object to determine the start address.
+## The Test is checking the following cases:
+## - For executable, check ELF e_entry is pathced.
+## - For shared object:
+##   - Bolt use DT_INIT for hooking runtime start function if that exists.
+##   - If it doesn't exists, DT_INIT_ARRAY takes its place.
+# REQUIRES: system-linux,bolt-runtime,target=aarch64{{.*}}
+
+## Check e_entry address is updated with ELF PIE executable.
+# RUN: %clang %cflags -pie %s -Wl,-q -o %t.exe
+# RUN: llvm-readelf -l %t.exe | FileCheck --check-prefix=CHECK-INTERP %s
+# RUN: llvm-readelf -r %t.exe | FileCheck --check-prefix=RELOC-PIE %s
+# RUN: llvm-readelf -hs %t.exe | FileCheck --check-prefix=CHECK-START %s
+# RUN: llvm-bolt %t.exe -o %t --instrument
+# RUN: llvm-readelf -dhs %t | FileCheck --check-prefix=CHECK-ENTRY %s
+
+## Create a shared library to use DT_INIT for the instrumentation.
+# RUN: %clang %cflags -fPIC -shared %s -Wl,-q -o %t-init.so
+# RUN: llvm-bolt %t-init.so -o %t-init --instrument
+# RUN: llvm-readelf -drs %t-init | FileCheck --check-prefix=CHECK-INIT %s
+
+# Create a shared library with no init to use DT_INIT_ARRAY for the instrumentation.
+# RUN: %clang %cflags -shared %s -Wl,-q,-init=0 -o %t-no-init.so
+# RUN: llvm-bolt %t-no-init.so -o %t-no-init --instrument
+# RUN: llvm-readelf -drs %t-no-init | FileCheck --check-prefix=CHECK-NO-INIT %s
+
+## Check the binary has InterP header
+# CHECK-INTERP: Program Headers:
+# CHECK-INTERP: INTERP
+
+## With PIE: binary should have relative relocations
+# RELOC-PIE: R_AARCH64_RELATIVE
+
+## ELF excecutable where e_entry is set to __bolt_runtime_start (PIE).
+## Check the input that e_entry points to _start by default.
+# CHECK-START: ELF Header:
+# CHECK-START-DAG:   Entry point address: 0x[[ENTRY:[[:xdigit:]]+]]
+# CHECK-START: Symbol table '.symtab' contains {{.*}} entries:
+# CHECK-START-DAG: {{0+}}[[ENTRY]] {{.*}} _start
+## Check that e_entry is set to __bolt_runtime_start after the instrumentation.
+# CHECK-ENTRY: ELF Header:
+# CHECK-ENTRY-DAG:   Entry point address: 0x[[ENTRY:[[:xdigit:]]+]]
+# CHECK-ENTRY: Symbol table '.symtab' contains {{.*}} entries:
+# CHECK-ENTRY-DAG: {{0+}}[[ENTRY]] {{.*}} __bolt_runtime_start
+
+## Check that DT_INIT is set to __bolt_runtime_start.
+# CHECK-INIT:     Dynamic section at offset {{.*}} contains {{.*}} entries:
+# CHECK-INIT-DAG: (INIT) 0x[[INIT:[[:xdigit:]]+]]
+# CHECK-INIT-DAG: (INIT_ARRAY) 0x[[INIT_ARRAY:[[:xdigit:]]+]]
+## Check that the dynamic relocation at .init_array was not patched
+# CHECK-INIT:     Relocation section '.rela.dyn' at offset {{.*}} contains {{.*}} entries
+# CHECK-INIT:     {{0+}}[[INIT_ARRAY]] {{.*}} R_AARCH64_RELATIVE [[MYINIT_ADDR:[[:xdigit:]]+]
+]
+# CHECK-INIT:     Symbol table '.symtab' contains {{.*}} entries:
+# CHECK-INIT-DAG: {{0+}}[[MYINIT_ADDR]] {{.*}} _myinit
+
+## Check that DT_INIT_ARRAY is set to __bolt_runtime_start.
+# CHECK-NO-INIT:     Dynamic section at offset {{.*}} contains {{.*}} entries:
+# CHECK-NO-INIT-NOT: (INIT)
+# CHECK-NO-INIT:     (INIT_ARRAY) 0x[[INIT_ARRAY:[a-f0-9]+]]
+# CHECK-NO-INIT:     Relocation section '.rela.dyn' at offset {{.*}} contains {{.*}} entries
+# CHECK-NO-INIT:     {{0+}}[[INIT_ARRAY]] {{.*}} R_AARCH64_RELATIVE [[INIT_ADDR:[[:xdigit:]]+]]
+# CHECK-NO-INIT:     Symbol table '.symtab' contains {{.*}} entries:
+# CHECK-NO-INIT-DAG: {{0+}}[[INIT_ADDR]] {{.*}} __bolt_runtime_start
+
+  .globl _start
+  .type _start, %function
+_start:
+  # Dummy relocation to force relocation mode.
+  .reloc 0, R_AARCH64_NONE
+  ret
+.size _start, .-_start
+
+  .globl _init
+  .type _init, %function
+_init:
+  ret
+  .size _init, .-_init
+
+  .globl _fini
+  .type _fini, %function
+_fini:
+  ret
+  .size _fini, .-_fini
+
+  .section .text
+_myinit:
+  ret
+  .size _myinit, .-_myinit
+
+  .section .init_array,"aw"
+  .align 3
+  .dword _myinit # For relative relocation