1//===-- AArch64BranchTargets.cpp -- Harden code using v8.5-A BTI extension -==//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This pass inserts BTI instructions at the start of every function and basic
10// block which could be indirectly called. The hardware will (when enabled)
11// trap when an indirect branch or call instruction targets an instruction
12// which is not a valid BTI instruction. This is intended to guard against
13// control-flow hijacking attacks. Note that this does not do anything for RET
14// instructions, as they can be more precisely protected by return address
15// signing.
16//
17//===----------------------------------------------------------------------===//
18
19#include "AArch64.h"
20#include "AArch64MachineFunctionInfo.h"
21#include "AArch64Subtarget.h"
22#include "Utils/AArch64BaseInfo.h"
23#include "llvm/CodeGen/MachineFunctionPass.h"
24#include "llvm/CodeGen/MachineInstrBuilder.h"
25#include "llvm/CodeGen/MachineJumpTableInfo.h"
26#include "llvm/CodeGen/MachineModuleInfo.h"
27#include "llvm/Support/Debug.h"
28
29using namespace llvm;
30
31#define DEBUG_TYPE "aarch64-branch-targets"
32#define AARCH64_BRANCH_TARGETS_NAME "AArch64 Branch Targets"
33
34namespace {
35// BTI HINT encoding: base (32) plus 'c' (2) and/or 'j' (4).
36enum : unsigned {
37 BTIBase = 32, // Base immediate for BTI HINT
38 BTIC = 1u << 1, // 2
39 BTIJ = 1u << 2, // 4
40 BTIMask = BTIC | BTIJ,
41};
42
43class AArch64BranchTargetsImpl {
44public:
45 bool run(MachineFunction &MF);
46
47private:
48 const AArch64Subtarget *Subtarget;
49
50 void addBTI(MachineBasicBlock &MBB, bool CouldCall, bool CouldJump,
51 bool NeedsWinCFI);
52};
53
54class AArch64BranchTargetsLegacy : public MachineFunctionPass {
55public:
56 static char ID;
57 AArch64BranchTargetsLegacy() : MachineFunctionPass(ID) {}
58 void getAnalysisUsage(AnalysisUsage &AU) const override;
59 bool runOnMachineFunction(MachineFunction &MF) override;
60 StringRef getPassName() const override { return AARCH64_BRANCH_TARGETS_NAME; }
61};
62
63} // end anonymous namespace
64
65char AArch64BranchTargetsLegacy::ID = 0;
66
67INITIALIZE_PASS(AArch64BranchTargetsLegacy, "aarch64-branch-targets",
68 AARCH64_BRANCH_TARGETS_NAME, false, false)
69
70void AArch64BranchTargetsLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
71 AU.setPreservesCFG();
72 MachineFunctionPass::getAnalysisUsage(AU);
73}
74
75FunctionPass *llvm::createAArch64BranchTargetsPass() {
76 return new AArch64BranchTargetsLegacy();
77}
78
79bool AArch64BranchTargetsLegacy::runOnMachineFunction(MachineFunction &MF) {
80 return AArch64BranchTargetsImpl().run(MF);
81}
82
83PreservedAnalyses
84AArch64BranchTargetsPass::run(MachineFunction &MF,
85 MachineFunctionAnalysisManager &MFAM) {
86 bool Changed = AArch64BranchTargetsImpl().run(MF);
87 if (!Changed)
88 return PreservedAnalyses::all();
89 PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
90 PA.preserveSet<CFGAnalyses>();
91 return PA;
92}
93
94bool AArch64BranchTargetsImpl::run(MachineFunction &MF) {
95 if (!MF.getInfo<AArch64FunctionInfo>()->branchTargetEnforcement())
96 return false;
97
98 LLVM_DEBUG(dbgs() << "********** AArch64 Branch Targets **********\n"
99 << "********** Function: " << MF.getName() << '\n');
100 const Function &F = MF.getFunction();
101
102 Subtarget = &MF.getSubtarget<AArch64Subtarget>();
103
104 // LLVM does not consider basic blocks which are the targets of jump tables
105 // to be address-taken (the address can't escape anywhere else), but they are
106 // used for indirect branches, so need BTI instructions.
107 SmallPtrSet<MachineBasicBlock *, 8> JumpTableTargets;
108 if (auto *JTI = MF.getJumpTableInfo())
109 for (auto &JTE : JTI->getJumpTables())
110 JumpTableTargets.insert_range(R: JTE.MBBs);
111
112 bool MadeChange = false;
113 bool HasWinCFI = MF.hasWinCFI();
114 for (MachineBasicBlock &MBB : MF) {
115 bool CouldCall = false, CouldJump = false;
116 // If the function is address-taken or externally-visible, it could be
117 // indirectly called. PLT entries and tail-calls use BR, but when they are
118 // are in guarded pages should all use x16 or x17 to hold the called
119 // address, so we don't need to set CouldJump here. BR instructions in
120 // non-guarded pages (which might be non-BTI-aware code) are allowed to
121 // branch to a "BTI c" using any register.
122 //
123 // For ELF targets, this is enough, because AAELF64 says that if the static
124 // linker later wants to use an indirect branch instruction in a
125 // long-branch thunk, it's also responsible for adding a 'landing pad' with
126 // a BTI, and pointing the indirect branch at that. For non-ELF targets we
127 // can't rely on that, so we assume that `CouldCall` is _always_ true due
128 // to the risk of long-branch thunks at link time.
129 if (&MBB == &*MF.begin() && (!Subtarget->isTargetELF() ||
130 (F.hasAddressTaken() || !F.hasLocalLinkage())))
131 CouldCall = true;
132
133 // If the block itself is address-taken, it could be indirectly branched
134 // to, but not called.
135 if (MBB.isMachineBlockAddressTaken() || MBB.isIRBlockAddressTaken() ||
136 JumpTableTargets.count(Ptr: &MBB))
137 CouldJump = true;
138
139 if (MBB.isEHPad()) {
140 if (HasWinCFI && (MBB.isEHFuncletEntry() || MBB.isCleanupFuncletEntry()))
141 CouldCall = true;
142 else
143 CouldJump = true;
144 }
145 if (CouldCall || CouldJump) {
146 addBTI(MBB, CouldCall, CouldJump, NeedsWinCFI: HasWinCFI);
147 MadeChange = true;
148 }
149 }
150
151 return MadeChange;
152}
153
154void AArch64BranchTargetsImpl::addBTI(MachineBasicBlock &MBB, bool CouldCall,
155 bool CouldJump, bool HasWinCFI) {
156 LLVM_DEBUG(dbgs() << "Adding BTI " << (CouldJump ? "j" : "")
157 << (CouldCall ? "c" : "") << " to " << MBB.getName()
158 << "\n");
159
160 unsigned HintNum = getBTIHintNum(CallTarget: CouldCall, JumpTarget: CouldJump);
161 auto MBBI = MBB.begin();
162
163 // If the block starts with EH_LABEL(s), skip them first.
164 while (MBBI != MBB.end() && MBBI->isEHLabel()) {
165 ++MBBI;
166 }
167
168 // Skip meta/CFI/etc. (and EMITBKEY) to reach the first executable insn.
169 for (; MBBI != MBB.end() &&
170 (MBBI->isMetaInstruction() || MBBI->getOpcode() == AArch64::EMITBKEY);
171 ++MBBI)
172 ;
173
174 // PACI[AB]SP are implicitly BTI c so insertion of a BTI can be skipped in
175 // this case. Depending on the runtime value of SCTLR_EL1.BT[01], they are not
176 // equivalent to a BTI jc, which still requires an additional BTI.
177 if (MBBI != MBB.end() && ((HintNum & BTIMask) == BTIC) &&
178 (MBBI->getOpcode() == AArch64::PACIASP ||
179 MBBI->getOpcode() == AArch64::PACIBSP))
180 return;
181
182 const AArch64InstrInfo *TII = Subtarget->getInstrInfo();
183
184 // Insert BTI exactly at the first executable instruction.
185 const DebugLoc DL = MBB.findDebugLoc(MBBI);
186 MachineInstr *BTI = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::HINT))
187 .addImm(Val: HintNum)
188 .getInstr();
189
190 // WinEH: put .seh_nop after BTI when the first real insn is FrameSetup.
191 if (HasWinCFI && MBBI != MBB.end() &&
192 MBBI->getFlag(Flag: MachineInstr::FrameSetup)) {
193 auto AfterBTI = std::next(x: MachineBasicBlock::iterator(BTI));
194 BuildMI(BB&: MBB, I: AfterBTI, MIMD: DL, MCID: TII->get(Opcode: AArch64::SEH_Nop));
195 }
196}
197