[BOLT][BTI] Refactor: move applyBTIFixup under MCPlusBuilder (#177164)

This patch moves the applyBTIFixup from LongJmp pass to MCPlusBuilder.
This refactor allows applyBTIFixup to be called from other passes
inserting indirect branches, such as:
- Hugify,
- PatchEntries.

As different passes have different information about their targets (e.g.
target BasicBlock, target Symbol, target Function), specialized versions
are created (applyBTIFixupToSymbol, applyBTIFixupToTarget), and each
calls
applyBTIFixupCommon, which implements the original logic from before.

Names of related lit tests are updated to have the "bti" prefix.

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -1787,6 +1787,21 @@ public:
     llvm_unreachable("not implemented");
   }
 
+  virtual void applyBTIFixupToTarget(BinaryBasicBlock &StubBB) {
+    llvm_unreachable("not implemented");
+  }
+
+  virtual void applyBTIFixupToSymbol(BinaryContext &BC, const MCSymbol *Symbol,
+                                     MCInst &Call) {
+    llvm_unreachable("not implemented");
+  }
+
+  virtual void applyBTIFixupCommon(const MCSymbol *RealTargetSym,
+                                   BinaryFunction *TgtFunction,
+                                   BinaryBasicBlock *TgtBB, MCInst &Call) {
+    llvm_unreachable("not implemented");
+  }
+
   virtual bool analyzeVirtualMethodCall(InstructionIterator Begin,
                                         InstructionIterator End,
                                         std::vector<MCInst *> &MethodFetchInsns,

bolt/lib/Passes/Hugify.cpp

@@ -42,8 +42,11 @@ Error HugePage::runOnFunctions(BinaryContext &BC) {
       BC.getBinaryFunctionAtAddress(*BC.StartFunctionAddress);
   assert(Start && "Entry point function not found");
   const MCSymbol *StartSym = Start->getSymbol();
-  createSimpleFunction("__bolt_hugify_start_program",
-                       BC.MIB->createSymbolTrampoline(StartSym, BC.Ctx.get()));
+  InstructionListType Insts =
+      BC.MIB->createSymbolTrampoline(StartSym, BC.Ctx.get());
+  createSimpleFunction("__bolt_hugify_start_program", Insts);
+  if (BC.usesBTI())
+    BC.MIB->applyBTIFixupToSymbol(BC, StartSym, *(Insts.end() - 1));
   return Error::success();
 }
 } // namespace bolt

bolt/lib/Passes/LongJmp.cpp

