113 using namespace llvm;
115 #define DEBUG_TYPE "aarch64-speculation-hardening"
117 #define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
120 cl::desc(
"Sanitize loads from memory."),
143 unsigned MisspeculatingTaintReg;
144 unsigned MisspeculatingTaintReg32Bit;
145 bool UseControlFlowSpeculationBarrier;
151 bool &UsesFullSpeculationBarrier);
161 unsigned TmpReg)
const;
171 bool UsesFullSpeculationBarrier);
174 bool UsesFullSpeculationBarrier);
186 bool AArch64SpeculationHardening::endsWithCondControlFlow(
194 if (analyzeBranchCondCode.empty())
212 assert(analyzeBranchCondCode.size() == 1 &&
"unknown Cond array format");
217 void AArch64SpeculationHardening::insertFullSpeculationBarrier(
225 void AArch64SpeculationHardening::insertTrackingCode(
228 if (UseControlFlowSpeculationBarrier) {
229 insertFullSpeculationBarrier(SplitEdgeBB, SplitEdgeBB.
begin(),
DL);
232 .
addDef(MisspeculatingTaintReg)
233 .
addUse(MisspeculatingTaintReg)
240 bool AArch64SpeculationHardening::instrumentControlFlow(
261 assert(SplitEdgeTBB !=
nullptr);
262 assert(SplitEdgeFBB !=
nullptr);
268 insertTrackingCode(*SplitEdgeTBB,
CondCode,
DL);
269 insertTrackingCode(*SplitEdgeFBB, InvCondCode,
DL);
286 bool TmpRegisterNotAvailableEverywhere =
false;
293 if (!
MI.isReturn() && !
MI.isCall())
309 << ((TmpReg == 0) ?
"no register " :
"register ");
311 dbgs() <<
"to be available at MI " <<
MI);
313 TmpRegisterNotAvailableEverywhere =
true;
315 ReturnInstructions.push_back({&
MI, TmpReg});
316 else if (
MI.isCall())
317 CallInstructions.push_back({&
MI, TmpReg});
320 if (TmpRegisterNotAvailableEverywhere) {
327 UsesFullSpeculationBarrier =
true;
330 for (
auto MI_Reg : ReturnInstructions) {
331 assert(MI_Reg.second != 0);
334 <<
" About to insert Reg to SP taint propagation with temp register "
336 <<
" on instruction: " << *MI_Reg.first);
337 insertRegToSPTaintPropagation(
MBB, MI_Reg.first, MI_Reg.second);
341 for (
auto MI_Reg : CallInstructions) {
342 assert(MI_Reg.second != 0);
344 "propagation with temp register "
346 <<
" around instruction: " << *MI_Reg.first);
348 insertSPToRegTaintPropagation(
351 insertRegToSPTaintPropagation(
MBB, MI_Reg.first, MI_Reg.second);
358 void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
363 if (UseControlFlowSpeculationBarrier) {
376 .
addDef(MisspeculatingTaintReg)
382 void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
384 unsigned TmpReg)
const {
388 if (UseControlFlowSpeculationBarrier)
411 bool AArch64SpeculationHardening::functionUsesHardeningRegister(
419 if (
MI.readsRegister(MisspeculatingTaintReg,
TRI) ||
420 MI.modifiesRegister(MisspeculatingTaintReg,
TRI))
430 bool AArch64SpeculationHardening::makeGPRSpeculationSafe(
441 if (
Reg == AArch64::SP ||
Reg == AArch64::WSP)
445 if (RegsAlreadyMasked[
Reg])
448 const bool Is64Bit = AArch64::GPR64allRegClass.contains(
Reg);
451 TII->get(Is64Bit ? AArch64::SpeculationSafeValueX
452 : AArch64::SpeculationSafeValueW))
455 RegsAlreadyMasked.set(
Reg);
464 RegsAlreadyMasked.reset();
470 NextMBBI = std::next(
MBBI);
485 return Op.isReg() && (AArch64::GPR32allRegClass.contains(Op.getReg()) ||
486 AArch64::GPR64allRegClass.contains(Op.getReg()));
492 bool HardenLoadedData = AllDefsAreGPR;
493 bool HardenAddressLoadedFrom = !HardenLoadedData;
500 RegsAlreadyMasked.reset(*AI);
508 if (HardenLoadedData)
509 for (
auto Def :
MI.defs()) {
521 if (HardenAddressLoadedFrom)
522 for (
auto Use :
MI.uses()) {
547 bool AArch64SpeculationHardening::expandSpeculationSafeValue(
549 bool UsesFullSpeculationBarrier) {
551 unsigned Opcode =
MI.getOpcode();
557 case AArch64::SpeculationSafeValueW:
560 case AArch64::SpeculationSafeValueX:
564 if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {
572 RegsNeedingCSDBBeforeUse.set(*AI);
576 Is64Bit ?
TII->get(AArch64::ANDXrs) :
TII->get(AArch64::ANDWrs))
579 .
addUse(Is64Bit ? MisspeculatingTaintReg
580 : MisspeculatingTaintReg32Bit)
583 MI.eraseFromParent();
592 assert(!UseControlFlowSpeculationBarrier &&
"No need to insert CSDBs when "
593 "control flow miss-speculation "
594 "is already blocked");
597 RegsNeedingCSDBBeforeUse.reset();
601 bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
605 RegsNeedingCSDBBeforeUse.reset();
621 DL =
MI.getDebugLoc();
627 bool NeedToEmitBarrier =
false;
628 if (RegsNeedingCSDBBeforeUse.any() && (
MI.isCall() ||
MI.isTerminator()))
629 NeedToEmitBarrier =
true;
630 if (!NeedToEmitBarrier)
632 if (
Op.isReg() && RegsNeedingCSDBBeforeUse[
Op.getReg()]) {
633 NeedToEmitBarrier =
true;
637 if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)
641 expandSpeculationSafeValue(
MBB,
MBBI, UsesFullSpeculationBarrier);
646 if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)
652 bool AArch64SpeculationHardening::runOnMachineFunction(
MachineFunction &MF) {
656 MisspeculatingTaintReg = AArch64::X16;
657 MisspeculatingTaintReg32Bit = AArch64::W16;
662 UseControlFlowSpeculationBarrier = functionUsesHardeningRegister(MF);
669 dbgs() <<
"***** AArch64SpeculationHardening - automatic insertion of "
670 "SpeculationSafeValue intrinsics *****\n");
678 <<
"***** AArch64SpeculationHardening - track control flow *****\n");
681 EntryBlocks.push_back(&MF.
front());
683 EntryBlocks.push_back(LPI.LandingPadBlock);
684 for (
auto *Entry : EntryBlocks)
685 insertSPToRegTaintPropagation(
686 *Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
689 for (
auto &
MBB : MF) {
690 bool UsesFullSpeculationBarrier =
false;
691 Modified |= instrumentControlFlow(
MBB, UsesFullSpeculationBarrier);
693 lowerSpeculationSafeValuePseudos(
MBB, UsesFullSpeculationBarrier);
701 return new AArch64SpeculationHardening();