MSan: handle llvm.lifetime.start intrinsic

[glider]

Summary

When a variable goes into scope several times within a single function
or when two variables from different scopes share a stack slot it may
be incorrect to poison such scoped locals at the beginning of the
function.
In the former case it may lead to false negatives (see
https://github.com/google/sanitizers/issues/590), in the latter - to
incorrect reports (because only one origin remains on the stack).

If Clang emits lifetime intrinsics with known size for such scoped
variables we insert code poisoning them after each call to
llvm.lifetime.start(). This is done on a best-effort basis.
If for a certain intrinsic we fail to find a corresponding alloca, we
fall back to poisoning allocas for the whole function, as it's now
impossible to tell which alloca was missed.

The new instrumentation may slow down hot loops containing local
variables with lifetime intrinsics, so we allow disabling it with
-mllvm -msan-handle-lifetime-intrinsics=false.

[glider]

Removed errs()

[eugenis]

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2580

IMHO failure to find the underlying alloca must result in poisoning of all allocas in the prologue - exactly because we don't know which one we've missed.
Alternatively, since we know the size, we could instrument this lifetime.start anyway using an invalid origin id.

3468

Let's call this function InstrumentAlloca, and rename instrumentAllocaUserspace to poisonAllocaUserspace.

[glider]

In D60617#1464971, @eugenis wrote:

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

Ok, will do. The repro in the bug doesn't produce IR containing a select statement for me, but the idea is clear.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Agreed. This also makes less sense without origins - do you think we shall hide it behind track-origins>0 ?

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

Sorry, I don't get it. We won't see any lifetime markers with that condition removed, right?
What's the point of this fix then?

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2580

I like your first proposal better.
Using an invalid origin is impractical, as we won't be able to even find the function this poisoned local belongs to.

3468

same for instrumentAllocaKmsan?

[eugenis]

In D60617#1467093, @glider wrote:

In D60617#1464971, @eugenis wrote:

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

Ok, will do. The repro in the bug doesn't produce IR containing a select statement for me, but the idea is clear.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Agreed. This also makes less sense without origins - do you think we shall hide it behind track-origins>0 ?

No, I think it is a good change. But it would be nice to have an llvm flag to back out just in case.

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

Sorry, I don't get it. We won't see any lifetime markers with that condition removed, right?
What's the point of this fix then?

We will see a lot more lifetime markers.
I'm talking about this:

// Disable lifetime markers in msan builds.
// FIXME: Remove this when msan works with lifetime markers.
if (LangOpts.Sanitize.has(SanitizerKind::Memory))
  return false;

http://llvm.org/viewvc/llvm-project?view=revision&revision=235613

[glider]

Added a test containing an intrinsic without a known alloca.
Handled this case by falling back to instrumenting only allocas.
Added a flag to disable llvm.lifetime.start handling.

[glider]

s/instrumentAllocaKmsan/poisonAllocaKmsan
s/instrumentAllocaUserspace/poisonAllocaUserspace

[glider]

Updated the patch, PTAL.

Re: shouldEmitLifetimeMarkers() - thanks, I wasn't aware we disable lifetime markers with -O0.
I don't need that for the kernel (which is built with -O2), but I can sure revert that CL if you think it's ok for the userspace.

FWIW those bits are in clang/lib/CodeGen/CodeGenFunction.cpp now, and removing just them doesn't break any tests.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1037

This does not have to be so complicated.
Get rid of FallbackAllocaSet; instead remove instructions from AllocaSet as they are being poisoned at lifetime calls, then poison remaining alloca instructions.

This way if untraceable lifetime was seen, you can simply skip the lifetime poisoning step.

Please rename InstrumentOnlyAllocas to something that does not suggest the we are instrumenting _only_ allocas (and nothing else).

1310

Move this logic to a separate function.

2574

What's wrong with variable length allocas?

[glider]

Fixed comments.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1037

Agreed, we don't need FallbackAllocaSet, we can just poison allocas from LifetimeStartList as if they were in AllocaSet (see the updated patch).

1301

Reused instrumentAlloca() here instead of calling KMSAN/userspace-specific functions.

2574

You're right, fixed this.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1300

What if InstrumentLifetimeStart == false, but there is an alloca with multiple lifetime.start ? It looks like you'll poison it several times in a row.

[glider]

Fixed the case when an alloca was initialized twice.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1300

Fixed this by adding alloca to AllocaSet instead of calling instrumentAlloca(*Item.second, nullptr).

[glider]

ptal

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2580

What if AllocaInst is visited after all of its lifetime.start calls? I'm not 100% that this is impossible, and then this code will poison both allocainst and lifetime.start.

IMHO, the visitor should simply collect allocas and lifetimes, and then the finalization code can erase allocas as it instruments their lifetime.starts.

[glider]

Simplified the logic per Evgeniy's request.
No test changes.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2580

You're right, we'd better not over-complicate the collection logic.
I however doubt a lifetime intrinsic can occur before an alloca, because it's using that alloca's result.

[glider]

PTAL

[eugenis]

LGTM
Thank you!

Please remove the condition in clang as well in a separate patch. This thing that you've just fixed is the only reason for MSan condition in shouldEmitLifetimeMarkers() to exist, so I'm afraid it's on you now to clean that up.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 ///  KernelMemorySanitizer only supports X86_64 at the moment.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"

 static cl::opt<bool> ClHandleICmp("msan-handle-icmp",
        cl::desc("propagate shadow through ICmpEQ and ICmpNE"),
        cl::Hidden, cl::init(true));
 
 static cl::opt<bool> ClHandleICmpExact("msan-handle-icmp-exact",
        cl::desc("exact handling of relational integer ICmp"),
        cl::Hidden, cl::init(false));
 
+static cl::opt<bool> ClHandleLifetimeIntrinsics(
+    "msan-handle-lifetime-intrinsics",
+    cl::desc(
+        "when possible, poison scoped variables at the beginning of the scope "
+        "(slower, but more precise)"),
+    cl::Hidden, cl::init(true));
+
 // When compiling the Linux kernel, we sometimes see false positives related to
 // MSan being unable to understand that inline assembly calls may initialize
 // local variables.
 // This flag makes the compiler conservatively unpoison every memory location
 // passed into an assembly call. Note that this may cause false positives.
 // Because it's impossible to figure out the array sizes, we can only unpoison
 // the first sizeof(type) bytes for each type* pointer.
 // The instrumentation is only enabled in KMSAN builds, and only if
struct ShadowOriginAndInsertPoint {
     Value *Shadow;
     Value *Origin;
     Instruction *OrigIns;
 
     ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I)
       : Shadow(S), Origin(O), OrigIns(I) {}
   };
   SmallVector<ShadowOriginAndInsertPoint, 16> InstrumentationList;