@@ -43,6 +43,8 @@ static void relaxStubToShortJmp(BinaryBasicBlock &StubBB, const MCSymbol *Tgt) {
   BC.MIB->createShortJmp(Seq, Tgt, BC.Ctx.get());
   StubBB.clear();
   StubBB.addInstructions(Seq.begin(), Seq.end());
+  if (BC.usesBTI())
+    BC.MIB->applyBTIFixupToTarget(StubBB);
 }
 
 static void relaxStubToLongJmp(BinaryBasicBlock &StubBB, const MCSymbol *Tgt) {
@@ -51,6 +53,8 @@ static void relaxStubToLongJmp(BinaryBasicBlock &StubBB, const MCSymbol *Tgt) {
   BC.MIB->createLongJmp(Seq, Tgt, BC.Ctx.get());
   StubBB.clear();
   StubBB.addInstructions(Seq.begin(), Seq.end());
+  if (BC.usesBTI())
+    BC.MIB->applyBTIFixupToTarget(StubBB);
 }
 
 static BinaryBasicBlock *getBBAtHotColdSplitPoint(BinaryFunction &Func) {
@@ -484,62 +488,12 @@ Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
       ~((1ULL << (RangeSingleInstr - 1)) - 1);
 
   const MCSymbol *RealTargetSym = BC.MIB->getTargetSymbol(*StubBB.begin());
-  BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(RealTargetSym);
-  BinaryFunction *TargetFunction = BC.getFunctionForSymbol(RealTargetSym);
+  const BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(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->isPLTFunction()) {
-      BC.MIB->patchPLTEntryForBTI(*TargetFunction,
-                                  *StubBB.getLastNonPseudoInstr());
-      return;
-    }
-    if (TargetFunction && TargetFunction->isIgnored()) {
-      // Includes PLT functions.
-      BC.errs() << "BOLT-ERROR: Cannot add BTI landing pad to ignored function "
-                << TargetFunction->getPrintName() << "\n";
-      exit(1);
-    }
-    if (TargetFunction && !TargetFunction->hasCFG()) {
-      if (TargetFunction->hasInstructions()) {
-        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();
@@ -554,8 +508,6 @@ 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();
@@ -571,8 +523,6 @@ 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/Passes/PatchEntries.cpp

@@ -120,6 +120,9 @@ Error PatchEntries::runOnFunctions(BinaryContext &BC) {
       BinaryFunction *PatchFunction = BC.createInstructionPatch(
           Patch.Address, Instructions,
           NameResolver::append(Patch.Symbol->getName(), ".org.0"));
+      if (BC.usesBTI())
+        BC.MIB->applyBTIFixupToSymbol(BC, Patch.Symbol,
+                                      *(Instructions.end() - 1));
 
       // Verify the size requirements.
       uint64_t HotSize, ColdSize;

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -1733,6 +1733,70 @@ public:
     BC.createInstructionPatch(PLTFunction.getAddress(), NewPLTSeq);
   }
 
+  void applyBTIFixupToSymbol(BinaryContext &BC, const MCSymbol *TargetSymbol,
+                             MCInst &Call) override {
+    BinaryFunction *TargetFunction = BC.getFunctionForSymbol(TargetSymbol);
+    applyBTIFixupCommon(TargetSymbol, TargetFunction, nullptr, Call);
+  }
+
+  void applyBTIFixupToTarget(BinaryBasicBlock &StubBB) override {
+    BinaryFunction &Func = *StubBB.getFunction();
+    BinaryContext &BC = Func.getBinaryContext();
+    const MCSymbol *RealTargetSym = BC.MIB->getTargetSymbol(*StubBB.begin());
+    BinaryFunction *TargetFunction = BC.getFunctionForSymbol(RealTargetSym);
+    BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(RealTargetSym);
+    applyBTIFixupCommon(RealTargetSym, TargetFunction, TgtBB,
+                        *StubBB.getLastNonPseudoInstr());
+  }
+
+  void applyBTIFixupCommon(const MCSymbol *RealTargetSym,
+                           BinaryFunction *TargetFunction,
+                           BinaryBasicBlock *TargetBB, MCInst &Call) override {
+    // TODO: add support for editing each type, and remove errors.
+    if (!TargetFunction && !TargetBB) {
+      errs() << "BOLT-ERROR: Cannot add BTI to function with symbol "
+             << RealTargetSym->getName() << "\n";
+      exit(1);
+    }
+    if (TargetFunction && TargetFunction->isPLTFunction()) {
+      patchPLTEntryForBTI(*TargetFunction, Call);
+      return;
+    }
+    if (TargetFunction && TargetFunction->isIgnored()) {
+      errs() << "BOLT-ERROR: Cannot add BTI landing pad to ignored function "
+             << TargetFunction->getPrintName() << "\n";
+      exit(1);
+    }
+    if (TargetFunction && !TargetFunction->hasCFG()) {
+      if (TargetFunction->hasInstructions()) {
+        auto FirstII = TargetFunction->instrs().begin();
+        MCInst FirstInst = FirstII->second;
+        if (isCallCoveredByBTI(Call, FirstInst))
+          return;
+      }
+      errs()
+          << "BOLT-ERROR: Cannot add BTI landing pad to function without CFG: "
+          << TargetFunction->getPrintName() << "\n";
+      exit(1);
+    }
+    if (!TargetBB)
+      // No need to check TargetFunction for nullptr, because
+      // !TargetBB &&!TargetFunction has already been checked.
+      TargetBB = &*TargetFunction->begin();
+    if (TargetBB) {
+      if (!TargetBB->hasParent()) {
+        errs() << "BOLT-ERROR: Cannot add BTI to block with no parent "
+                  "function. Targeted symbol: "
+               << RealTargetSym->getName() << "\n";
+        exit(1);
+      }
+      insertBTI(*TargetBB, Call);
+      return;
+    }
+    errs() << "BOLT-ERROR: unhandled case when applying BTI fixup\n";
+    exit(1);
+  }
+
   unsigned getInvertedBranchOpcode(unsigned Opcode) const {
     switch (Opcode) {
     default:

bolt/test/AArch64/bti-patch-entries.s

@@ -0,0 +1,38 @@
+# This test checks that BOLT can add BTI to targets at patched entries.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: %clang %cflags -pie %t.o -o %t.exe -nostdlib -Wl,-q,-z,force-bti
+# RUN: llvm-bolt %t.exe -o %t.bolt --use-old-text=0 --lite=0 --force-patch | FileCheck %s --check-prefix=CHECK-BOLT
+# RUN: llvm-objdump -dz %t.bolt | FileCheck %s
+
+# CHECK-BOLT: binary is using BTI
+
+# CHECK: <pathedEntries.org.0>:
+# CHECK-NEXT: adrp x16, 0x[[#%x,ADRP:]]
+# CHECK-NEXT: add x16, x16, #0x[[#%x,ADD:]]
+# CHECK-NEXT: br x16
+
+# CHECK: [[#ADRP + ADD]] <pathedEntries>:
+# CHECK-NEXT: [[#ADRP + ADD]]: {{.*}} bti c
+# CHECK-NEXT: ret
+
+.text
+.balign 4
+.global pathedEntries
+.type pathedEntries, %function
+pathedEntries:
+  .rept 32
+  nop
+  .endr
+  ret
+.size pathedEntries, .-pathedEntries
+
+.global _start
+.type _start, %function
+_start:
+  bl pathedEntries
+  ret
+.size _start, .-_start

bolt/test/runtime/AArch64/bti-hugify.c

@@ -0,0 +1,36 @@
+// Make sure that during the Hugify pass BOLT adds a BTI instruction to _start.
+
+#include <stdio.h>
+
+int main(int argc, char **argv) {
+  printf("Hello world\n");
+  return 0;
+}
+
+/*
+REQUIRES: system-linux,bolt-runtime
+
+RUN: %clang %cflags -mbranch-protection=standard -no-pie %s -o %t.nopie.exe \
+RUN: -Wl,-q,-z,force-bti
+RUN: %clang %cflags -mbranch-protection=standard -fpic %s -o %t.pie.exe \
+RUN: -Wl,-q,-z,force-bti
+
+RUN: llvm-bolt %t.nopie.exe --lite=0 -o %t.nopie.bolt --hugify | FileCheck %s \
+RUN: -check-prefix=CHECK-BOLT
+RUN: llvm-bolt %t.pie.exe --lite=0 -o %t.pie.bolt --hugify | FileCheck %s \
+RUN: -check-prefix=CHECK-BOLT
+
+CHECK-BOLT: binary is using BTI
+
+RUN: llvm-objdump -D %t.nopie.bolt | FileCheck %s -check-prefix=CHECK-OBJDUMP
+RUN: llvm-objdump -D %t.pie.bolt | FileCheck %s -check-prefix=CHECK-OBJDUMP
+
+CHECK-OBJDUMP: <_start>:
+CHECK-OBJDUMP-NEXT: bti
+
+CHECK-OBJDUMP: <__bolt_hugify_start_program>:
+CHECK-OBJDUMP-NEXT: adrp x16, 0x[[#%x,ADRP:]]
+CHECK-OBJDUMP-NEXT: add x16, x16, #0x[[#%x,ADD:]]
+CHECK-OBJDUMP-NEXT: br x16
+
+*/