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