[Function Specialisation] Fix use after free

[labrinea]

Summary

This is a fix for a use-after-free found by the address sanitizer when compiling GCC: https://github.com/llvm/llvm-project/issues/52821

The Function Specialization pass may remove instructions, cached inside the PredicateBase class, which are later being dereferenced from the SCCPInstVisitor class. To prevent the dangling references I am lazily deleting the dead instructions after the Solver has run.

[nikic]

Can you please explain in more detail where exactly an instruction referenced by PredicateInfo gets removed, and why other passes using PredicateInfo (IPSCCP and NewGVN) do not encounter this issue?

[labrinea]

In D118591#3283559, @nikic wrote:

Can you please explain in more detail where exactly an instruction referenced by PredicateInfo gets removed, and why other passes using PredicateInfo (IPSCCP and NewGVN) do not encounter this issue?

Hi. If you follow the github link there's a dump from the address sanitizer and some explanation on the description of the issue. The instruction is removed by the function specialization pass inside tryToReplaceWithConstant after all its uses have been replaced. As far as I understand, the dangling pointer is then being dereferenced inside PredicateBase::getConstraint() (the function I've patched), which is being called from SCCPInstVisitor::handleCallResult. Both FunctionSpecialization and IPSCCP are using the SCCPSolver, which in turn is using the SCCPInstVisitor. The bug was exposed with ASAN when compilig GCC, so if the NewGVN is using the same machinery as the other passes, then it is potentially affected. I made a bootstrap build of clang with ASAN and FunctionSpecialization enabled and it succeeded. Also the reproducer is not triggering the use-after-free anymore.

[fhahn]

In D118591#3283811, @labrinea wrote:

In D118591#3283559, @nikic wrote:

Can you please explain in more detail where exactly an instruction referenced by PredicateInfo gets removed, and why other passes using PredicateInfo (IPSCCP and NewGVN) do not encounter this issue?

Hi. If you follow the github link there's a dump from the address sanitizer and some explanation on the description of the issue. The instruction is removed by the function specialization pass inside tryToReplaceWithConstant after all its uses have been replaced. As far as I understand, the dangling pointer is then being dereferenced inside PredicateBase::getConstraint() (the function I've patched), which is being called from SCCPInstVisitor::handleCallResult. Both FunctionSpecialization and IPSCCP are using the SCCPSolver, which in turn is using the SCCPInstVisitor. The bug was exposed with ASAN when compilig GCC, so if the NewGVN is using the same machinery as the other passes, then it is potentially affected. I made a bootstrap build of clang with ASAN and FunctionSpecialization enabled and it succeeded. Also the reproducer is not triggering the use-after-free anymore.

It sounds like the real problem may be that FunctionSpecialization removes instructions before the solver is finished and this lead to issues if the removed instructions has an associated PredicateInfo. WeakVH seems like a very heavy hammer to fix the underlying issue. Could FunctionSpecialization just delay removing instructions until all solving is done? If that's not possible, we may have to add a way to remove stale predicate info entries.

[labrinea]

As suggested by Florian, instead of using a WeakVH I am lazily removing the replaced instructions after the Solver has run. None of the existing tests actually covers this code path I am afraid. Examining the debug output of the reproducer I found two PhiNode instructions being replaced with null.

[SjoerdMeijer]

Agreed, lazily removing the instructions in this way is what I was expecting.

But this is a bit odd, a.k.a. not what I was expecting:

None of the existing tests actually covers this code path I am afraid.

Do you mean this:

if (I->isSafeToRemove()) {
   I->eraseFromParent();
   Solver.removeLatticeValueFor(I);
 }

Only triggers in the reproducer attached to the bugreport and not in the regression tests? I think we should first try hard to get a test for that, which I hope is doable. Otherwise, we might as well remove this code and rely on later clean up passes to delete dead code which I think is a valid strategy and alternative if it makes things really complicated here. But again, better would be to see if we can trigger this and a test for it.

[labrinea]

Added some debug output:

• when replacing a value
• when removing dead instructions

[labrinea]

In D118591#3286671, @SjoerdMeijer wrote:

Agreed, lazily removing the instructions in this way is what I was expecting.

