[lld] [lld-macho] Ensure all sections in __TEXT get thunks if necessary (PR #87818)

[Leonard Grey via llvm-commits]

https://github.com/speednoisemovement created https://github.com/llvm/llvm-project/pull/87818

This is a proposed fix for issue #52767. It has some flaws, but I wanted to put it out for discussion.

We decide whether a section has thunk callsites by comparing its size with the branch range. This can fail in the case of sections that are themselves smaller than the branch range, but need to branch to addresses in other sections that are farther than the branch range away. We ran into this in Chrome ASAN builds with the `__lcxx_override` section (see https://crbug.com/326898585). Putting the section closer to `__text`, or before `__stubs` doesn't help since `__text` is large enough on its own to require thunks for branches into it.

Probably the right way to do this is to check if the total size of text sections that can serve as branch targets is bigger than the range, but I couldn't think of a nice way to do this. This change sets a bit on the segment if one section is found to need thunks. Sections processed afterwards in the same segment will take the thunks branch regardless of their size.

This makes the assumption that the biggest section comes first (probably true almost always) and that a single section is bigger than the branch range (probably not true at the margins).


>From d92a3fbe49858d3873e6f2ae757df63fa174af19 Mon Sep 17 00:00:00 2001
From: Leonard Grey <lgrey at chromium.org>
Date: Fri, 5 Apr 2024 13:49:27 -0400
Subject: [PATCH] [lld-macho] Ensure all sections in __TEXT get thunks if
 necessary

---
 lld/MachO/ConcatOutputSection.cpp | 14 ++++++++++++--
 lld/MachO/OutputSegment.h         |  1 +
 lld/test/MachO/arm64-thunks.s     | 14 ++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/lld/MachO/ConcatOutputSection.cpp b/lld/MachO/ConcatOutputSection.cpp
index 279423720be9d5..cd87f9cc4c3caf 100644
--- a/lld/MachO/ConcatOutputSection.cpp
+++ b/lld/MachO/ConcatOutputSection.cpp
@@ -127,10 +127,20 @@ bool TextOutputSection::needsThunks() const {
   uint64_t isecAddr = addr;
   for (ConcatInputSection *isec : inputs)
     isecAddr = alignToPowerOf2(isecAddr, isec->align) + isec->getSize();
-  if (isecAddr - addr + in.stubs->getSize() <=
-      std::min(target->backwardBranchRange, target->forwardBranchRange))
+  // Other sections besides __text might be small enough to pass this
+  // test but nevertheless need thunks for calling into oher sections.
+  // An imperfect heuristic to use in this case is that if a section
+  // we've already processed in this segment needs thunks, so do the
+  // rest.
+  bool needsThunks = parent && parent->needsThunks;
+  if (!needsThunks &&
+      isecAddr - addr + in.stubs->getSize() <=
+          std::min(target->backwardBranchRange, target->forwardBranchRange))
     return false;
   // Yes, this program is large enough to need thunks.
+  if (parent) {
+    parent->needsThunks = true;
+  }
   for (ConcatInputSection *isec : inputs) {
     for (Reloc &r : isec->relocs) {
       if (!target->hasAttr(r.type, RelocAttrBits::BRANCH))
diff --git a/lld/MachO/OutputSegment.h b/lld/MachO/OutputSegment.h
index 7a0c4a2065a15b..5f39734ffdbd7a 100644
--- a/lld/MachO/OutputSegment.h
+++ b/lld/MachO/OutputSegment.h
@@ -57,6 +57,7 @@ class OutputSegment {
   uint32_t initProt = 0;
   uint32_t flags = 0;
   uint8_t index;
+  bool needsThunks = false;
 
   llvm::TinyPtrVector<Defined *> segmentStartSymbols;
   llvm::TinyPtrVector<Defined *> segmentEndSymbols;
diff --git a/lld/test/MachO/arm64-thunks.s b/lld/test/MachO/arm64-thunks.s
index 0e4157ece72de2..ea40effd696f59 100644
--- a/lld/test/MachO/arm64-thunks.s
+++ b/lld/test/MachO/arm64-thunks.s
@@ -7,6 +7,7 @@
 ## (3) a second thunk is created when the first one goes out of range
 ## (4) early calls to a dylib stub use a thunk, and later calls the stub
 ##     directly
+## (5) Thunks are created for all sections in the text segment with branches.
 ## Notes:
 ## 0x4000000 = 64 Mi = half the magnitude of the forward-branch range
 
@@ -168,6 +169,10 @@
 
 # CHECK: [[#%x, NAN_PAGE + NAN_OFFSET]] <__stubs>:
 
+# CHECK: Disassembly of section __TEXT,__lcxx_override:
+# CHECK: <_z>:
+# CHECK:  bl 0x[[#%x, A_THUNK_0]] <_a.thunk.0>
+
 .subsections_via_symbols
 .text
 
@@ -300,3 +305,12 @@ _main:
   bl _h
   bl ___nan
   ret
+
+.section __TEXT,__lcxx_override,regular,pure_instructions
+
+.globl _z
+.no_dead_strip _z
+.p2align 2
+_z:
+  bl _a
+  ret