X86LoadValueInjectionLoadHardening.cpp revision 363496 
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
67using namespace llvm;
68
69#define PASS_KEY "x86-lvi-load"
70#define DEBUG_TYPE PASS_KEY
71
72STATISTIC(NumFences, "Number of LFENCEs inserted for LVI mitigation");
73STATISTIC(NumFunctionsConsidered, "Number of functions analyzed");
74STATISTIC(NumFunctionsMitigated, "Number of functions for which mitigations "
75                                 "were deployed");
76STATISTIC(NumGadgets, "Number of LVI gadgets detected during analysis");
77
78static cl::opt<std::string> OptimizePluginPath(
79    PASS_KEY "-opt-plugin",
80    cl::desc("Specify a plugin to optimize LFENCE insertion"), cl::Hidden);
81
82static 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
88static 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
94static 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
100static 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
106static llvm::sys::DynamicLibrary OptimizeDL;
107typedef 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);
110static OptimizeCutT OptimizeCut = nullptr;
111
112namespace {
113
114struct 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
137class X86LoadValueInjectionLoadHardeningPass : public MachineFunctionPass {
138public:
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
149private:
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
186namespace llvm {
187
188template <>
189struct GraphTraits<MachineGadgetGraph *>
190    : GraphTraits<ImmutableGraph<MachineInstr *, int> *> {};
191
192template <>
193struct 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
232constexpr MachineInstr *MachineGadgetGraph::ArgNodeSentinel;
233constexpr int MachineGadgetGraph::GadgetEdgeSentinel;
234
235char X86LoadValueInjectionLoadHardeningPass::ID = 0;
236
237void 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
246static 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
252bool 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
325std::unique_ptr<MachineGadgetGraph>
326X86LoadValueInjectionLoadHardeningPass::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
542int 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
598std::unique_ptr<MachineGadgetGraph>
599X86LoadValueInjectionLoadHardeningPass::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
615int 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
661int 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
733int 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
777bool 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
804bool 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
814INITIALIZE_PASS_BEGIN(X86LoadValueInjectionLoadHardeningPass, PASS_KEY,
815                      "X86 LVI load hardening", false, false)
816INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
817INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
818INITIALIZE_PASS_DEPENDENCY(MachineDominanceFrontier)
819INITIALIZE_PASS_END(X86LoadValueInjectionLoadHardeningPass, PASS_KEY,
820                    "X86 LVI load hardening", false, false)
821
822FunctionPass *llvm::createX86LoadValueInjectionLoadHardeningPass() {
823  return new X86LoadValueInjectionLoadHardeningPass();
824}
825
826namespace {
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.
837class X86LoadValueInjectionLoadHardeningUnoptimizedPass
838    : public MachineFunctionPass {
839public:
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
852char X86LoadValueInjectionLoadHardeningUnoptimizedPass::ID = 0;
853
854bool 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
895INITIALIZE_PASS(X86LoadValueInjectionLoadHardeningUnoptimizedPass, PASS_KEY,
896                "X86 LVI load hardening", false, false)
897
898FunctionPass *llvm::createX86LoadValueInjectionLoadHardeningUnoptimizedPass() {
899  return new X86LoadValueInjectionLoadHardeningUnoptimizedPass();
900}
901