AArch64SpeculationHardening.cpp Source File (original) (raw)

108#include

109

110using namespace llvm;

111

112#define DEBUG_TYPE "aarch64-speculation-hardening"

113

114#define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"

115

117 cl::desc("Sanitize loads from memory."),

119

120namespace {

121

123public:

126

127 static char ID;

128

130

131 bool runOnMachineFunction(MachineFunction &Fn) override;

132

133 StringRef getPassName() const override {

135 }

136

137private:

138 unsigned MisspeculatingTaintReg;

139 unsigned MisspeculatingTaintReg32Bit;

140 bool UseControlFlowSpeculationBarrier;

141 BitVector RegsNeedingCSDBBeforeUse;

142 BitVector RegsAlreadyMasked;

143

144 bool functionUsesHardeningRegister(MachineFunction &MF) const;

145 bool instrumentControlFlow(MachineBasicBlock &MBB,

146 bool &UsesFullSpeculationBarrier);

147 bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,

148 MachineBasicBlock *&FBB,

150 void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,

152 void insertSPToRegTaintPropagation(MachineBasicBlock &MBB,

154 void insertRegToSPTaintPropagation(MachineBasicBlock &MBB,

156 unsigned TmpReg) const;

157 void insertFullSpeculationBarrier(MachineBasicBlock &MBB,

160

161 bool slhLoads(MachineBasicBlock &MBB);

162 bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,

164 MachineInstr &MI, unsigned Reg);

165 bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB,

166 bool UsesFullSpeculationBarrier);

167 bool expandSpeculationSafeValue(MachineBasicBlock &MBB,

169 bool UsesFullSpeculationBarrier);

172};

173

174}

175

176char AArch64SpeculationHardening::ID = 0;

177

178 INITIALIZE_PASS(AArch64SpeculationHardening, "aarch64-speculation-hardening",

180

181bool AArch64SpeculationHardening::endsWithCondControlFlow(

183 AArch64CC::CondCode &CondCode) const {

185 if (TII->analyzeBranch(MBB, TBB, FBB, analyzeBranchCondCode, false))

186 return false;

187

188

189 if (analyzeBranchCondCode.empty())

190 return false;

191

192

193

194

196 if (FBB == nullptr)

197 FBB = MBB.getFallThrough();

198

199

200

201

203 return false;

204

206

207 assert(analyzeBranchCondCode.size() == 1 && "unknown Cond array format");

209 return true;

210}

211

212void AArch64SpeculationHardening::insertFullSpeculationBarrier(

215

218}

219

220void AArch64SpeculationHardening::insertTrackingCode(

223 if (UseControlFlowSpeculationBarrier) {

224 insertFullSpeculationBarrier(SplitEdgeBB, SplitEdgeBB.begin(), DL);

225 } else {

227 .addDef(MisspeculatingTaintReg)

228 .addUse(MisspeculatingTaintReg)

229 .addUse(AArch64::XZR)

231 SplitEdgeBB.addLiveIn(AArch64::NZCV);

232 }

233}

234

235bool AArch64SpeculationHardening::instrumentControlFlow(

236 MachineBasicBlock &MBB, bool &UsesFullSpeculationBarrier) {

237 LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);

238

240 MachineBasicBlock *TBB = nullptr;

241 MachineBasicBlock *FBB = nullptr;

243

244 if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {

245 LLVM_DEBUG(dbgs() << "... doesn't end with CondControlFlow\n");

246 } else {

247

248

249

250

252

255

256 assert(SplitEdgeTBB != nullptr);

257 assert(SplitEdgeFBB != nullptr);

258

262

263 insertTrackingCode(*SplitEdgeTBB, CondCode, DL);

264 insertTrackingCode(*SplitEdgeFBB, InvCondCode, DL);

265

266 LLVM_DEBUG(dbgs() << "SplitEdgeTBB: " << *SplitEdgeTBB << "\n");

267 LLVM_DEBUG(dbgs() << "SplitEdgeFBB: " << *SplitEdgeFBB << "\n");

269 }

270

271

272

273

274

277

278

279

280

281 bool TmpRegisterNotAvailableEverywhere = false;

282

283 RegScavenger RS;

284 RS.enterBasicBlockEnd(MBB);

285

287 MachineInstr &MI = *--I;

288 if (.isReturn() && .isCall())

289 continue;

290

291

292

293

295 RS.enterBasicBlock(MBB);

296 else

297 RS.backward(I);

298

299

300

301

302

303

304 Register TmpReg = RS.FindUnusedReg(&AArch64::GPR64commonRegClass);

306 << ((TmpReg == 0) ? "no register " : "register ");

308 dbgs() << "to be available at MI " << MI);

