[BOLT] Search section based on relocation symbol

We need to search referenced section based on relocations symbol section
to properly match end section symbols. For example on some binaries we
can observe that init_array_end/fini_array_end might be "placed" in to
the gap and since no section could be found for address the relocation
would be skipped resulting in wrong ADRP imm after emitting new text
resulting in binary sigsegv.

Credits for the test to Vladislav Khmelevskii aka yota9.

(cherry picked from commit 0776fc32b15dc76c6b43c41fc4471552833265de)

3 files changed, 50 insertions, 2 deletions

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 77323adc0b38..9f63a69365cc 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -2462,8 +2462,21 @@ void RewriteInstance::handleRelocation(const SectionRef &RelocatedSection,
     if (BinaryData *BD = BC->getBinaryDataByName(SymbolName))
       ReferencedSymbol = BD->getSymbol();
 
-  ErrorOr<BinarySection &> ReferencedSection =
-      BC->getSectionForAddress(SymbolAddress);
+  ErrorOr<BinarySection &> ReferencedSection{std::errc::bad_address};
+  symbol_iterator SymbolIter = Rel.getSymbol();
+  if (SymbolIter != InputFile->symbol_end()) {
+    SymbolRef Symbol = *SymbolIter;
+    section_iterator Section =
+        cantFail(Symbol.getSection(), "cannot get symbol section");
+    if (Section != InputFile->section_end()) {
+      Expected<StringRef> SectionName = Section->getName();
+      if (SectionName && !SectionName->empty())
+        ReferencedSection = BC->getUniqueSectionByName(*SectionName);
+    }
+  }
+
+  if (!ReferencedSection)
+    ReferencedSection = BC->getSectionForAddress(SymbolAddress);
 
   const bool IsToCode = ReferencedSection && ReferencedSection->isText();
 
diff --git a/bolt/test/AArch64/Inputs/array_end.lld_script b/bolt/test/AArch64/Inputs/array_end.lld_script
new file mode 100644
index 000000000000..182c13d370a3
--- /dev/null
+++ b/bolt/test/AArch64/Inputs/array_end.lld_script
@@ -0,0 +1,12 @@
+SECTIONS {
+  .fini_array    :
+  {
+    PROVIDE_HIDDEN (__fini_array_start = .);
+    KEEP (*(SORT_BY_INIT_PRIORITY(.fini_array.*) SORT_BY_INIT_PRIORITY(.dtors.*)))
+    PROVIDE_HIDDEN (__fini_array_end = .);
+  }
+
+  . = . + 128;
+
+  .text : { *(.text) }
+}
diff --git a/bolt/test/AArch64/array_end.c b/bolt/test/AArch64/array_end.c
new file mode 100644
index 000000000000..6f4f9800c390
--- /dev/null
+++ b/bolt/test/AArch64/array_end.c
@@ -0,0 +1,23 @@
+// Test checks that bolt properly finds end section label.
+// Linker script contains gap after destructor array, so
+// __init_array_end address would not be owned by any section.
+
+// REQUIRES: system-linux
+// RUN: %clang %cflags -no-pie %s -o %t.exe -Wl,-q \
+// RUN:   -Wl,-T %S/Inputs/array_end.lld_script
+// RUN: llvm-bolt %t.exe -o %t.bolt --print-disasm \
+// RUN:  --print-only="callFini" | FileCheck %s
+
+// CHECK: adr [[REG:x[0-28]+]], "__fini_array_end/1"
+
+__attribute__((destructor)) void destr() {}
+
+__attribute__((noinline)) void callFini() {
+  extern void (*__fini_array_start[])();
+  extern void (*__fini_array_end[])();
+  unsigned long Count = __fini_array_end - __fini_array_start;
+  for (unsigned long I = 0; I < Count; ++I)
+    (*__fini_array_start[I])();
+}
+
+void _start() { callFini(); }