10b57cec5SDimitry Andric //===- AArch64SpeculationHardening.cpp - Harden Against Missspeculation --===//
20b57cec5SDimitry Andric //
30b57cec5SDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
40b57cec5SDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
50b57cec5SDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
60b57cec5SDimitry Andric //
70b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
80b57cec5SDimitry Andric //
90b57cec5SDimitry Andric // This file contains a pass to insert code to mitigate against side channel
100b57cec5SDimitry Andric // vulnerabilities that may happen under control flow miss-speculation.
110b57cec5SDimitry Andric //
120b57cec5SDimitry Andric // The pass implements tracking of control flow miss-speculation into a "taint"
130b57cec5SDimitry Andric // register. That taint register can then be used to mask off registers with
140b57cec5SDimitry Andric // sensitive data when executing under miss-speculation, a.k.a. "transient
150b57cec5SDimitry Andric // execution".
160b57cec5SDimitry Andric // This pass is aimed at mitigating against SpectreV1-style vulnarabilities.
170b57cec5SDimitry Andric //
180b57cec5SDimitry Andric // It also implements speculative load hardening, i.e. using the taint register
190b57cec5SDimitry Andric // to automatically mask off loaded data.
200b57cec5SDimitry Andric //
210b57cec5SDimitry Andric // As a possible follow-on improvement, also an intrinsics-based approach as
220b57cec5SDimitry Andric // explained at https://lwn.net/Articles/759423/ could be implemented on top of
230b57cec5SDimitry Andric // the current design.
240b57cec5SDimitry Andric //
250b57cec5SDimitry Andric // For AArch64, the following implementation choices are made to implement the
260b57cec5SDimitry Andric // tracking of control flow miss-speculation into a taint register:
270b57cec5SDimitry Andric // Some of these are different than the implementation choices made in
280b57cec5SDimitry Andric // the similar pass implemented in X86SpeculativeLoadHardening.cpp, as
290b57cec5SDimitry Andric // the instruction set characteristics result in different trade-offs.
300b57cec5SDimitry Andric // - The speculation hardening is done after register allocation. With a
310b57cec5SDimitry Andric // relative abundance of registers, one register is reserved (X16) to be
320b57cec5SDimitry Andric // the taint register. X16 is expected to not clash with other register
330b57cec5SDimitry Andric // reservation mechanisms with very high probability because:
340b57cec5SDimitry Andric // . The AArch64 ABI doesn't guarantee X16 to be retained across any call.
350b57cec5SDimitry Andric // . The only way to request X16 to be used as a programmer is through
360b57cec5SDimitry Andric // inline assembly. In the rare case a function explicitly demands to
370b57cec5SDimitry Andric // use X16/W16, this pass falls back to hardening against speculation
380b57cec5SDimitry Andric // by inserting a DSB SYS/ISB barrier pair which will prevent control
390b57cec5SDimitry Andric // flow speculation.
400b57cec5SDimitry Andric // - It is easy to insert mask operations at this late stage as we have
410b57cec5SDimitry Andric // mask operations available that don't set flags.
420b57cec5SDimitry Andric // - The taint variable contains all-ones when no miss-speculation is detected,
430b57cec5SDimitry Andric // and contains all-zeros when miss-speculation is detected. Therefore, when
440b57cec5SDimitry Andric // masking, an AND instruction (which only changes the register to be masked,
450b57cec5SDimitry Andric // no other side effects) can easily be inserted anywhere that's needed.
460b57cec5SDimitry Andric // - The tracking of miss-speculation is done by using a data-flow conditional
470b57cec5SDimitry Andric // select instruction (CSEL) to evaluate the flags that were also used to
480b57cec5SDimitry Andric // make conditional branch direction decisions. Speculation of the CSEL
490b57cec5SDimitry Andric // instruction can be limited with a CSDB instruction - so the combination of
500b57cec5SDimitry Andric // CSEL + a later CSDB gives the guarantee that the flags as used in the CSEL
510b57cec5SDimitry Andric // aren't speculated. When conditional branch direction gets miss-speculated,
520b57cec5SDimitry Andric // the semantics of the inserted CSEL instruction is such that the taint
530b57cec5SDimitry Andric // register will contain all zero bits.
540b57cec5SDimitry Andric // One key requirement for this to work is that the conditional branch is
550b57cec5SDimitry Andric // followed by an execution of the CSEL instruction, where the CSEL
560b57cec5SDimitry Andric // instruction needs to use the same flags status as the conditional branch.
570b57cec5SDimitry Andric // This means that the conditional branches must not be implemented as one
580b57cec5SDimitry Andric // of the AArch64 conditional branches that do not use the flags as input
590b57cec5SDimitry Andric // (CB(N)Z and TB(N)Z). This is implemented by ensuring in the instruction
600b57cec5SDimitry Andric // selectors to not produce these instructions when speculation hardening
610b57cec5SDimitry Andric // is enabled. This pass will assert if it does encounter such an instruction.
620b57cec5SDimitry Andric // - On function call boundaries, the miss-speculation state is transferred from
630b57cec5SDimitry Andric // the taint register X16 to be encoded in the SP register as value 0.
640b57cec5SDimitry Andric //
650b57cec5SDimitry Andric // For the aspect of automatically hardening loads, using the taint register,
660b57cec5SDimitry Andric // (a.k.a. speculative load hardening, see
670b57cec5SDimitry Andric // https://llvm.org/docs/SpeculativeLoadHardening.html), the following
680b57cec5SDimitry Andric // implementation choices are made for AArch64:
690b57cec5SDimitry Andric // - Many of the optimizations described at
700b57cec5SDimitry Andric // https://llvm.org/docs/SpeculativeLoadHardening.html to harden fewer
710b57cec5SDimitry Andric // loads haven't been implemented yet - but for some of them there are
720b57cec5SDimitry Andric // FIXMEs in the code.
730b57cec5SDimitry Andric // - loads that load into general purpose (X or W) registers get hardened by
740b57cec5SDimitry Andric // masking the loaded data. For loads that load into other registers, the
750b57cec5SDimitry Andric // address loaded from gets hardened. It is expected that hardening the
760b57cec5SDimitry Andric // loaded data may be more efficient; but masking data in registers other
770b57cec5SDimitry Andric // than X or W is not easy and may result in being slower than just
780b57cec5SDimitry Andric // hardening the X address register loaded from.
790b57cec5SDimitry Andric // - On AArch64, CSDB instructions are inserted between the masking of the
800b57cec5SDimitry Andric // register and its first use, to ensure there's no non-control-flow
810b57cec5SDimitry Andric // speculation that might undermine the hardening mechanism.
820b57cec5SDimitry Andric //
830b57cec5SDimitry Andric // Future extensions/improvements could be:
840b57cec5SDimitry Andric // - Implement this functionality using full speculation barriers, akin to the
850b57cec5SDimitry Andric // x86-slh-lfence option. This may be more useful for the intrinsics-based
860b57cec5SDimitry Andric // approach than for the SLH approach to masking.
870b57cec5SDimitry Andric // Note that this pass already inserts the full speculation barriers if the
880b57cec5SDimitry Andric // function for some niche reason makes use of X16/W16.
890b57cec5SDimitry Andric // - no indirect branch misprediction gets protected/instrumented; but this
900b57cec5SDimitry Andric // could be done for some indirect branches, such as switch jump tables.
910b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
920b57cec5SDimitry Andric
930b57cec5SDimitry Andric #include "AArch64InstrInfo.h"
940b57cec5SDimitry Andric #include "AArch64Subtarget.h"
950b57cec5SDimitry Andric #include "Utils/AArch64BaseInfo.h"
960b57cec5SDimitry Andric #include "llvm/ADT/BitVector.h"
970b57cec5SDimitry Andric #include "llvm/ADT/SmallVector.h"
980b57cec5SDimitry Andric #include "llvm/CodeGen/MachineBasicBlock.h"
990b57cec5SDimitry Andric #include "llvm/CodeGen/MachineFunction.h"
1000b57cec5SDimitry Andric #include "llvm/CodeGen/MachineFunctionPass.h"
1010b57cec5SDimitry Andric #include "llvm/CodeGen/MachineInstr.h"
1020b57cec5SDimitry Andric #include "llvm/CodeGen/MachineInstrBuilder.h"
1030b57cec5SDimitry Andric #include "llvm/CodeGen/MachineOperand.h"
1040b57cec5SDimitry Andric #include "llvm/CodeGen/MachineRegisterInfo.h"
1050b57cec5SDimitry Andric #include "llvm/CodeGen/RegisterScavenging.h"
1060b57cec5SDimitry Andric #include "llvm/IR/DebugLoc.h"
1070b57cec5SDimitry Andric #include "llvm/Pass.h"
1080b57cec5SDimitry Andric #include "llvm/Support/CodeGen.h"
1098bcb0991SDimitry Andric #include "llvm/Support/Debug.h"
1100b57cec5SDimitry Andric #include "llvm/Target/TargetMachine.h"
1110b57cec5SDimitry Andric #include <cassert>
1120b57cec5SDimitry Andric
1130b57cec5SDimitry Andric using namespace llvm;
1140b57cec5SDimitry Andric
1150b57cec5SDimitry Andric #define DEBUG_TYPE "aarch64-speculation-hardening"
1160b57cec5SDimitry Andric
1170b57cec5SDimitry Andric #define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
1180b57cec5SDimitry Andric
1198bcb0991SDimitry Andric static cl::opt<bool> HardenLoads("aarch64-slh-loads", cl::Hidden,
1200b57cec5SDimitry Andric cl::desc("Sanitize loads from memory."),
1210b57cec5SDimitry Andric cl::init(true));
1220b57cec5SDimitry Andric
1230b57cec5SDimitry Andric namespace {
1240b57cec5SDimitry Andric
1250b57cec5SDimitry Andric class AArch64SpeculationHardening : public MachineFunctionPass {
1260b57cec5SDimitry Andric public:
1270b57cec5SDimitry Andric const TargetInstrInfo *TII;
1280b57cec5SDimitry Andric const TargetRegisterInfo *TRI;
1290b57cec5SDimitry Andric
1300b57cec5SDimitry Andric static char ID;
1310b57cec5SDimitry Andric
AArch64SpeculationHardening()1320b57cec5SDimitry Andric AArch64SpeculationHardening() : MachineFunctionPass(ID) {
1330b57cec5SDimitry Andric initializeAArch64SpeculationHardeningPass(*PassRegistry::getPassRegistry());
1340b57cec5SDimitry Andric }
1350b57cec5SDimitry Andric
1360b57cec5SDimitry Andric bool runOnMachineFunction(MachineFunction &Fn) override;
1370b57cec5SDimitry Andric
getPassName() const1380b57cec5SDimitry Andric StringRef getPassName() const override {
1390b57cec5SDimitry Andric return AARCH64_SPECULATION_HARDENING_NAME;
1400b57cec5SDimitry Andric }
1410b57cec5SDimitry Andric
1420b57cec5SDimitry Andric private:
1430b57cec5SDimitry Andric unsigned MisspeculatingTaintReg;
1440b57cec5SDimitry Andric unsigned MisspeculatingTaintReg32Bit;
1450b57cec5SDimitry Andric bool UseControlFlowSpeculationBarrier;
1460b57cec5SDimitry Andric BitVector RegsNeedingCSDBBeforeUse;
1470b57cec5SDimitry Andric BitVector RegsAlreadyMasked;
1480b57cec5SDimitry Andric
1490b57cec5SDimitry Andric bool functionUsesHardeningRegister(MachineFunction &MF) const;
1500b57cec5SDimitry Andric bool instrumentControlFlow(MachineBasicBlock &MBB,
1510b57cec5SDimitry Andric bool &UsesFullSpeculationBarrier);
1520b57cec5SDimitry Andric bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
1530b57cec5SDimitry Andric MachineBasicBlock *&FBB,
1540b57cec5SDimitry Andric AArch64CC::CondCode &CondCode) const;
1550b57cec5SDimitry Andric void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
1560b57cec5SDimitry Andric AArch64CC::CondCode &CondCode, DebugLoc DL) const;
1570b57cec5SDimitry Andric void insertSPToRegTaintPropagation(MachineBasicBlock &MBB,
1580b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI) const;
1590b57cec5SDimitry Andric void insertRegToSPTaintPropagation(MachineBasicBlock &MBB,
1600b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI,
1610b57cec5SDimitry Andric unsigned TmpReg) const;
1620b57cec5SDimitry Andric void insertFullSpeculationBarrier(MachineBasicBlock &MBB,
1630b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI,
1640b57cec5SDimitry Andric DebugLoc DL) const;
1650b57cec5SDimitry Andric
1660b57cec5SDimitry Andric bool slhLoads(MachineBasicBlock &MBB);
1670b57cec5SDimitry Andric bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,
1680b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI,
1690b57cec5SDimitry Andric MachineInstr &MI, unsigned Reg);
1700b57cec5SDimitry Andric bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB,
1710b57cec5SDimitry Andric bool UsesFullSpeculationBarrier);
1720b57cec5SDimitry Andric bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
1730b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI,
1740b57cec5SDimitry Andric bool UsesFullSpeculationBarrier);
1750b57cec5SDimitry Andric bool insertCSDB(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
1760b57cec5SDimitry Andric DebugLoc DL);
1770b57cec5SDimitry Andric };
1780b57cec5SDimitry Andric
1790b57cec5SDimitry Andric } // end anonymous namespace
1800b57cec5SDimitry Andric
1810b57cec5SDimitry Andric char AArch64SpeculationHardening::ID = 0;
1820b57cec5SDimitry Andric
1830b57cec5SDimitry Andric INITIALIZE_PASS(AArch64SpeculationHardening, "aarch64-speculation-hardening",
1840b57cec5SDimitry Andric AARCH64_SPECULATION_HARDENING_NAME, false, false)
1850b57cec5SDimitry Andric
endsWithCondControlFlow(MachineBasicBlock & MBB,MachineBasicBlock * & TBB,MachineBasicBlock * & FBB,AArch64CC::CondCode & CondCode) const1860b57cec5SDimitry Andric bool AArch64SpeculationHardening::endsWithCondControlFlow(
1870b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
1880b57cec5SDimitry Andric AArch64CC::CondCode &CondCode) const {
1890b57cec5SDimitry Andric SmallVector<MachineOperand, 1> analyzeBranchCondCode;
1900b57cec5SDimitry Andric if (TII->analyzeBranch(MBB, TBB, FBB, analyzeBranchCondCode, false))
1910b57cec5SDimitry Andric return false;
1920b57cec5SDimitry Andric
1930b57cec5SDimitry Andric // Ignore if the BB ends in an unconditional branch/fall-through.
1940b57cec5SDimitry Andric if (analyzeBranchCondCode.empty())
1950b57cec5SDimitry Andric return false;
1960b57cec5SDimitry Andric
1970b57cec5SDimitry Andric // If the BB ends with a single conditional branch, FBB will be set to
1980b57cec5SDimitry Andric // nullptr (see API docs for TII->analyzeBranch). For the rest of the
1990b57cec5SDimitry Andric // analysis we want the FBB block to be set always.
2000b57cec5SDimitry Andric assert(TBB != nullptr);
2010b57cec5SDimitry Andric if (FBB == nullptr)
2020b57cec5SDimitry Andric FBB = MBB.getFallThrough();
2030b57cec5SDimitry Andric
2040b57cec5SDimitry Andric // If both the true and the false condition jump to the same basic block,
2050b57cec5SDimitry Andric // there isn't need for any protection - whether the branch is speculated
2060b57cec5SDimitry Andric // correctly or not, we end up executing the architecturally correct code.
2070b57cec5SDimitry Andric if (TBB == FBB)
2080b57cec5SDimitry Andric return false;
2090b57cec5SDimitry Andric
2100b57cec5SDimitry Andric assert(MBB.succ_size() == 2);
2110b57cec5SDimitry Andric // translate analyzeBranchCondCode to CondCode.
2120b57cec5SDimitry Andric assert(analyzeBranchCondCode.size() == 1 && "unknown Cond array format");
2130b57cec5SDimitry Andric CondCode = AArch64CC::CondCode(analyzeBranchCondCode[0].getImm());
2140b57cec5SDimitry Andric return true;
2150b57cec5SDimitry Andric }
2160b57cec5SDimitry Andric
insertFullSpeculationBarrier(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI,DebugLoc DL) const2170b57cec5SDimitry Andric void AArch64SpeculationHardening::insertFullSpeculationBarrier(
2180b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
2190b57cec5SDimitry Andric DebugLoc DL) const {
2200b57cec5SDimitry Andric // A full control flow speculation barrier consists of (DSB SYS + ISB)
2210b57cec5SDimitry Andric BuildMI(MBB, MBBI, DL, TII->get(AArch64::DSB)).addImm(0xf);
2220b57cec5SDimitry Andric BuildMI(MBB, MBBI, DL, TII->get(AArch64::ISB)).addImm(0xf);
2230b57cec5SDimitry Andric }
2240b57cec5SDimitry Andric
insertTrackingCode(MachineBasicBlock & SplitEdgeBB,AArch64CC::CondCode & CondCode,DebugLoc DL) const2250b57cec5SDimitry Andric void AArch64SpeculationHardening::insertTrackingCode(
2260b57cec5SDimitry Andric MachineBasicBlock &SplitEdgeBB, AArch64CC::CondCode &CondCode,
2270b57cec5SDimitry Andric DebugLoc DL) const {
2280b57cec5SDimitry Andric if (UseControlFlowSpeculationBarrier) {
2290b57cec5SDimitry Andric insertFullSpeculationBarrier(SplitEdgeBB, SplitEdgeBB.begin(), DL);
2300b57cec5SDimitry Andric } else {
2310b57cec5SDimitry Andric BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::CSELXr))
2320b57cec5SDimitry Andric .addDef(MisspeculatingTaintReg)
2330b57cec5SDimitry Andric .addUse(MisspeculatingTaintReg)
2340b57cec5SDimitry Andric .addUse(AArch64::XZR)
2350b57cec5SDimitry Andric .addImm(CondCode);
2360b57cec5SDimitry Andric SplitEdgeBB.addLiveIn(AArch64::NZCV);
2370b57cec5SDimitry Andric }
2380b57cec5SDimitry Andric }
2390b57cec5SDimitry Andric
instrumentControlFlow(MachineBasicBlock & MBB,bool & UsesFullSpeculationBarrier)2400b57cec5SDimitry Andric bool AArch64SpeculationHardening::instrumentControlFlow(
2410b57cec5SDimitry Andric MachineBasicBlock &MBB, bool &UsesFullSpeculationBarrier) {
2420b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);
2430b57cec5SDimitry Andric
2440b57cec5SDimitry Andric bool Modified = false;
2450b57cec5SDimitry Andric MachineBasicBlock *TBB = nullptr;
2460b57cec5SDimitry Andric MachineBasicBlock *FBB = nullptr;
2470b57cec5SDimitry Andric AArch64CC::CondCode CondCode;
2480b57cec5SDimitry Andric
2490b57cec5SDimitry Andric if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {
2500b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "... doesn't end with CondControlFlow\n");
2510b57cec5SDimitry Andric } else {
2520b57cec5SDimitry Andric // Now insert:
2530b57cec5SDimitry Andric // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, cond" on the True edge and
2540b57cec5SDimitry Andric // "CSEL MisSpeculatingR, MisSpeculatingR, XZR, Invertcond" on the False
2550b57cec5SDimitry Andric // edge.
2560b57cec5SDimitry Andric AArch64CC::CondCode InvCondCode = AArch64CC::getInvertedCondCode(CondCode);
2570b57cec5SDimitry Andric
2580b57cec5SDimitry Andric MachineBasicBlock *SplitEdgeTBB = MBB.SplitCriticalEdge(TBB, *this);
2590b57cec5SDimitry Andric MachineBasicBlock *SplitEdgeFBB = MBB.SplitCriticalEdge(FBB, *this);
2600b57cec5SDimitry Andric
2610b57cec5SDimitry Andric assert(SplitEdgeTBB != nullptr);
2620b57cec5SDimitry Andric assert(SplitEdgeFBB != nullptr);
2630b57cec5SDimitry Andric
2640b57cec5SDimitry Andric DebugLoc DL;
2650b57cec5SDimitry Andric if (MBB.instr_end() != MBB.instr_begin())
2660b57cec5SDimitry Andric DL = (--MBB.instr_end())->getDebugLoc();
2670b57cec5SDimitry Andric
2680b57cec5SDimitry Andric insertTrackingCode(*SplitEdgeTBB, CondCode, DL);
2690b57cec5SDimitry Andric insertTrackingCode(*SplitEdgeFBB, InvCondCode, DL);
2700b57cec5SDimitry Andric
2710b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "SplitEdgeTBB: " << *SplitEdgeTBB << "\n");
2720b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "SplitEdgeFBB: " << *SplitEdgeFBB << "\n");
2730b57cec5SDimitry Andric Modified = true;
2740b57cec5SDimitry Andric }
2750b57cec5SDimitry Andric
2760b57cec5SDimitry Andric // Perform correct code generation around function calls and before returns.
2770b57cec5SDimitry Andric // The below variables record the return/terminator instructions and the call
2780b57cec5SDimitry Andric // instructions respectively; including which register is available as a
2790b57cec5SDimitry Andric // temporary register just before the recorded instructions.
2800b57cec5SDimitry Andric SmallVector<std::pair<MachineInstr *, unsigned>, 4> ReturnInstructions;
2810b57cec5SDimitry Andric SmallVector<std::pair<MachineInstr *, unsigned>, 4> CallInstructions;
2820b57cec5SDimitry Andric // if a temporary register is not available for at least one of the
2830b57cec5SDimitry Andric // instructions for which we need to transfer taint to the stack pointer, we
2840b57cec5SDimitry Andric // need to insert a full speculation barrier.
2850b57cec5SDimitry Andric // TmpRegisterNotAvailableEverywhere tracks that condition.
2860b57cec5SDimitry Andric bool TmpRegisterNotAvailableEverywhere = false;
2870b57cec5SDimitry Andric
2880b57cec5SDimitry Andric RegScavenger RS;
28906c3fb27SDimitry Andric RS.enterBasicBlockEnd(MBB);
2900b57cec5SDimitry Andric
29106c3fb27SDimitry Andric for (MachineBasicBlock::iterator I = MBB.end(); I != MBB.begin(); ) {
29206c3fb27SDimitry Andric MachineInstr &MI = *--I;
2930b57cec5SDimitry Andric if (!MI.isReturn() && !MI.isCall())
2940b57cec5SDimitry Andric continue;
2950b57cec5SDimitry Andric
2960b57cec5SDimitry Andric // The RegScavenger represents registers available *after* the MI
2970b57cec5SDimitry Andric // instruction pointed to by RS.getCurrentPosition().
2980b57cec5SDimitry Andric // We need to have a register that is available *before* the MI is executed.
29906c3fb27SDimitry Andric if (I == MBB.begin())
30006c3fb27SDimitry Andric RS.enterBasicBlock(MBB);
30106c3fb27SDimitry Andric else
3025f757f3fSDimitry Andric RS.backward(I);
3030b57cec5SDimitry Andric // FIXME: The below just finds *a* unused register. Maybe code could be
3040b57cec5SDimitry Andric // optimized more if this looks for the register that isn't used for the
3050b57cec5SDimitry Andric // longest time around this place, to enable more scheduling freedom. Not
3060b57cec5SDimitry Andric // sure if that would actually result in a big performance difference
3070b57cec5SDimitry Andric // though. Maybe RegisterScavenger::findSurvivorBackwards has some logic
3080b57cec5SDimitry Andric // already to do this - but it's unclear if that could easily be used here.
30904eeddc0SDimitry Andric Register TmpReg = RS.FindUnusedReg(&AArch64::GPR64commonRegClass);
3100b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "RS finds "
3110b57cec5SDimitry Andric << ((TmpReg == 0) ? "no register " : "register ");
3120b57cec5SDimitry Andric if (TmpReg != 0) dbgs() << printReg(TmpReg, TRI) << " ";
3130b57cec5SDimitry Andric dbgs() << "to be available at MI " << MI);
3140b57cec5SDimitry Andric if (TmpReg == 0)
3150b57cec5SDimitry Andric TmpRegisterNotAvailableEverywhere = true;
3160b57cec5SDimitry Andric if (MI.isReturn())
3170b57cec5SDimitry Andric ReturnInstructions.push_back({&MI, TmpReg});
3180b57cec5SDimitry Andric else if (MI.isCall())
3190b57cec5SDimitry Andric CallInstructions.push_back({&MI, TmpReg});
3200b57cec5SDimitry Andric }
3210b57cec5SDimitry Andric
3220b57cec5SDimitry Andric if (TmpRegisterNotAvailableEverywhere) {
3230b57cec5SDimitry Andric // When a temporary register is not available everywhere in this basic
3240b57cec5SDimitry Andric // basic block where a propagate-taint-to-sp operation is needed, just
3250b57cec5SDimitry Andric // emit a full speculation barrier at the start of this basic block, which
3260b57cec5SDimitry Andric // renders the taint/speculation tracking in this basic block unnecessary.
3270b57cec5SDimitry Andric insertFullSpeculationBarrier(MBB, MBB.begin(),
3280b57cec5SDimitry Andric (MBB.begin())->getDebugLoc());
3290b57cec5SDimitry Andric UsesFullSpeculationBarrier = true;
3300b57cec5SDimitry Andric Modified = true;
3310b57cec5SDimitry Andric } else {
3320b57cec5SDimitry Andric for (auto MI_Reg : ReturnInstructions) {
3330b57cec5SDimitry Andric assert(MI_Reg.second != 0);
3340b57cec5SDimitry Andric LLVM_DEBUG(
3350b57cec5SDimitry Andric dbgs()
3360b57cec5SDimitry Andric << " About to insert Reg to SP taint propagation with temp register "
3370b57cec5SDimitry Andric << printReg(MI_Reg.second, TRI)
3380b57cec5SDimitry Andric << " on instruction: " << *MI_Reg.first);
3390b57cec5SDimitry Andric insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
3400b57cec5SDimitry Andric Modified = true;
3410b57cec5SDimitry Andric }
3420b57cec5SDimitry Andric
3430b57cec5SDimitry Andric for (auto MI_Reg : CallInstructions) {
3440b57cec5SDimitry Andric assert(MI_Reg.second != 0);
3450b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << " About to insert Reg to SP and back taint "
3460b57cec5SDimitry Andric "propagation with temp register "
3470b57cec5SDimitry Andric << printReg(MI_Reg.second, TRI)
3480b57cec5SDimitry Andric << " around instruction: " << *MI_Reg.first);
3490b57cec5SDimitry Andric // Just after the call:
3500b57cec5SDimitry Andric insertSPToRegTaintPropagation(
3510b57cec5SDimitry Andric MBB, std::next((MachineBasicBlock::iterator)MI_Reg.first));
3520b57cec5SDimitry Andric // Just before the call:
3530b57cec5SDimitry Andric insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
3540b57cec5SDimitry Andric Modified = true;
3550b57cec5SDimitry Andric }
3560b57cec5SDimitry Andric }
3570b57cec5SDimitry Andric return Modified;
3580b57cec5SDimitry Andric }
3590b57cec5SDimitry Andric
insertSPToRegTaintPropagation(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI) const3600b57cec5SDimitry Andric void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
3610b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const {
3620b57cec5SDimitry Andric // If full control flow speculation barriers are used, emit a control flow
3630b57cec5SDimitry Andric // barrier to block potential miss-speculation in flight coming in to this
3640b57cec5SDimitry Andric // function.
3650b57cec5SDimitry Andric if (UseControlFlowSpeculationBarrier) {
3660b57cec5SDimitry Andric insertFullSpeculationBarrier(MBB, MBBI, DebugLoc());
3670b57cec5SDimitry Andric return;
3680b57cec5SDimitry Andric }
3690b57cec5SDimitry Andric
3700b57cec5SDimitry Andric // CMP SP, #0 === SUBS xzr, SP, #0
3710b57cec5SDimitry Andric BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
3720b57cec5SDimitry Andric .addDef(AArch64::XZR)
3730b57cec5SDimitry Andric .addUse(AArch64::SP)
3740b57cec5SDimitry Andric .addImm(0)
3750b57cec5SDimitry Andric .addImm(0); // no shift
3760b57cec5SDimitry Andric // CSETM x16, NE === CSINV x16, xzr, xzr, EQ
3770b57cec5SDimitry Andric BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
3780b57cec5SDimitry Andric .addDef(MisspeculatingTaintReg)
3790b57cec5SDimitry Andric .addUse(AArch64::XZR)
3800b57cec5SDimitry Andric .addUse(AArch64::XZR)
3810b57cec5SDimitry Andric .addImm(AArch64CC::EQ);
3820b57cec5SDimitry Andric }
3830b57cec5SDimitry Andric
insertRegToSPTaintPropagation(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI,unsigned TmpReg) const3840b57cec5SDimitry Andric void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
3850b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
3860b57cec5SDimitry Andric unsigned TmpReg) const {
3870b57cec5SDimitry Andric // If full control flow speculation barriers are used, there will not be
3880b57cec5SDimitry Andric // miss-speculation when returning from this function, and therefore, also
3890b57cec5SDimitry Andric // no need to encode potential miss-speculation into the stack pointer.
3900b57cec5SDimitry Andric if (UseControlFlowSpeculationBarrier)
3910b57cec5SDimitry Andric return;
3920b57cec5SDimitry Andric
3930b57cec5SDimitry Andric // mov Xtmp, SP === ADD Xtmp, SP, #0
3940b57cec5SDimitry Andric BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
3950b57cec5SDimitry Andric .addDef(TmpReg)
3960b57cec5SDimitry Andric .addUse(AArch64::SP)
3970b57cec5SDimitry Andric .addImm(0)
3980b57cec5SDimitry Andric .addImm(0); // no shift
3990b57cec5SDimitry Andric // and Xtmp, Xtmp, TaintReg === AND Xtmp, Xtmp, TaintReg, #0
4000b57cec5SDimitry Andric BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
4010b57cec5SDimitry Andric .addDef(TmpReg, RegState::Renamable)
4020b57cec5SDimitry Andric .addUse(TmpReg, RegState::Kill | RegState::Renamable)
4030b57cec5SDimitry Andric .addUse(MisspeculatingTaintReg, RegState::Kill)
4040b57cec5SDimitry Andric .addImm(0);
4050b57cec5SDimitry Andric // mov SP, Xtmp === ADD SP, Xtmp, #0
4060b57cec5SDimitry Andric BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
4070b57cec5SDimitry Andric .addDef(AArch64::SP)
4080b57cec5SDimitry Andric .addUse(TmpReg, RegState::Kill)
4090b57cec5SDimitry Andric .addImm(0)
4100b57cec5SDimitry Andric .addImm(0); // no shift
4110b57cec5SDimitry Andric }
4120b57cec5SDimitry Andric
functionUsesHardeningRegister(MachineFunction & MF) const4130b57cec5SDimitry Andric bool AArch64SpeculationHardening::functionUsesHardeningRegister(
4140b57cec5SDimitry Andric MachineFunction &MF) const {
4150b57cec5SDimitry Andric for (MachineBasicBlock &MBB : MF) {
4160b57cec5SDimitry Andric for (MachineInstr &MI : MBB) {
4170b57cec5SDimitry Andric // treat function calls specially, as the hardening register does not
4180b57cec5SDimitry Andric // need to remain live across function calls.
4190b57cec5SDimitry Andric if (MI.isCall())
4200b57cec5SDimitry Andric continue;
4210b57cec5SDimitry Andric if (MI.readsRegister(MisspeculatingTaintReg, TRI) ||
4220b57cec5SDimitry Andric MI.modifiesRegister(MisspeculatingTaintReg, TRI))
4230b57cec5SDimitry Andric return true;
4240b57cec5SDimitry Andric }
4250b57cec5SDimitry Andric }
4260b57cec5SDimitry Andric return false;
4270b57cec5SDimitry Andric }
4280b57cec5SDimitry Andric
4290b57cec5SDimitry Andric // Make GPR register Reg speculation-safe by putting it through the
4300b57cec5SDimitry Andric // SpeculationSafeValue pseudo instruction, if we can't prove that
4310b57cec5SDimitry Andric // the value in the register has already been hardened.
makeGPRSpeculationSafe(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI,MachineInstr & MI,unsigned Reg)4320b57cec5SDimitry Andric bool AArch64SpeculationHardening::makeGPRSpeculationSafe(
4330b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, MachineInstr &MI,
4340b57cec5SDimitry Andric unsigned Reg) {
4350b57cec5SDimitry Andric assert(AArch64::GPR32allRegClass.contains(Reg) ||
4360b57cec5SDimitry Andric AArch64::GPR64allRegClass.contains(Reg));
4370b57cec5SDimitry Andric
4380b57cec5SDimitry Andric // Loads cannot directly load a value into the SP (nor WSP).
4390b57cec5SDimitry Andric // Therefore, if Reg is SP or WSP, it is because the instruction loads from
4400b57cec5SDimitry Andric // the stack through the stack pointer.
4410b57cec5SDimitry Andric //
4420b57cec5SDimitry Andric // Since the stack pointer is never dynamically controllable, don't harden it.
4430b57cec5SDimitry Andric if (Reg == AArch64::SP || Reg == AArch64::WSP)
4440b57cec5SDimitry Andric return false;
4450b57cec5SDimitry Andric
4460b57cec5SDimitry Andric // Do not harden the register again if already hardened before.
4470b57cec5SDimitry Andric if (RegsAlreadyMasked[Reg])
4480b57cec5SDimitry Andric return false;
4490b57cec5SDimitry Andric
4500b57cec5SDimitry Andric const bool Is64Bit = AArch64::GPR64allRegClass.contains(Reg);
4510b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "About to harden register : " << Reg << "\n");
4520b57cec5SDimitry Andric BuildMI(MBB, MBBI, MI.getDebugLoc(),
4530b57cec5SDimitry Andric TII->get(Is64Bit ? AArch64::SpeculationSafeValueX
4540b57cec5SDimitry Andric : AArch64::SpeculationSafeValueW))
4550b57cec5SDimitry Andric .addDef(Reg)
4560b57cec5SDimitry Andric .addUse(Reg);
4570b57cec5SDimitry Andric RegsAlreadyMasked.set(Reg);
4580b57cec5SDimitry Andric return true;
4590b57cec5SDimitry Andric }
4600b57cec5SDimitry Andric
slhLoads(MachineBasicBlock & MBB)4610b57cec5SDimitry Andric bool AArch64SpeculationHardening::slhLoads(MachineBasicBlock &MBB) {
4620b57cec5SDimitry Andric bool Modified = false;
4630b57cec5SDimitry Andric
4640b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "slhLoads running on MBB: " << MBB);
4650b57cec5SDimitry Andric
4660b57cec5SDimitry Andric RegsAlreadyMasked.reset();
4670b57cec5SDimitry Andric
4680b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
4690b57cec5SDimitry Andric MachineBasicBlock::iterator NextMBBI;
4700b57cec5SDimitry Andric for (; MBBI != E; MBBI = NextMBBI) {
4710b57cec5SDimitry Andric MachineInstr &MI = *MBBI;
4720b57cec5SDimitry Andric NextMBBI = std::next(MBBI);
4730b57cec5SDimitry Andric // Only harden loaded values or addresses used in loads.
4740b57cec5SDimitry Andric if (!MI.mayLoad())
4750b57cec5SDimitry Andric continue;
4760b57cec5SDimitry Andric
4770b57cec5SDimitry Andric LLVM_DEBUG(dbgs() << "About to harden: " << MI);
4780b57cec5SDimitry Andric
4790b57cec5SDimitry Andric // For general purpose register loads, harden the registers loaded into.
4800b57cec5SDimitry Andric // For other loads, harden the address loaded from.
4810b57cec5SDimitry Andric // Masking the loaded value is expected to result in less performance
4820b57cec5SDimitry Andric // overhead, as the load can still execute speculatively in comparison to
4830b57cec5SDimitry Andric // when the address loaded from gets masked. However, masking is only
4840b57cec5SDimitry Andric // easy to do efficiently on GPR registers, so for loads into non-GPR
4850b57cec5SDimitry Andric // registers (e.g. floating point loads), mask the address loaded from.
4860b57cec5SDimitry Andric bool AllDefsAreGPR = llvm::all_of(MI.defs(), [&](MachineOperand &Op) {
4870b57cec5SDimitry Andric return Op.isReg() && (AArch64::GPR32allRegClass.contains(Op.getReg()) ||
4880b57cec5SDimitry Andric AArch64::GPR64allRegClass.contains(Op.getReg()));
4890b57cec5SDimitry Andric });
4900b57cec5SDimitry Andric // FIXME: it might be a worthwhile optimization to not mask loaded
4910b57cec5SDimitry Andric // values if all the registers involved in address calculation are already
4920b57cec5SDimitry Andric // hardened, leading to this load not able to execute on a miss-speculated
4930b57cec5SDimitry Andric // path.
4940b57cec5SDimitry Andric bool HardenLoadedData = AllDefsAreGPR;
4950b57cec5SDimitry Andric bool HardenAddressLoadedFrom = !HardenLoadedData;
4960b57cec5SDimitry Andric
4970b57cec5SDimitry Andric // First remove registers from AlreadyMaskedRegisters if their value is
4980b57cec5SDimitry Andric // updated by this instruction - it makes them contain a new value that is
4990b57cec5SDimitry Andric // not guaranteed to already have been masked.
5000b57cec5SDimitry Andric for (MachineOperand Op : MI.defs())
5010b57cec5SDimitry Andric for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)
5020b57cec5SDimitry Andric RegsAlreadyMasked.reset(*AI);
5030b57cec5SDimitry Andric
5040b57cec5SDimitry Andric // FIXME: loads from the stack with an immediate offset from the stack
5050b57cec5SDimitry Andric // pointer probably shouldn't be hardened, which could result in a
5060b57cec5SDimitry Andric // significant optimization. See section "Don’t check loads from
5070b57cec5SDimitry Andric // compile-time constant stack offsets", in
5080b57cec5SDimitry Andric // https://llvm.org/docs/SpeculativeLoadHardening.html
5090b57cec5SDimitry Andric
5100b57cec5SDimitry Andric if (HardenLoadedData)
5110b57cec5SDimitry Andric for (auto Def : MI.defs()) {
5120b57cec5SDimitry Andric if (Def.isDead())
5130b57cec5SDimitry Andric // Do not mask a register that is not used further.
5140b57cec5SDimitry Andric continue;
5150b57cec5SDimitry Andric // FIXME: For pre/post-increment addressing modes, the base register
5160b57cec5SDimitry Andric // used in address calculation is also defined by this instruction.
5170b57cec5SDimitry Andric // It might be a worthwhile optimization to not harden that
5180b57cec5SDimitry Andric // base register increment/decrement when the increment/decrement is
5190b57cec5SDimitry Andric // an immediate.
5200b57cec5SDimitry Andric Modified |= makeGPRSpeculationSafe(MBB, NextMBBI, MI, Def.getReg());
5210b57cec5SDimitry Andric }
5220b57cec5SDimitry Andric
5230b57cec5SDimitry Andric if (HardenAddressLoadedFrom)
5240b57cec5SDimitry Andric for (auto Use : MI.uses()) {
5250b57cec5SDimitry Andric if (!Use.isReg())
5260b57cec5SDimitry Andric continue;
5278bcb0991SDimitry Andric Register Reg = Use.getReg();
5280b57cec5SDimitry Andric // Some loads of floating point data have implicit defs/uses on a
5290b57cec5SDimitry Andric // super register of that floating point data. Some examples:
5300b57cec5SDimitry Andric // $s0 = LDRSui $sp, 22, implicit-def $q0
5310b57cec5SDimitry Andric // $q0 = LD1i64 $q0, 1, renamable $x0
5320b57cec5SDimitry Andric // We need to filter out these uses for non-GPR register which occur
5330b57cec5SDimitry Andric // because the load partially fills a non-GPR register with the loaded
5340b57cec5SDimitry Andric // data. Just skipping all non-GPR registers is safe (for now) as all
5350b57cec5SDimitry Andric // AArch64 load instructions only use GPR registers to perform the
5360b57cec5SDimitry Andric // address calculation. FIXME: However that might change once we can
5370b57cec5SDimitry Andric // produce SVE gather instructions.
5380b57cec5SDimitry Andric if (!(AArch64::GPR32allRegClass.contains(Reg) ||
5390b57cec5SDimitry Andric AArch64::GPR64allRegClass.contains(Reg)))
5400b57cec5SDimitry Andric continue;
5410b57cec5SDimitry Andric Modified |= makeGPRSpeculationSafe(MBB, MBBI, MI, Reg);
5420b57cec5SDimitry Andric }
5430b57cec5SDimitry Andric }
5440b57cec5SDimitry Andric return Modified;
5450b57cec5SDimitry Andric }
5460b57cec5SDimitry Andric
5470b57cec5SDimitry Andric /// \brief If MBBI references a pseudo instruction that should be expanded
5480b57cec5SDimitry Andric /// here, do the expansion and return true. Otherwise return false.
expandSpeculationSafeValue(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI,bool UsesFullSpeculationBarrier)5490b57cec5SDimitry Andric bool AArch64SpeculationHardening::expandSpeculationSafeValue(
5500b57cec5SDimitry Andric MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
5510b57cec5SDimitry Andric bool UsesFullSpeculationBarrier) {
5520b57cec5SDimitry Andric MachineInstr &MI = *MBBI;
5530b57cec5SDimitry Andric unsigned Opcode = MI.getOpcode();
5540b57cec5SDimitry Andric bool Is64Bit = true;
5550b57cec5SDimitry Andric
5560b57cec5SDimitry Andric switch (Opcode) {
5570b57cec5SDimitry Andric default:
5580b57cec5SDimitry Andric break;
5590b57cec5SDimitry Andric case AArch64::SpeculationSafeValueW:
5600b57cec5SDimitry Andric Is64Bit = false;
561bdd1243dSDimitry Andric [[fallthrough]];
5620b57cec5SDimitry Andric case AArch64::SpeculationSafeValueX:
5630b57cec5SDimitry Andric // Just remove the SpeculationSafe pseudo's if control flow
5640b57cec5SDimitry Andric // miss-speculation isn't happening because we're already inserting barriers
5650b57cec5SDimitry Andric // to guarantee that.
5660b57cec5SDimitry Andric if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {
5678bcb0991SDimitry Andric Register DstReg = MI.getOperand(0).getReg();
5688bcb0991SDimitry Andric Register SrcReg = MI.getOperand(1).getReg();
5690b57cec5SDimitry Andric // Mark this register and all its aliasing registers as needing to be
5700b57cec5SDimitry Andric // value speculation hardened before its next use, by using a CSDB
5710b57cec5SDimitry Andric // barrier instruction.
5720b57cec5SDimitry Andric for (MachineOperand Op : MI.defs())
5730b57cec5SDimitry Andric for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)
5740b57cec5SDimitry Andric RegsNeedingCSDBBeforeUse.set(*AI);
5750b57cec5SDimitry Andric
5760b57cec5SDimitry Andric // Mask off with taint state.
5770b57cec5SDimitry Andric BuildMI(MBB, MBBI, MI.getDebugLoc(),
5780b57cec5SDimitry Andric Is64Bit ? TII->get(AArch64::ANDXrs) : TII->get(AArch64::ANDWrs))
5790b57cec5SDimitry Andric .addDef(DstReg)
5800b57cec5SDimitry Andric .addUse(SrcReg, RegState::Kill)
5810b57cec5SDimitry Andric .addUse(Is64Bit ? MisspeculatingTaintReg
5820b57cec5SDimitry Andric : MisspeculatingTaintReg32Bit)
5830b57cec5SDimitry Andric .addImm(0);
5840b57cec5SDimitry Andric }
5850b57cec5SDimitry Andric MI.eraseFromParent();
5860b57cec5SDimitry Andric return true;
5870b57cec5SDimitry Andric }
5880b57cec5SDimitry Andric return false;
5890b57cec5SDimitry Andric }
5900b57cec5SDimitry Andric
insertCSDB(MachineBasicBlock & MBB,MachineBasicBlock::iterator MBBI,DebugLoc DL)5910b57cec5SDimitry Andric bool AArch64SpeculationHardening::insertCSDB(MachineBasicBlock &MBB,
5920b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI,
5930b57cec5SDimitry Andric DebugLoc DL) {
5940b57cec5SDimitry Andric assert(!UseControlFlowSpeculationBarrier && "No need to insert CSDBs when "
5950b57cec5SDimitry Andric "control flow miss-speculation "
5960b57cec5SDimitry Andric "is already blocked");
5970b57cec5SDimitry Andric // insert data value speculation barrier (CSDB)
5980b57cec5SDimitry Andric BuildMI(MBB, MBBI, DL, TII->get(AArch64::HINT)).addImm(0x14);
5990b57cec5SDimitry Andric RegsNeedingCSDBBeforeUse.reset();
6000b57cec5SDimitry Andric return true;
6010b57cec5SDimitry Andric }
6020b57cec5SDimitry Andric
lowerSpeculationSafeValuePseudos(MachineBasicBlock & MBB,bool UsesFullSpeculationBarrier)6030b57cec5SDimitry Andric bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
6040b57cec5SDimitry Andric MachineBasicBlock &MBB, bool UsesFullSpeculationBarrier) {
6050b57cec5SDimitry Andric bool Modified = false;
6060b57cec5SDimitry Andric
6070b57cec5SDimitry Andric RegsNeedingCSDBBeforeUse.reset();
6080b57cec5SDimitry Andric
6090b57cec5SDimitry Andric // The following loop iterates over all instructions in the basic block,
6100b57cec5SDimitry Andric // and performs 2 operations:
6110b57cec5SDimitry Andric // 1. Insert a CSDB at this location if needed.
6120b57cec5SDimitry Andric // 2. Expand the SpeculationSafeValuePseudo if the current instruction is
6130b57cec5SDimitry Andric // one.
6140b57cec5SDimitry Andric //
6150b57cec5SDimitry Andric // The insertion of the CSDB is done as late as possible (i.e. just before
6160b57cec5SDimitry Andric // the use of a masked register), in the hope that that will reduce the
6170b57cec5SDimitry Andric // total number of CSDBs in a block when there are multiple masked registers
6180b57cec5SDimitry Andric // in the block.
6190b57cec5SDimitry Andric MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
6200b57cec5SDimitry Andric DebugLoc DL;
6210b57cec5SDimitry Andric while (MBBI != E) {
6220b57cec5SDimitry Andric MachineInstr &MI = *MBBI;
6230b57cec5SDimitry Andric DL = MI.getDebugLoc();
6240b57cec5SDimitry Andric MachineBasicBlock::iterator NMBBI = std::next(MBBI);
6250b57cec5SDimitry Andric
6260b57cec5SDimitry Andric // First check if a CSDB needs to be inserted due to earlier registers
6270b57cec5SDimitry Andric // that were masked and that are used by the next instruction.
6280b57cec5SDimitry Andric // Also emit the barrier on any potential control flow changes.
6290b57cec5SDimitry Andric bool NeedToEmitBarrier = false;
6300b57cec5SDimitry Andric if (RegsNeedingCSDBBeforeUse.any() && (MI.isCall() || MI.isTerminator()))
6310b57cec5SDimitry Andric NeedToEmitBarrier = true;
6320b57cec5SDimitry Andric if (!NeedToEmitBarrier)
6330b57cec5SDimitry Andric for (MachineOperand Op : MI.uses())
6340b57cec5SDimitry Andric if (Op.isReg() && RegsNeedingCSDBBeforeUse[Op.getReg()]) {
6350b57cec5SDimitry Andric NeedToEmitBarrier = true;
6360b57cec5SDimitry Andric break;
6370b57cec5SDimitry Andric }
6380b57cec5SDimitry Andric
6390b57cec5SDimitry Andric if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)
6400b57cec5SDimitry Andric Modified |= insertCSDB(MBB, MBBI, DL);
6410b57cec5SDimitry Andric
6420b57cec5SDimitry Andric Modified |=
6430b57cec5SDimitry Andric expandSpeculationSafeValue(MBB, MBBI, UsesFullSpeculationBarrier);
6440b57cec5SDimitry Andric
6450b57cec5SDimitry Andric MBBI = NMBBI;
6460b57cec5SDimitry Andric }
6470b57cec5SDimitry Andric
6480b57cec5SDimitry Andric if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)
6490b57cec5SDimitry Andric Modified |= insertCSDB(MBB, MBBI, DL);
6500b57cec5SDimitry Andric
6510b57cec5SDimitry Andric return Modified;
6520b57cec5SDimitry Andric }
6530b57cec5SDimitry Andric
runOnMachineFunction(MachineFunction & MF)6540b57cec5SDimitry Andric bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
6550b57cec5SDimitry Andric if (!MF.getFunction().hasFnAttribute(Attribute::SpeculativeLoadHardening))
6560b57cec5SDimitry Andric return false;
6570b57cec5SDimitry Andric
6580b57cec5SDimitry Andric MisspeculatingTaintReg = AArch64::X16;
6590b57cec5SDimitry Andric MisspeculatingTaintReg32Bit = AArch64::W16;
6600b57cec5SDimitry Andric TII = MF.getSubtarget().getInstrInfo();
6610b57cec5SDimitry Andric TRI = MF.getSubtarget().getRegisterInfo();
6620b57cec5SDimitry Andric RegsNeedingCSDBBeforeUse.resize(TRI->getNumRegs());
6630b57cec5SDimitry Andric RegsAlreadyMasked.resize(TRI->getNumRegs());
6640b57cec5SDimitry Andric UseControlFlowSpeculationBarrier = functionUsesHardeningRegister(MF);
6650b57cec5SDimitry Andric
6660b57cec5SDimitry Andric bool Modified = false;
6670b57cec5SDimitry Andric
6680b57cec5SDimitry Andric // Step 1: Enable automatic insertion of SpeculationSafeValue.
6690b57cec5SDimitry Andric if (HardenLoads) {
6700b57cec5SDimitry Andric LLVM_DEBUG(
6710b57cec5SDimitry Andric dbgs() << "***** AArch64SpeculationHardening - automatic insertion of "
6720b57cec5SDimitry Andric "SpeculationSafeValue intrinsics *****\n");
6730b57cec5SDimitry Andric for (auto &MBB : MF)
6740b57cec5SDimitry Andric Modified |= slhLoads(MBB);
6750b57cec5SDimitry Andric }
6760b57cec5SDimitry Andric
6770b57cec5SDimitry Andric // 2. Add instrumentation code to function entry and exits.
6780b57cec5SDimitry Andric LLVM_DEBUG(
6790b57cec5SDimitry Andric dbgs()
6800b57cec5SDimitry Andric << "***** AArch64SpeculationHardening - track control flow *****\n");
6810b57cec5SDimitry Andric
6820b57cec5SDimitry Andric SmallVector<MachineBasicBlock *, 2> EntryBlocks;
6830b57cec5SDimitry Andric EntryBlocks.push_back(&MF.front());
6840b57cec5SDimitry Andric for (const LandingPadInfo &LPI : MF.getLandingPads())
6850b57cec5SDimitry Andric EntryBlocks.push_back(LPI.LandingPadBlock);
686bdd1243dSDimitry Andric for (auto *Entry : EntryBlocks)
6870b57cec5SDimitry Andric insertSPToRegTaintPropagation(
6880b57cec5SDimitry Andric *Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
6890b57cec5SDimitry Andric
6900b57cec5SDimitry Andric // 3. Add instrumentation code to every basic block.
6910b57cec5SDimitry Andric for (auto &MBB : MF) {
6920b57cec5SDimitry Andric bool UsesFullSpeculationBarrier = false;
6930b57cec5SDimitry Andric Modified |= instrumentControlFlow(MBB, UsesFullSpeculationBarrier);
6940b57cec5SDimitry Andric Modified |=
6950b57cec5SDimitry Andric lowerSpeculationSafeValuePseudos(MBB, UsesFullSpeculationBarrier);
6960b57cec5SDimitry Andric }
6970b57cec5SDimitry Andric
6980b57cec5SDimitry Andric return Modified;
6990b57cec5SDimitry Andric }
7000b57cec5SDimitry Andric
7010b57cec5SDimitry Andric /// \brief Returns an instance of the pseudo instruction expansion pass.
createAArch64SpeculationHardeningPass()7020b57cec5SDimitry Andric FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
7030b57cec5SDimitry Andric return new AArch64SpeculationHardening();
7040b57cec5SDimitry Andric }
705