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

Wed Apr 15 01:36:08 PDT 2020

mattdr requested changes to this revision.
mattdr marked an inline comment as done.
mattdr added a comment.
This revision now requires changes to proceed.

Reviewed the algorithm, have to confess I skipped the tests.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:351
+
+    Graph = GraphBuilder::trim(
+        *ElimGraph, MachineGadgetGraph::NodeSet{*ElimGraph}, CutEdges);
----------------
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?

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

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:537
+      const MachineGadgetGraph::Node *Dest = E.getDest();
+      if (isFence(Dest->getValue())) {
+        ElimNodes.insert(*Dest);
----------------
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?

================
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()) {
----------------
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.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:565
+      continue;
+    std::function<void(const MachineGadgetGraph::Node *, bool)> TraverseDFS =
+        [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
----------------
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.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:567
+        [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
+          if (!FirstNode) {
+            Visited.insert(*N);
----------------
We already capture `RootN` later on. We can remove the boolean conditional and make this more readable by writing `if (*N != RootN)` (maybe `if (N != &RootN)`? I've sort of lost track of the API for nodes and edges.)

Pushing further: maybe we can just remove the tiny optimization we get by special-casing the root node in the interest of simplicity?

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:591
+
+  if (!(ElimEdges.empty() && ElimNodes.empty())) {
+    int NumRemainingGadgets = NumGadgets - ElimGadgets.count();
----------------
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()`

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:602
+
+void X86LoadValueInjectionLoadHardeningPass::cutEdges(
+    MachineGadgetGraph &G,
----------------
likewise `findEdgesToCut` or similar

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:605
+    MachineGadgetGraph::EdgeSet &CutEdges /* out */) const {
+  if (!OptimizePluginPath.empty()) {
+    if (!OptimizeDL.isValid()) {
----------------
I'm still not convinced it's worth the extra (untestable) complexity to add this plugin point, but I defer to the committer.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:607
+    if (!OptimizeDL.isValid()) {
+      std::string ErrorMsg{};
+      OptimizeDL = llvm::sys::DynamicLibrary::getPermanentLibrary(
----------------
Extra `{}` should be unnecessary

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:641
+        if (MachineGadgetGraph::isGadgetEdge(E)) {
+          // NI is a SOURCE node. Look for a cheap egress edge
+          for (const auto &EE : N.edges()) {
----------------
Seems like `NI` is a vestige of a previous iteration

================
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())
----------------
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.

================
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
----------------
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.

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:651
+          if (GadgetSinks.contains(*E.getDest())) {
+            // The dest is a SINK node. Hence EI is an ingress edge
+            if (!CheapestSoFar || E.getValue() < CheapestSoFar->getValue())
----------------
`NI` also seems like it comes from a previous version

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:668
+  int FencesInserted = 0, AdditionalEdgesCut = 0;
+  auto CutAllCFGEdges =
+      [&CutEdges, &AdditionalEdgesCut](const MachineGadgetGraph::Node *N) {
----------------
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

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:681
+        MachineInstr *MI = N.getValue(), *Prev;
+        MachineBasicBlock *MBB;
+        MachineBasicBlock::iterator InsertionPt;
----------------
`// Insert an LFENCE in this basic block`

================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:682
+        MachineBasicBlock *MBB;
+        MachineBasicBlock::iterator InsertionPt;
+        if (MI == MachineGadgetGraph::ArgNodeSentinel) {
----------------
`// ... at this point in the basic block`

================
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));
----------------
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

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

https://reviews.llvm.org/D75937