[BOLT][BTI] Add needed BTIs in LongJmp or refuse to optimize binary

This patch adds BTI landing pads to ShortJmp/LongJmp targets in the
LongJmp pass when optimizing BTI binaries.

BOLT does not have the ability to add BTI to all types of functions.
This patch aims to insert the landing pad where possible, and emit an
error where it currently is not.

BOLT cannot insert BTIs into several function "types", including:
- ignored functions,
- PLT functions,
- other functions without a CFG.

Additional context:

In #161206, BOLT gained the ability to decode the .note.gnu.property
section, and warn about lack of BTI support for BOLT. However, this
warning is misleading: the emitted binary may not need extra BTI landing
pads.

With this patch, the emitted binary will be "BTI-safe".

Author: bgergely0

bolt/include/bolt/Core/BinaryBasicBlock.h

@@ -890,6 +890,8 @@ class BinaryBasicBlock {
   /// Needed by graph traits.
   BinaryFunction *getParent() const { return getFunction(); }
 
+  bool hasParent() const { return getFunction() != nullptr; }
+
   /// Return true if the containing function is in CFG state.
   bool hasCFG() const;
 

bolt/lib/Passes/LongJmp.cpp

@@ -469,8 +469,8 @@ uint64_t LongJmpPass::getSymbolAddress(const BinaryContext &BC,
 }
 
 Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
-  const BinaryFunction &Func = *StubBB.getFunction();
-  const BinaryContext &BC = Func.getBinaryContext();
+  BinaryFunction &Func = *StubBB.getFunction();
+  BinaryContext &BC = Func.getBinaryContext();
   const int Bits = StubBits[&StubBB];
   // Already working with the largest range?
   if (Bits == static_cast<int>(BC.AsmInfo->getCodePointerSize() * 8))
@@ -483,11 +483,54 @@ Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
       ~((1ULL << (RangeSingleInstr - 1)) - 1);
 
   const MCSymbol *RealTargetSym = BC.MIB->getTargetSymbol(*StubBB.begin());
-  const BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(RealTargetSym);
+  BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(RealTargetSym);
+  BinaryFunction *TargetFunction = BC.getFunctionForSymbol(RealTargetSym);
   uint64_t TgtAddress = getSymbolAddress(BC, RealTargetSym, TgtBB);
   uint64_t DotAddress = BBAddresses[&StubBB];
   uint64_t PCRelTgtAddress = DotAddress > TgtAddress ? DotAddress - TgtAddress
                                                      : TgtAddress - DotAddress;
+
+  auto applyBTIFixup = [&](BinaryFunction *TargetFunction,
+                           BinaryBasicBlock *RealTgtBB) {
+    // TODO: add support for editing each type, and remove errors.
+    if (!TargetFunction && !RealTgtBB) {
+      BC.errs() << "BOLT-ERROR: Cannot add BTI to function with symbol "
+                << RealTargetSym->getName() << "\n";
+      exit(1);
+    }
+    if (TargetFunction && TargetFunction->isIgnored()) {
+      BC.errs() << "BOLT-ERROR: Cannot add BTI landing pad to ignored function "
+                << TargetFunction->getPrintName() << "\n";
+      exit(1);
+    }
+    if (TargetFunction && !TargetFunction->hasCFG()) {
+      auto FirstII = TargetFunction->instrs().begin();
+      MCInst FirstInst = FirstII->second;
+      if (BC.MIB->isBTIVariantCoveringCall(FirstInst,
+                                           *StubBB.getLastNonPseudoInstr()))
+        return;
+      BC.errs()
+          << "BOLT-ERROR: Cannot add BTI landing pad to function without CFG: "
+          << TargetFunction->getPrintName() << "\n";
+      exit(1);
+    }
+    if (!RealTgtBB)
+      // !RealTgtBB -> TargetFunction is not a nullptr
+      RealTgtBB = &*TargetFunction->begin();
+    if (RealTgtBB) {
+      if (!RealTgtBB->hasParent()) {
+        BC.errs() << "BOLT-ERROR: Cannot add BTI to block with no parent "
+                     "function. Targeted symbol: "
+                  << RealTargetSym->getName() << "\n";
+        exit(1);
+      }
+      // The BR is the last inst of the StubBB.
+      BC.MIB->insertBTI(*RealTgtBB, *StubBB.getLastNonPseudoInstr());
+      return;
+    }
+    BC.errs() << "BOLT-ERROR: unhandled case when applying BTI fixup\n";
+    exit(1);
+  };
   // If it fits in one instruction, do not relax
   if (!(PCRelTgtAddress & SingleInstrMask))
     return Error::success();
@@ -502,6 +545,8 @@ Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
                       << " RealTargetSym = " << RealTargetSym->getName()
                       << "\n");
     relaxStubToShortJmp(StubBB, RealTargetSym);
+    if (BC.usesBTI())
+      applyBTIFixup(TargetFunction, TgtBB);
     StubBits[&StubBB] = RangeShortJmp;
     Modified = true;
     return Error::success();
@@ -517,6 +562,8 @@ Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
                     << Twine::utohexstr(PCRelTgtAddress)
                     << " RealTargetSym = " << RealTargetSym->getName() << "\n");
   relaxStubToLongJmp(StubBB, RealTargetSym);
+  if (BC.usesBTI())
+    applyBTIFixup(TargetFunction, TgtBB);
   StubBits[&StubBB] = static_cast<int>(BC.AsmInfo->getCodePointerSize() * 8);
   Modified = true;
   return Error::success();

