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

[bgergely0]

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.

[bgergely0]

• [BOLT][BTI] Disassemble PLT entries when processing BTI binaries #169663 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Gergely Bálint (bgergely0)

Changes

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 --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@PLT" after disassembly {
+// CHECK-BTI: adrp
+// CHECK-BTI: ldr
+// CHECK-BTI: add
+// CHECK-BTI: 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);
+}

[bgergely0]

Note: there is an assert checking that PLT functions' Size must be 0. This never triggered during my local testing. I'm not sure why the Size (meaning the original size of the function) was set to 0 though.

[paschalis-mpeis]

Thanks Gergely, looks good!

This is well guarded under BTI, so it shouldn't affect othes, but give it a bit of time in case there are any comments.

[maksfb]

Note: there is an assert checking that PLT functions' Size must be 0. This never triggered during my local testing. I'm not sure why the Size (meaning the original size of the function) was set to 0 though.

Since PLT functions are not marked in the symbol table, they were considered pseudo functions with zero size when the support was added. BinaryFunction::IsPseudo was added later, and after that I believe it became safe to remove the zero-size assertion.

[maksfb]

Currently, there's a bespoke support for PLT functions where they are disassembled in RewriteInstance and then analyzed in MCPlusBuilder::analyzePLTEntry(). I don't know the exact analysis you plan on the first instruction, but maybe it's possible within the existing analysis code by extending analyzePLTEntry() interface?

[bgergely0]

Hi @maksfb! Thanks for taking a look.

Currently, there's a bespoke support for PLT functions where they are disassembled in RewriteInstance and then analyzed in MCPlusBuilder::analyzePLTEntry(). I don't know the exact analysis you plan on the first instruction, but maybe it's possible within the existing analysis code by extending analyzePLTEntry() interface?

It is likely that in the near future we might need to make PLTs emittable as well, so we can add extra BTIs.
For now, the AnalyzePLTEntry should be good as well.

I will mark this draft until we really need it.

[maksfb]

Sounds good. Thanks!