[asan] Avoid lifetime analysis for allocas with can be in ambiguous state

[vitalybuka]

Summary

C allows to jump over variables declaration so lifetime.start can be
avoid before variable usage. To avoid false-positives on such rare cases
we detect them and remove from lifetime analysis.

PR27453
PR28267

[vitalybuka]

comment

[eugenis]

LGTM

lib/Transforms/Instrumentation/AddressSanitizer.cpp

838 ↗

(On Diff #70612)

please add short comments for these two maps

845 ↗

(On Diff #70612)

succ_range

[vitalybuka]

addressed comments

llvm/trunk/lib/Transforms/Instrumentation/AddressSanitizer.cpp

struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   void poisonAlloca(Value *V, uint64_t Size, IRBuilder<> &IRB, bool DoPoison);
 
   Value *createAllocaForLayout(IRBuilder<> &IRB, const ASanStackFrameLayout &L,
                                bool Dynamic);
   PHINode *createPHI(IRBuilder<> &IRB, Value *Cond, Value *ValueIfTrue,
                      Instruction *ThenTerm, Value *ValueIfFalse);
 };
 
+// Performs depth-first search on the control flow graph of block and checks if
+// we can get into the same block with different lifetime state.
+class AllocaLifetimeChecker {
+  // Contains values of the last lifetime intrinsics in the block.
+  // true: llvm.lifetime.start, false: llvm.lifetime.end
+  DenseMap<const BasicBlock *, bool> Markers;
+  // Contains the lifetime state we detected doing depth-first search on the
+  // control flow graph. We expect all future hits will have the same value.
+  // true: llvm.lifetime.start, false: llvm.lifetime.end
+  DenseMap<const BasicBlock *, bool> InboundState;
+  bool Processed = false;
+  bool CollisionDetected = false;
+
+  bool FindCollision(const std::pair<const BasicBlock *, bool> &BlockState) {
+    auto Ins = InboundState.insert(BlockState);
+    if (!Ins.second) {
+      // Already there. Return collision if they are different.
+      return BlockState.second != Ins.first->second;
+    }
+
+    // Use marker for successors if block contains any.
+    auto M = Markers.find(BlockState.first);
+    bool NewState = (M != Markers.end() ? M->second : BlockState.second);
+    for (const BasicBlock *SB : successors(BlockState.first))
+      if (FindCollision({SB, NewState}))
+        return true;
+
+    return false;
+  }
+
+public:
+  // Assume that markers of the same block will be added in the same order as
+  // the order of corresponding intrinsics, so in the end we will keep only
+  // value of the last intrinsic.
+  void AddMarker(const BasicBlock *BB, bool start) {
+    assert(!Processed);
+    Markers[BB] = start;
+  }
+
+  bool HasAmbiguousLifetime() {
+    if (!Processed) {
+      Processed = true;
+      const Function *F = Markers.begin()->first->getParent();
+      CollisionDetected = FindCollision({&F->getEntryBlock(), false});
+    }
+    return CollisionDetected;
+  }
+};
+
+// Removes allocas for which exists at least one block simultaneously
+// reachable in both states: allocas is inside the scope, and alloca is outside
+// of the scope. We don't have enough information to validate access to such
+// variable, so we just remove such allocas from lifetime analysis.
+// This is workaround for PR28267.
+void removeAllocasWithAmbiguousLifetime(
+    SmallVectorImpl<FunctionStackPoisoner::AllocaPoisonCall> &PoisonCallVec) {
+  DenseMap<const AllocaInst *, AllocaLifetimeChecker> Checkers;
+  for (const auto &APC : PoisonCallVec)
+    Checkers[APC.AI].AddMarker(APC.InsBefore->getParent(), !APC.DoPoison);
+
+  auto IsAmbiguous =
+      [&Checkers](const FunctionStackPoisoner::AllocaPoisonCall &APC) {
+        return Checkers[APC.AI].HasAmbiguousLifetime();
+      };
+
+  PoisonCallVec.erase(
+      std::remove_if(PoisonCallVec.begin(), PoisonCallVec.end(), IsAmbiguous),
+      PoisonCallVec.end());
+}
+
 } // anonymous namespace
 
 char AddressSanitizer::ID = 0;
 INITIALIZE_PASS_BEGIN(
     AddressSanitizer, "asan",
     "AddressSanitizer: detects use-after-free and out-of-bounds bugs.", false,
     false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)

 }
 
 void FunctionStackPoisoner::processDynamicAllocas() {
   if (!ClInstrumentDynamicAllocas || DynamicAllocaVec.empty()) {
     assert(DynamicAllocaPoisonCallVec.empty());
     return;
   }
 
+  removeAllocasWithAmbiguousLifetime(DynamicAllocaPoisonCallVec);
+
   // Insert poison calls for lifetime intrinsics for dynamic allocas.
   for (const auto &APC : DynamicAllocaPoisonCallVec) {
     assert(APC.InsBefore);
     assert(APC.AI);
     assert(ASan.isInterestingAlloca(*APC.AI));
     assert(!APC.AI->isStaticAlloca());
 
     IRBuilder<> IRB(APC.InsBefore);

 }
 
 void FunctionStackPoisoner::processStaticAllocas() {
   if (AllocaVec.empty()) {
     assert(StaticAllocaPoisonCallVec.empty());
     return;
   }
 
+  removeAllocasWithAmbiguousLifetime(StaticAllocaPoisonCallVec);
+
   int StackMallocIdx = -1;
   DebugLoc EntryDebugLocation;
   if (auto SP = F.getSubprogram())
     EntryDebugLocation = DebugLoc::get(SP->getScopeLine(), 0, SP);
 
   Instruction *InsBefore = AllocaVec[0];
   IRBuilder<> IRB(InsBefore);
   IRB.SetCurrentDebugLocation(EntryDebugLocation);

llvm/trunk/test/Instrumentation/AddressSanitizer/lifetime.ll

define void @lifetime() sanitize_address {
   ; CHECK-NEXT: llvm.lifetime.end
 
   ; Memory is unpoisoned at function exit (only once).
   ; CHECK: store i64 0, i64* %{{[0-9]+}}
   ; CHECK-NEXT: ret void
   ret void
 }
 
+; Case of lifetime depends on how we get into the block.
+define void @ambiguous_lifetime(i1 %x) sanitize_address {
+  ; CHECK-LABEL: define void @ambiguous_lifetime
+
+entry:
+  ; Regular variable lifetime intrinsics.
+  %i = alloca i8, align 4  ; Good
+  %j = alloca i8, align 4  ; Bad
+
+  call void @llvm.lifetime.start(i64 1, i8* %i)
+  ; CHECK: store i8 1, i8* %{{[0-9]+}}
+  ; CHECK-NEXT: call void @llvm.lifetime.start
+
+  br i1 %x, label %bb0, label %bb1
+
+bb0:
+  ; CHECK-LABEL: bb0:
+
+  call void @llvm.lifetime.start(i64 1, i8* %j)
+  ; CHECK-NOT: store i8 1, i8* %{{[0-9]+}}
+  ; CHECK-NEXT: call void @llvm.lifetime.start
+
+  br label %bb1
+
+bb1:
+  ; CHECK-LABEL: bb1:
+
+  store volatile i8 0, i8* %i
+  store volatile i8 0, i8* %j
+
+  call void @llvm.lifetime.end(i64 1, i8* %i)
+  ; CHECK: store i8 -8, i8* %{{[0-9]+}}
+  ; CHECK-NEXT: call void @llvm.lifetime.end
+
+  call void @llvm.lifetime.end(i64 1, i8* %j)
+  ; CHECK-NOT: store i8 -8, i8* %{{[0-9]+}}
+  ; CHECK-NEXT: call void @llvm.lifetime.end
+
+  ret void
+}
+
 ; Check that arguments of lifetime may come from phi nodes.
 define void @phi_args(i1 %x) sanitize_address {
   ; CHECK-LABEL: define void @phi_args(i1 %x)
 
 entry:
   %i = alloca i64, align 4
   %i.ptr = bitcast i64* %i to i8*