[PATCH] D109221: [LowerConstantIntrinsics] Fix heap-use-after-free bug in worklist

[David Stenberg via Phabricator via llvm-commits]

dstenb created this revision.
dstenb added reviewers: lebedev.ri, joerg, nickdesaulniers.
Herald added a subscriber: hiraditya.
dstenb requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

This fixes PR51730, a heap-use-after-free bug in
replaceConditionalBranchesOnConstant().

With the attached reproducer we were left with a function looking
something like this after replaceAndRecursivelySimplify():

  [...]
  
  cont2.i:
    br i1 %.not1.i, label %handler.type_mismatch3.i, label %cont4.i
  
  handler.type_mismatch3.i:
    %3 = phi i1 [ %2, %cont2.thread.i ], [ false, %cont2.i ]
    unreachable
  
  cont4.i:
    unreachable
  
  [...]

with both the branch instruction and PHI node being in the worklist. As
a result of replacing the branch instruction with an unconditional
branch, the PHI node in %handler.type_mismatch3.i would be removed. This
then resulted in a heap-use-after-free bug due to accessing that removed
PHI node in the next worklist iteration.

This is solved by using a value handle worklist. I am a unsure if this
is the most idiomatic solution. Another solution could have been to
produce a worklist just containing the interesting branch instructions,
but I thought that it perhaps was a bit cleaner to keep all worklist
filtering in the loop that does the rewrites.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D109221

Files:
  llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
  llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll


Index: llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll
@@ -0,0 +1,48 @@
+; RUN: opt -lower-constant-intrinsics -S < %s | FileCheck %s
+
+; This is a reproducer for a heap-use-after-free bug that occured due to trying
+; to process a PHI node that was removed in a preceding worklist iteration. The
+; conditional branch in %cont2.i will be replaced with an unconditional branch
+; to %cont4.i. As a result of that, the PHI node in %handler.type_mismatch3.i
+; will be left with one predecessor and will thus be removed in that iteration.
+
+; CHECK-NOT: phi
+; CHECK: cont2.i:
+; CHECK-NEXT: br label %cont4.i
+; CHECK-NOT: phi
+
+%s = type { [2 x i16] }
+
+define fastcc void @foo(%s* %p) unnamed_addr {
+entry:
+  %0 = bitcast %s* %p to i8*
+  %1 = tail call i32 @llvm.objectsize.i32.p0i8(i8* %0, i1 false, i1 false, i1 false) #2
+  %2 = icmp ne i32 %1, 0
+  %.not1.i = icmp eq i32 %1, 0
+  br label %for.cond
+
+for.cond:                                         ; preds = %entry
+  br label %cont.i
+
+cont.i:                                           ; preds = %for.cond
+  br i1 undef, label %cont2.i, label %cont2.thread.i
+
+cont2.thread.i:                                   ; preds = %cont.i
+  br label %handler.type_mismatch3.i
+
+cont2.i:                                          ; preds = %cont.i
+  br i1 %.not1.i, label %handler.type_mismatch3.i, label %cont4.i
+
+handler.type_mismatch3.i:                         ; preds = %cont2.i, %cont2.thread.i
+  %3 = phi i1 [ %2, %cont2.thread.i ], [ false, %cont2.i ]
+  unreachable
+
+cont4.i:                                          ; preds = %cont2.i
+  unreachable
+}
+
+; Function Attrs: nofree nosync nounwind readnone speculatable willreturn
+declare i32 @llvm.objectsize.i32.p0i8(i8*, i1 immarg, i1 immarg, i1 immarg) #1
+
+attributes #1 = { nofree nosync nounwind readnone speculatable willreturn }
+attributes #2 = { nounwind }
Index: llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
+++ llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
@@ -55,11 +55,17 @@
                                                  Value *NewValue,
                                                  DomTreeUpdater *DTU) {
   bool HasDeadBlocks = false;
-  SmallSetVector<Instruction *, 8> Worklist;
+  SmallSetVector<Instruction *, 8> UnsimplifiedUsers;
   replaceAndRecursivelySimplify(II, NewValue, nullptr, nullptr, nullptr,
-                                &Worklist);
-  for (auto I : Worklist) {
-    BranchInst *BI = dyn_cast<BranchInst>(I);
+                                &UnsimplifiedUsers);
+  // UnsimplifiedUsers can contain PHI nodes that may be removed when
+  // replacing the branch instructions, so use a value handle worklist
+  // to handle those possibly removed instructions.
+  SmallVector<WeakVH, 8> Worklist(UnsimplifiedUsers.begin(),
+                                  UnsimplifiedUsers.end());
+
+  for (auto &VH : Worklist) {
+    BranchInst *BI = dyn_cast_or_null<BranchInst>(VH);
     if (!BI)
       continue;
     if (BI->isUnconditional())


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D109221.370531.patch
Type: text/x-patch
Size: 3328 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210903/8de816f3/attachment-0001.bin>