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