1 //==-- X86LoadValueInjectionLoadHardening.cpp - LVI load hardening for x86 --=// 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 /// Description: This pass finds Load Value Injection (LVI) gadgets consisting 10 /// of a load from memory (i.e., SOURCE), and any operation that may transmit 11 /// the value loaded from memory over a covert channel, or use the value loaded 12 /// from memory to determine a branch/call target (i.e., SINK). After finding 13 /// all such gadgets in a given function, the pass minimally inserts LFENCE 14 /// instructions in such a manner that the following property is satisfied: for 15 /// all SOURCE+SINK pairs, all paths in the CFG from SOURCE to SINK contain at 16 /// least one LFENCE instruction. The algorithm that implements this minimal 17 /// insertion is influenced by an academic paper that minimally inserts memory 18 /// fences for high-performance concurrent programs: 19 /// http://www.cs.ucr.edu/~lesani/companion/oopsla15/OOPSLA15.pdf 20 /// The algorithm implemented in this pass is as follows: 21 /// 1. Build a condensed CFG (i.e., a GadgetGraph) consisting only of the 22 /// following components: 23 /// - SOURCE instructions (also includes function arguments) 24 /// - SINK instructions 25 /// - Basic block entry points 26 /// - Basic block terminators 27 /// - LFENCE instructions 28 /// 2. Analyze the GadgetGraph to determine which SOURCE+SINK pairs (i.e., 29 /// gadgets) are already mitigated by existing LFENCEs. If all gadgets have been 30 /// mitigated, go to step 6. 31 /// 3. Use a heuristic or plugin to approximate minimal LFENCE insertion. 32 /// 4. Insert one LFENCE along each CFG edge that was cut in step 3. 33 /// 5. Go to step 2. 34 /// 6. If any LFENCEs were inserted, return `true` from runOnMachineFunction() 35 /// to tell LLVM that the function was modified. 36 /// 37 //===----------------------------------------------------------------------===// 38 39 #include "ImmutableGraph.h" 40 #include "X86.h" 41 #include "X86Subtarget.h" 42 #include "X86TargetMachine.h" 43 #include "llvm/ADT/DenseMap.h" 44 #include "llvm/ADT/DenseSet.h" 45 #include "llvm/ADT/STLExtras.h" 46 #include "llvm/ADT/SmallSet.h" 47 #include "llvm/ADT/Statistic.h" 48 #include "llvm/ADT/StringRef.h" 49 #include "llvm/CodeGen/MachineBasicBlock.h" 50 #include "llvm/CodeGen/MachineDominanceFrontier.h" 51 #include "llvm/CodeGen/MachineDominators.h" 52 #include "llvm/CodeGen/MachineFunction.h" 53 #include "llvm/CodeGen/MachineFunctionPass.h" 54 #include "llvm/CodeGen/MachineInstr.h" 55 #include "llvm/CodeGen/MachineInstrBuilder.h" 56 #include "llvm/CodeGen/MachineLoopInfo.h" 57 #include "llvm/CodeGen/MachineRegisterInfo.h" 58 #include "llvm/CodeGen/RDFGraph.h" 59 #include "llvm/CodeGen/RDFLiveness.h" 60 #include "llvm/InitializePasses.h" 61 #include "llvm/Support/CommandLine.h" 62 #include "llvm/Support/DOTGraphTraits.h" 63 #include "llvm/Support/Debug.h" 64 #include "llvm/Support/DynamicLibrary.h" 65 #include "llvm/Support/GraphWriter.h" 66 #include "llvm/Support/raw_ostream.h" 67 68 using namespace llvm; 69 70 #define PASS_KEY "x86-lvi-load" 71 #define DEBUG_TYPE PASS_KEY 72 73 STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation"); 74 STATISTIC(NumFunctionsConsidered, "Number of functions analyzed"); 75 STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations " 76 "were deployed"); 77 STATISTIC(NumGadgets, "Number of LVI gadgets detected during analysis"); 78 79 static cl::opt<std::string> OptimizePluginPath( 80 PASS_KEY "-opt-plugin", 81 cl::desc("Specify a plugin to optimize LFENCE insertion"), cl::Hidden); 82 83 static cl::opt<bool> NoConditionalBranches( 84 PASS_KEY "-no-cbranch", 85 cl::desc("Don't treat conditional branches as disclosure gadgets. This " 86 "may improve performance, at the cost of security."), 87 cl::init(false), cl::Hidden); 88 89 static cl::opt<bool> EmitDot( 90 PASS_KEY "-dot", 91 cl::desc( 92 "For each function, emit a dot graph depicting potential LVI gadgets"), 93 cl::init(false), cl::Hidden); 94 95 static cl::opt<bool> EmitDotOnly( 96 PASS_KEY "-dot-only", 97 cl::desc("For each function, emit a dot graph depicting potential LVI " 98 "gadgets, and do not insert any fences"), 99 cl::init(false), cl::Hidden); 100 101 static cl::opt<bool> EmitDotVerify( 102 PASS_KEY "-dot-verify", 103 cl::desc("For each function, emit a dot graph to stdout depicting " 104 "potential LVI gadgets, used for testing purposes only"), 105 cl::init(false), cl::Hidden); 106 107 static llvm::sys::DynamicLibrary OptimizeDL; 108 typedef int (*OptimizeCutT)(unsigned int *Nodes, unsigned int NodesSize, 109 unsigned int *Edges, int *EdgeValues, 110 int *CutEdges /* out */, unsigned int EdgesSize); 111 static OptimizeCutT OptimizeCut = nullptr; 112 113 namespace { 114 115 struct MachineGadgetGraph : ImmutableGraph<MachineInstr *, int> { 116 static constexpr int GadgetEdgeSentinel = -1; 117 static constexpr MachineInstr *const ArgNodeSentinel = nullptr; 118 119 using GraphT = ImmutableGraph<MachineInstr *, int>; 120 using Node = typename GraphT::Node; 121 using Edge = typename GraphT::Edge; 122 using size_type = typename GraphT::size_type; 123 MachineGadgetGraph(std::unique_ptr<Node[]> Nodes, 124 std::unique_ptr<Edge[]> Edges, size_type NodesSize, 125 size_type EdgesSize, int NumFences = 0, int NumGadgets = 0) 126 : GraphT(std::move(Nodes), std::move(Edges), NodesSize, EdgesSize), 127 NumFences(NumFences), NumGadgets(NumGadgets) {} 128 static inline bool isCFGEdge(const Edge &E) { 129 return E.getValue() != GadgetEdgeSentinel; 130 } 131 static inline bool isGadgetEdge(const Edge &E) { 132 return E.getValue() == GadgetEdgeSentinel; 133 } 134 int NumFences; 135 int NumGadgets; 136 }; 137 138 class X86LoadValueInjectionLoadHardeningPass : public MachineFunctionPass { 139 public: 140 X86LoadValueInjectionLoadHardeningPass() : MachineFunctionPass(ID) {} 141 142 StringRef getPassName() const override { 143 return "X86 Load Value Injection (LVI) Load Hardening"; 144 } 145 void getAnalysisUsage(AnalysisUsage &AU) const override; 146 bool runOnMachineFunction(MachineFunction &MF) override; 147 148 static char ID; 149 150 private: 151 using GraphBuilder = ImmutableGraphBuilder<MachineGadgetGraph>; 152 using Edge = MachineGadgetGraph::Edge; 153 using Node = MachineGadgetGraph::Node; 154 using EdgeSet = MachineGadgetGraph::EdgeSet; 155 using NodeSet = MachineGadgetGraph::NodeSet; 156 157 const X86Subtarget *STI; 158 const TargetInstrInfo *TII; 159 const TargetRegisterInfo *TRI; 160 161 std::unique_ptr<MachineGadgetGraph> 162 getGadgetGraph(MachineFunction &MF, const MachineLoopInfo &MLI, 163 const MachineDominatorTree &MDT, 164 const MachineDominanceFrontier &MDF) const; 165 int hardenLoadsWithPlugin(MachineFunction &MF, 166 std::unique_ptr<MachineGadgetGraph> Graph) const; 167 int hardenLoadsWithHeuristic(MachineFunction &MF, 168 std::unique_ptr<MachineGadgetGraph> Graph) const; 169 int elimMitigatedEdgesAndNodes(MachineGadgetGraph &G, 170 EdgeSet &ElimEdges /* in, out */, 171 NodeSet &ElimNodes /* in, out */) const; 172 std::unique_ptr<MachineGadgetGraph> 173 trimMitigatedEdges(std::unique_ptr<MachineGadgetGraph> Graph) const; 174 int insertFences(MachineFunction &MF, MachineGadgetGraph &G, 175 EdgeSet &CutEdges /* in, out */) const; 176 bool instrUsesRegToAccessMemory(const MachineInstr &I, unsigned Reg) const; 177 bool instrUsesRegToBranch(const MachineInstr &I, unsigned Reg) const; 178 inline bool isFence(const MachineInstr *MI) const { 179 return MI && (MI->getOpcode() == X86::LFENCE || 180 (STI->useLVIControlFlowIntegrity() && MI->isCall())); 181 } 182 }; 183 184 } // end anonymous namespace 185 186 namespace llvm { 187 188 template <> 189 struct GraphTraits<MachineGadgetGraph *> 190 : GraphTraits<ImmutableGraph<MachineInstr *, int> *> {}; 191 192 template <> 193 struct DOTGraphTraits<MachineGadgetGraph *> : DefaultDOTGraphTraits { 194 using GraphType = MachineGadgetGraph; 195 using Traits = llvm::GraphTraits<GraphType *>; 196 using NodeRef = typename Traits::NodeRef; 197 using EdgeRef = typename Traits::EdgeRef; 198 using ChildIteratorType = typename Traits::ChildIteratorType; 199 using ChildEdgeIteratorType = typename Traits::ChildEdgeIteratorType; 200 201 DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {} 202 203 std::string getNodeLabel(NodeRef Node, GraphType *) { 204 if (Node->getValue() == MachineGadgetGraph::ArgNodeSentinel) 205 return "ARGS"; 206 207 std::string Str; 208 raw_string_ostream OS(Str); 209 OS << *Node->getValue(); 210 return OS.str(); 211 } 212 213 static std::string getNodeAttributes(NodeRef Node, GraphType *) { 214 MachineInstr *MI = Node->getValue(); 215 if (MI == MachineGadgetGraph::ArgNodeSentinel) 216 return "color = blue"; 217 if (MI->getOpcode() == X86::LFENCE) 218 return "color = green"; 219 return ""; 220 } 221 222 static std::string getEdgeAttributes(NodeRef, ChildIteratorType E, 223 GraphType *) { 224 int EdgeVal = (*E.getCurrent()).getValue(); 225 return EdgeVal >= 0 ? "label = " + std::to_string(EdgeVal) 226 : "color = red, style = \"dashed\""; 227 } 228 }; 229 230 } // end namespace llvm 231 232 constexpr MachineInstr *MachineGadgetGraph::ArgNodeSentinel; 233 constexpr int MachineGadgetGraph::GadgetEdgeSentinel; 234 235 char X86LoadValueInjectionLoadHardeningPass::ID = 0; 236 237 void X86LoadValueInjectionLoadHardeningPass::getAnalysisUsage( 238 AnalysisUsage &AU) const { 239 MachineFunctionPass::getAnalysisUsage(AU); 240 AU.addRequired<MachineLoopInfo>(); 241 AU.addRequired<MachineDominatorTree>(); 242 AU.addRequired<MachineDominanceFrontier>(); 243 AU.setPreservesCFG(); 244 } 245 246 static void writeGadgetGraph(raw_ostream &OS, MachineFunction &MF, 247 MachineGadgetGraph *G) { 248 WriteGraph(OS, G, /*ShortNames*/ false, 249 "Speculative gadgets for \"" + MF.getName() + "\" function"); 250 } 251 252 bool X86LoadValueInjectionLoadHardeningPass::runOnMachineFunction( 253 MachineFunction &MF) { 254 LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName() 255 << " *****\n"); 256 STI = &MF.getSubtarget<X86Subtarget>(); 257 if (!STI->useLVILoadHardening()) 258 return false; 259 260 // FIXME: support 32-bit 261 if (!STI->is64Bit()) 262 report_fatal_error("LVI load hardening is only supported on 64-bit", false); 263 264 // Don't skip functions with the "optnone" attr but participate in opt-bisect. 265 const Function &F = MF.getFunction(); 266 if (!F.hasOptNone() && skipFunction(F)) 267 return false; 268 269 ++NumFunctionsConsidered; 270 TII = STI->getInstrInfo(); 271 TRI = STI->getRegisterInfo(); 272 LLVM_DEBUG(dbgs() << "Building gadget graph...\n"); 273 const auto &MLI = getAnalysis<MachineLoopInfo>(); 274 const auto &MDT = getAnalysis<MachineDominatorTree>(); 275 const auto &MDF = getAnalysis<MachineDominanceFrontier>(); 276 std::unique_ptr<MachineGadgetGraph> Graph = getGadgetGraph(MF, MLI, MDT, MDF); 277 LLVM_DEBUG(dbgs() << "Building gadget graph... Done\n"); 278 if (Graph == nullptr) 279 return false; // didn't find any gadgets 280 281 if (EmitDotVerify) { 282 writeGadgetGraph(outs(), MF, Graph.get()); 283 return false; 284 } 285 286 if (EmitDot || EmitDotOnly) { 287 LLVM_DEBUG(dbgs() << "Emitting gadget graph...\n"); 288 std::error_code FileError; 289 std::string FileName = "lvi."; 290 FileName += MF.getName(); 291 FileName += ".dot"; 292 raw_fd_ostream FileOut(FileName, FileError); 293 if (FileError) 294 errs() << FileError.message(); 295 writeGadgetGraph(FileOut, MF, Graph.get()); 296 FileOut.close(); 297 LLVM_DEBUG(dbgs() << "Emitting gadget graph... Done\n"); 298 if (EmitDotOnly) 299 return false; 300 } 301 302 int FencesInserted; 303 if (!OptimizePluginPath.empty()) { 304 if (!OptimizeDL.isValid()) { 305 std::string ErrorMsg; 306 OptimizeDL = llvm::sys::DynamicLibrary::getPermanentLibrary( 307 OptimizePluginPath.c_str(), &ErrorMsg); 308 if (!ErrorMsg.empty()) 309 report_fatal_error("Failed to load opt plugin: \"" + ErrorMsg + '\"'); 310 OptimizeCut = (OptimizeCutT)OptimizeDL.getAddressOfSymbol("optimize_cut"); 311 if (!OptimizeCut) 312 report_fatal_error("Invalid optimization plugin"); 313 } 314 FencesInserted = hardenLoadsWithPlugin(MF, std::move(Graph)); 315 } else { // Use the default greedy heuristic 316 FencesInserted = hardenLoadsWithHeuristic(MF, std::move(Graph)); 317 } 318 319 if (FencesInserted > 0) 320 ++NumFunctionsMitigated; 321 NumFences += FencesInserted; 322 return (FencesInserted > 0); 323 } 324 325 std::unique_ptr<MachineGadgetGraph> 326 X86LoadValueInjectionLoadHardeningPass::getGadgetGraph( 327 MachineFunction &MF, const MachineLoopInfo &MLI, 328 const MachineDominatorTree &MDT, 329 const MachineDominanceFrontier &MDF) const { 330 using namespace rdf; 331 332 // Build the Register Dataflow Graph using the RDF framework 333 TargetOperandInfo TOI{*TII}; 334 DataFlowGraph DFG{MF, *TII, *TRI, MDT, MDF, TOI}; 335 DFG.build(); 336 Liveness L{MF.getRegInfo(), DFG}; 337 L.computePhiInfo(); 338 339 GraphBuilder Builder; 340 using GraphIter = typename GraphBuilder::BuilderNodeRef; 341 DenseMap<MachineInstr *, GraphIter> NodeMap; 342 int FenceCount = 0, GadgetCount = 0; 343 auto MaybeAddNode = [&NodeMap, &Builder](MachineInstr *MI) { 344 auto Ref = NodeMap.find(MI); 345 if (Ref == NodeMap.end()) { 346 auto I = Builder.addVertex(MI); 347 NodeMap[MI] = I; 348 return std::pair<GraphIter, bool>{I, true}; 349 } 350 return std::pair<GraphIter, bool>{Ref->getSecond(), false}; 351 }; 352 353 // The `Transmitters` map memoizes transmitters found for each def. If a def 354 // has not yet been analyzed, then it will not appear in the map. If a def 355 // has been analyzed and was determined not to have any transmitters, then 356 // its list of transmitters will be empty. 357 DenseMap<NodeId, std::vector<NodeId>> Transmitters; 358 359 // Analyze all machine instructions to find gadgets and LFENCEs, adding 360 // each interesting value to `Nodes` 361 auto AnalyzeDef = [&](NodeAddr<DefNode *> SourceDef) { 362 SmallSet<NodeId, 8> UsesVisited, DefsVisited; 363 std::function<void(NodeAddr<DefNode *>)> AnalyzeDefUseChain = 364 [&](NodeAddr<DefNode *> Def) { 365 if (Transmitters.find(Def.Id) != Transmitters.end()) 366 return; // Already analyzed `Def` 367 368 // Use RDF to find all the uses of `Def` 369 rdf::NodeSet Uses; 370 RegisterRef DefReg = Def.Addr->getRegRef(DFG); 371 for (auto UseID : L.getAllReachedUses(DefReg, Def)) { 372 auto Use = DFG.addr<UseNode *>(UseID); 373 if (Use.Addr->getFlags() & NodeAttrs::PhiRef) { // phi node 374 NodeAddr<PhiNode *> Phi = Use.Addr->getOwner(DFG); 375 for (auto I : L.getRealUses(Phi.Id)) { 376 if (DFG.getPRI().alias(RegisterRef(I.first), DefReg)) { 377 for (auto UA : I.second) 378 Uses.emplace(UA.first); 379 } 380 } 381 } else { // not a phi node 382 Uses.emplace(UseID); 383 } 384 } 385 386 // For each use of `Def`, we want to know whether: 387 // (1) The use can leak the Def'ed value, 388 // (2) The use can further propagate the Def'ed value to more defs 389 for (auto UseID : Uses) { 390 if (!UsesVisited.insert(UseID).second) 391 continue; // Already visited this use of `Def` 392 393 auto Use = DFG.addr<UseNode *>(UseID); 394 assert(!(Use.Addr->getFlags() & NodeAttrs::PhiRef)); 395 MachineOperand &UseMO = Use.Addr->getOp(); 396 MachineInstr &UseMI = *UseMO.getParent(); 397 assert(UseMO.isReg()); 398 399 // We naively assume that an instruction propagates any loaded 400 // uses to all defs unless the instruction is a call, in which 401 // case all arguments will be treated as gadget sources during 402 // analysis of the callee function. 403 if (UseMI.isCall()) 404 continue; 405 406 // Check whether this use can transmit (leak) its value. 407 if (instrUsesRegToAccessMemory(UseMI, UseMO.getReg()) || 408 (!NoConditionalBranches && 409 instrUsesRegToBranch(UseMI, UseMO.getReg()))) { 410 Transmitters[Def.Id].push_back(Use.Addr->getOwner(DFG).Id); 411 if (UseMI.mayLoad()) 412 continue; // Found a transmitting load -- no need to continue 413 // traversing its defs (i.e., this load will become 414 // a new gadget source anyways). 415 } 416 417 // Check whether the use propagates to more defs. 418 NodeAddr<InstrNode *> Owner{Use.Addr->getOwner(DFG)}; 419 rdf::NodeList AnalyzedChildDefs; 420 for (auto &ChildDef : 421 Owner.Addr->members_if(DataFlowGraph::IsDef, DFG)) { 422 if (!DefsVisited.insert(ChildDef.Id).second) 423 continue; // Already visited this def 424 if (Def.Addr->getAttrs() & NodeAttrs::Dead) 425 continue; 426 if (Def.Id == ChildDef.Id) 427 continue; // `Def` uses itself (e.g., increment loop counter) 428 429 AnalyzeDefUseChain(ChildDef); 430 431 // `Def` inherits all of its child defs' transmitters. 432 for (auto TransmitterId : Transmitters[ChildDef.Id]) 433 Transmitters[Def.Id].push_back(TransmitterId); 434 } 435 } 436 437 // Note that this statement adds `Def.Id` to the map if no 438 // transmitters were found for `Def`. 439 auto &DefTransmitters = Transmitters[Def.Id]; 440 441 // Remove duplicate transmitters 442 llvm::sort(DefTransmitters); 443 DefTransmitters.erase( 444 std::unique(DefTransmitters.begin(), DefTransmitters.end()), 445 DefTransmitters.end()); 446 }; 447 448 // Find all of the transmitters 449 AnalyzeDefUseChain(SourceDef); 450 auto &SourceDefTransmitters = Transmitters[SourceDef.Id]; 451 if (SourceDefTransmitters.empty()) 452 return; // No transmitters for `SourceDef` 453 454 MachineInstr *Source = SourceDef.Addr->getFlags() & NodeAttrs::PhiRef 455 ? MachineGadgetGraph::ArgNodeSentinel 456 : SourceDef.Addr->getOp().getParent(); 457 auto GadgetSource = MaybeAddNode(Source); 458 // Each transmitter is a sink for `SourceDef`. 459 for (auto TransmitterId : SourceDefTransmitters) { 460 MachineInstr *Sink = DFG.addr<StmtNode *>(TransmitterId).Addr->getCode(); 461 auto GadgetSink = MaybeAddNode(Sink); 462 // Add the gadget edge to the graph. 463 Builder.addEdge(MachineGadgetGraph::GadgetEdgeSentinel, 464 GadgetSource.first, GadgetSink.first); 465 ++GadgetCount; 466 } 467 }; 468 469 LLVM_DEBUG(dbgs() << "Analyzing def-use chains to find gadgets\n"); 470 // Analyze function arguments 471 NodeAddr<BlockNode *> EntryBlock = DFG.getFunc().Addr->getEntryBlock(DFG); 472 for (NodeAddr<PhiNode *> ArgPhi : 473 EntryBlock.Addr->members_if(DataFlowGraph::IsPhi, DFG)) { 474 NodeList Defs = ArgPhi.Addr->members_if(DataFlowGraph::IsDef, DFG); 475 llvm::for_each(Defs, AnalyzeDef); 476 } 477 // Analyze every instruction in MF 478 for (NodeAddr<BlockNode *> BA : DFG.getFunc().Addr->members(DFG)) { 479 for (NodeAddr<StmtNode *> SA : 480 BA.Addr->members_if(DataFlowGraph::IsCode<NodeAttrs::Stmt>, DFG)) { 481 MachineInstr *MI = SA.Addr->getCode(); 482 if (isFence(MI)) { 483 MaybeAddNode(MI); 484 ++FenceCount; 485 } else if (MI->mayLoad()) { 486 NodeList Defs = SA.Addr->members_if(DataFlowGraph::IsDef, DFG); 487 llvm::for_each(Defs, AnalyzeDef); 488 } 489 } 490 } 491 LLVM_DEBUG(dbgs() << "Found " << FenceCount << " fences\n"); 492 LLVM_DEBUG(dbgs() << "Found " << GadgetCount << " gadgets\n"); 493 if (GadgetCount == 0) 494 return nullptr; 495 NumGadgets += GadgetCount; 496 497 // Traverse CFG to build the rest of the graph 498 SmallSet<MachineBasicBlock *, 8> BlocksVisited; 499 std::function<void(MachineBasicBlock *, GraphIter, unsigned)> TraverseCFG = 500 [&](MachineBasicBlock *MBB, GraphIter GI, unsigned ParentDepth) { 501 unsigned LoopDepth = MLI.getLoopDepth(MBB); 502 if (!MBB->empty()) { 503 // Always add the first instruction in each block 504 auto NI = MBB->begin(); 505 auto BeginBB = MaybeAddNode(&*NI); 506 Builder.addEdge(ParentDepth, GI, BeginBB.first); 507 if (!BlocksVisited.insert(MBB).second) 508 return; 509 510 // Add any instructions within the block that are gadget components 511 GI = BeginBB.first; 512 while (++NI != MBB->end()) { 513 auto Ref = NodeMap.find(&*NI); 514 if (Ref != NodeMap.end()) { 515 Builder.addEdge(LoopDepth, GI, Ref->getSecond()); 516 GI = Ref->getSecond(); 517 } 518 } 519 520 // Always add the terminator instruction, if one exists 521 auto T = MBB->getFirstTerminator(); 522 if (T != MBB->end()) { 523 auto EndBB = MaybeAddNode(&*T); 524 if (EndBB.second) 525 Builder.addEdge(LoopDepth, GI, EndBB.first); 526 GI = EndBB.first; 527 } 528 } 529 for (MachineBasicBlock *Succ : MBB->successors()) 530 TraverseCFG(Succ, GI, LoopDepth); 531 }; 532 // ArgNodeSentinel is a pseudo-instruction that represents MF args in the 533 // GadgetGraph 534 GraphIter ArgNode = MaybeAddNode(MachineGadgetGraph::ArgNodeSentinel).first; 535 TraverseCFG(&MF.front(), ArgNode, 0); 536 std::unique_ptr<MachineGadgetGraph> G{Builder.get(FenceCount, GadgetCount)}; 537 LLVM_DEBUG(dbgs() << "Found " << G->nodes_size() << " nodes\n"); 538 return G; 539 } 540 541 // Returns the number of remaining gadget edges that could not be eliminated 542 int X86LoadValueInjectionLoadHardeningPass::elimMitigatedEdgesAndNodes( 543 MachineGadgetGraph &G, EdgeSet &ElimEdges /* in, out */, 544 NodeSet &ElimNodes /* in, out */) const { 545 if (G.NumFences > 0) { 546 // Eliminate fences and CFG edges that ingress and egress the fence, as 547 // they are trivially mitigated. 548 for (const Edge &E : G.edges()) { 549 const Node *Dest = E.getDest(); 550 if (isFence(Dest->getValue())) { 551 ElimNodes.insert(*Dest); 552 ElimEdges.insert(E); 553 for (const Edge &DE : Dest->edges()) 554 ElimEdges.insert(DE); 555 } 556 } 557 } 558 559 // Find and eliminate gadget edges that have been mitigated. 560 int MitigatedGadgets = 0, RemainingGadgets = 0; 561 NodeSet ReachableNodes{G}; 562 for (const Node &RootN : G.nodes()) { 563 if (llvm::none_of(RootN.edges(), MachineGadgetGraph::isGadgetEdge)) 564 continue; // skip this node if it isn't a gadget source 565 566 // Find all of the nodes that are CFG-reachable from RootN using DFS 567 ReachableNodes.clear(); 568 std::function<void(const Node *, bool)> FindReachableNodes = 569 [&](const Node *N, bool FirstNode) { 570 if (!FirstNode) 571 ReachableNodes.insert(*N); 572 for (const Edge &E : N->edges()) { 573 const Node *Dest = E.getDest(); 574 if (MachineGadgetGraph::isCFGEdge(E) && !ElimEdges.contains(E) && 575 !ReachableNodes.contains(*Dest)) 576 FindReachableNodes(Dest, false); 577 } 578 }; 579 FindReachableNodes(&RootN, true); 580 581 // Any gadget whose sink is unreachable has been mitigated 582 for (const Edge &E : RootN.edges()) { 583 if (MachineGadgetGraph::isGadgetEdge(E)) { 584 if (ReachableNodes.contains(*E.getDest())) { 585 // This gadget's sink is reachable 586 ++RemainingGadgets; 587 } else { // This gadget's sink is unreachable, and therefore mitigated 588 ++MitigatedGadgets; 589 ElimEdges.insert(E); 590 } 591 } 592 } 593 } 594 return RemainingGadgets; 595 } 596 597 std::unique_ptr<MachineGadgetGraph> 598 X86LoadValueInjectionLoadHardeningPass::trimMitigatedEdges( 599 std::unique_ptr<MachineGadgetGraph> Graph) const { 600 NodeSet ElimNodes{*Graph}; 601 EdgeSet ElimEdges{*Graph}; 602 int RemainingGadgets = 603 elimMitigatedEdgesAndNodes(*Graph, ElimEdges, ElimNodes); 604 if (ElimEdges.empty() && ElimNodes.empty()) { 605 Graph->NumFences = 0; 606 Graph->NumGadgets = RemainingGadgets; 607 } else { 608 Graph = GraphBuilder::trim(*Graph, ElimNodes, ElimEdges, 0 /* NumFences */, 609 RemainingGadgets); 610 } 611 return Graph; 612 } 613 614 int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithPlugin( 615 MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const { 616 int FencesInserted = 0; 617 618 do { 619 LLVM_DEBUG(dbgs() << "Eliminating mitigated paths...\n"); 620 Graph = trimMitigatedEdges(std::move(Graph)); 621 LLVM_DEBUG(dbgs() << "Eliminating mitigated paths... Done\n"); 622 if (Graph->NumGadgets == 0) 623 break; 624 625 LLVM_DEBUG(dbgs() << "Cutting edges...\n"); 626 EdgeSet CutEdges{*Graph}; 627 auto Nodes = std::make_unique<unsigned int[]>(Graph->nodes_size() + 628 1 /* terminator node */); 629 auto Edges = std::make_unique<unsigned int[]>(Graph->edges_size()); 630 auto EdgeCuts = std::make_unique<int[]>(Graph->edges_size()); 631 auto EdgeValues = std::make_unique<int[]>(Graph->edges_size()); 632 for (const Node &N : Graph->nodes()) { 633 Nodes[Graph->getNodeIndex(N)] = Graph->getEdgeIndex(*N.edges_begin()); 634 } 635 Nodes[Graph->nodes_size()] = Graph->edges_size(); // terminator node 636 for (const Edge &E : Graph->edges()) { 637 Edges[Graph->getEdgeIndex(E)] = Graph->getNodeIndex(*E.getDest()); 638 EdgeValues[Graph->getEdgeIndex(E)] = E.getValue(); 639 } 640 OptimizeCut(Nodes.get(), Graph->nodes_size(), Edges.get(), EdgeValues.get(), 641 EdgeCuts.get(), Graph->edges_size()); 642 for (int I = 0; I < Graph->edges_size(); ++I) 643 if (EdgeCuts[I]) 644 CutEdges.set(I); 645 LLVM_DEBUG(dbgs() << "Cutting edges... Done\n"); 646 LLVM_DEBUG(dbgs() << "Cut " << CutEdges.count() << " edges\n"); 647 648 LLVM_DEBUG(dbgs() << "Inserting LFENCEs...\n"); 649 FencesInserted += insertFences(MF, *Graph, CutEdges); 650 LLVM_DEBUG(dbgs() << "Inserting LFENCEs... Done\n"); 651 LLVM_DEBUG(dbgs() << "Inserted " << FencesInserted << " fences\n"); 652 653 Graph = GraphBuilder::trim(*Graph, NodeSet{*Graph}, CutEdges); 654 } while (true); 655 656 return FencesInserted; 657 } 658 659 int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithHeuristic( 660 MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const { 661 // If `MF` does not have any fences, then no gadgets would have been 662 // mitigated at this point. 663 if (Graph->NumFences > 0) { 664 LLVM_DEBUG(dbgs() << "Eliminating mitigated paths...\n"); 665 Graph = trimMitigatedEdges(std::move(Graph)); 666 LLVM_DEBUG(dbgs() << "Eliminating mitigated paths... Done\n"); 667 } 668 669 if (Graph->NumGadgets == 0) 670 return 0; 671 672 LLVM_DEBUG(dbgs() << "Cutting edges...\n"); 673 EdgeSet CutEdges{*Graph}; 674 675 // Begin by collecting all ingress CFG edges for each node 676 DenseMap<const Node *, SmallVector<const Edge *, 2>> IngressEdgeMap; 677 for (const Edge &E : Graph->edges()) 678 if (MachineGadgetGraph::isCFGEdge(E)) 679 IngressEdgeMap[E.getDest()].push_back(&E); 680 681 // For each gadget edge, make cuts that guarantee the gadget will be 682 // mitigated. A computationally efficient way to achieve this is to either: 683 // (a) cut all egress CFG edges from the gadget source, or 684 // (b) cut all ingress CFG edges to the gadget sink. 685 // 686 // Moreover, the algorithm tries not to make a cut into a loop by preferring 687 // to make a (b)-type cut if the gadget source resides at a greater loop depth 688 // than the gadget sink, or an (a)-type cut otherwise. 689 for (const Node &N : Graph->nodes()) { 690 for (const Edge &E : N.edges()) { 691 if (!MachineGadgetGraph::isGadgetEdge(E)) 692 continue; 693 694 SmallVector<const Edge *, 2> EgressEdges; 695 SmallVector<const Edge *, 2> &IngressEdges = IngressEdgeMap[E.getDest()]; 696 for (const Edge &EgressEdge : N.edges()) 697 if (MachineGadgetGraph::isCFGEdge(EgressEdge)) 698 EgressEdges.push_back(&EgressEdge); 699 700 int EgressCutCost = 0, IngressCutCost = 0; 701 for (const Edge *EgressEdge : EgressEdges) 702 if (!CutEdges.contains(*EgressEdge)) 703 EgressCutCost += EgressEdge->getValue(); 704 for (const Edge *IngressEdge : IngressEdges) 705 if (!CutEdges.contains(*IngressEdge)) 706 IngressCutCost += IngressEdge->getValue(); 707 708 auto &EdgesToCut = 709 IngressCutCost < EgressCutCost ? IngressEdges : EgressEdges; 710 for (const Edge *E : EdgesToCut) 711 CutEdges.insert(*E); 712 } 713 } 714 LLVM_DEBUG(dbgs() << "Cutting edges... Done\n"); 715 LLVM_DEBUG(dbgs() << "Cut " << CutEdges.count() << " edges\n"); 716 717 LLVM_DEBUG(dbgs() << "Inserting LFENCEs...\n"); 718 int FencesInserted = insertFences(MF, *Graph, CutEdges); 719 LLVM_DEBUG(dbgs() << "Inserting LFENCEs... Done\n"); 720 LLVM_DEBUG(dbgs() << "Inserted " << FencesInserted << " fences\n"); 721 722 return FencesInserted; 723 } 724 725 int X86LoadValueInjectionLoadHardeningPass::insertFences( 726 MachineFunction &MF, MachineGadgetGraph &G, 727 EdgeSet &CutEdges /* in, out */) const { 728 int FencesInserted = 0; 729 for (const Node &N : G.nodes()) { 730 for (const Edge &E : N.edges()) { 731 if (CutEdges.contains(E)) { 732 MachineInstr *MI = N.getValue(), *Prev; 733 MachineBasicBlock *MBB; // Insert an LFENCE in this MBB 734 MachineBasicBlock::iterator InsertionPt; // ...at this point 735 if (MI == MachineGadgetGraph::ArgNodeSentinel) { 736 // insert LFENCE at beginning of entry block 737 MBB = &MF.front(); 738 InsertionPt = MBB->begin(); 739 Prev = nullptr; 740 } else if (MI->isBranch()) { // insert the LFENCE before the branch 741 MBB = MI->getParent(); 742 InsertionPt = MI; 743 Prev = MI->getPrevNode(); 744 // Remove all egress CFG edges from this branch because the inserted 745 // LFENCE prevents gadgets from crossing the branch. 746 for (const Edge &E : N.edges()) { 747 if (MachineGadgetGraph::isCFGEdge(E)) 748 CutEdges.insert(E); 749 } 750 } else { // insert the LFENCE after the instruction 751 MBB = MI->getParent(); 752 InsertionPt = MI->getNextNode() ? MI->getNextNode() : MBB->end(); 753 Prev = InsertionPt == MBB->end() 754 ? (MBB->empty() ? nullptr : &MBB->back()) 755 : InsertionPt->getPrevNode(); 756 } 757 // Ensure this insertion is not redundant (two LFENCEs in sequence). 758 if ((InsertionPt == MBB->end() || !isFence(&*InsertionPt)) && 759 (!Prev || !isFence(Prev))) { 760 BuildMI(*MBB, InsertionPt, DebugLoc(), TII->get(X86::LFENCE)); 761 ++FencesInserted; 762 } 763 } 764 } 765 } 766 return FencesInserted; 767 } 768 769 bool X86LoadValueInjectionLoadHardeningPass::instrUsesRegToAccessMemory( 770 const MachineInstr &MI, unsigned Reg) const { 771 if (!MI.mayLoadOrStore() || MI.getOpcode() == X86::MFENCE || 772 MI.getOpcode() == X86::SFENCE || MI.getOpcode() == X86::LFENCE) 773 return false; 774 775 // FIXME: This does not handle pseudo loading instruction like TCRETURN* 776 const MCInstrDesc &Desc = MI.getDesc(); 777 int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags); 778 if (MemRefBeginIdx < 0) { 779 LLVM_DEBUG(dbgs() << "Warning: unable to obtain memory operand for loading " 780 "instruction:\n"; 781 MI.print(dbgs()); dbgs() << '\n';); 782 return false; 783 } 784 MemRefBeginIdx += X86II::getOperandBias(Desc); 785 786 const MachineOperand &BaseMO = 787 MI.getOperand(MemRefBeginIdx + X86::AddrBaseReg); 788 const MachineOperand &IndexMO = 789 MI.getOperand(MemRefBeginIdx + X86::AddrIndexReg); 790 return (BaseMO.isReg() && BaseMO.getReg() != X86::NoRegister && 791 TRI->regsOverlap(BaseMO.getReg(), Reg)) || 792 (IndexMO.isReg() && IndexMO.getReg() != X86::NoRegister && 793 TRI->regsOverlap(IndexMO.getReg(), Reg)); 794 } 795 796 bool X86LoadValueInjectionLoadHardeningPass::instrUsesRegToBranch( 797 const MachineInstr &MI, unsigned Reg) const { 798 if (!MI.isConditionalBranch()) 799 return false; 800 for (const MachineOperand &Use : MI.uses()) 801 if (Use.isReg() && Use.getReg() == Reg) 802 return true; 803 return false; 804 } 805 806 INITIALIZE_PASS_BEGIN(X86LoadValueInjectionLoadHardeningPass, PASS_KEY, 807 "X86 LVI load hardening", false, false) 808 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo) 809 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree) 810 INITIALIZE_PASS_DEPENDENCY(MachineDominanceFrontier) 811 INITIALIZE_PASS_END(X86LoadValueInjectionLoadHardeningPass, PASS_KEY, 812 "X86 LVI load hardening", false, false) 813 814 FunctionPass *llvm::createX86LoadValueInjectionLoadHardeningPass() { 815 return new X86LoadValueInjectionLoadHardeningPass(); 816 } 817