But this is a bit odd, a.k.a. not what I was expecting:

None of the existing tests actually covers this code path I am afraid.

Do you mean this:

if (I->isSafeToRemove()) {
   I->eraseFromParent();
   Solver.removeLatticeValueFor(I);
 }

Only triggers in the reproducer attached to the bugreport and not in the regression tests? I think we should first try hard to get a test for that, which I hope is doable. Otherwise, we might as well remove this code and rely on later clean up passes to delete dead code which I think is a valid strategy and alternative if it makes things really complicated here. But again, better would be to see if we can trigger this and a test for it.

I was wrong when I said we are not testing the removal of dead instructions. I was looking at the tests under the original commit D93838, but we have actually added more tests since. For example function-specialization-constant-expression.ll is exercising this code path.

[SjoerdMeijer]

In D118591#3286952, @labrinea wrote:

I was wrong when I said we are not testing the removal of dead instructions. I was looking at the tests under the original commit D93838, but we have actually added more tests since. For example function-specialization-constant-expression.ll is exercising this code path.

Ok, very good, that's what I thought.

One more question about tests: does the sanitzer failure also happens/happened on these existing tests? If so, I think we are done. If not, we are possibly missing something and in this case, do we need to add a sanitizer test? I have never looked that, so don't even know if there are sanitizers tests that run opt.

[fhahn]

In D118591#3287222, @SjoerdMeijer wrote:

In D118591#3286952, @labrinea wrote:

I was wrong when I said we are not testing the removal of dead instructions. I was looking at the tests under the original commit D93838, but we have actually added more tests since. For example function-specialization-constant-expression.ll is exercising this code path.

Ok, very good, that's what I thought.

One more question about tests: does the sanitzer failure also happens/happened on these existing tests? If so, I think we are done. If not, we are possibly missing something and in this case, do we need to add a sanitizer test? I have never looked that, so don't even know if there are sanitizers tests that run opt.

To expose the issue, we might have to add a test where a compare that is used in a predicate gets simplified to a constant and removed and the solver needs to run again on the function.

[labrinea]

I have heavily modified the original reproducer, but the problem is that we don't have buildbots running the llvm regression tests with asan. I can only see two builders for libc with asan here https://lab.llvm.org/buildbot/#/builders. Do you think there's still value in adding this testcase?

[fhahn]

In D118591#3290145, @labrinea wrote:

I have heavily modified the original reproducer, but the problem is that we don't have buildbots running the llvm regression tests with asan. I can only see two builders for libc with asan here https://lab.llvm.org/buildbot/#/builders. Do you think there's still value in adding this testcase?

There should be multiple bots that run with various sanitizers, e.g. https://lab.llvm.org/buildbot/#/builders/168

[SjoerdMeijer]

In D118591#3290167, @fhahn wrote:

In D118591#3290145, @labrinea wrote:

I have heavily modified the original reproducer, but the problem is that we don't have buildbots running the llvm regression tests with asan. I can only see two builders for libc with asan here https://lab.llvm.org/buildbot/#/builders. Do you think there's still value in adding this testcase?

There should be multiple bots that run with various sanitizers, e.g. https://lab.llvm.org/buildbot/#/builders/168

In that case, let's do it; we have got a fix that looks reasonable and we have a test case. Before committing, the test case could do with a comment and checking a bit more though. Let's at least check the statements that the solver is changing and/or perhaps check for absence of the statements that are deleted (if that is still the case here).

[labrinea]

Ah, right, I didn't notice there was another builder page. I'll generate some check lines and commit. Thanks both :)

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

class FunctionSpecializer {
   SCCPSolver &Solver;
 
   /// Analyses used to help determine if a function should be specialized.
   std::function<AssumptionCache &(Function &)> GetAC;
   std::function<TargetTransformInfo &(Function &)> GetTTI;
   std::function<TargetLibraryInfo &(Function &)> GetTLI;
 
   SmallPtrSet<Function *, 2> SpecializedFuncs;
+  SmallVector<Instruction *> ReplacedWithConstant;
 
