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

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 getParent() != 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->isCallCoveredByBTI(*StubBB.getLastNonPseudoInstr(),
+                                     FirstInst))
+        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,38 @@
+# 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
+
+# This is skipped, so it stays in the .bolt.org.text.
+# The .text produced by BOLT is aligned to 0x10000000,
+# so _start will need a stub to jump here.
+  .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.