Skip to content

Commit

Permalink
Merging r351930 and r351932:
Browse files Browse the repository at this point in the history 
------------------------------------------------------------------------
r351930 | kbeyls | 2019-01-23 09:18:39 +0100 (Wed, 23 Jan 2019) | 30 lines

[SLH] AArch64: correctly pick temporary register to mask SP

As part of speculation hardening, the stack pointer gets masked with the
taint register (X16) before a function call or before a function return.
Since there are no instructions that can directly mask writing to the
stack pointer, the stack pointer must first be transferred to another
register, where it can be masked, before that value is transferred back
to the stack pointer.
Before, that temporary register was always picked to be x17, since the
ABI allows clobbering x17 on any function call, resulting in the
following instruction pattern being inserted before function calls and
returns/tail calls:

mov x17, sp
and x17, x17, x16
mov sp, x17
However, x17 can be live in those locations, for example when the call
is an indirect call, using x17 as the target address (blr x17).

To fix this, this patch looks for an available register just before the
call or terminator instruction and uses that.

In the rare case when no register turns out to be available (this
situation is only encountered twice across the whole test-suite), just
insert a full speculation barrier at the start of the basic block where
this occurs.

Differential Revision: https://reviews.llvm.org/D56717


------------------------------------------------------------------------

------------------------------------------------------------------------
r351932 | kbeyls | 2019-01-23 10:10:12 +0100 (Wed, 23 Jan 2019) | 3 lines

[SLH][AArch64] Remove accidentally retained -debug-only line from test.


------------------------------------------------------------------------

llvm-svn: 352115
  • Loading branch information
@zmodem
zmodem committed Jan 24, 2019
1 parent 17c9824 commit 22c2b2b
Show file tree
Hide file tree
Showing 4 changed files with 258 additions and 104 deletions.
175 changes: 118 additions & 57 deletions llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp
View file
Expand Up @@ -103,6 +103,7 @@
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegisterScavenging.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/Pass.h"
#include "llvm/Support/CodeGen.h"
Expand Down Expand Up @@ -146,25 +147,31 @@ class AArch64SpeculationHardening : public MachineFunctionPass {
BitVector RegsAlreadyMasked;

bool functionUsesHardeningRegister(MachineFunction &MF) const;
bool instrumentControlFlow(MachineBasicBlock &MBB);
bool instrumentControlFlow(MachineBasicBlock &MBB,
bool &UsesFullSpeculationBarrier);
bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
MachineBasicBlock *&FBB,
AArch64CC::CondCode &CondCode) const;
void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
AArch64CC::CondCode &CondCode, DebugLoc DL) const;
void insertSPToRegTaintPropagation(MachineBasicBlock *MBB,
void insertSPToRegTaintPropagation(MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI) const;
void insertRegToSPTaintPropagation(MachineBasicBlock *MBB,
void insertRegToSPTaintPropagation(MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI,
unsigned TmpReg) const;
void insertFullSpeculationBarrier(MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI,
DebugLoc DL) const;

bool slhLoads(MachineBasicBlock &MBB);
bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI,
MachineInstr &MI, unsigned Reg);
bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB);
bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB,
bool UsesFullSpeculationBarrier);
bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
MachineBasicBlock::iterator MBBI);
MachineBasicBlock::iterator MBBI,
bool UsesFullSpeculationBarrier);
bool insertCSDB(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
DebugLoc DL);
};
Expand Down Expand Up @@ -207,15 +214,19 @@ bool AArch64SpeculationHardening::endsWithCondControlFlow(
return true;
}

void AArch64SpeculationHardening::insertFullSpeculationBarrier(
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
DebugLoc DL) const {
// A full control flow speculation barrier consists of (DSB SYS + ISB)
BuildMI(MBB, MBBI, DL, TII->get(AArch64::DSB)).addImm(0xf);
BuildMI(MBB, MBBI, DL, TII->get(AArch64::ISB)).addImm(0xf);
}

void AArch64SpeculationHardening::insertTrackingCode(
MachineBasicBlock &SplitEdgeBB, AArch64CC::CondCode &CondCode,
DebugLoc DL) const {
if (UseControlFlowSpeculationBarrier) {
// insert full control flow speculation barrier (DSB SYS + ISB)
BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::ISB))
.addImm(0xf);
BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::DSB))
.addImm(0xf);
insertFullSpeculationBarrier(SplitEdgeBB, SplitEdgeBB.begin(), DL);
} else {
BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::CSELXr))
.addDef(MisspeculatingTaintReg)
Expand All @@ -227,7 +238,7 @@ void AArch64SpeculationHardening::insertTrackingCode(
}

bool AArch64SpeculationHardening::instrumentControlFlow(
MachineBasicBlock &MBB) {
MachineBasicBlock &MBB, bool &UsesFullSpeculationBarrier) {
LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);