+  SmallSet<AllocaInst *, 16> FallbackAllocaSet;
+  bool InstrumentOnlyAllocas = !ClHandleLifetimeIntrinsics;
+  SmallSet<AllocaInst *, 16> AllocaSet;
+  SmallVector<std::pair<IntrinsicInst *, AllocaInst *>, 16> LifetimeStartList;
   SmallVector<StoreInst *, 16> StoreList;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
     bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeMemory);
     InsertChecks = SanitizeFunction;
     PropagateShadow = SanitizeFunction;
for (PHINode *PN : ShadowPHINodes) {
       for (size_t v = 0; v < NumValues; v++) {
         PNS->addIncoming(getShadow(PN, v), PN->getIncomingBlock(v));
         if (PNO) PNO->addIncoming(getOrigin(PN, v), PN->getIncomingBlock(v));
       }
     }
 
     VAHelper->finalizeInstrumentation();
 
+    // Poison allocas that weren't instrumented at llvm.lifetime.start.
+    // If we failed to find an alloca for one of llvm.lifetime.start intrinsics,
+    // fall back to instrumenting all the allocas.
+    auto Allocas = InstrumentOnlyAllocas ? FallbackAllocaSet : AllocaSet;
+    for (AllocaInst *AI : Allocas)
+      instrumentAlloca(*AI);
+    if (!InstrumentOnlyAllocas)
+      for (auto Item : LifetimeStartList) {
+        auto Intr = Item.first;
+        auto Alloca = Item.second;
+        ConstantInt *Len = dyn_cast<ConstantInt>(Intr->getArgOperand(0));
+        IRBuilder<> IRB(Intr->getNextNode());
+        if (MS.CompileKernel)
+          instrumentAllocaKmsan(*Alloca, IRB, Len);
+        else
+          instrumentAllocaUserspace(*Alloca, IRB, Len);
+      }
+
     bool InstrumentWithCalls = ClInstrumentationWithCallThreshold >= 0 &&
                                InstrumentationList.size() + StoreList.size() >
                                    (unsigned)ClInstrumentationWithCallThreshold;
 
     // Insert shadow value checks.
     materializeChecks(InstrumentWithCalls);
 
     // Delayed instrumentation of StoreInst.
