[llvm] 7b4cc09 - [LowerConstantIntrinsics] Fix heap-use-after-free bug in worklist

[David Stenberg via llvm-commits]

Author: David Stenberg
Date: 2021-09-21T11:33:07+02:00
New Revision: 7b4cc09b1424c7f53051f971347c00d5f27fbb4e

URL: https://github.com/llvm/llvm-project/commit/7b4cc09b1424c7f53051f971347c00d5f27fbb4e
DIFF: https://github.com/llvm/llvm-project/commit/7b4cc09b1424c7f53051f971347c00d5f27fbb4e.diff

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

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.

Reviewed By: lebedev.ri

Differential Revision: https://reviews.llvm.org/D109221

Added: 
    llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll

Modified: 
    llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp b/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
index bd30019883696..186065db327e6 100644
--- a/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
@@ -55,11 +55,17 @@ static bool replaceConditionalBranchesOnConstant(Instruction *II,
                                                  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())

diff  --git a/llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll b/llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll
new file mode 100644
index 0000000000000..304e86584862b
--- /dev/null
+++ b/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 }