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