bool handleUnknownIntrinsic(IntrinsicInst &I) {
     if (I.doesNotAccessMemory())
       if (maybeHandleSimpleNomemIntrinsic(I))
         return true;
 
     // FIXME: detect and handle SSE maskstore/maskload
     return false;
   }
 
+  void handleLifetimeStart(IntrinsicInst &I) {
+    if (!PoisonStack || InstrumentOnlyAllocas)
+      return;
+    ConstantInt *Len = dyn_cast<ConstantInt>(I.getArgOperand(0));
+    if (Len->isMinusOne())
+      return;
+
+    DenseMap<Value *, AllocaInst *> AllocaForValue;
+    AllocaInst *AI =
+        llvm::findAllocaForValue(I.getArgOperand(1), AllocaForValue);
+    if (!AI) {
+      InstrumentOnlyAllocas = true;
+      return;
+    }
+    if (AllocaSet.count(AI))
+      AllocaSet.erase(AI);
+    LifetimeStartList.push_back(std::make_pair(&I, AI));
+  }
+
   void handleBswap(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *Op = I.getArgOperand(0);
     Type *OpType = Op->getType();
     Function *BswapFunc = Intrinsic::getDeclaration(
       F.getParent(), Intrinsic::bswap, makeArrayRef(&OpType, 1));
     setShadow(&I, IRB.CreateCall(BswapFunc, getShadow(Op)));
     setOrigin(&I, getOrigin(Op));
Value *S = IRB.CreateCall(I.getCalledFunction(),
                               {getShadow(&I, 0), I.getOperand(1)});
     S = IRB.CreateOr(SMask, S);
     setShadow(&I, S);
     setOriginForNaryOp(I);
   }
 
   void visitIntrinsicInst(IntrinsicInst &I) {
     switch (I.getIntrinsicID()) {
+    case Intrinsic::lifetime_start:
+      handleLifetimeStart(I);
+      break;
     case Intrinsic::bswap:
       handleBswap(I);
       break;
     case Intrinsic::masked_store:
       handleMaskedStore(I);
       break;
     case Intrinsic::masked_load:
       handleMaskedLoad(I);
if (PoisonStack) {
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
                       IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy())});
     } else {
       IRB.CreateCall(MS.MsanUnpoisonAllocaFn,
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
     }
   }
 
