[SCCPSolver] Fix use-after-free in markArgInFuncSpecialization.

[duck-37]

Summary

In SCCPSolver::markArgInFuncSpecialization, the ValueState map may be reallocated *after* the initial ValueLatticeElement reference is grabbed, but *before* its use in copy initialization. This causes a use-after-free.
To fix this, this commit changes the behavior to create the new ValueLatticeElement before assigning the old one to it.

[SjoerdMeijer]

Thanks for looking into this!

I think I may have been bitten by this before (or something similar), but it's easy to forget the details here.... So perhaps you can help me a bit by expanding a on this:

In SCCPSolver::markArgInFuncSpecialization, the ValueState map may be reallocated *after* the initial ValueLatticeElement reference is grabbed, but *before* its use in copy initialization.

I know the operator[] is kind of considered harmful for maps, but in this case that seemed fine and doesn't seem the cause of the problem, is that right? So perhaps you can expand on the reallocation (of the map?) that I don't quite get.

llvm/lib/Transforms/Utils/SCCPSolver.cpp

543–550

This is probably non-obvious enough that a little comment would be useful why we are doing this.

[duck-37]

In D111112#3041838, @SjoerdMeijer wrote:

Thanks for looking into this!

I think I may have been bitten by this before (or something similar), but it's easy to forget the details here.... So perhaps you can help me a bit by expanding a on this:

In SCCPSolver::markArgInFuncSpecialization, the ValueState map may be reallocated *after* the initial ValueLatticeElement reference is grabbed, but *before* its use in copy initialization.

I know the operator[] is kind of considered harmful for maps, but in this case that seemed fine and doesn't seem the cause of the problem, is that right? So perhaps you can expand on the reallocation (of the map?) that I don't quite get.

Sure! The issue here isn't operator[] itself, it's the order in which the operations here are done. The DenseMap class differs from std::unordered_map in a few ways; the important one here is that references can be invalidated if you insert a new element, due to the fact that everything is in one block of memory.
What happens here is that the code takes a reference to ValueState[I], then tries to index into ValueState[J]. Since J doesn't exist in the ValueState map, the DenseMap may be resized to accommodate it. After this, the reference to ValueState[I] points to invalid or garbage memory.
The copy assignment operator for ValueState[J] is then called with the now-invalid reference to ValueState[I], causing all sorts of fun things to happen.

The fix here is just to make sure that ValueState[J] exists in the map *before* indexing into ValueState[I] so that no references get invalidated.

[duck-37]

Add some comments.

[SjoerdMeijer]

Wowsers, this is quite something, and I've learned something today!
But now reading things, and after asking my colleague @momchil.velikov for a second opinion, we agree with your analysis + fix, so thanks for fixing, LGTM.
Do you have commit rights? Would you like this to be committed on your behalf?

[duck-37]

In D111112#3042115, @SjoerdMeijer wrote:

Wowsers, this is quite something, and I've learned something today!
But now reading things, and after asking my colleague @momchil.velikov for a second opinion, we agree with your analysis + fix, so thanks for fixing, LGTM.
Do you have commit rights? Would you like this to be committed on your behalf?

I don't have commit rights, so that would be nice, yes. Thanks for the fast review!

[SjoerdMeijer]

No worries, and I will leave a "Patch by: <your-name>" note in the commit message. How would you like to be referenced, what should I be using for <your-name>?

[duck-37]

In D111112#3042178, @SjoerdMeijer wrote:

No worries, and I will leave a "Patch by: <your-name>" note in the commit message. How would you like to be referenced, what should I be using for <your-name>?

Just linking my GitHub account is fine :)
https://github.com/duck-37/

[chill]

In D111112#3042178, @SjoerdMeijer wrote:

No worries, and I will leave a "Patch by: <your-name>" note in the commit message. How would you like to be referenced, what should I be using for <your-name>?

Oh well, a bit too late, but nowadays the policy seems to be to use GIT for attribution https://llvm.org/docs/DeveloperPolicy.html#commit-messages

[SjoerdMeijer]

In D111112#3042303, @chill wrote:

In D111112#3042178, @SjoerdMeijer wrote:

No worries, and I will leave a "Patch by: <your-name>" note in the commit message. How would you like to be referenced, what should I be using for <your-name>?

Oh well, a bit too late, but nowadays the policy seems to be to use GIT for attribution https://llvm.org/docs/DeveloperPolicy.html#commit-messages

Oh, didn't know, next time!

llvm/lib/Transforms/Utils/SCCPSolver.cpp

void SCCPInstVisitor::markArgInFuncSpecialization(Function *F, Argument *A,
   markConstant(A, C);
 
   // For the remaining arguments in the new function, copy the lattice state
   // over from the old function.
   for (auto I = F->arg_begin(), J = A->getParent()->arg_begin(),
             E = F->arg_end();
        I != E; ++I, ++J)
     if (J != A && ValueState.count(I)) {
-      ValueState[J] = ValueState[I];
-      pushToWorkList(ValueState[J], J);
+      // Note: This previously looked like this:
+      // ValueState[J] = ValueState[I];
+      // This is incorrect because the DenseMap class may resize the underlying
+      // memory when inserting `J`, which will invalidate the reference to `I`.
+      // Instead, we make sure `J` exists, then set it to `I` afterwards.
+      auto &NewValue = ValueState[J];
+      NewValue = ValueState[I];
+      pushToWorkList(NewValue, J);
     }
 }
 
 void SCCPInstVisitor::visitInstruction(Instruction &I) {
   // All the instructions we don't do any special handling for just
   // go to overdefined.
   LLVM_DEBUG(dbgs() << "SCCP: Don't know how to handle: " << I << '\n');
   markOverdefined(&I);