bolt/lib/Rewrite/GNUPropertyRewriter.cpp

@@ -75,8 +75,7 @@ Error GNUPropertyRewriter::sectionInitializer() {
   if (BC.isAArch64()) {
     BC.setUsesBTI(FeaturesAcc & llvm::ELF::GNU_PROPERTY_AARCH64_FEATURE_1_BTI);
     if (BC.usesBTI())
-      BC.outs() << "BOLT-WARNING: binary is using BTI. Optimized binary may be "
-                   "corrupted\n";
+      BC.outs() << "BOLT-INFO: binary is using BTI.\n";
   }
 
   return Error::success();

bolt/test/AArch64/bti-note.test

@@ -1,10 +1,10 @@
 // This test checks that the GNUPropertyRewriter can decode the BTI feature flag.
-// It decodes an executable with BTI, and checks for the warning.
+// It decodes an executable with BTI, and checks for the message.
 
 RUN: yaml2obj %p/Inputs/property-note-bti.yaml &> %t.exe
 
 RUN: llvm-readelf -n %t.exe | FileCheck %s
 CHECK: BTI
 
 RUN: llvm-bolt %t.exe -o %t.exe.bolt | FileCheck %s -check-prefix=CHECK-BOLT
-CHECK-BOLT: BOLT-WARNING: binary is using BTI. Optimized binary may be corrupted
+CHECK-BOLT: BOLT-INFO: binary is using BTI.

bolt/test/AArch64/long-jmp-bti-ignored.s

@@ -0,0 +1,35 @@
+# This test checks the situation where LongJmp adds a stub targeting an ignored (skipped) function.
+# The problem is that by default BOLT cannot modify ignored functions, so it cannot add the needed BTI.
+
+# Current behaviour is to emit an error.
+
+# REQUIRES: system-linux, asserts
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
+# RUN:   -mattr=+bti -aarch64-mark-bti-property %s -o %t.o
+# RUN: %clang %cflags -O0 %t.o -o %t.exe -Wl,-q -Wl,-z,force-bti
+# RUN: not llvm-bolt %t.exe -o %t.bolt \
+# RUN:   --align-text=0x10000000 --skip-funcs=far_away_func 2>&1 | FileCheck %s
+
+# CHECK: BOLT-ERROR: Cannot add BTI landing pad to ignored function far_away_func
+
+  .section .text
+  .global _start
+  .global far_away_func
+
+  .align 4
+  .global _start
+  .type _start, %function
+_start:
+    bti c
+    bl far_away_func
+    ret
+
+  .global far_away_func
+  .type far_away_func, %function
+far_away_func:
+    add x0, x0, #1
+    ret
+
+.reloc 0, R_AARCH64_NONE
+

bolt/test/AArch64/long-jmp-bti.s

@@ -0,0 +1,46 @@
+# This test checks that BOLT can generate BTI landing pads for targets of stubs inserted in LongJmp.
+
+# REQUIRES: system-linux
+
+# RUN: %clang %s %cflags -Wl,-q -o %t -mbranch-protection=bti -Wl,-z,force-bti
+# RUN: link_fdata --no-lbr %s %t %t.fdata
+# RUN: llvm-bolt %t -o %t.bolt --data %t.fdata -split-functions \
+# RUN: --print-split --print-only foo --print-longjmp 2>&1 | FileCheck %s
+
+#CHECK: BOLT-INFO: Starting stub-insertion pass
+#CHECK: Binary Function "foo" after long-jmp
+
+#CHECK:      cmp     x0, #0x0
+#CHECK-NEXT: Successors: .LStub0
+
+#CHECK:      adrp    x16, .Ltmp0
+#CHECK-NEXT: add     x16, x16, :lo12:.Ltmp0
+#CHECK-NEXT: br      x16 # UNKNOWN CONTROL FLOW
+
+#CHECK: -------   HOT-COLD SPLIT POINT   -------
+
+#CHECK:      bti     c
+#CHECK-NEXT: mov     x0, #0x2
+#CHECK-NEXT: ret
+
+  .text
+  .globl  foo
+  .type foo, %function
+foo:
+.cfi_startproc
+.entry_bb:
+# FDATA: 1 foo #.entry_bb# 10
+    cmp x0, #0
+    b .Lcold_bb1
+.Lcold_bb1:
+    mov x0, #2
+    ret
+.cfi_endproc
+  .size foo, .-foo
+
+# empty space, so the splitting needs short stubs
+.data
+.space 0x8000000
+
+## Force relocation mode.
+.reloc 0, R_AARCH64_NONE

bolt/test/AArch64/no-bti-note.test

@@ -1,10 +1,10 @@
 // This test checks that the GNUPropertyRewriter can decode the BTI feature flag.
-// It decodes an executable without BTI, and checks for the warning.
+// It decodes an executable without BTI, and checks that the BTI message is not emitted.
 
 RUN: yaml2obj %p/Inputs/property-note-nobti.yaml &> %t.exe
 
 RUN: llvm-readelf -n %t.exe | FileCheck %s
 CHECK-NOT: BTI
 
 RUN: llvm-bolt %t.exe -o %t.exe.bolt | FileCheck %s -check-prefix=CHECK-BOLT
-CHECK-BOLT-NOT: BOLT-WARNING: binary is using BTI. Optimized binary may be corrupted
+CHECK-BOLT-NOT: BOLT-INFO: binary is using BTI.