[PATCH] D75937: Add Support to X86 for Load Hardening to Mitigate Load Value Injection (LVI) [5/6]

Thu Apr 23 15:15:13 PDT 2020

sconstab added a comment.

@mattdr Thanks for the very scrupulous review!

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:351
+
+    Graph = GraphBuilder::trim(
+        *ElimGraph, MachineGadgetGraph::NodeSet{*ElimGraph}, CutEdges);
----------------
mattdr wrote:
> I know there was a lot of effort to make `ImmutableGraph` efficient, but this is an `O(N)` reallocation on every loop, right?
> 
> In practice, how many times do we end up looping before we hit a fixed point?
The answers to your questions are "yes" and "until all of the gadgets have been mitigated". The intent here is to optimize for the MILP plugin:
https://github.com/intel/lvi-llvm-optimization-plugin
The plugin eliminates lots of gadgets all at once, and therefore rebuilding the graph is a relatively rare event. I'm not sure it would be worthwhile to use an entirely different data structure and algorithm for the greedy mitigation strategy.

What I have done is rewrite the Greedy heuristic to keep track of which edges have been cut and eliminated, thus never actually having to rebuild the graph.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:529
+std::unique_ptr<MachineGadgetGraph>
+X86LoadValueInjectionLoadHardeningPass::elimEdges(
+    std::unique_ptr<MachineGadgetGraph> Graph) const {
----------------
mattdr wrote:
> Could be more descriptive to call this, say, `trimMitigatedEdges`
Yes, done.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:537
+      const MachineGadgetGraph::Node *Dest = E.getDest();
+      if (isFence(Dest->getValue())) {
+        ElimNodes.insert(*Dest);
----------------
mattdr wrote:
> In what circumstance will we add fences to the graph as sources or sinks? Can we just avoid that at the point of insertion, rather than running this extra culling pass on every iteration?
> 
Fences are never added as sources or sinks. We need to know where the fences are so that we can eliminate gadget edges for which all CFG paths from source to sink cross a fence. We //could// perform this analysis at the point of insertion, i.e., during the `getGadgetGraph()` function, BUT this analysis would be much more expensive there because it would be performed on the actual `MachineFunction`, instead of on the gadget graph.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:555
+  for (const auto &RootN : Graph->nodes()) {
+    // collect the gadgets for this node
+    for (const auto &E : RootN.edges()) {
----------------
mattdr wrote:
> I would write something stronger here:
> 
> Start off assuming all gadgets originating at this source node have been mitigated and can be removed. Later we will add back unmitigated gadgets by erasing them from the removal set.
> 
Well, I substantially rewrote this algorithm to build a set of mitigated edges, instead of trimming down a set of mitigated edges. I think that the new set of comments should be much easier to follow.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:565
+      continue;
+    std::function<void(const MachineGadgetGraph::Node *, bool)> TraverseDFS =
+        [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
----------------
mattdr wrote:
> This is really, really hard to read and understand. I think that's in large part because we have this one graph that represents _both_ control-flow _and_ source/sink pairs.
> 
> Given that it's the load-bearing part of the whole stack, let me offer the best way I've come up with to explain it, then a suggestion for making it simpler to follow.
> 
> My summary based on a few readings:
> 
> First we compute `GadgetSinks`, the set of gadget sinks whose source is the current root. Then we do a depth-first search of the //control-flow graph// to find all gadget sinks that are control-flow-reachable from the given root. When we find such a sink, we look to see if it is **also** in `GadgetSinks` -- again, a sink whose source is the current root -- at which point we know we have found an unmitigated gadget sink. We iterate through the root's gadget edges to find the specific edge that points to the current DFS node -- the unmitigated gadget sink -- and remove that edge from `ElimGadgets`, where we had previously added it on the presumption it was mitigated.
> 
> One idea for making this easier to follow:
> 
> 1. Start off assuming no edges are mitigated. Don't pre-populate the `EdgeSet` with every edge.
> 2. Use depth-first search to create `ReachableSinks`, the set of all gadget sinks (regardless of source) that are control-flow-reachable from `RootN`. (I don't think this adds more work than `TraverseDFS` already does.)
> 3. Iterate over gadget edges from `RootN`. For each edge, ask: is the destination reachable? If so, add this to a new set, `GadgetEdges`.
> 4. After iterating over all nodes, let `TrimEdges` be `AllEdges - GadgetEdges`. Then trim those.
> 
You of course are correct that this was terribly difficult to follow. I actually found an even simpler solution that what you had suggested. For each `RootN` that is the source of at least one gadget, I DFS to find all of the `Visited` nodes, and then check to see which node members of `Visited` are also the destinations of a gadget edge rooted at `RootN`.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:591
+
+  if (!(ElimEdges.empty() && ElimNodes.empty())) {
+    int NumRemainingGadgets = NumGadgets - ElimGadgets.count();
----------------
mattdr wrote:
> It's weird to have a negated predicate like this in a conditional with an `else`. Let's either rewrite it or un-negate it and flip the `if`/`else` blocks.
> 
> Rewrite: `!ElimEdges.empty() || !ElimNodes.empty()`
> 
Right. By the way, I decomposed `elimEdges()` into two functions:
- `elimEdges()` just runs the edge elimination algorithm without rebuilding the gadget graph.
- `trimMitigatedEdges()` wraps `elimEdges()` and does rebuild (i.e., "trim") the gadget graph.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:643
+          for (const auto &EE : N.edges()) {
+            if (MachineGadgetGraph::isCFGEdge(EE)) {
+              if (!CheapestSoFar || EE.getValue() < CheapestSoFar->getValue())
----------------
mattdr wrote:
> Every edge touched by this loop will //also// be touched by the other leg of this loop at one point or another ("E is a CFG edge"). Can we find a way to avoid the `O(E^2)` inner loop?
> 
> For example: what if we made the outermost loop edge-major? Keep sets of mitigated source-nodes and sink-nodes. Create a list of all edges by weight, sort it, then go through the edges from lowest to highest weight and use them to try to mitigate a new source or sink node.
> 
> There's a good chance I'm missing something fundamental there, and I thank you in advance for your patience explaining it to me. But if I can I'd really like to avoid repeating the `CheapestSoFar` comparison on edge weights in two places.
Yes... in hindsight this whole algorithm was more than a bit sloppy on my part. I completely revamped the algorithm and I think that now it is O(N + E). Please check and make sure you agree.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:650
+        } else { // E is a CFG edge
+          if (GadgetSinks.contains(*E.getDest())) {
+            // The dest is a SINK node. Hence EI is an ingress edge
----------------
mattdr wrote:
> What guarantee do we have here that all of the gadget sinks have, in fact, been added to `GadgetSinks`? It seems to be up to the order in which we iterate through nodes and edges.
Yikes! Can't believe I overlooked this. I fixed the issue, and it looks like the algorithm had been cutting a few more edges in some cases than necessary.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:668
+  int FencesInserted = 0, AdditionalEdgesCut = 0;
+  auto CutAllCFGEdges =
+      [&CutEdges, &AdditionalEdgesCut](const MachineGadgetGraph::Node *N) {
----------------
mattdr wrote:
> A comment for what this lambda does, and how it's intended to be used, is much appreciated for the top-to-bottom reader
> 
> Probably something like:
> 
> When we add an `LFENCE` before an instruction, remove any CFG edges that point to that instruction because they all now refer to a mitigated codepath
I simplified this by inlining lambda with a descriptive comment.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:700-701
+        }
+        if ((InsertionPt == MBB->end() || !isFence(&*InsertionPt)) &&
+            (!Prev || !isFence(Prev))) {
+          BuildMI(*MBB, InsertionPt, DebugLoc(), TII->get(X86::LFENCE));
----------------
mattdr wrote:
> This one took me a bit. It seems like the summary is: add a fence unless it would be redundant, i.e. if we can see the next instruction and see it is itself a fence
Right. I added a clarifying comment.

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D75937/new/

https://reviews.llvm.org/D75937