bool Modified = false;
Expand Down Expand Up @@ -263,63 +274,113 @@ bool AArch64SpeculationHardening::instrumentControlFlow(
}

// Perform correct code generation around function calls and before returns.
{
SmallVector<MachineInstr *, 4> ReturnInstructions;
SmallVector<MachineInstr *, 4> CallInstructions;
// The below variables record the return/terminator instructions and the call
// instructions respectively; including which register is available as a
// temporary register just before the recorded instructions.
SmallVector<std::pair<MachineInstr *, unsigned>, 4> ReturnInstructions;
SmallVector<std::pair<MachineInstr *, unsigned>, 4> CallInstructions;
// if a temporary register is not available for at least one of the
// instructions for which we need to transfer taint to the stack pointer, we
// need to insert a full speculation barrier.
// TmpRegisterNotAvailableEverywhere tracks that condition.
bool TmpRegisterNotAvailableEverywhere = false;

RegScavenger RS;
RS.enterBasicBlock(MBB);

for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); I++) {
MachineInstr &MI = *I;
if (!MI.isReturn() && !MI.isCall())
continue;

for (MachineInstr &MI : MBB) {
if (MI.isReturn())
ReturnInstructions.push_back(&MI);
else if (MI.isCall())
CallInstructions.push_back(&MI);
}
// The RegScavenger represents registers available *after* the MI
// instruction pointed to by RS.getCurrentPosition().
// We need to have a register that is available *before* the MI is executed.
if (I != MBB.begin())
RS.forward(std::prev(I));
// FIXME: The below just finds *a* unused register. Maybe code could be
// optimized more if this looks for the register that isn't used for the
// longest time around this place, to enable more scheduling freedom. Not
// sure if that would actually result in a big performance difference
// though. Maybe RegisterScavenger::findSurvivorBackwards has some logic
// already to do this - but it's unclear if that could easily be used here.
unsigned TmpReg = RS.FindUnusedReg(&AArch64::GPR64commonRegClass);
LLVM_DEBUG(dbgs() << "RS finds "
<< ((TmpReg == 0) ? "no register " : "register ");
if (TmpReg != 0) dbgs() << printReg(TmpReg, TRI) << " ";
dbgs() << "to be available at MI " << MI);
if (TmpReg == 0)
TmpRegisterNotAvailableEverywhere = true;
if (MI.isReturn())
ReturnInstructions.push_back({&MI, TmpReg});
else if (MI.isCall())
CallInstructions.push_back({&MI, TmpReg});
}

Modified |=
(ReturnInstructions.size() > 0) || (CallInstructions.size() > 0);
if (TmpRegisterNotAvailableEverywhere) {
// When a temporary register is not available everywhere in this basic
// basic block where a propagate-taint-to-sp operation is needed, just
// emit a full speculation barrier at the start of this basic block, which
// renders the taint/speculation tracking in this basic block unnecessary.
insertFullSpeculationBarrier(MBB, MBB.begin(),
(MBB.begin())->getDebugLoc());
UsesFullSpeculationBarrier = true;
Modified = true;
} else {
for (auto MI_Reg : ReturnInstructions) {
assert(MI_Reg.second != 0);
LLVM_DEBUG(
dbgs()
<< " About to insert Reg to SP taint propagation with temp register "
<< printReg(MI_Reg.second, TRI)
<< " on instruction: " << *MI_Reg.first);
insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
Modified = true;
}

for (MachineInstr *Return : ReturnInstructions)
insertRegToSPTaintPropagation(Return->getParent(), Return, AArch64::X17);
for (MachineInstr *Call : CallInstructions) {
for (auto MI_Reg : CallInstructions) {
assert(MI_Reg.second != 0);
LLVM_DEBUG(dbgs() << " About to insert Reg to SP and back taint "
"propagation with temp register "
<< printReg(MI_Reg.second, TRI)
<< " around instruction: " << *MI_Reg.first);
// Just after the call:
MachineBasicBlock::iterator i = Call;
i++;
insertSPToRegTaintPropagation(Call->getParent(), i);
insertSPToRegTaintPropagation(
MBB, std::next((MachineBasicBlock::iterator)MI_Reg.first));
// Just before the call:
insertRegToSPTaintPropagation(Call->getParent(), Call, AArch64::X17);
insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
Modified = true;
}
}

return Modified;
}

void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI) const {
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const {
// If full control flow speculation barriers are used, emit a control flow
// barrier to block potential miss-speculation in flight coming in to this
// function.
if (UseControlFlowSpeculationBarrier) {
// insert full control flow speculation barrier (DSB SYS + ISB)
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::DSB)).addImm(0xf);
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ISB)).addImm(0xf);
insertFullSpeculationBarrier(MBB, MBBI, DebugLoc());
return;
}

