From f7c5316468ec077157c70252c5bf97ab50b3dbd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gergely=20B=C3=A1lint?= Date: Thu, 12 Feb 2026 08:29:16 +0100 Subject: [PATCH] [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 | 15 +++++ bolt/lib/Passes/Hugify.cpp | 7 +- bolt/lib/Passes/LongJmp.cpp | 60 ++--------------- bolt/lib/Passes/PatchEntries.cpp | 3 + .../Target/AArch64/AArch64MCPlusBuilder.cpp | 64 +++++++++++++++++++ .../{inline-bti-dbg.s => bti-inline-dbg.s} | 0 .../AArch64/{inline-bti.s => bti-inline.s} | 0 ...p-bti-ignored.s => bti-long-jmp-ignored.s} | 0 .../{long-jmp-bti.s => bti-long-jmp.s} | 0 ...{no-bti-note.test => bti-note-unused.test} | 0 bolt/test/AArch64/bti-patch-entries.s | 38 +++++++++++ bolt/test/runtime/AArch64/bti-hugify.c | 36 +++++++++++ ...l-bti.c => bti-instrumentation-ind-call.c} | 0 ...{long-jmp-bti-plt.c => bti-long-jmp-plt.c} | 0 14 files changed, 166 insertions(+), 57 deletions(-) rename bolt/test/AArch64/{inline-bti-dbg.s => bti-inline-dbg.s} (100%) rename bolt/test/AArch64/{inline-bti.s => bti-inline.s} (100%) rename bolt/test/AArch64/{long-jmp-bti-ignored.s => bti-long-jmp-ignored.s} (100%) rename bolt/test/AArch64/{long-jmp-bti.s => bti-long-jmp.s} (100%) rename bolt/test/AArch64/{no-bti-note.test => bti-note-unused.test} (100%) create mode 100644 bolt/test/AArch64/bti-patch-entries.s create mode 100644 bolt/test/runtime/AArch64/bti-hugify.c rename bolt/test/runtime/AArch64/{instrumentation-ind-call-bti.c => bti-instrumentation-ind-call.c} (100%) rename bolt/test/runtime/AArch64/{long-jmp-bti-plt.c => bti-long-jmp-plt.c} (100%) diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h index e571e91d8513..b78270b31298 100644 --- a/bolt/include/bolt/Core/MCPlusBuilder.h +++ b/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 &MethodFetchInsns, diff --git a/bolt/lib/Passes/Hugify.cpp b/bolt/lib/Passes/Hugify.cpp index 1ac1b08573b8..2f240b7cdeb1 100644 --- a/bolt/lib/Passes/Hugify.cpp +++ b/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 diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp index 798e1ba08918..5b8cbe623dea 100644 --- a/bolt/lib/Passes/LongJmp.cpp +++ b/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(BC.AsmInfo->getCodePointerSize() * 8); Modified = true; return Error::success(); diff --git a/bolt/lib/Passes/PatchEntries.cpp b/bolt/lib/Passes/PatchEntries.cpp index 1b79c122c5f9..9af9e2ca3bde 100644 --- a/bolt/lib/Passes/PatchEntries.cpp +++ b/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; diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp index aa5cf3c671cd..ea043c12583c 100644 --- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp +++ b/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: diff --git a/bolt/test/AArch64/inline-bti-dbg.s b/bolt/test/AArch64/bti-inline-dbg.s similarity index 100% rename from bolt/test/AArch64/inline-bti-dbg.s rename to bolt/test/AArch64/bti-inline-dbg.s diff --git a/bolt/test/AArch64/inline-bti.s b/bolt/test/AArch64/bti-inline.s similarity index 100% rename from bolt/test/AArch64/inline-bti.s rename to bolt/test/AArch64/bti-inline.s diff --git a/bolt/test/AArch64/long-jmp-bti-ignored.s b/bolt/test/AArch64/bti-long-jmp-ignored.s similarity index 100% rename from bolt/test/AArch64/long-jmp-bti-ignored.s rename to bolt/test/AArch64/bti-long-jmp-ignored.s diff --git a/bolt/test/AArch64/long-jmp-bti.s b/bolt/test/AArch64/bti-long-jmp.s similarity index 100% rename from bolt/test/AArch64/long-jmp-bti.s rename to bolt/test/AArch64/bti-long-jmp.s diff --git a/bolt/test/AArch64/no-bti-note.test b/bolt/test/AArch64/bti-note-unused.test similarity index 100% rename from bolt/test/AArch64/no-bti-note.test rename to bolt/test/AArch64/bti-note-unused.test diff --git a/bolt/test/AArch64/bti-patch-entries.s b/bolt/test/AArch64/bti-patch-entries.s new file mode 100644 index 000000000000..005b811ff4d7 --- /dev/null +++ b/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: : +# CHECK-NEXT: adrp x16, 0x[[#%x,ADRP:]] +# CHECK-NEXT: add x16, x16, #0x[[#%x,ADD:]] +# CHECK-NEXT: br x16 + +# CHECK: [[#ADRP + ADD]] : +# 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 diff --git a/bolt/test/runtime/AArch64/bti-hugify.c b/bolt/test/runtime/AArch64/bti-hugify.c new file mode 100644 index 000000000000..9887a9746f11 --- /dev/null +++ b/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 + +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 + +*/ diff --git a/bolt/test/runtime/AArch64/instrumentation-ind-call-bti.c b/bolt/test/runtime/AArch64/bti-instrumentation-ind-call.c similarity index 100% rename from bolt/test/runtime/AArch64/instrumentation-ind-call-bti.c rename to bolt/test/runtime/AArch64/bti-instrumentation-ind-call.c diff --git a/bolt/test/runtime/AArch64/long-jmp-bti-plt.c b/bolt/test/runtime/AArch64/bti-long-jmp-plt.c similarity index 100% rename from bolt/test/runtime/AArch64/long-jmp-bti-plt.c rename to bolt/test/runtime/AArch64/bti-long-jmp-plt.c