309 if (TmpReg == 0)

310 TmpRegisterNotAvailableEverywhere = true;

311 if (MI.isReturn())

312 ReturnInstructions.push_back({&MI, TmpReg});

313 else if (MI.isCall())

314 CallInstructions.push_back({&MI, TmpReg});

315 }

316

317 if (TmpRegisterNotAvailableEverywhere) {

318

319

320

321

322 insertFullSpeculationBarrier(MBB, MBB.begin(),

323 (MBB.begin())->getDebugLoc());

324 UsesFullSpeculationBarrier = true;

326 } else {

327 for (auto MI_Reg : ReturnInstructions) {

328 assert(MI_Reg.second != 0);

331 << " About to insert Reg to SP taint propagation with temp register "

333 << " on instruction: " << *MI_Reg.first);

334 insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);

336 }

337

338 for (auto MI_Reg : CallInstructions) {

339 assert(MI_Reg.second != 0);

340 LLVM_DEBUG(dbgs() << " About to insert Reg to SP and back taint "

341 "propagation with temp register "

343 << " around instruction: " << *MI_Reg.first);

344

345 insertSPToRegTaintPropagation(

347

348 insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);

350 }

351 }

353}

354

355void AArch64SpeculationHardening::insertSPToRegTaintPropagation(

357

358

359

360 if (UseControlFlowSpeculationBarrier) {

362 return;

363 }

364

365

367 .addDef(AArch64::XZR)

370 .addImm(0);

371

373 .addDef(MisspeculatingTaintReg)

374 .addUse(AArch64::XZR)

375 .addUse(AArch64::XZR)

377}

378

379void AArch64SpeculationHardening::insertRegToSPTaintPropagation(

381 unsigned TmpReg) const {

382

383

384

385 if (UseControlFlowSpeculationBarrier)

386 return;

387

388

393 .addImm(0);

394

400

405 .addImm(0);

406}

407

408bool AArch64SpeculationHardening::functionUsesHardeningRegister(

409 MachineFunction &MF) const {

410 for (MachineBasicBlock &MBB : MF) {

411 for (MachineInstr &MI : MBB) {

412

413

414 if (MI.isCall())

415 continue;

416 if (MI.readsRegister(MisspeculatingTaintReg, TRI) ||

417 MI.modifiesRegister(MisspeculatingTaintReg, TRI))

418 return true;

419 }

420 }

421 return false;

422}

423

424

425

426

427bool AArch64SpeculationHardening::makeGPRSpeculationSafe(

429 unsigned Reg) {

431 AArch64::GPR64allRegClass.contains(Reg));

432

433

434

435

436

437

438 if (Reg == AArch64::SP || Reg == AArch64::WSP)

439 return false;

440

441

442 if (RegsAlreadyMasked[Reg])

443 return false;

444

445 const bool Is64Bit = AArch64::GPR64allRegClass.contains(Reg);

446 LLVM_DEBUG(dbgs() << "About to harden register : " << Reg << "\n");

448 TII->get(Is64Bit ? AArch64::SpeculationSafeValueX

449 : AArch64::SpeculationSafeValueW))

452 RegsAlreadyMasked.set(Reg);

453 return true;

454}

455

456bool AArch64SpeculationHardening::slhLoads(MachineBasicBlock &MBB) {

458

460

461 RegsAlreadyMasked.reset();

462

465 for (; MBBI != E; MBBI = NextMBBI) {

466 MachineInstr &MI = *MBBI;

467 NextMBBI = std::next(MBBI);

468

469 if (.mayLoad())

470 continue;

471

473

474

475

476

477

478

479

480

481 bool AllDefsAreGPR = llvm::all_of(MI.defs(), [&](MachineOperand &Op) {

482 return Op.isReg() && (AArch64::GPR32allRegClass.contains(Op.getReg()) ||

483 AArch64::GPR64allRegClass.contains(Op.getReg()));

484 });

485

486

487

488

489 bool HardenLoadedData = AllDefsAreGPR;

490 bool HardenAddressLoadedFrom = !HardenLoadedData;

491

492

493

494

495 for (MachineOperand Op : MI.defs())

496 for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)

497 RegsAlreadyMasked.reset(*AI);

498

499

500

501

502

503

504

505 if (HardenLoadedData)

506 for (auto Def : MI.defs()) {

507 if (Def.isDead())

508

509 continue;

510

511

512

513

514

515 Modified |= makeGPRSpeculationSafe(MBB, NextMBBI, MI, Def.getReg());

516 }

517

518 if (HardenAddressLoadedFrom)

519 for (auto Use : MI.uses()) {

520 if ( Use .isReg())

521 continue;

523

524

525

526

527

528

529

530

531

532

533 if (!(AArch64::GPR32allRegClass.contains(Reg) ||

534 AArch64::GPR64allRegClass.contains(Reg)))

535 continue;

537 }

538 }