// CMP SP, #0 === SUBS xzr, SP, #0
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
.addDef(AArch64::XZR)
.addUse(AArch64::SP)
.addImm(0)
.addImm(0); // no shift
// CSETM x16, NE === CSINV x16, xzr, xzr, EQ
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
.addDef(MisspeculatingTaintReg)
.addUse(AArch64::XZR)
.addUse(AArch64::XZR)
.addImm(AArch64CC::EQ);
}

void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI,
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
unsigned TmpReg) const {
// If full control flow speculation barriers are used, there will not be
// miss-speculation when returning from this function, and therefore, also
Expand All @@ -328,19 +389,19 @@ void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
return;

// mov Xtmp, SP === ADD Xtmp, SP, #0
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
.addDef(TmpReg)
.addUse(AArch64::SP)
.addImm(0)
.addImm(0); // no shift
// and Xtmp, Xtmp, TaintReg === AND Xtmp, Xtmp, TaintReg, #0
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
.addDef(TmpReg, RegState::Renamable)
.addUse(TmpReg, RegState::Kill | RegState::Renamable)
.addUse(MisspeculatingTaintReg, RegState::Kill)
.addImm(0);
// mov SP, Xtmp === ADD SP, Xtmp, #0
BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
.addDef(AArch64::SP)
.addUse(TmpReg, RegState::Kill)
.addImm(0)
Expand Down Expand Up @@ -484,7 +545,8 @@ bool AArch64SpeculationHardening::slhLoads(MachineBasicBlock &MBB) {
/// \brief If MBBI references a pseudo instruction that should be expanded
/// here, do the expansion and return true. Otherwise return false.
bool AArch64SpeculationHardening::expandSpeculationSafeValue(
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
bool UsesFullSpeculationBarrier) {
MachineInstr &MI = *MBBI;
unsigned Opcode = MI.getOpcode();
bool Is64Bit = true;
Expand All @@ -499,7 +561,7 @@ bool AArch64SpeculationHardening::expandSpeculationSafeValue(
// Just remove the SpeculationSafe pseudo's if control flow
// miss-speculation isn't happening because we're already inserting barriers
// to guarantee that.
if (!UseControlFlowSpeculationBarrier) {
if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {
unsigned DstReg = MI.getOperand(0).getReg();
unsigned SrcReg = MI.getOperand(1).getReg();
// Mark this register and all its aliasing registers as needing to be
Expand Down Expand Up @@ -537,7 +599,7 @@ bool AArch64SpeculationHardening::insertCSDB(MachineBasicBlock &MBB,
}

bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
MachineBasicBlock &MBB) {
MachineBasicBlock &MBB, bool UsesFullSpeculationBarrier) {
bool Modified = false;

RegsNeedingCSDBBeforeUse.reset();
Expand Down Expand Up @@ -572,15 +634,16 @@ bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
break;
}

if (NeedToEmitBarrier)
if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)
Modified |= insertCSDB(MBB, MBBI, DL);

Modified |= expandSpeculationSafeValue(MBB, MBBI);
Modified |=
expandSpeculationSafeValue(MBB, MBBI, UsesFullSpeculationBarrier);

MBBI = NMBBI;
}

if (RegsNeedingCSDBBeforeUse.any())
if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)
Modified |= insertCSDB(MBB, MBBI, DL);

return Modified;
Expand Down Expand Up @@ -609,7 +672,7 @@ bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
Modified |= slhLoads(MBB);
}

// 2.a Add instrumentation code to function entry and exits.
// 2. Add instrumentation code to function entry and exits.
LLVM_DEBUG(
dbgs()
<< "***** AArch64SpeculationHardening - track control flow *****\n");
Expand All @@ -620,17 +683,15 @@ bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
EntryBlocks.push_back(LPI.LandingPadBlock);
for (auto Entry : EntryBlocks)
insertSPToRegTaintPropagation(
Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));

// 2.b Add instrumentation code to every basic block.
for (auto &MBB : MF)
Modified |= instrumentControlFlow(MBB);
*Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));

LLVM_DEBUG(dbgs() << "***** AArch64SpeculationHardening - Lowering "
"SpeculationSafeValue Pseudos *****\n");
// Step 3: Lower SpeculationSafeValue pseudo instructions.
for (auto &MBB : MF)
Modified |= lowerSpeculationSafeValuePseudos(MBB);
// 3. Add instrumentation code to every basic block.
for (auto &MBB : MF) {
bool UsesFullSpeculationBarrier = false;
Modified |= instrumentControlFlow(MBB, UsesFullSpeculationBarrier);
Modified |=
lowerSpeculationSafeValuePseudos(MBB, UsesFullSpeculationBarrier);
}

return Modified;
}
Expand Down

0 comments on commit 22c2b2b

Please sign in to comment.