[ASan] fix lifetime intrinsics handling

[samsonov]

Summary

Change the way we handle lifetime intrinsics: for each intrinsic check
if it describes one of the "interesting" allocas. Assume that allocas can go through
casts and phi-nodes before appearing as llvm.lifetime arguments.

[kcc]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

328

do you really need this map?
w/o it things may become simpler

1353

return 0?

[samsonov]

Address kcc's comments

lib/Transforms/Instrumentation/AddressSanitizer.cpp

328

Replaced map with a vector of AllocaPoisonCall

1353

Done

[kcc]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

328

Maybe like this?
Maps Value to an AllocaInst from which the Value is originated

329

no need for this typedef anymore

1238

do you need this if?
If it's false, the vector will be empty.

1341

Res = AllocaForValue[V]; // Creates a 0 value if there is not such element.
if (Res)

return Res;

[samsonov]

Address kcc's comments.

lib/Transforms/Instrumentation/AddressSanitizer.cpp

328

Done

329

See below

1238

Done

1341

This is not equivalent: we should also do an early return if the map contains "0" as the value for key V.

[kcc]

LGTM

[samsonov]

r171153, thanks!

lib/Transforms/Instrumentation/AddressSanitizer.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "asan"
 
 #include "llvm/Transforms/Instrumentation.h"
 #include "BlackList.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/OwningPtr.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/DataLayout.h"
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   SmallVector<AllocaInst*, 16> AllocaVec;
   SmallVector<Instruction*, 8> RetVec;
   uint64_t TotalStackSize;
   unsigned StackAlignment;
 
   Function *AsanStackMallocFunc, *AsanStackFreeFunc;
   Function *AsanPoisonStackMemoryFunc, *AsanUnpoisonStackMemoryFunc;
 
+  // Stores a place and arguments of poisoning/unpoisoning call for alloca.
+  struct AllocaPoisonCall {
+    IntrinsicInst *InsBefore;
+    uint64_t Size;
+    bool DoPoison;
+  };
+  SmallVector<AllocaPoisonCall, 8> AllocaPoisonCallVec;
+
+  // Maps Value to an AllocaInst from which the Value is originated.
+  typedef DenseMap<Value*, AllocaInst*> AllocaForValueMapTy;
+  AllocaForValueMapTy AllocaForValue;
+
   FunctionStackPoisoner(Function &F, AddressSanitizer &ASan)
       : F(F), ASan(ASan), DIB(*F.getParent()), C(ASan.C),
         IntptrTy(ASan.IntptrTy), IntptrPtrTy(PointerType::get(IntptrTy, 0)),
         TotalStackSize(0), StackAlignment(1 << MappingScale()) {}
 
   bool runOnFunction() {
     if (!ClStack) return false;
     // Collect alloca, ret, lifetime instructions etc.
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   // ----------------------- Visitors.
   /// \brief Collect all Ret instructions.
   void visitReturnInst(ReturnInst &RI) {
     RetVec.push_back(&RI);
   }
 
   /// \brief Collect Alloca instructions we want (and can) handle.
   void visitAllocaInst(AllocaInst &AI) {
-    if (AI.isArrayAllocation()) return;
-    if (!AI.isStaticAlloca()) return;
-    if (!AI.getAllocatedType()->isSized()) return;
+    if (!isInterestingAlloca(AI)) return;
 
     StackAlignment = std::max(StackAlignment, AI.getAlignment());
     AllocaVec.push_back(&AI);
     uint64_t AlignedSize =  getAlignedAllocaSize(&AI);
     TotalStackSize += AlignedSize;
   }
 
+  /// \brief Collect lifetime intrinsic calls to check for use-after-scope
+  /// errors.
+  void visitIntrinsicInst(IntrinsicInst &II) {
+    if (!ASan.CheckLifetime) return;
+    Intrinsic::ID ID = II.getIntrinsicID();
+    if (ID != Intrinsic::lifetime_start &&
+        ID != Intrinsic::lifetime_end)
+      return;
+    // Found lifetime intrinsic, add ASan instrumentation if necessary.
+    ConstantInt *Size = dyn_cast<ConstantInt>(II.getArgOperand(0));
+    // If size argument is undefined, don't do anything.
+    if (Size->isMinusOne()) return;
+    // Check that size doesn't saturate uint64_t and can
+    // be stored in IntptrTy.
+    const uint64_t SizeValue = Size->getValue().getLimitedValue();
+    if (SizeValue == ~0ULL ||
+        !ConstantInt::isValueValidForType(IntptrTy, SizeValue))
+      return;
+    // Find alloca instruction that corresponds to llvm.lifetime argument.
+    AllocaInst *AI = findAllocaForValue(II.getArgOperand(1));
+    if (!AI) return;
+    bool DoPoison = (ID == Intrinsic::lifetime_end);
+    AllocaPoisonCall APC = {&II, SizeValue, DoPoison};
+    AllocaPoisonCallVec.push_back(APC);
+  }
+
   // ---------------------- Helpers.
   void initializeCallbacks(Module &M);
 
+  // Check if we want (and can) handle this alloca.
+  bool isInterestingAlloca(AllocaInst &AI) {
+    return (!AI.isArrayAllocation() &&
+            AI.isStaticAlloca() &&
+            AI.getAllocatedType()->isSized());
+  }
+
   uint64_t getAllocaSizeInBytes(AllocaInst *AI) {
     Type *Ty = AI->getAllocatedType();
     uint64_t SizeInBytes = ASan.TD->getTypeAllocSize(Ty);
     return SizeInBytes;
   }
   uint64_t getAlignedSize(uint64_t SizeInBytes) {
     size_t RZ = RedzoneSize();
     return ((SizeInBytes + RZ - 1) / RZ) * RZ;
   }
   uint64_t getAlignedAllocaSize(AllocaInst *AI) {
     uint64_t SizeInBytes = getAllocaSizeInBytes(AI);
     return getAlignedSize(SizeInBytes);
   }
+  /// Finds alloca where the value comes from.
+  AllocaInst *findAllocaForValue(Value *V);
   void poisonRedZones(const ArrayRef<AllocaInst*> &AllocaVec, IRBuilder<> IRB,
                       Value *ShadowBase, bool DoPoison);
   void poisonAlloca(Value *V, uint64_t Size, IRBuilder<> IRB, bool DoPoison);
-  /// Analyze lifetime intrinsics for given alloca. Use Value* instead of
-  /// AllocaInst* here, as we call this method after we merge all allocas into a
-  /// single one. Returns true if ASan added some instrumentation.
-  bool handleAllocaLifetime(Value *Alloca);
-  /// Analyze lifetime intrinsics for a specific value, casted from alloca.
-  /// Returns true if if ASan added some instrumentation.
-  bool handleValueLifetime(Value *V);
 };
 
 }  // namespace
 
 char AddressSanitizer::ID = 0;
 INITIALIZE_PASS(AddressSanitizer, "asan",
     "AddressSanitizer: detects use-after-free and out-of-bounds bugs.",
     false, false)