 public:
   FunctionSpecializer(SCCPSolver &Solver,
                       std::function<AssumptionCache &(Function &)> GetAC,
                       std::function<TargetTransformInfo &(Function &)> GetTTI,
                       std::function<TargetLibraryInfo &(Function &)> GetTLI)
       : Solver(Solver), GetAC(GetAC), GetTTI(GetTTI), GetTLI(GetTLI) {}
 
for (auto *F : FuncDecls) {
       }
     }
 
     updateSpecializedFuncs(FuncDecls, CurrentSpecializations);
     NumFuncSpecialized += NbFunctionsSpecialized;
     return Changed;
   }
 
+  void removeDeadInstructions() {
+    for (auto *I : ReplacedWithConstant) {
+      LLVM_DEBUG(dbgs() << "FnSpecialization: Removing dead instruction "
+                        << *I << "\n");
+      I->eraseFromParent();
+    }
+    ReplacedWithConstant.clear();
+  }
+
   bool tryToReplaceWithConstant(Value *V) {
     if (!V->getType()->isSingleValueType() || isa<CallBase>(V) ||
         V->user_empty())
       return false;
 
     const ValueLatticeElement &IV = Solver.getLatticeValueFor(V);
     if (isOverdefined(IV))
       return false;
     auto *Const =
         isConstant(IV) ? Solver.getConstant(IV) : UndefValue::get(V->getType());
+
+    LLVM_DEBUG(dbgs() << "FnSpecialization: Replacing " << *V
+                      << "\nFnSpecialization: with " << *Const << "\n");
+
     V->replaceAllUsesWith(Const);
 
     for (auto *U : Const->users())
       if (auto *I = dyn_cast<Instruction>(U))
         if (Solver.isBlockExecutable(I->getParent()))
           Solver.visit(I);
 
     // Remove the instruction from Block and Solver.
     if (auto *I = dyn_cast<Instruction>(V)) {
       if (I->isSafeToRemove()) {
-        I->eraseFromParent();
+        ReplacedWithConstant.push_back(I);
         Solver.removeLatticeValueFor(I);
       }
     }
     return true;
   }
 
 private:
   // The number of functions specialised, used for collecting statistics and
while (FuncSpecializationMaxIters != I++ &&
 
     // Replace some unresolved constant arguments.
     constantArgPropagation(FuncDecls, M, Solver);
 
     CurrentSpecializations.clear();
     Changed = true;
   }
 
-  // Clean up the IR by removing ssa_copy intrinsics.
+  // Clean up the IR by removing dead instructions and ssa_copy intrinsics.
+  FS.removeDeadInstructions();
   removeSSACopy(M);
   return Changed;
 }

llvm/test/Transforms/FunctionSpecialization/bug52821-use-after-free.ll

+; RUN: opt -function-specialization -S < %s | FileCheck %s
+
+%mystruct = type { i32, [2 x i64] }
+
+;CHECK-LABEL: @myfunc
+
+define internal %mystruct* @myfunc(%mystruct* %arg) {
+entry:
+  br label %for.cond
+
+for.cond:                                     ; preds = %for.body, %entry
+  %phi = phi %mystruct* [ undef, %for.body ], [ null, %entry ]
+  %cond = icmp eq %mystruct* %phi, null
+  br i1 %cond, label %for.cond2, label %for.body
+
+for.body:                                     ; preds = %for.cond
+  call void @callee(%mystruct* nonnull %phi)
+  br label %for.cond
+
+for.cond2:                                     ; preds = %for.body2, %for.cond
+  %phi2 = phi %mystruct* [ undef, %for.body2 ], [ null, %for.cond ]
+  br i1 undef, label %for.end, label %for.body2
+
+for.body2:                                     ; preds = %for.cond2
+  %arrayidx = getelementptr inbounds %mystruct, %mystruct* %phi2, i64 0, i32 1, i64 3
+  br label %for.cond2
+
+for.end:                                      ; preds = %for.cond2
+  ret %mystruct* %arg
+}
+
+define %mystruct* @caller() {
+entry:
+  %call = call %mystruct* @myfunc(%mystruct* undef)
+  ret %mystruct* %call
+}
+
+declare void @callee(%mystruct*)