[hwasan] Support more complicated lifetimes.

[fmayer]

Summary

This is important as with exceptions enabled, non-POD allocas often have
two lifetime ends: the exception handler, and the normal one.

[fmayer]

formatting

[fmayer]

add comment

[fmayer]

Ping :)

[eugenis]

LGTM

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

50

Update the comment.

74

This does not handle diamond CFG, right? A case where two ends dominate a return as a group. I don't know if LLVM has an implementation of that.

I think it is fine, that must be a relatively uncommon case.

[fmayer]

update comment

[fmayer]

add todo

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

74

Correct. I added a TODO.

compiler-rt/test/hwasan/TestCases/use-after-scope-types.cpp

 // This is the ASAN test of the same name ported to HWAsan.
 
 // RUN: %clangxx_hwasan -mllvm -hwasan-use-after-scope -std=c++11 -O0 %s -o %t
 // RUN: not %run %t 0 2>&1 | FileCheck %s
 // RUN: not %run %t 1 2>&1 | FileCheck %s
 // RUN: not %run %t 2 2>&1 | FileCheck %s
 // RUN: not %run %t 3 2>&1 | FileCheck %s
 // RUN: not %run %t 4 2>&1 | FileCheck %s
 // RUN: not %run %t 5 2>&1 | FileCheck %s
-// The std::vector case is broken because of limited lifetime tracking.
-// TODO(fmayer): Fix and enable.
+// RUN: not %run %t 6 2>&1 | FileCheck %s
 // RUN: not %run %t 7 2>&1 | FileCheck %s
 // RUN: not %run %t 8 2>&1 | FileCheck %s
 // RUN: not %run %t 9 2>&1 | FileCheck %s
 // RUN: not %run %t 10 2>&1 | FileCheck %s
 
 // REQUIRES: aarch64-target-arch
 // REQUIRES: stable-runtime
 

llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h

 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
 #define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H
 
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
 
 class InterestingMemoryOperand {
 public:
   Use *PtrUse;
   bool IsWrite;
InterestingMemoryOperand(Instruction *I, unsigned OperandNo, bool IsWrite,
     PtrUse = &I->getOperandUse(OperandNo);
   }
 
   Instruction *getInsn() { return cast<Instruction>(PtrUse->getUser()); }
 
   Value *getPtr() { return PtrUse->get(); }
 };
 
-// For an alloca valid between lifetime markers Start and End, call the
+// For an alloca valid between lifetime markers Start and Ends, call the
 // Callback for all possible exits out of the lifetime in the containing
 // function, which can return from the instructions in RetVec.
 //
-// Returns whether End was the only possible exit. If it wasn't, the caller
-// should remove End to ensure that work done at the other exits does not
-// happen outside of the lifetime.
+// Returns whether Ends covered all possible exits. If they did not,
+// the caller should remove Ends to ensure that work done at the other
+// exits does not happen outside of the lifetime.
 template <typename F>
 bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
-                          const Instruction *Start, Instruction *End,
+                          const Instruction *Start,
+                          const SmallVectorImpl<IntrinsicInst *> &Ends,
                           const SmallVectorImpl<Instruction *> &RetVec,
                           F Callback) {
-  // We need to ensure that if we tag some object, we certainly untag it
-  // before the function exits.
-  if (PDT.dominates(End, Start)) {
-    Callback(End);
-  } else {
-    SmallVector<Instruction *, 8> ReachableRetVec;
-    unsigned NumCoveredExits = 0;
-    for (auto &RI : RetVec) {
-      if (!isPotentiallyReachable(Start, RI, nullptr, &DT))
-        continue;
-      ReachableRetVec.push_back(RI);
-      if (DT.dominates(End, RI))
-        ++NumCoveredExits;
-    }
-    // If there's a mix of covered and non-covered exits, just put the untag
-    // on exits, so we avoid the redundancy of untagging twice.
-    if (NumCoveredExits == ReachableRetVec.size()) {
+  if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
+    Callback(Ends[0]);
+    return true;
+  }
+    // TODO(fmayer): We don't support diamond shapes, where multiple lifetime
+    // ends together dominate the RI, but none of them does by itself.
+    // Check how often this happens and decide whether to support this here.
+    if (std::any_of(Ends.begin(), Ends.end(),
+                    [&](Instruction *End) { return DT.dominates(End, RI); }))
+    for (auto *End : Ends)
       Callback(End);
-    } else {
-      for (auto &RI : ReachableRetVec)
-        Callback(RI);
-      // We may have inserted untag outside of the lifetime interval.
-      // Signal the caller to remove the lifetime end call for this alloca.
-      return false;
-    }
-  }
   return true;
 }
 
 // Get AddressSanitizer parameters.
 void getAddressSanitizerParams(const Triple &TargetTriple, int LongSize,
                                bool IsKasan, uint64_t *ShadowBase,
                                int *MappingScale, bool *OrShadowOffset);
 
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

