Fix use-after-free in DestroyHoisting

meg-gupta · meg-gupta · commit 4c9cb0c76858 · 2021-08-25T13:07:21.000-07:00
Due to mismatch in the instructions handled in DestroyHoisting::getUsedLocationsOfInst
and MemoryLocations::analyzeLocationUsesRecursively, certain users of addresses
were not considered and the destroys were hoisted before valid uses causing use-after-frees
diff --git a/lib/SILOptimizer/Transforms/DestroyHoisting.cpp b/lib/SILOptimizer/Transforms/DestroyHoisting.cpp
@@ -333,16 +333,21 @@ void DestroyHoisting::getUsedLocationsOfOperands(Bits &bits, SILInstruction *I)
   }
 }
 
-// Set all bits of locations which instruction \p I is using. It's including
-// parent and sub-locations (see comment in getUsedLocationsOfAddr).
+// NOTE: All instructions handled in
+// MemoryLocations::analyzeLocationUsesRecursively should also be handled
+// explicitly here.
+// Set all bits of locations which instruction \p I is using.
+// It's including parent and sub-locations (see comment in
+// getUsedLocationsOfAddr).
 void DestroyHoisting::getUsedLocationsOfInst(Bits &bits, SILInstruction *I) {
   switch (I->getKind()) {
-    case SILInstructionKind::EndBorrowInst:
-      if (auto *LBI = dyn_cast<LoadBorrowInst>(
-                                        cast<EndBorrowInst>(I)->getOperand())) {
+    case SILInstructionKind::EndBorrowInst: {
+      auto op = cast<EndBorrowInst>(I)->getOperand();
+      if (auto *LBI = dyn_cast<LoadBorrowInst>(op)) {
         getUsedLocationsOfAddr(bits, LBI->getOperand());
       }
       break;
+    }
     case SILInstructionKind::EndApplyInst:
       // Operands passed to begin_apply are alive throughout an end_apply ...
       getUsedLocationsOfOperands(bits, cast<EndApplyInst>(I)->getBeginApply());
@@ -351,6 +356,19 @@ void DestroyHoisting::getUsedLocationsOfInst(Bits &bits, SILInstruction *I) {
       // ... or abort_apply.
       getUsedLocationsOfOperands(bits, cast<AbortApplyInst>(I)->getBeginApply());
       break;
+    case SILInstructionKind::EndAccessInst:
+      getUsedLocationsOfOperands(bits,
+                                 cast<EndAccessInst>(I)->getBeginAccess());
+      break;
+    // These instructions do not access the memory location for read/write
+    case SILInstructionKind::StructElementAddrInst:
+    case SILInstructionKind::TupleElementAddrInst:
+    case SILInstructionKind::WitnessMethodInst:
+    case SILInstructionKind::OpenExistentialAddrInst:
+      break;
+    case SILInstructionKind::InitExistentialAddrInst:
+    case SILInstructionKind::InitEnumDataAddrInst:
+    case SILInstructionKind::UncheckedTakeEnumDataAddrInst:
     case SILInstructionKind::SelectEnumAddrInst:
     case SILInstructionKind::ExistentialMetatypeInst:
     case SILInstructionKind::ValueMetatypeInst:
@@ -368,6 +386,7 @@ void DestroyHoisting::getUsedLocationsOfInst(Bits &bits, SILInstruction *I) {
     case SILInstructionKind::ApplyInst:
     case SILInstructionKind::TryApplyInst:
     case SILInstructionKind::YieldInst:
+    case SILInstructionKind::SwitchEnumAddrInst:
       getUsedLocationsOfOperands(bits, I);
       break;
     case SILInstructionKind::DebugValueAddrInst:
diff --git a/test/SILOptimizer/destroy_hoisting.sil b/test/SILOptimizer/destroy_hoisting.sil
@@ -348,3 +348,104 @@ bb0(%0 :  $*X):
   return %res : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_switch_enum_addr :
+// CHECK-NOT:     destroy_addr
+// CHECK:         switch_enum_addr
+// CHECK:       } // end sil function 'test_switch_enum_addr'
+sil [ossa] @test_switch_enum_addr : $@convention(thin) (@in TwoCases) -> () {
+bb0(%0 : $*TwoCases):
+  %2 = alloc_stack $TwoCases
+  copy_addr %0 to [initialization] %2 : $*TwoCases
+  switch_enum_addr %2 : $*TwoCases, case #TwoCases.A!enumelt: bb1, case #TwoCases.B!enumelt: bb2
+
+bb1:
+  destroy_addr %2 : $*TwoCases
+  dealloc_stack %2 : $*TwoCases
+  br bb3
+
+bb2:
+  destroy_addr %2 : $*TwoCases
+  dealloc_stack %2 : $*TwoCases
+  br bb3
+
+bb3:
+  destroy_addr %0 : $*TwoCases
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_load_borrow1 :
+// CHECK-NOT:         destroy_addr
+// CHECK:             load_borrow
+// CHECK-LABEL: } // end sil function 'test_load_borrow1'
+sil [ossa] @test_load_borrow1 : $@convention(thin) (@in TwoCases) -> () {
+bb0(%0 : $*TwoCases):
+  %2 = alloc_stack $TwoCases
+  copy_addr %0 to [initialization] %2 : $*TwoCases
+  %ld = load_borrow %2 : $*TwoCases
+  br bb1(%ld : $TwoCases)
+
+bb1(%newld : @guaranteed $TwoCases):
+  end_borrow %newld : $TwoCases
+  br bb2
+
+bb2:
+  destroy_addr %2 : $*TwoCases
+  dealloc_stack %2 : $*TwoCases
+  br bb3
+
+bb3:
+  destroy_addr %0 : $*TwoCases
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_load_borrow2 :
+// CHECK-NOT:         destroy_addr
+// CHECK:             load_borrow
+// CHECK-LABEL: } // end sil function 'test_load_borrow2'
+sil [ossa] @test_load_borrow2 : $@convention(thin) (@in TwoCases) -> () {
+bb0(%0 : $*TwoCases):
+  %2 = alloc_stack $TwoCases
+  copy_addr %0 to [initialization] %2 : $*TwoCases
+  %ld = load_borrow %2 : $*TwoCases
+  cond_br undef, bb1, bb2
+
+bb1:
+  end_borrow %ld : $TwoCases
+  br bb3
+
+bb2:
+  br bb2a(%ld : $TwoCases)
+
+bb2a(%newld : @guaranteed $TwoCases):
+  end_borrow %newld : $TwoCases
+  br bb3
+
+bb3:
+  destroy_addr %2 : $*TwoCases
+  dealloc_stack %2 : $*TwoCases
+  br bb3
+
+bb4:
+  destroy_addr %0 : $*TwoCases
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_begin_access :
+// CHECK-NOT:         destroy_addr
+// CHECK:             begin_access
+// CHECK-LABEL: } // end sil function 'test_begin_access'
+sil [ossa] @test_begin_access : $@convention(thin) (@in X) -> () {
+bb0(%0 : $*X):
+  %a = begin_access [read] [dynamic] %0 : $*X
+  br bb1
+
+bb1:
+  end_access %a : $*X
+  destroy_addr %0 : $*X
+  %r = tuple ()
+  return %r : $()
+}
+
diff --git a/test/SILOptimizer/pointer_conversion.swift b/test/SILOptimizer/pointer_conversion.swift
@@ -56,11 +56,12 @@ public func testMutableArray() {
   // CHECK: [[POINTER:%.+]] = struct $UnsafeMutableRawPointer (
   // CHECK-NEXT: [[DEP_POINTER:%.+]] = mark_dependence [[POINTER]] : $UnsafeMutableRawPointer on {{.*}} : $__ContiguousArrayStorageBase
   // CHECK: [[FN:%.+]] = function_ref @takesMutableRawPointer
-  // CHECK-NEXT: apply [[FN]]([[DEP_POINTER]])
-  // CHECK-NOT: release
+  // CHECK: strong_retain {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
+  // CHECK: apply [[FN]]([[DEP_POINTER]])
   // CHECK-NOT: {{^bb[0-9]+:}}
   // CHECK: strong_release {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
-  // CHECK-NEXT: dealloc_stack {{%.+}} : $*Array<Int>
+  // CHECK: strong_release {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
+  // CHECK: dealloc_stack {{%.+}} : $*Array<Int>
   // CHECK-NEXT: [[EMPTY:%.+]] = tuple ()
   // CHECK-NEXT: return [[EMPTY]]
 }
@@ -73,11 +74,12 @@ public func testMutableArrayToOptional() {
   // CHECK-NEXT: [[DEP_POINTER:%.+]] = mark_dependence [[POINTER]] : $UnsafeMutableRawPointer on {{.*}} : $__ContiguousArrayStorageBase
   // CHECK-NEXT: [[OPT_POINTER:%.+]] = enum $Optional<UnsafeMutableRawPointer>, #Optional.some!enumelt, [[DEP_POINTER]]
   // CHECK: [[FN:%.+]] = function_ref @takesOptMutableRawPointer
-  // CHECK-NEXT: apply [[FN]]([[OPT_POINTER]])
-  // CHECK-NOT: release
+  // CHECK: strong_retain {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
+  // CHECK: apply [[FN]]([[OPT_POINTER]])
   // CHECK-NOT: {{^bb[0-9]+:}}
   // CHECK: strong_release {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
-  // CHECK-NEXT: dealloc_stack {{%.+}} : $*Array<Int>
+  // CHECK: strong_release {{%.+}} : ${{Builtin[.]BridgeObject|__ContiguousArrayStorageBase}}
+  // CHECK: dealloc_stack {{%.+}} : $*Array<Int>
   // CHECK-NEXT: [[EMPTY:%.+]] = tuple ()
   // CHECK-NEXT: return [[EMPTY]]
 }
