[BOLT][BTI] Disassemble PLT entries when processing BTI binaries

bolt/lib/Core/BinaryFunction.cpp

@@ -147,6 +147,11 @@ static cl::opt<bool> TrapOnAVX512(
     cl::init(false), cl::ZeroOrMore, cl::Hidden, cl::cat(BoltCategory));
 
 bool shouldPrint(const BinaryFunction &Function) {
+  // PLT stubs are disassembled for BTI binaries, therefore they should be
+  // printed.
+  if (Function.getBinaryContext().usesBTI() && Function.isPLTFunction())
+    return true;
+
   if (Function.isIgnored())
     return false;
 

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -461,6 +461,13 @@ Error RewriteInstance::setProfile(StringRef Filename) {
 
 /// Return true if the function \p BF should be disassembled.
 static bool shouldDisassemble(const BinaryFunction &BF) {
+
+  const BinaryContext &BC = BF.getBinaryContext();
+  // Disassemble PLT functions for BTI binaries to check if they need landing
+  // pads when targeting them in LongJmp.
+  if (BC.usesBTI() && BF.isPLTFunction())
+    return true;
+
   if (BF.isPseudo())
     return false;
 

bolt/test/runtime/AArch64/disassemble-plts.c

@@ -0,0 +1,31 @@
+// This test checks that BOLT disassembles PLT stubs in binaries using BTI,
+// while keeping them not disassembled in non-BTI binaries.
+
+// RUN: %clang -fuse-ld=lld --target=aarch64-unknown-linux-gnu %s -o %t.exe \
+// RUN: -Wl,-q
+// RUN: llvm-bolt %t.exe -o %t.bolt --print-disasm | FileCheck %s
+
+// RUN: %clang -fuse-ld=lld --target=aarch64-unknown-linux-gnu \
+// RUN: -mbranch-protection=standard %s -o %t.bti.exe -Wl,-q -Wl,-z,force-bti
+// RUN: llvm-bolt %t.bti.exe -o %t.bolt --print-disasm | FileCheck %s \
+// RUN: --check-prefix=CHECK-BTI
+
+// For the non-BTI binary, PLTs should not be disassembled.
+// CHECK-NOT: Binary Function "{{.*}}@PLT" after disassembly {
+
+// Check that PLTs are disassembled for the BTI binary.
+// CHECK-BTI: Binary Function "__libc_start_main@PLT" after disassembly {
+// CHECK-BTI: adrp
+// CHECK-BTI-NEXT: ldr
+// CHECK-BTI-NEXT: add
+// CHECK-BTI-NEXT: br
+// CHECK-BTI: End of Function "__libc_start_main@PLT"
+
+#include <stdio.h>
+#include <stdlib.h>
+int main(int argc, char **argv) {
+  if (argc > 3)
+    exit(42);
+  else
+    printf("Number of args: %d\n", argc);
+}