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;
MachineGadgetGraph__anon2656d5e80111::MachineGadgetGraph123 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) {}
isCFGEdge__anon2656d5e80111::MachineGadgetGraph128 static inline bool isCFGEdge(const Edge &E) {
129 return E.getValue() != GadgetEdgeSentinel;
130 }
isGadgetEdge__anon2656d5e80111::MachineGadgetGraph131 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:
X86LoadValueInjectionLoadHardeningPass()140 X86LoadValueInjectionLoadHardeningPass() : MachineFunctionPass(ID) {}
141
getPassName() const142 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;
isFence(const MachineInstr * MI) const178 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
DOTGraphTraitsllvm::DOTGraphTraits201 DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
202
getNodeLabelllvm::DOTGraphTraits203 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
getNodeAttributesllvm::DOTGraphTraits213 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
getEdgeAttributesllvm::DOTGraphTraits222 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
getAnalysisUsage(AnalysisUsage & AU) const237 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
writeGadgetGraph(raw_ostream & OS,MachineFunction & MF,MachineGadgetGraph * G)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
runOnMachineFunction(MachineFunction & MF)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>
getGadgetGraph(MachineFunction & MF,const MachineLoopInfo & MLI,const MachineDominatorTree & MDT,const MachineDominanceFrontier & MDF) const326 X86LoadValueInjectionLoadHardeningPass::getGadgetGraph(
327 MachineFunction &MF, const MachineLoopInfo &MLI,
328 const MachineDominatorTree &MDT,
329 const MachineDominanceFrontier &MDF) const {
330 using namespace rdf;
331
332 // Build the Register Dataflow Graph using the RDF framework
333 TargetOperandInfo TOI{*TII};
334 DataFlowGraph DFG{MF, *TII, *TRI, MDT, MDF, TOI};
335 DFG.build();
336 Liveness L{MF.getRegInfo(), DFG};
337 L.computePhiInfo();
338
339 GraphBuilder Builder;
340 using GraphIter = typename GraphBuilder::BuilderNodeRef;
341 DenseMap<MachineInstr *, GraphIter> NodeMap;
342 int FenceCount = 0, GadgetCount = 0;
343 auto MaybeAddNode = [&NodeMap, &Builder](MachineInstr *MI) {
344 auto Ref = NodeMap.find(MI);
345 if (Ref == NodeMap.end()) {
346 auto I = Builder.addVertex(MI);
347 NodeMap[MI] = I;
348 return std::pair<GraphIter, bool>{I, true};
349 }
350 return std::pair<GraphIter, bool>{Ref->getSecond(), false};
351 };
352
353 // The `Transmitters` map memoizes transmitters found for each def. If a def
354 // has not yet been analyzed, then it will not appear in the map. If a def
355 // has been analyzed and was determined not to have any transmitters, then
356 // its list of transmitters will be empty.
357 DenseMap<NodeId, std::vector<NodeId>> Transmitters;
358
359 // Analyze all machine instructions to find gadgets and LFENCEs, adding
360 // each interesting value to `Nodes`
361 auto AnalyzeDef = [&](NodeAddr<DefNode *> SourceDef) {
362 SmallSet<NodeId, 8> UsesVisited, DefsVisited;
363 std::function<void(NodeAddr<DefNode *>)> AnalyzeDefUseChain =
364 [&](NodeAddr<DefNode *> Def) {
365 if (Transmitters.find(Def.Id) != Transmitters.end())
366 return; // Already analyzed `Def`
367
368 // Use RDF to find all the uses of `Def`
369 rdf::NodeSet Uses;
370 RegisterRef DefReg = Def.Addr->getRegRef(DFG);
371 for (auto UseID : L.getAllReachedUses(DefReg, Def)) {
372 auto Use = DFG.addr<UseNode *>(UseID);
373 if (Use.Addr->getFlags() & NodeAttrs::PhiRef) { // phi node
374 NodeAddr<PhiNode *> Phi = Use.Addr->getOwner(DFG);
375 for (auto I : L.getRealUses(Phi.Id)) {
376 if (DFG.getPRI().alias(RegisterRef(I.first), DefReg)) {
377 for (auto UA : I.second)
378 Uses.emplace(UA.first);
379 }
380 }
381 } else { // not a phi node
382 Uses.emplace(UseID);
383 }
384 }
385
386 // For each use of `Def`, we want to know whether:
387 // (1) The use can leak the Def'ed value,
388 // (2) The use can further propagate the Def'ed value to more defs
389 for (auto UseID : Uses) {
390 if (!UsesVisited.insert(UseID).second)
391 continue; // Already visited this use of `Def`
392
393 auto Use = DFG.addr<UseNode *>(UseID);
394 assert(!(Use.Addr->getFlags() & NodeAttrs::PhiRef));
395 MachineOperand &UseMO = Use.Addr->getOp();
396 MachineInstr &UseMI = *UseMO.getParent();
397 assert(UseMO.isReg());
398
399 // We naively assume that an instruction propagates any loaded
400 // uses to all defs unless the instruction is a call, in which
401 // case all arguments will be treated as gadget sources during
402 // analysis of the callee function.
403 if (UseMI.isCall())
404 continue;
405
406 // Check whether this use can transmit (leak) its value.
407 if (instrUsesRegToAccessMemory(UseMI, UseMO.getReg()) ||
408 (!NoConditionalBranches &&
409 instrUsesRegToBranch(UseMI, UseMO.getReg()))) {
410 Transmitters[Def.Id].push_back(Use.Addr->getOwner(DFG).Id);
411 if (UseMI.mayLoad())
412 continue; // Found a transmitting load -- no need to continue
413 // traversing its defs (i.e., this load will become
414 // a new gadget source anyways).
415 }
416
417 // Check whether the use propagates to more defs.
418 NodeAddr<InstrNode *> Owner{Use.Addr->getOwner(DFG)};
419 rdf::NodeList AnalyzedChildDefs;
420 for (auto &ChildDef :
421 Owner.Addr->members_if(DataFlowGraph::IsDef, DFG)) {
422 if (!DefsVisited.insert(ChildDef.Id).second)
423 continue; // Already visited this def
424 if (Def.Addr->getAttrs() & NodeAttrs::Dead)
425 continue;
426 if (Def.Id == ChildDef.Id)
427 continue; // `Def` uses itself (e.g., increment loop counter)
428
429 AnalyzeDefUseChain(ChildDef);
430
431 // `Def` inherits all of its child defs' transmitters.
432 for (auto TransmitterId : Transmitters[ChildDef.Id])
433 Transmitters[Def.Id].push_back(TransmitterId);
434 }
435 }
436
437 // Note that this statement adds `Def.Id` to the map if no
438 // transmitters were found for `Def`.
439 auto &DefTransmitters = Transmitters[Def.Id];
440
441 // Remove duplicate transmitters
442 llvm::sort(DefTransmitters);
443 DefTransmitters.erase(
444 std::unique(DefTransmitters.begin(), DefTransmitters.end()),
445 DefTransmitters.end());
446 };
447
448 // Find all of the transmitters
449 AnalyzeDefUseChain(SourceDef);
450 auto &SourceDefTransmitters = Transmitters[SourceDef.Id];
451 if (SourceDefTransmitters.empty())
452 return; // No transmitters for `SourceDef`
453
454 MachineInstr *Source = SourceDef.Addr->getFlags() & NodeAttrs::PhiRef
455 ? MachineGadgetGraph::ArgNodeSentinel
456 : SourceDef.Addr->getOp().getParent();
457 auto GadgetSource = MaybeAddNode(Source);
458 // Each transmitter is a sink for `SourceDef`.
459 for (auto TransmitterId : SourceDefTransmitters) {
460 MachineInstr *Sink = DFG.addr<StmtNode *>(TransmitterId).Addr->getCode();
461 auto GadgetSink = MaybeAddNode(Sink);
462 // Add the gadget edge to the graph.
463 Builder.addEdge(MachineGadgetGraph::GadgetEdgeSentinel,
464 GadgetSource.first, GadgetSink.first);
465 ++GadgetCount;
466 }
467 };
468
469 LLVM_DEBUG(dbgs() << "Analyzing def-use chains to find gadgets\n");
470 // Analyze function arguments
471 NodeAddr<BlockNode *> EntryBlock = DFG.getFunc().Addr->getEntryBlock(DFG);
472 for (NodeAddr<PhiNode *> ArgPhi :
473 EntryBlock.Addr->members_if(DataFlowGraph::IsPhi, DFG)) {
474 NodeList Defs = ArgPhi.Addr->members_if(DataFlowGraph::IsDef, DFG);
475 llvm::for_each(Defs, AnalyzeDef);
476 }
477 // Analyze every instruction in MF
478 for (NodeAddr<BlockNode *> BA : DFG.getFunc().Addr->members(DFG)) {
479 for (NodeAddr<StmtNode *> SA :
480 BA.Addr->members_if(DataFlowGraph::IsCode<NodeAttrs::Stmt>, DFG)) {
481 MachineInstr *MI = SA.Addr->getCode();
482 if (isFence(MI)) {
483 MaybeAddNode(MI);
484 ++FenceCount;
485 } else if (MI->mayLoad()) {
486 NodeList Defs = SA.Addr->members_if(DataFlowGraph::IsDef, DFG);
487 llvm::for_each(Defs, AnalyzeDef);
488 }
489 }
490 }
491 LLVM_DEBUG(dbgs() << "Found " << FenceCount << " fences\n");
492 LLVM_DEBUG(dbgs() << "Found " << GadgetCount << " gadgets\n");
493 if (GadgetCount == 0)
494 return nullptr;
495 NumGadgets += GadgetCount;
496
497 // Traverse CFG to build the rest of the graph
498 SmallSet<MachineBasicBlock *, 8> BlocksVisited;
499 std::function<void(MachineBasicBlock *, GraphIter, unsigned)> TraverseCFG =
500 [&](MachineBasicBlock *MBB, GraphIter GI, unsigned ParentDepth) {
501 unsigned LoopDepth = MLI.getLoopDepth(MBB);
502 if (!MBB->empty()) {
503 // Always add the first instruction in each block
504 auto NI = MBB->begin();
505 auto BeginBB = MaybeAddNode(&*NI);
506 Builder.addEdge(ParentDepth, GI, BeginBB.first);
507 if (!BlocksVisited.insert(MBB).second)
508 return;
509
510 // Add any instructions within the block that are gadget components
511 GI = BeginBB.first;
512 while (++NI != MBB->end()) {
513 auto Ref = NodeMap.find(&*NI);
514 if (Ref != NodeMap.end()) {
515 Builder.addEdge(LoopDepth, GI, Ref->getSecond());
516 GI = Ref->getSecond();
517 }
518 }
519
520 // Always add the terminator instruction, if one exists
521 auto T = MBB->getFirstTerminator();
522 if (T != MBB->end()) {
523 auto EndBB = MaybeAddNode(&*T);
524 if (EndBB.second)
525 Builder.addEdge(LoopDepth, GI, EndBB.first);
526 GI = EndBB.first;
527 }
528 }
529 for (MachineBasicBlock *Succ : MBB->successors())
530 TraverseCFG(Succ, GI, LoopDepth);
531 };
532 // ArgNodeSentinel is a pseudo-instruction that represents MF args in the
533 // GadgetGraph
534 GraphIter ArgNode = MaybeAddNode(MachineGadgetGraph::ArgNodeSentinel).first;
535 TraverseCFG(&MF.front(), ArgNode, 0);
536 std::unique_ptr<MachineGadgetGraph> G{Builder.get(FenceCount, GadgetCount)};
537 LLVM_DEBUG(dbgs() << "Found " << G->nodes_size() << " nodes\n");
538 return G;
539 }
540
541 // Returns the number of remaining gadget edges that could not be eliminated
elimMitigatedEdgesAndNodes(MachineGadgetGraph & G,EdgeSet & ElimEdges,NodeSet & ElimNodes) const542 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>
trimMitigatedEdges(std::unique_ptr<MachineGadgetGraph> Graph) const598 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
hardenLoadsWithPlugin(MachineFunction & MF,std::unique_ptr<MachineGadgetGraph> Graph) const614 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
hardenLoadsWithHeuristic(MachineFunction & MF,std::unique_ptr<MachineGadgetGraph> Graph) const659 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
insertFences(MachineFunction & MF,MachineGadgetGraph & G,EdgeSet & CutEdges) const725 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
instrUsesRegToAccessMemory(const MachineInstr & MI,unsigned Reg) const769 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
instrUsesRegToBranch(const MachineInstr & MI,unsigned Reg) const796 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)
INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)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