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