[llvm] [BOLT][BTI] Disassemble PLT entries when processing BTI binaries (PR #169663)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-bolt

Author: Gergely Bálint (bgergely0)

<details>
<summary>Changes</summary>

PLT entries are PseudoFunctions, and are not disassembled or emitted.
For BTI, we need to check the first MCInst of PLT entries, to see
if indirectly calling them is safe or not.

This patch disassembles PLTs for binaries using BTI, while not changing
the behaviour for binaries without BTI.

The PLTs are only disassembled, not emitted.

---
Full diff: https://github.com/llvm/llvm-project/pull/169663.diff


3 Files Affected:

- (modified) bolt/lib/Core/BinaryFunction.cpp (+5) 
- (modified) bolt/lib/Rewrite/RewriteInstance.cpp (+6) 
- (added) bolt/test/runtime/AArch64/disassemble-plts.c (+30) 


``````````diff
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index a5fdf79a737f5..8cb0ccf7bd396 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/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;
 
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 8a5bbe28e5f66..1be5fc4e2a707 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -461,6 +461,12 @@ 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 on AArch64 to check BTI landing pads.
+  if (BC.usesBTI() && BF.isPLTFunction())
+    return true;
+
   if (BF.isPseudo())
     return false;
 
diff --git a/bolt/test/runtime/AArch64/disassemble-plts.c b/bolt/test/runtime/AArch64/disassemble-plts.c
new file mode 100644
index 0000000000000..5bb0e8615eeab
--- /dev/null
+++ b/bolt/test/runtime/AArch64/disassemble-plts.c
@@ -0,0 +1,30 @@
+// 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
+// -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
+// --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 at PLT" after disassembly {
+// CHECK-BTI: adrp
+// CHECK-BTI: ldr
+// CHECK-BTI: add
+// CHECK-BTI: br
+// CHECK-BTI: End of Function "__libc_start_main at 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);
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/169663