if (SizeInBytes < AlignedSize) {
       }
       Value *PartialPoison = ConstantInt::get(RZTy, Poison);
       IRB.CreateStore(PartialPoison, IRB.CreateIntToPtr(Ptr, RZPtrTy));
     }
 
     // Poison the full redzone at right.
     Ptr = IRB.CreateAdd(ShadowBase,
                         ConstantInt::get(IntptrTy, Pos >> MappingScale()));
-    Value *Poison = i == AllocaVec.size() - 1 ? PoisonRight : PoisonMid;
+    bool LastAlloca = (i == AllocaVec.size() - 1);
+    Value *Poison = LastAlloca ? PoisonRight : PoisonMid;
     IRB.CreateStore(Poison, IRB.CreateIntToPtr(Ptr, RZPtrTy));
 
     Pos += RedzoneSize();
   }
 }
 
 void FunctionStackPoisoner::poisonStack() {
-  bool HavePoisonedAllocas = false;
   uint64_t LocalStackSize = TotalStackSize +
                             (AllocaVec.size() + 1) * RedzoneSize();
 
   bool DoStackMalloc = ASan.CheckUseAfterReturn
       && LocalStackSize <= kMaxStackMallocSize;
 
+  assert(AllocaVec.size() > 0);
   Instruction *InsBefore = AllocaVec[0];
   IRBuilder<> IRB(InsBefore);
 
 
   Type *ByteArrayTy = ArrayType::get(IRB.getInt8Ty(), LocalStackSize);
   AllocaInst *MyAlloca =
       new AllocaInst(ByteArrayTy, "MyAlloca", InsBefore);
   if (ClRealignStack && StackAlignment < RedzoneSize())
     StackAlignment = RedzoneSize();
   MyAlloca->setAlignment(StackAlignment);
   assert(MyAlloca->isStaticAlloca());
   Value *OrigStackBase = IRB.CreatePointerCast(MyAlloca, IntptrTy);
   Value *LocalStackBase = OrigStackBase;
 
   if (DoStackMalloc) {
     LocalStackBase = IRB.CreateCall2(AsanStackMallocFunc,
         ConstantInt::get(IntptrTy, LocalStackSize), OrigStackBase);
   }
 
   // This string will be parsed by the run-time (DescribeStackAddress).
   SmallString<2048> StackDescriptionStorage;
   raw_svector_ostream StackDescription(StackDescriptionStorage);
   StackDescription << F.getName() << " " << AllocaVec.size() << " ";
 
+  // Insert poison calls for lifetime intrinsics for alloca.
+  bool HavePoisonedAllocas = false;
+  for (size_t i = 0, n = AllocaPoisonCallVec.size(); i < n; i++) {
+    const AllocaPoisonCall &APC = AllocaPoisonCallVec[i];
+    IntrinsicInst *II = APC.InsBefore;
+    AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
+    assert(AI);
+    IRBuilder<> IRB(II);
+    poisonAlloca(AI, APC.Size, IRB, APC.DoPoison);
+    HavePoisonedAllocas |= APC.DoPoison;
+  }
+
   uint64_t Pos = RedzoneSize();
   // Replace Alloca instructions with base+offset.
   for (size_t i = 0, n = AllocaVec.size(); i < n; i++) {
     AllocaInst *AI = AllocaVec[i];
     uint64_t SizeInBytes = getAllocaSizeInBytes(AI);
     StringRef Name = AI->getName();
     StackDescription << Pos << " " << SizeInBytes << " "
                      << Name.size() << " " << Name << " ";
     uint64_t AlignedSize = getAlignedAllocaSize(AI);
     assert((AlignedSize % RedzoneSize()) == 0);
     Value *NewAllocaPtr = IRB.CreateIntToPtr(
             IRB.CreateAdd(LocalStackBase, ConstantInt::get(IntptrTy, Pos)),
             AI->getType());
     replaceDbgDeclareForAlloca(AI, NewAllocaPtr, DIB);
     AI->replaceAllUsesWith(NewAllocaPtr);
-    // Analyze lifetime intrinsics only for static allocas we handle.
-    if (ASan.CheckLifetime)
-      HavePoisonedAllocas |= handleAllocaLifetime(NewAllocaPtr);
     Pos += AlignedSize + RedzoneSize();
   }
   assert(Pos == LocalStackSize);
 
   // Write the Magic value and the frame description constant to the redzone.
   Value *BasePlus0 = IRB.CreateIntToPtr(LocalStackBase, IntptrPtrTy);
   IRB.CreateStore(ConstantInt::get(IntptrTy, kCurrentStackFrameMagic),
                   BasePlus0);

 // Handling llvm.lifetime intrinsics for a given %alloca:
 // (1) collect all llvm.lifetime.xxx(%size, %value) describing the alloca.
 // (2) if %size is constant, poison memory for llvm.lifetime.end (to detect
 //     invalid accesses) and unpoison it for llvm.lifetime.start (the memory
 //     could be poisoned by previous llvm.lifetime.end instruction, as the
 //     variable may go in and out of scope several times, e.g. in loops).
 // (3) if we poisoned at least one %alloca in a function,
 //     unpoison the whole stack frame at function exit.
-bool FunctionStackPoisoner::handleAllocaLifetime(Value *Alloca) {
-  assert(ASan.CheckLifetime);
-  Type *AllocaType = Alloca->getType();
-  Type *Int8PtrTy = Type::getInt8PtrTy(AllocaType->getContext());
-
-  bool Res = false;
-  // Typical code looks like this:
-  // %alloca = alloca <type>, <alignment>
-  // ... some code ...
-  // %val1 = bitcast <type>* %alloca to i8*
-  // call void @llvm.lifetime.start(i64 <size>, i8* %val1)
-  // ... more code ...
-  // %val2 = bitcast <type>* %alloca to i8*
-  // call void @llvm.lifetime.start(i64 <size>, i8* %val2)
-  // That is, to handle %alloca we must find all its casts to
-  // i8* values, and find lifetime instructions for these values.
-  if (AllocaType == Int8PtrTy)
-    Res |= handleValueLifetime(Alloca);
-  for (Value::use_iterator UI = Alloca->use_begin(), UE = Alloca->use_end();
-       UI != UE; ++UI) {
-    if (UI->getType() != Int8PtrTy) continue;
-    if (UI->stripPointerCasts() != Alloca) continue;
-    Res |= handleValueLifetime(*UI);
-  }
-  return Res;
-}
 
-bool FunctionStackPoisoner::handleValueLifetime(Value *V) {
-  assert(ASan.CheckLifetime);
-  bool Res = false;
-  for (Value::use_iterator UI = V->use_begin(), UE = V->use_end(); UI != UE;
-       ++UI) {
-    IntrinsicInst *II = dyn_cast<IntrinsicInst>(*UI);
-    if (!II) continue;
-    Intrinsic::ID ID = II->getIntrinsicID();
-    if (ID != Intrinsic::lifetime_start &&
-        ID != Intrinsic::lifetime_end)
-      continue;
-    if (V != II->getArgOperand(1))
-      continue;
-    // Found lifetime intrinsic, add ASan instrumentation if necessary.
-    ConstantInt *Size = dyn_cast<ConstantInt>(II->getArgOperand(0));
-    // If size argument is undefined, don't do anything.
-    if (Size->isMinusOne())
-      continue;
-    // Check that size doesn't saturate uint64_t and can
-    // be stored in IntptrTy.
-    const uint64_t SizeValue = Size->getValue().getLimitedValue();
-    if (SizeValue == ~0ULL ||
-        !ConstantInt::isValueValidForType(IntptrTy, SizeValue)) {
-      continue;
+AllocaInst *FunctionStackPoisoner::findAllocaForValue(Value *V) {
+  if (AllocaInst *AI = dyn_cast<AllocaInst>(V))
+    // We're intested only in allocas we can handle.
+    return isInterestingAlloca(*AI) ? AI : 0;
+  // See if we've already calculated (or started to calculate) alloca for a
+  // given value.
+  AllocaForValueMapTy::iterator I = AllocaForValue.find(V);
+  if (I != AllocaForValue.end())
+    return I->second;
+  // Store 0 while we're calculating alloca for value V to avoid
+  // infinite recursion if the value references itself.
+  AllocaForValue[V] = 0;
+  AllocaInst *Res = 0;
+  if (CastInst *CI = dyn_cast<CastInst>(V))
+    Res = findAllocaForValue(CI->getOperand(0));
+  else if (PHINode *PN = dyn_cast<PHINode>(V)) {
+    for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
+      Value *IncValue = PN->getIncomingValue(i);
+      // Allow self-referencing phi-nodes.
+      if (IncValue == PN) continue;
+      AllocaInst *IncValueAI = findAllocaForValue(IncValue);
+      // AI for incoming values should exist and should all be equal.
+      if (IncValueAI == 0 || (Res != 0 && IncValueAI != Res))
+        return 0;
+      Res = IncValueAI;
     }
-    IRBuilder<> IRB(II);
-    bool DoPoison = (ID == Intrinsic::lifetime_end);
-    poisonAlloca(V, SizeValue, IRB, DoPoison);
-    Res = true;
   }
+  if (Res != 0)
+    AllocaForValue[V] = Res;
   return Res;
 }

test/Instrumentation/AddressSanitizer/lifetime.ll

 define void @lifetime() address_safety {
   ; CHECK: @lifetime
 
   ; Regular variable lifetime intrinsics.
   %i = alloca i32, align 4
   %i.ptr = bitcast i32* %i to i8*
   call void @llvm.lifetime.start(i64 3, i8* %i.ptr)
   ; Memory is unpoisoned at llvm.lifetime.start
-  ; CHECK: %[[VAR:[^ ]*]] = ptrtoint i8* %i.ptr to i64
+  ; CHECK: %[[VAR:[^ ]*]] = ptrtoint i32* %{{[^ ]+}} to i64
   ; CHECK-NEXT: call void @__asan_unpoison_stack_memory(i64 %[[VAR]], i64 3)
   call void @llvm.lifetime.end(i64 4, i8* %i.ptr)
   call void @llvm.lifetime.end(i64 2, i8* %i.ptr)
   ; Memory is poisoned at every call to llvm.lifetime.end
   ; CHECK: call void @__asan_poison_stack_memory(i64 %{{[^ ]+}}, i64 4)
   ; CHECK: call void @__asan_poison_stack_memory(i64 %{{[^ ]+}}, i64 2)
 
   ; Lifetime intrinsics for array.
define void @lifetime() address_safety {
   ; CHECK: call void @__asan_poison_stack_memory(i64 %{{[^ ]+}}, i64 4)
 
   ; Memory is unpoisoned at function exit (only once).
   ; CHECK: call void @__asan_unpoison_stack_memory(i64 %{{[^ ]+}}, i64 {{.*}})
   ; CHECK-NOT: @__asan_unpoison_stack_memory
   ; CHECK: ret void
   ret void
 }
+
+; Check that arguments of lifetime may come from phi nodes.
+define void @phi_args(i1 %x) address_safety {
+  ; CHECK: @phi_args
+
+entry:
+  %i = alloca i64, align 4
+  %i.ptr = bitcast i64* %i to i8*
+  call void @llvm.lifetime.start(i64 8, i8* %i.ptr)
+  ; CHECK: __asan_unpoison_stack_memory
+  br i1 %x, label %bb0, label %bb1
+
+bb0:
+  %i.ptr2 = bitcast i64* %i to i8*
+  br label %bb1
+
+bb1:
+  %i.phi = phi i8* [ %i.ptr, %entry ], [ %i.ptr2, %bb0 ]
+  call void @llvm.lifetime.end(i64 8, i8* %i.phi)
+  ; CHECK: __asan_poison_stack_memory
+  ; CHECK: ret void
+  ret void
+}