if (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 &&
       IntrinsicInst *End = Info.LifetimeEnd[0];
       uint64_t Size =
           cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
       Size = alignTo(Size, kTagGranuleSize);
       tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
 
       auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
       if (!DT || !PDT ||
-          !forAllReachableExits(*DT, *PDT, Start, End, RetVec, TagEnd))
+          !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec,
+                                TagEnd))
         End->eraseFromParent();
     } else {
       uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
       tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);
       for (auto &RI : RetVec) {
         untagAlloca(AI, RI, Size);
       }

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
static cl::opt<bool> ClInstrumentStack("hwasan-instrument-stack",
                                        cl::desc("instrument stack (allocas)"),
                                        cl::Hidden, cl::init(true));
 
 static cl::opt<bool>
     ClUseStackSafety("hwasan-use-stack-safety", cl::Hidden, cl::init(true),
                      cl::Hidden, cl::desc("Use Stack Safety analysis results"),
                      cl::Optional);
 
+static cl::opt<size_t> ClMaxLifetimes(
+    "hwasan-max-lifetimes-for-alloca", cl::Hidden, cl::init(3),
+    cl::ReallyHidden,
+    cl::desc("How many lifetime ends to handle for a single alloca."),
+    cl::Optional);
+
 static cl::opt<bool>
     ClUseAfterScope("hwasan-use-after-scope",
                     cl::desc("detect use after scope within function"),
                     cl::Hidden, cl::init(false));
 
 static cl::opt<bool> ClUARRetagToZero(
     "hwasan-uar-retag-to-zero",
     cl::desc("Clear alloca tags before returning from the function to allow "
public:
   bool ignoreAccess(Value *Ptr);
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
   bool isInterestingAlloca(const AllocaInst &AI);
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
+  static bool isStandardLifetime(const AllocaInfo &AllocaInfo,
+                                 const DominatorTree &DT);
   bool instrumentStack(
       MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
       SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
       DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
       SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
       llvm::function_ref<const DominatorTree &()> GetDT,
       llvm::function_ref<const PostDominatorTree &()> GetPDT);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
for (auto *LP : LandingPadVec) {
     IRB.CreateCall(
         HWAsanHandleVfork,
         {readRegister(IRB, (TargetTriple.getArch() == Triple::x86_64) ? "rsp"
                                                                       : "sp")});
   }
   return true;
 }
 
+static bool
+maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
+                            const DominatorTree &DT) {
+  // If we have too many lifetime ends, give up, as the algorithm below is N^2.
+  if (Insts.size() > ClMaxLifetimes)
+    return true;
+  for (size_t I = 0; I < Insts.size(); ++I) {
+    for (size_t J = 0; J < Insts.size(); ++J) {
+      if (I == J)
+        continue;
+      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
+        return true;
+    }
+  }
+  return false;
+}
+
+// static
+bool HWAddressSanitizer::isStandardLifetime(const AllocaInfo &AllocaInfo,
+                                            const DominatorTree &DT) {
+  // An alloca that has exactly one start and end in every possible execution.
+  // If it has multiple ends, they have to be unreachable from each other, so
+  // at most one of them is actually used for each execution of the function.
+  return AllocaInfo.LifetimeStart.size() == 1 &&
+         (AllocaInfo.LifetimeEnd.size() == 1 ||
+          (AllocaInfo.LifetimeEnd.size() > 0 &&
+           !maybeReachableFromEachOther(AllocaInfo.LifetimeEnd, DT)));
+}
+
 bool HWAddressSanitizer::instrumentStack(
     MapVector<AllocaInst *, AllocaInfo> &AllocasToInstrument,
     SmallVector<Instruction *, 4> &UnrecognizedLifetimes,
     DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
     SmallVectorImpl<Instruction *> &RetVec, Value *StackTag,
     llvm::function_ref<const DominatorTree &()> GetDT,
     llvm::function_ref<const PostDominatorTree &()> GetPDT) {
   // Ideally, we want to calculate tagged stack base pointer, and rewrite all
for (auto *DDI : AllocaDbgMap.lookup(AI)) {
       for (size_t LocNo = 0; LocNo < DDI->getNumVariableLocationOps(); ++LocNo)
         if (DDI->getVariableLocationOp(LocNo) == AI)
           DDI->setExpression(DIExpression::appendOpsToArg(DDI->getExpression(),
                                                           NewOps, LocNo));
     }
 
     size_t Size = getAllocaSizeInBytes(*AI);
     size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
-    bool StandardLifetime = UnrecognizedLifetimes.empty() &&
-                            Info.LifetimeStart.size() == 1 &&
-                            Info.LifetimeEnd.size() == 1;
+    bool StandardLifetime =
+        UnrecognizedLifetimes.empty() && isStandardLifetime(Info, GetDT());
     if (DetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
-      IntrinsicInst *End = Info.LifetimeEnd[0];
       IRB.SetInsertPoint(Start->getNextNode());
       auto TagEnd = [&](Instruction *Node) {
         IRB.SetInsertPoint(Node);
         Value *UARTag = getUARTag(IRB, StackTag);
         tagAlloca(IRB, AI, UARTag, AlignedSize);
       };
       tagAlloca(IRB, AI, Tag, Size);
-      if (!forAllReachableExits(GetDT(), GetPDT(), Start, End, RetVec, TagEnd))
+      if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd,
+                                RetVec, TagEnd)) {
+        for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
+      }
     } else {
       tagAlloca(IRB, AI, Tag, Size);
       for (auto *RI : RetVec) {
         IRB.SetInsertPoint(RI);
         Value *UARTag = getUARTag(IRB, StackTag);
         tagAlloca(IRB, AI, UARTag, AlignedSize);
       }
       if (!StandardLifetime) {

llvm/test/Instrumentation/HWAddressSanitizer/exception-lifetime.ll

+; Test allocas with multiple lifetime ends, as frequently seen for exception
+; handling.
+;
+; RUN: opt -hwasan -hwasan-use-after-scope -S -o - %s | FileCheck %s --check-prefix=CHECK
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-android"
+
+declare void @mayFail(i32* %x) sanitize_hwaddress
+declare void @onExcept(i32* %x) sanitize_hwaddress
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) nounwind
+declare i32 @__gxx_personality_v0(...)
+
+define void @test() sanitize_hwaddress personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+entry:
+  %x = alloca i32, align 8
+  %exn.slot = alloca i8*, align 8
+  %ehselector.slot = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @llvm.lifetime.start.p0i8(i64 8, i8* %0)
+  invoke void @mayFail(i32* %x) to label %invoke.cont unwind label %lpad
+
+invoke.cont:                                      ; preds = %entry
+; CHECK: invoke.cont:
+; CHECK:  call void @llvm.memset.p0i8.i64(i8* align 1 %31, i8 0, i64 1, i1 false)
+; CHECK:  call void @llvm.lifetime.end.p0i8(i64 8, i8* %28)
+; CHECK:  ret void
+
+  %1 = bitcast i32* %x to i8*
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* %1)
+  ret void
+
+lpad:                                             ; preds = %entry
+; CHECK: lpad
+; CHECK:  %41 = getelementptr i8, i8* %17, i64 %40
+; CHECK:  call void @llvm.memset.p0i8.i64(i8* align 1 %41, i8 0, i64 1, i1 false)
+; CHECK:  call void @llvm.lifetime.end.p0i8(i64 8, i8* %38)
+; CHECK:  br label %eh.resume
+
+  %2 = landingpad { i8*, i32 }
+  cleanup
+  %3 = extractvalue { i8*, i32 } %2, 0
+  store i8* %3, i8** %exn.slot, align 8
+  %4 = extractvalue { i8*, i32 } %2, 1
+  store i32 %4, i32* %ehselector.slot, align 4
+  call void @onExcept(i32* %x) #18
+  %5 = bitcast i32* %x to i8*
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* %5)
+  br label %eh.resume
+
+eh.resume:                                        ; preds = %lpad
+  %exn = load i8*, i8** %exn.slot, align 8
+  %sel = load i32, i32* %ehselector.slot, align 4
+  %lpad.val = insertvalue { i8*, i32 } undef, i8* %exn, 0
+  %lpad.val1 = insertvalue { i8*, i32 } %lpad.val, i32 %sel, 1
+  resume { i8*, i32 } %lpad.val1
+}