540}

541

542

543

544bool AArch64SpeculationHardening::expandSpeculationSafeValue(

546 bool UsesFullSpeculationBarrier) {

547 MachineInstr &MI = *MBBI;

548 unsigned Opcode = MI.getOpcode();

549 bool Is64Bit = true;

550

551 switch (Opcode) {

552 default:

553 break;

554 case AArch64::SpeculationSafeValueW:

555 Is64Bit = false;

556 [[fallthrough]];

557 case AArch64::SpeculationSafeValueX:

558

559

560

561 if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {

562 Register DstReg = MI.getOperand(0).getReg();

563 Register SrcReg = MI.getOperand(1).getReg();

564

565

566

567 for (MachineOperand Op : MI.defs())

568 for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)

569 RegsNeedingCSDBBeforeUse.set(*AI);

570

571

573 Is64Bit ? TII->get(AArch64::ANDXrs) : TII->get(AArch64::ANDWrs))

576 .addUse(Is64Bit ? MisspeculatingTaintReg

577 : MisspeculatingTaintReg32Bit)

579 }

580 MI.eraseFromParent();

581 return true;

582 }

583 return false;

584}

585

586bool AArch64SpeculationHardening::insertCSDB(MachineBasicBlock &MBB,

589 assert(!UseControlFlowSpeculationBarrier && "No need to insert CSDBs when "

590 "control flow miss-speculation "

591 "is already blocked");

592

594 RegsNeedingCSDBBeforeUse.reset();

595 return true;

596}

597

598bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(

599 MachineBasicBlock &MBB, bool UsesFullSpeculationBarrier) {

601

602 RegsNeedingCSDBBeforeUse.reset();

603

604

605

606

607

608

609

610

611

612

613

616 while (MBBI != E) {

617 MachineInstr &MI = *MBBI;

618 DL = MI.getDebugLoc();

620

621

622

623

624 bool NeedToEmitBarrier = false;

625 if (RegsNeedingCSDBBeforeUse.any() && (MI.isCall() || MI.isTerminator()))

626 NeedToEmitBarrier = true;

627 if (!NeedToEmitBarrier)

628 for (MachineOperand Op : MI.uses())

629 if (Op.isReg() && RegsNeedingCSDBBeforeUse[Op.getReg()]) {

630 NeedToEmitBarrier = true;

631 break;

632 }

633

634 if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)

636

638 expandSpeculationSafeValue(MBB, MBBI, UsesFullSpeculationBarrier);

639

640 MBBI = NMBBI;

641 }

642

643 if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)

645

647}

648

649bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {

651 return false;

652

653 MisspeculatingTaintReg = AArch64::X16;

654 MisspeculatingTaintReg32Bit = AArch64::W16;

657 RegsNeedingCSDBBeforeUse.resize(TRI->getNumRegs());

658 RegsAlreadyMasked.resize(TRI->getNumRegs());

659 UseControlFlowSpeculationBarrier = functionUsesHardeningRegister(MF);

660

662

663

666 dbgs() << "***** AArch64SpeculationHardening - automatic insertion of "

667 "SpeculationSafeValue intrinsics *****\n");

668 for (auto &MBB : MF)

670 }

671

672

675 << "***** AArch64SpeculationHardening - track control flow *****\n");

676

679 for (const LandingPadInfo &LPI : MF.getLandingPads())

680 EntryBlocks.push_back(LPI.LandingPadBlock);

681 for (auto *Entry : EntryBlocks)

682 insertSPToRegTaintPropagation(

683 *Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));

684

685

686 for (auto &MBB : MF) {

687 bool UsesFullSpeculationBarrier = false;

688 Modified |= instrumentControlFlow(MBB, UsesFullSpeculationBarrier);

690 lowerSpeculationSafeValuePseudos(MBB, UsesFullSpeculationBarrier);

691 }

692

694}

695

696

698 return new AArch64SpeculationHardening();

699}

assert(UImm &&(UImm !=~static_cast< T >(0)) &&"Invalid immediate!")

const TargetInstrInfo & TII

#define AARCH64_SPECULATION_HARDENING_NAME

Definition AArch64SpeculationHardening.cpp:114

