diff --git a/compiler-rt/test/msan/Linux/fopencookie.cpp b/compiler-rt/test/msan/Linux/fopencookie.cpp --- a/compiler-rt/test/msan/Linux/fopencookie.cpp +++ b/compiler-rt/test/msan/Linux/fopencookie.cpp @@ -36,7 +36,10 @@ return 0; } -void PoisonStack() { char a[8192]; } +void PoisonStack() { + char a[8192]; + __asm__ volatile("" ::"r"(a)); +} void TestPoisonStack() { // Verify that PoisonStack has poisoned the stack - otherwise this test is not diff --git a/compiler-rt/test/msan/getaddrinfo.cpp b/compiler-rt/test/msan/getaddrinfo.cpp --- a/compiler-rt/test/msan/getaddrinfo.cpp +++ b/compiler-rt/test/msan/getaddrinfo.cpp @@ -7,6 +7,7 @@ void poison_stack_ahead() { char buf[100000]; + __asm__ volatile("" ::"r"(x)); // With -O0 this poisons a large chunk of stack. } diff --git a/compiler-rt/test/msan/qsort.cpp b/compiler-rt/test/msan/qsort.cpp --- a/compiler-rt/test/msan/qsort.cpp +++ b/compiler-rt/test/msan/qsort.cpp @@ -19,6 +19,7 @@ void poison_stack_and_param() { char x[10000]; + __asm__ volatile("" ::"r"(x)); #ifndef TEST_MSAN_EAGER_CHECKS int y; dummy(y, y, y, y, y); diff --git a/compiler-rt/test/msan/vararg.cpp b/compiler-rt/test/msan/vararg.cpp --- a/compiler-rt/test/msan/vararg.cpp +++ b/compiler-rt/test/msan/vararg.cpp @@ -42,6 +42,7 @@ __attribute__((noinline)) void poison_stack_and_param() { char x[10000]; + __asm__ volatile("" ::"r"(x)); #ifndef TEST_MSAN_EAGER_CHECKS int y; dummy(y, y, y, y, y); diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp @@ -3870,6 +3870,73 @@ } } + // Determine whether all paths from this alloca lead to storing into it. + // If true, we can omit poisoning the alloca because it'll just be + // overwritten anyway. + bool firstUsesAreStore(SmallPtrSet &TraversedSet, + SmallPtrSet &Users, + SmallPtrSet &Defs, + Instruction *CurrInst) { + if (TraversedSet.size() == 100) { + // Have to cut it off somewhere + return false; + } + BasicBlock *BB = CurrInst->getParent(); + TraversedSet.insert(BB); + + while (CurrInst) { + if (Users.count(CurrInst)) { + if (StoreInst *Store = dyn_cast(CurrInst)) { + // The store could be a use if it's storing the alloca + // pointer somewhere. But we don't want that. + return Defs.count(Store->getPointerOperand()); + } else if (IntrinsicInst *Intrinsic = + dyn_cast(CurrInst)) { + // Ignore lifetime intrinsics, count others as a use. + Intrinsic::ID ID = Intrinsic->getIntrinsicID(); + if (ID != Intrinsic::lifetime_start && ID != Intrinsic::lifetime_end) + return false; + } else { + return false; + } + } + CurrInst = CurrInst->getNextNonDebugInstruction(); + if (Defs.count(CurrInst)) { + // Append chained uses and defs to their relevant sets + populateRelevantAllocaUsers(CurrInst, Defs, Users); + } + } + + // Traverse to the Node's successors if we haven't yet + for (auto Node = succ_begin(BB); Node != succ_end(BB); Node++) { + if (TraversedSet.count(*Node) == 0) { + Instruction *NextInst = Node->getFirstNonPHIOrDbgOrLifetime(); + if (!firstUsesAreStore(TraversedSet, Users, Defs, NextInst)) + return false; + } + } + return true; + }; + + void populateRelevantAllocaUsers(Instruction *I, + SmallPtrSet &Defs, + SmallPtrSet &Users) { + Defs.insert(I); + + for (User *User : I->users()) { + if (isa(User)) { + Defs.insert(User); + continue; + } else if (GetElementPtrInst *GEP = dyn_cast(User)) { + if (GEP->getNumIndices() == 1 && GEP->hasAllZeroIndices()) { + Defs.insert(GEP); + continue; + } + } + Users.insert(User); + } + } + void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) { if (!InsPoint) InsPoint = &I; @@ -3877,9 +3944,15 @@ const DataLayout &DL = F.getParent()->getDataLayout(); uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType()); Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize); - if (I.isArrayAllocation()) + if (I.isArrayAllocation()) { Len = IRB.CreateMul(Len, I.getArraySize()); - + } else { + SmallPtrSet Users, Defs; + SmallPtrSet TraversedSet; + populateRelevantAllocaUsers(&I, Defs, Users); + if (firstUsesAreStore(TraversedSet, Users, Defs, &I)) + return; + } if (MS.CompileKernel) poisonAllocaKmsan(I, IRB, Len); else diff --git a/llvm/test/Instrumentation/MemorySanitizer/alloca-store.ll b/llvm/test/Instrumentation/MemorySanitizer/alloca-store.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Instrumentation/MemorySanitizer/alloca-store.ll @@ -0,0 +1,38 @@ +; Tests if MSAN can optimize away dead-code alloca poisons + +; RUN: opt < %s -msan-track-origins=2 -msan-check-access-address=0 -msan-poison-stack-with-call=1 -S -passes=msan 2>&1 | FileCheck \ +; RUN: %s "--check-prefixes=CHECK" + +target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +define i32* @simple(i32 %val) sanitize_memory { + ; CHECK: define {{.*}} @simple + %p = alloca i32 + store i32 %val, i32* %p + ; CHECK-NOT: call void @__msan_poison_stack + ret i32* %p +} + +define i32* @bitcast({ i32 } %val) sanitize_memory { + ; CHECK: define {{.*}} @bitcast + %p = alloca i32 + %ps = bitcast i32* %p to { i32 }* + store { i32 } %val, { i32 }* %ps + ; CHECK-NOT: call void @__msan_poison_stack + ret i32* %p +} + +define i32* @lifetime(i32 %val) sanitize_memory { + ; CHECK: define {{.*}} @lifetime + %p = alloca i32 + %p8 = bitcast i32* %p to i8* + call void @llvm.lifetime.start(i64 4, i8* %p8) + store i32 %val, i32* %p + ; CHECK-NOT: call void @__msan_poison_stack + call void @llvm.lifetime.end(i64 4, i8* %p8) + ret i32* %p +} + +declare void @llvm.lifetime.start(i64, i8* nocapture) +declare void @llvm.lifetime.end(i64, i8* nocapture) diff --git a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll --- a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll +++ b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll @@ -22,6 +22,7 @@ define void @static() sanitize_memory { entry: %x = alloca i32, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -38,6 +39,7 @@ br label %l l: %x = alloca i32, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -51,6 +53,7 @@ define void @array() sanitize_memory { entry: %x = alloca i32, i64 5, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -64,6 +67,7 @@ define void @array_non_const(i64 %cnt) sanitize_memory { entry: %x = alloca i32, i64 %cnt, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -79,6 +83,7 @@ define void @unpoison_local() { entry: %x = alloca i32, i64 5, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } diff --git a/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll b/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll --- a/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll +++ b/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll @@ -58,9 +58,9 @@ ret i32 1 } -; %nr is first poisoned, then unpoisoned (written to). Need to optimize this in the future. -; CHECK: [[NRC1:%.*]] = bitcast i64* %nr to i8* -; CHECK: call void @__msan_poison_alloca(i8* [[NRC1]]{{.*}}) +; %nr is first poisoned, then unpoisoned (written to). No need to poison. +; CHECK-NOT: [[NRC1:%.*]] = bitcast i64* %nr to i8* +; CHECK-NOT: call void @__msan_poison_alloca(i8* [[NRC1]]{{.*}}) ; CHECK: [[NRC2:%.*]] = bitcast i64* %nr to i8* ; CHECK: call { i8*, i32* } @__msan_metadata_ptr_for_store_8(i8* [[NRC2]])