-  void visitAllocaInst(AllocaInst &I) {
-    setShadow(&I, getCleanShadow(&I));
-    setOrigin(&I, getCleanOrigin());
+  void instrumentAlloca(AllocaInst &I) {
     IRBuilder<> IRB(I.getNextNode());
     const DataLayout &DL = F.getParent()->getDataLayout();
     uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
     Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
     if (I.isArrayAllocation())
       Len = IRB.CreateMul(Len, I.getArraySize());
 
     if (MS.CompileKernel)
       instrumentAllocaKmsan(I, IRB, Len);
     else
       instrumentAllocaUserspace(I, IRB, Len);
   }
 
+  void visitAllocaInst(AllocaInst &I) {
+    setShadow(&I, getCleanShadow(&I));
+    setOrigin(&I, getCleanOrigin());
+    // We'll get to this alloca later unless it's poisoned at the corresponding
+    // llvm.lifetime.start.
+    AllocaSet.insert(&I);
+    FallbackAllocaSet.insert(&I);
+  }
+
   void visitSelectInst(SelectInst& I) {
     IRBuilder<> IRB(&I);
     // a = select b, c, d
     Value *B = I.getCondition();
     Value *C = I.getTrueValue();
     Value *D = I.getFalseValue();
     Value *Sb = getShadow(B);
     Value *Sc = getShadow(C);

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 
 ; CHECK-LABEL: define void @unpoison_local(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
 ; CALL: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
 ; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
 ; KMSAN: call void @__msan_unpoison_alloca(i8* {{.*}}, i64 20)
 ; CHECK: ret void
 
+; Check that every llvm.lifetime.start() causes poisoning of locals.
+define void @lifetime_start() sanitize_memory {
+entry:
+  %x = alloca i32, align 4
+  %c = bitcast i32* %x to i8*
+  br label %another_bb
+
+another_bb:
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %c)
+  store i32 7, i32* %x
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %c)
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %c)
+  store i32 8, i32* %x
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %c)
+  ret void
+}
+
+; CHECK-LABEL: define void @lifetime_start(
+; CHECK-LABEL: entry:
+; CHECK: %x = alloca i32
+; CHECK-LABEL: another_bb:
+
+; CHECK: call void @llvm.lifetime.start
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+
+; CHECK: call void @llvm.lifetime.start
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK: ret void
+
+; If an object is variable sized, ignore llvm.lifetime.start.
+define void @lifetime_start_var(i64 %cnt) sanitize_memory {
+entry:
+  %x = alloca i32, i64 %cnt, align 4
+  %c = bitcast i32* %x to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %c)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %c)
+  ret void
+}
+
+; CHECK-LABEL: define void @lifetime_start_var(
+; CHECK-LABEL: entry:
+; CHECK: %x = alloca i32, i64 %cnt
+; CHECK: %[[A:.*]] = mul i64 4, %cnt
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 %[[A]])
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 %[[A]],
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 %[[A]],
+; CHECK: call void @llvm.lifetime.start
+; CHECK: call void @llvm.lifetime.end
+; CHECK: ret void
+
+
+; If we can't trace one of the lifetime markers to a single alloca, fall back
+; to poisoning allocas at the beginning of the function.
+define void @lifetime_no_alloca(i8 %v) sanitize_memory {
+entry:
+  %x = alloca i32, align 4
+  %y = alloca i32, align 4
+  %z = alloca i32, align 4
+  %cx = bitcast i32* %x to i8*
+  %cy = bitcast i32* %y to i8*
+  %cz = bitcast i32* %z to i8*
+  %tobool = icmp eq i8 %v, 0
+  %xy = select i1 %tobool, i32* %x, i32* %y
+  %cxcy = select i1 %tobool, i8* %cx, i8* %cy
+  br label %another_bb
+
+another_bb:
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %cz)
+  store i32 7, i32* %z
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %cz)
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %cxcy)
+  store i32 8, i32* %xy
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %cxcy)
+  ret void
+}
+
+; CHECK-LABEL: define void @lifetime_no_alloca(
+; CHECK-LABEL: entry:
+; CHECK: %x = alloca i32
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK: %y = alloca i32
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK: %z = alloca i32
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK-LABEL: another_bb:
+
+; CHECK: call void @llvm.lifetime.start
+; INLINE-NOT: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL-NOT: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN-NOT: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK: call void @llvm.lifetime.end
+; CHECK: call void @llvm.lifetime.start
+; INLINE-NOT: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
+; CALL-NOT: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN-NOT: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
+; CHECK: call void @llvm.lifetime.end
+
+
+
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)