contrib/llvm-project/llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.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-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/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/MachineRegisterInfo.h"
  57 #include "llvm/CodeGen/RDFGraph.h"
  58 #include "llvm/CodeGen/RDFLiveness.h"
  59 #include "llvm/InitializePasses.h"
  60 #include "llvm/Support/CommandLine.h"
  61 #include "llvm/Support/DOTGraphTraits.h"
  62 #include "llvm/Support/Debug.h"
  63 #include "llvm/Support/DynamicLibrary.h"
  64 #include "llvm/Support/GraphWriter.h"
  65 #include "llvm/Support/raw_ostream.h"
  66
  67 using namespace llvm;
  68
  69 #define PASS_KEY "x86-lvi-load"
  70 #define DEBUG_TYPE PASS_KEY
  71
  72 STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation");
  73 STATISTIC(NumFunctionsConsidered, "Number of functions analyzed");
  74 STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations "
  75                                  "were deployed");
  76 STATISTIC(NumGadgets, "Number of LVI gadgets detected during analysis");
  77
  78 static cl::opt<std::string> OptimizePluginPath(
  79     PASS_KEY "-opt-plugin",
  80     cl::desc("Specify a plugin to optimize LFENCE insertion"), cl::Hidden);
  81
  82 static cl::opt<bool> NoConditionalBranches(
  83     PASS_KEY "-no-cbranch",
  84     cl::desc("Don't treat conditional branches as disclosure gadgets. This "
  85              "may improve performance, at the cost of security."),
  86     cl::init(false), cl::Hidden);
  87
  88 static cl::opt<bool> EmitDot(
  89     PASS_KEY "-dot",
  90     cl::desc(
  91         "For each function, emit a dot graph depicting potential LVI gadgets"),
  92     cl::init(false), cl::Hidden);
  93
  94 static cl::opt<bool> EmitDotOnly(
  95     PASS_KEY "-dot-only",
  96     cl::desc("For each function, emit a dot graph depicting potential LVI "
  97              "gadgets, and do not insert any fences"),
  98     cl::init(false), cl::Hidden);
  99
 100 static cl::opt<bool> EmitDotVerify(
 101     PASS_KEY "-dot-verify",
 102     cl::desc("For each function, emit a dot graph to stdout depicting "
 103              "potential LVI gadgets, used for testing purposes only"),
 104     cl::init(false), cl::Hidden);
 105
 106 static llvm::sys::DynamicLibrary OptimizeDL;
 107 typedef int (*OptimizeCutT)(unsigned int *nodes, unsigned int nodes_size,
 108                             unsigned int *edges, int *edge_values,
 109                             int *cut_edges /* out */, unsigned int edges_size);
 110 static OptimizeCutT OptimizeCut = nullptr;
 111
 112 namespace {
 113
 114 struct MachineGadgetGraph : ImmutableGraph<MachineInstr *, int> {
 115   static constexpr int GadgetEdgeSentinel = -1;
 116   static constexpr MachineInstr *const ArgNodeSentinel = nullptr;
 117
 118   using GraphT = ImmutableGraph<MachineInstr *, int>;
 119   using Node = typename GraphT::Node;
 120   using Edge = typename GraphT::Edge;
 121   using size_type = typename GraphT::size_type;
 122   MachineGadgetGraph(std::unique_ptr<Node[]> Nodes,
 123                      std::unique_ptr<Edge[]> Edges, size_type NodesSize,
 124                      size_type EdgesSize, int NumFences = 0, int NumGadgets = 0)
 125       : GraphT(std::move(Nodes), std::move(Edges), NodesSize, EdgesSize),
 126         NumFences(NumFences), NumGadgets(NumGadgets) {}
 127   static inline bool isCFGEdge(const Edge &E) {
 128     return E.getValue() != GadgetEdgeSentinel;
 129   }
 130   static inline bool isGadgetEdge(const Edge &E) {
 131     return E.getValue() == GadgetEdgeSentinel;
 132   }
 133   int NumFences;
 134   int NumGadgets;
 135 };
 136
 137 class X86LoadValueInjectionLoadHardeningPass : public MachineFunctionPass {
 138 public:
 139   X86LoadValueInjectionLoadHardeningPass() : MachineFunctionPass(ID) {}
 140
 141   StringRef getPassName() const override {
 142     return "X86 Load Value Injection (LVI) Load Hardening";
 143   }
 144   void getAnalysisUsage(AnalysisUsage &AU) const override;
 145   bool runOnMachineFunction(MachineFunction &MF) override;
 146
 147   static char ID;
 148
 149 private:
 150   using GraphBuilder = ImmutableGraphBuilder<MachineGadgetGraph>;
 151   using EdgeSet = MachineGadgetGraph::EdgeSet;
 152   using NodeSet = MachineGadgetGraph::NodeSet;
 153   using Gadget = std::pair<MachineInstr *, MachineInstr *>;
 154
 155   const X86Subtarget *STI;
 156   const TargetInstrInfo *TII;
 157   const TargetRegisterInfo *TRI;
 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 hardenLoadsWithGreedyHeuristic(
 166       MachineFunction &MF, 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   void findAndCutEdges(MachineGadgetGraph &G,
 173                        EdgeSet &CutEdges /* out */) 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 = hardenLoadsWithGreedyHeuristic(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 = DFG.getPRI().normalize(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, MachineGadgetGraph::EdgeSet &ElimEdges /* in, out */,
 544     MachineGadgetGraph::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 auto &E : G.edges()) {
 549       const MachineGadgetGraph::Node *Dest = E.getDest();
 550       if (isFence(Dest->getValue())) {
 551         ElimNodes.insert(*Dest);
 552         ElimEdges.insert(E);
 553         for (const auto &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   MachineGadgetGraph::NodeSet ReachableNodes{G};
 562   for (const auto &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 MachineGadgetGraph::Node *, bool)>
 569         FindReachableNodes =
 570             [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
 571               if (!FirstNode)
 572                 ReachableNodes.insert(*N);
 573               for (const auto &E : N->edges()) {
 574                 const MachineGadgetGraph::Node *Dest = E.getDest();
 575                 if (MachineGadgetGraph::isCFGEdge(E) &&
 576                     !ElimEdges.contains(E) && !ReachableNodes.contains(*Dest))
 577                   FindReachableNodes(Dest, false);
 578               }
 579             };
 580     FindReachableNodes(&RootN, true);
 581
 582     // Any gadget whose sink is unreachable has been mitigated
 583     for (const auto &E : RootN.edges()) {
 584       if (MachineGadgetGraph::isGadgetEdge(E)) {
 585         if (ReachableNodes.contains(*E.getDest())) {
 586           // This gadget's sink is reachable
 587           ++RemainingGadgets;
 588         } else { // This gadget's sink is unreachable, and therefore mitigated
 589           ++MitigatedGadgets;
 590           ElimEdges.insert(E);
 591         }
 592       }
 593     }
 594   }
 595   return RemainingGadgets;
 596 }
 597
 598 std::unique_ptr<MachineGadgetGraph>
 599 X86LoadValueInjectionLoadHardeningPass::trimMitigatedEdges(
 600     std::unique_ptr<MachineGadgetGraph> Graph) const {
 601   MachineGadgetGraph::NodeSet ElimNodes{*Graph};
 602   MachineGadgetGraph::EdgeSet ElimEdges{*Graph};
 603   int RemainingGadgets =
 604       elimMitigatedEdgesAndNodes(*Graph, ElimEdges, ElimNodes);
 605   if (ElimEdges.empty() && ElimNodes.empty()) {
 606     Graph->NumFences = 0;
 607     Graph->NumGadgets = RemainingGadgets;
 608   } else {
 609     Graph = GraphBuilder::trim(*Graph, ElimNodes, ElimEdges, 0 /* NumFences */,
 610                                RemainingGadgets);
 611   }
 612   return Graph;
 613 }
 614
 615 int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithPlugin(
 616     MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const {
 617   int FencesInserted = 0;
 618
 619   do {
 620     LLVM_DEBUG(dbgs() << "Eliminating mitigated paths...\n");
 621     Graph = trimMitigatedEdges(std::move(Graph));
 622     LLVM_DEBUG(dbgs() << "Eliminating mitigated paths... Done\n");
 623     if (Graph->NumGadgets == 0)
 624       break;
 625
 626     LLVM_DEBUG(dbgs() << "Cutting edges...\n");
 627     EdgeSet CutEdges{*Graph};
 628     auto Nodes = std::make_unique<unsigned int[]>(Graph->nodes_size() +
 629                                                   1 /* terminator node */);
 630     auto Edges = std::make_unique<unsigned int[]>(Graph->edges_size());
 631     auto EdgeCuts = std::make_unique<int[]>(Graph->edges_size());
 632     auto EdgeValues = std::make_unique<int[]>(Graph->edges_size());
 633     for (const auto &N : Graph->nodes()) {
 634       Nodes[Graph->getNodeIndex(N)] = Graph->getEdgeIndex(*N.edges_begin());
 635     }
 636     Nodes[Graph->nodes_size()] = Graph->edges_size(); // terminator node
 637     for (const auto &E : Graph->edges()) {
 638       Edges[Graph->getEdgeIndex(E)] = Graph->getNodeIndex(*E.getDest());
 639       EdgeValues[Graph->getEdgeIndex(E)] = E.getValue();
 640     }
 641     OptimizeCut(Nodes.get(), Graph->nodes_size(), Edges.get(), EdgeValues.get(),
 642                 EdgeCuts.get(), Graph->edges_size());
 643     for (int I = 0; I < Graph->edges_size(); ++I)
 644       if (EdgeCuts[I])
 645         CutEdges.set(I);
 646     LLVM_DEBUG(dbgs() << "Cutting edges... Done\n");
 647     LLVM_DEBUG(dbgs() << "Cut " << CutEdges.count() << " edges\n");
 648
 649     LLVM_DEBUG(dbgs() << "Inserting LFENCEs...\n");
 650     FencesInserted += insertFences(MF, *Graph, CutEdges);
 651     LLVM_DEBUG(dbgs() << "Inserting LFENCEs... Done\n");
 652     LLVM_DEBUG(dbgs() << "Inserted " << FencesInserted << " fences\n");
 653
 654     Graph = GraphBuilder::trim(*Graph, MachineGadgetGraph::NodeSet{*Graph},
 655                                CutEdges);
 656   } while (true);
 657
 658   return FencesInserted;
 659 }
 660
 661 int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithGreedyHeuristic(
 662     MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const {
 663   LLVM_DEBUG(dbgs() << "Eliminating mitigated paths...\n");
 664   Graph = trimMitigatedEdges(std::move(Graph));
 665   LLVM_DEBUG(dbgs() << "Eliminating mitigated paths... Done\n");
 666   if (Graph->NumGadgets == 0)
 667     return 0;
 668
 669   LLVM_DEBUG(dbgs() << "Cutting edges...\n");
 670   MachineGadgetGraph::NodeSet ElimNodes{*Graph}, GadgetSinks{*Graph};
 671   MachineGadgetGraph::EdgeSet ElimEdges{*Graph}, CutEdges{*Graph};
 672   auto IsCFGEdge = [&ElimEdges, &CutEdges](const MachineGadgetGraph::Edge &E) {
 673     return !ElimEdges.contains(E) && !CutEdges.contains(E) &&
 674            MachineGadgetGraph::isCFGEdge(E);
 675   };
 676   auto IsGadgetEdge = [&ElimEdges,
 677                        &CutEdges](const MachineGadgetGraph::Edge &E) {
 678     return !ElimEdges.contains(E) && !CutEdges.contains(E) &&
 679            MachineGadgetGraph::isGadgetEdge(E);
 680   };
 681
 682   // FIXME: this is O(E^2), we could probably do better.
 683   do {
 684     // Find the cheapest CFG edge that will eliminate a gadget (by being
 685     // egress from a SOURCE node or ingress to a SINK node), and cut it.
 686     const MachineGadgetGraph::Edge *CheapestSoFar = nullptr;
 687
 688     // First, collect all gadget source and sink nodes.
 689     MachineGadgetGraph::NodeSet GadgetSources{*Graph}, GadgetSinks{*Graph};
 690     for (const auto &N : Graph->nodes()) {
 691       if (ElimNodes.contains(N))
 692         continue;
 693       for (const auto &E : N.edges()) {
 694         if (IsGadgetEdge(E)) {
 695           GadgetSources.insert(N);
 696           GadgetSinks.insert(*E.getDest());
 697         }
 698       }
 699     }
 700
 701     // Next, look for the cheapest CFG edge which, when cut, is guaranteed to
 702     // mitigate at least one gadget by either:
 703     // (a) being egress from a gadget source, or
 704     // (b) being ingress to a gadget sink.
 705     for (const auto &N : Graph->nodes()) {
 706       if (ElimNodes.contains(N))
 707         continue;
 708       for (const auto &E : N.edges()) {
 709         if (IsCFGEdge(E)) {
 710           if (GadgetSources.contains(N) || GadgetSinks.contains(*E.getDest())) {
 711             if (!CheapestSoFar || E.getValue() < CheapestSoFar->getValue())
 712               CheapestSoFar = &E;
 713           }
 714         }
 715       }
 716     }
 717
 718     assert(CheapestSoFar && "Failed to cut an edge");
 719     CutEdges.insert(*CheapestSoFar);
 720     ElimEdges.insert(*CheapestSoFar);
 721   } while (elimMitigatedEdgesAndNodes(*Graph, ElimEdges, ElimNodes));
 722   LLVM_DEBUG(dbgs() << "Cutting edges... Done\n");
 723   LLVM_DEBUG(dbgs() << "Cut " << CutEdges.count() << " edges\n");
 724
 725   LLVM_DEBUG(dbgs() << "Inserting LFENCEs...\n");
 726   int FencesInserted = insertFences(MF, *Graph, CutEdges);
 727   LLVM_DEBUG(dbgs() << "Inserting LFENCEs... Done\n");
 728   LLVM_DEBUG(dbgs() << "Inserted " << FencesInserted << " fences\n");
 729
 730   return FencesInserted;
 731 }
 732
 733 int X86LoadValueInjectionLoadHardeningPass::insertFences(
 734     MachineFunction &MF, MachineGadgetGraph &G,
 735     EdgeSet &CutEdges /* in, out */) const {
 736   int FencesInserted = 0;
 737   for (const auto &N : G.nodes()) {
 738     for (const auto &E : N.edges()) {
 739       if (CutEdges.contains(E)) {
 740         MachineInstr *MI = N.getValue(), *Prev;
 741         MachineBasicBlock *MBB;                  // Insert an LFENCE in this MBB
 742         MachineBasicBlock::iterator InsertionPt; // ...at this point
 743         if (MI == MachineGadgetGraph::ArgNodeSentinel) {
 744           // insert LFENCE at beginning of entry block
 745           MBB = &MF.front();
 746           InsertionPt = MBB->begin();
 747           Prev = nullptr;
 748         } else if (MI->isBranch()) { // insert the LFENCE before the branch
 749           MBB = MI->getParent();
 750           InsertionPt = MI;
 751           Prev = MI->getPrevNode();
 752           // Remove all egress CFG edges from this branch because the inserted
 753           // LFENCE prevents gadgets from crossing the branch.
 754           for (const auto &E : N.edges()) {
 755             if (MachineGadgetGraph::isCFGEdge(E))
 756               CutEdges.insert(E);
 757           }
 758         } else { // insert the LFENCE after the instruction
 759           MBB = MI->getParent();
 760           InsertionPt = MI->getNextNode() ? MI->getNextNode() : MBB->end();
 761           Prev = InsertionPt == MBB->end()
 762                      ? (MBB->empty() ? nullptr : &MBB->back())
 763                      : InsertionPt->getPrevNode();
 764         }
 765         // Ensure this insertion is not redundant (two LFENCEs in sequence).
 766         if ((InsertionPt == MBB->end() || !isFence(&*InsertionPt)) &&
 767             (!Prev || !isFence(Prev))) {
 768           BuildMI(*MBB, InsertionPt, DebugLoc(), TII->get(X86::LFENCE));
 769           ++FencesInserted;
 770         }
 771       }
 772     }
 773   }
 774   return FencesInserted;
 775 }
 776
 777 bool X86LoadValueInjectionLoadHardeningPass::instrUsesRegToAccessMemory(
 778     const MachineInstr &MI, unsigned Reg) const {
 779   if (!MI.mayLoadOrStore() || MI.getOpcode() == X86::MFENCE ||
 780       MI.getOpcode() == X86::SFENCE || MI.getOpcode() == X86::LFENCE)
 781     return false;
 782
 783   // FIXME: This does not handle pseudo loading instruction like TCRETURN*
 784   const MCInstrDesc &Desc = MI.getDesc();
 785   int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
 786   if (MemRefBeginIdx < 0) {
 787     LLVM_DEBUG(dbgs() << "Warning: unable to obtain memory operand for loading "
 788                          "instruction:\n";
 789                MI.print(dbgs()); dbgs() << '\n';);
 790     return false;
 791   }
 792   MemRefBeginIdx += X86II::getOperandBias(Desc);
 793
 794   const MachineOperand &BaseMO =
 795       MI.getOperand(MemRefBeginIdx + X86::AddrBaseReg);
 796   const MachineOperand &IndexMO =
 797       MI.getOperand(MemRefBeginIdx + X86::AddrIndexReg);
 798   return (BaseMO.isReg() && BaseMO.getReg() != X86::NoRegister &&
 799           TRI->regsOverlap(BaseMO.getReg(), Reg)) ||
 800          (IndexMO.isReg() && IndexMO.getReg() != X86::NoRegister &&
 801           TRI->regsOverlap(IndexMO.getReg(), Reg));
 802 }
 803
 804 bool X86LoadValueInjectionLoadHardeningPass::instrUsesRegToBranch(
 805     const MachineInstr &MI, unsigned Reg) const {
 806   if (!MI.isConditionalBranch())
 807     return false;
 808   for (const MachineOperand &Use : MI.uses())
 809     if (Use.isReg() && Use.getReg() == Reg)
 810       return true;
 811   return false;
 812 }
 813
 814 INITIALIZE_PASS_BEGIN(X86LoadValueInjectionLoadHardeningPass, PASS_KEY,
 815                       "X86 LVI load hardening", false, false)
 816 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 817 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
 818 INITIALIZE_PASS_DEPENDENCY(MachineDominanceFrontier)
 819 INITIALIZE_PASS_END(X86LoadValueInjectionLoadHardeningPass, PASS_KEY,
 820                     "X86 LVI load hardening", false, false)
 821
 822 FunctionPass *llvm::createX86LoadValueInjectionLoadHardeningPass() {
 823   return new X86LoadValueInjectionLoadHardeningPass();
 824 }
 825
 826 namespace {
 827
 828 /// The `X86LoadValueInjectionLoadHardeningPass` above depends on expensive
 829 /// analysis passes that add complexity to the pipeline. This complexity
 830 /// can cause noticable overhead when no optimizations are enabled, i.e., -O0.
 831 /// The purpose of `X86LoadValueInjectionLoadHardeningUnoptimizedPass` is to
 832 /// provide the same security as the optimized pass, but without adding
 833 /// unnecessary complexity to the LLVM pipeline.
 834 ///
 835 /// The behavior of this pass is simply to insert an LFENCE after every load
 836 /// instruction.
 837 class X86LoadValueInjectionLoadHardeningUnoptimizedPass
 838     : public MachineFunctionPass {
 839 public:
 840   X86LoadValueInjectionLoadHardeningUnoptimizedPass()
 841       : MachineFunctionPass(ID) {}
 842
 843   StringRef getPassName() const override {
 844     return "X86 Load Value Injection (LVI) Load Hardening (Unoptimized)";
 845   }
 846   bool runOnMachineFunction(MachineFunction &MF) override;
 847   static char ID;
 848 };
 849
 850 } // end anonymous namespace
 851
 852 char X86LoadValueInjectionLoadHardeningUnoptimizedPass::ID = 0;
 853
 854 bool X86LoadValueInjectionLoadHardeningUnoptimizedPass::runOnMachineFunction(
 855     MachineFunction &MF) {
 856   LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName()
 857                     << " *****\n");
 858   const X86Subtarget *STI = &MF.getSubtarget<X86Subtarget>();
 859   if (!STI->useLVILoadHardening())
 860     return false;
 861
 862   // FIXME: support 32-bit
 863   if (!STI->is64Bit())
 864     report_fatal_error("LVI load hardening is only supported on 64-bit", false);
 865
 866   // Don't skip functions with the "optnone" attr but participate in opt-bisect.
 867   const Function &F = MF.getFunction();
 868   if (!F.hasOptNone() && skipFunction(F))
 869     return false;
 870
 871   bool Modified = false;
 872   ++NumFunctionsConsidered;
 873
 874   const TargetInstrInfo *TII = STI->getInstrInfo();
 875   for (auto &MBB : MF) {
 876     for (auto &MI : MBB) {
 877       if (!MI.mayLoad() || MI.getOpcode() == X86::LFENCE ||
 878           MI.getOpcode() == X86::MFENCE)
 879         continue;
 880
 881       MachineBasicBlock::iterator InsertionPt =
 882           MI.getNextNode() ? MI.getNextNode() : MBB.end();
 883       BuildMI(MBB, InsertionPt, DebugLoc(), TII->get(X86::LFENCE));
 884       ++NumFences;
 885       Modified = true;
 886     }
 887   }
 888
 889   if (Modified)
 890     ++NumFunctionsMitigated;
 891
 892   return Modified;
 893 }
 894
 895 INITIALIZE_PASS(X86LoadValueInjectionLoadHardeningUnoptimizedPass, PASS_KEY,
 896                 "X86 LVI load hardening", false, false)
 897
 898 FunctionPass *llvm::createX86LoadValueInjectionLoadHardeningUnoptimizedPass() {
 899   return new X86LoadValueInjectionLoadHardeningUnoptimizedPass();
 900 }