[lld][AArch64] Add BTI landing pad to PLT when it is accessed by a range extension thunk.

Adding BTI to those PLT's which accessed with by a range extension thunk due to those preform an indirect call.
Fixes: #62140

Reviewed By: MaskRay

Differential Revision: https://reviews.llvm.org/D148704

4 files changed, 62 insertions, 2 deletions

diff --git a/lld/ELF/Arch/AArch64.cpp b/lld/ELF/Arch/AArch64.cpp
index 0f5c5933fd60..0301508fea23 100644
--- a/lld/ELF/Arch/AArch64.cpp
+++ b/lld/ELF/Arch/AArch64.cpp
@@ -912,7 +912,8 @@ void AArch64BtiPac::writePlt(uint8_t *buf, const Symbol &sym,
   // escape to shared objects. isInIplt indicates a non-preemptible ifunc. Its
   // address may escape if referenced by a direct relocation. The condition is
   // conservative.
-  bool hasBti = btiHeader && (sym.hasFlag(NEEDS_COPY) || sym.isInIplt);
+  bool hasBti = btiHeader &&
+                (sym.hasFlag(NEEDS_COPY) || sym.isInIplt || sym.thunkAccessed);
   if (hasBti) {
     memcpy(buf, btiData, sizeof(btiData));
     buf += sizeof(btiData);
diff --git a/lld/ELF/Symbols.h b/lld/ELF/Symbols.h
index e0b74faafeca..bb440530b4df 100644
--- a/lld/ELF/Symbols.h
+++ b/lld/ELF/Symbols.h
@@ -292,6 +292,9 @@ public:
   // True if defined in a DSO as protected visibility.
   uint8_t dsoProtected : 1;
 
+  // True if targeted by a range extension thunk.
+  uint8_t thunkAccessed : 1;
+
   // Temporary flags used to communicate which symbol entries need PLT and GOT
   // entries during postScanRelocations();
   std::atomic<uint16_t> flags;
diff --git a/lld/ELF/Thunks.cpp b/lld/ELF/Thunks.cpp
index 37896d9483d1..5bfcf6218084 100644
--- a/lld/ELF/Thunks.cpp
+++ b/lld/ELF/Thunks.cpp
@@ -1134,7 +1134,9 @@ bool PPC64LongBranchThunk::isCompatibleWith(const InputSection &isec,
   return rel.type == R_PPC64_REL24 || rel.type == R_PPC64_REL14;
 }
 
-Thunk::Thunk(Symbol &d, int64_t a) : destination(d), addend(a), offset(0) {}
+Thunk::Thunk(Symbol &d, int64_t a) : destination(d), addend(a), offset(0) {
+  destination.thunkAccessed = true;
+}
 
 Thunk::~Thunk() = default;
 
diff --git a/lld/test/ELF/aarch64-feature-bti-plt.s b/lld/test/ELF/aarch64-feature-bti-plt.s
new file mode 100644
index 000000000000..12025c4d4166
--- /dev/null
+++ b/lld/test/ELF/aarch64-feature-bti-plt.s
@@ -0,0 +1,54 @@
+# REQUIRES: aarch64
+
+# RUN: rm -rf %t && split-file %s %t
+
+# RUN: llvm-mc --triple=aarch64 --filetype=obj -o %t.o %t/a.s
+# RUN: ld.lld --shared -T %t/largegap.lds -z force-bti %t.o -o %t.elf
+# RUN: llvm-objdump -d %t.elf | FileCheck %s
+
+#--- largegap.lds
+SECTIONS {
+  .plt : { *(.plt) }
+  .text.near 0x1000 : AT(0x1000) { *(.text.near) }
+  .text.far 0xf0000000 : AT(0xf0000000) { *(.text.far) }
+}
+
+#--- a.s
+# CHECK:        <.plt>:
+# CHECK-NEXT:     bti     c
+
+## foo@plt is targeted by a range extension thunk with an indirect branch.
+## Add a bti c instruction.
+# CHECK:        <foo@plt>:
+# CHECK-NEXT:     bti     c
+
+## biz is not targeted by a thunk using an indirect branch, so no need for bti c.
+# CHECK:        <biz@plt>:
+# CHECK-NEXT:     adrp    x16, {{.*}} <func>
+
+# CHECK:         <bar>:
+# CHECK-NEXT:      bl   {{.*}} <foo@plt>
+# CHECK-NEXT:      bl   {{.*}} <biz@plt>
+
+# CHECK:         <func>:
+# CHECK-NEXT:      bl   {{.*}} <__AArch64ADRPThunk_foo>
+
+# CHECK:         <__AArch64ADRPThunk_foo>:
+# CHECK-NEXT:      adrp    x16, 0x0 <foo>
+# CHECK-NEXT:      add     x16, x16, {{.*}}
+# CHECK-NEXT:      br      x16
+
+        .global foo
+        .global biz
+        .section .text.near, "ax", %progbits
+bar:
+        .type bar, %function
+        bl foo
+        bl biz
+        ret
+
+        .section .text.far, "ax", %progbits
+func:
+        .type func, %function
+        bl foo
+        ret