static cl::opt< bool > HardenLoads("aarch64-slh-loads", cl::Hidden, cl::desc("Sanitize loads from memory."), cl::init(true))

MachineBasicBlock MachineBasicBlock::iterator DebugLoc DL

MachineBasicBlock MachineBasicBlock::iterator MBBI

This file implements the BitVector class.

static GCRegistry::Add< CoreCLRGC > E("coreclr", "CoreCLR-compatible GC")

Promote Memory to Register

#define INITIALIZE_PASS(passName, arg, name, cfg, analysis)

const SmallVectorImpl< MachineOperand > MachineBasicBlock * TBB

This file declares the machine register scavenger class.

static bool contains(SmallPtrSetImpl< ConstantExpr * > &Cache, ConstantExpr *Expr, Constant *C)

This file defines the SmallVector class.

void resize(unsigned N, bool t=false)

resize - Grow or shrink the bitvector.

bool any() const

any - Returns true if any bit is set.

FunctionPass class - This class is used to implement most global optimizations.

bool hasFnAttribute(Attribute::AttrKind Kind) const

Return true if the function has the attribute.

const MCInstrDesc & get(unsigned Opcode) const

Return the machine instruction descriptor that corresponds to the specified instruction opcode.

instr_iterator instr_begin()

MachineBasicBlock * SplitCriticalEdge(MachineBasicBlock *Succ, Pass &P, std::vector< SparseBitVector<> > *LiveInSets=nullptr, MachineDomTreeUpdater *MDTU=nullptr)

instr_iterator instr_end()

void addLiveIn(MCRegister PhysReg, LaneBitmask LaneMask=LaneBitmask::getAll())

Adds the specified register as a live in.

MachineInstrBundleIterator< MachineInstr > iterator

MachineFunctionPass - This class adapts the FunctionPass interface to allow convenient creation of pa...

const TargetSubtargetInfo & getSubtarget() const

getSubtarget - Return the subtarget for which this machine code is being compiled.

Function & getFunction()

Return the LLVM function that this machine code represents.

const std::vector< LandingPadInfo > & getLandingPads() const

Return a reference to the landing pad info for the current function.

const MachineBasicBlock & front() const

const MachineInstrBuilder & addImm(int64_t Val) const

Add a new immediate operand.

const MachineInstrBuilder & addUse(Register RegNo, unsigned Flags=0, unsigned SubReg=0) const

Add a virtual register use operand.

const MachineInstrBuilder & addDef(Register RegNo, unsigned Flags=0, unsigned SubReg=0) const

Add a virtual register definition operand.

void push_back(const T &Elt)

This is a 'vector' (really, a variable-sized array), optimized for the case when the array is small.

TargetInstrInfo - Interface to description of machine instruction set.

TargetRegisterInfo base class - We assume that the target defines a static array of TargetRegisterDes...

virtual const TargetInstrInfo * getInstrInfo() const

virtual const TargetRegisterInfo * getRegisterInfo() const =0

Return the target's register information.

static CondCode getInvertedCondCode(CondCode Code)

unsigned ID

LLVM IR allows to use arbitrary numbers as calling convention identifiers.

CondCode

ISD::CondCode enum - These are ordered carefully to make the bitfields below work out,...

@ Renamable

@ Kill

The last use of a register.

initializer< Ty > init(const Ty &Val)

NodeAddr< DefNode * > Def

NodeAddr< UseNode * > Use

This is an optimization pass for GlobalISel generic memory operations.

bool all_of(R &&range, UnaryPredicate P)

Provide wrappers to std::all_of which take ranges instead of having to pass begin/end explicitly.

MachineInstrBuilder BuildMI(MachineFunction &MF, const MIMetadata &MIMD, const MCInstrDesc &MCID)

Builder interface. Specify how to create the initial instruction itself.

MachineInstr * getImm(const MachineOperand &MO, const MachineRegisterInfo *MRI)

LLVM_ABI raw_ostream & dbgs()

dbgs() - This returns a reference to a raw_ostream for debugging messages.

class LLVM_GSL_OWNER SmallVector

Forward declaration of SmallVector so that calculateSmallVectorDefaultInlinedElements can reference s...

DWARFExpression::Operation Op

FunctionPass * createAArch64SpeculationHardeningPass()

Returns an instance of the pseudo instruction expansion pass.

Definition AArch64SpeculationHardening.cpp:697

LLVM_ABI Printable printReg(Register Reg, const TargetRegisterInfo *TRI=nullptr, unsigned SubIdx=0, const MachineRegisterInfo *MRI=nullptr)

Prints virtual and physical registers with or without a TRI instance.