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"(buf)); // 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 @@ -145,6 +145,7 @@ #include "llvm/Transforms/Instrumentation/MemorySanitizer.h" #include "llvm/ADT/APInt.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/BitVector.h" #include "llvm/ADT/DepthFirstIterator.h" #include "llvm/ADT/SmallSet.h" #include "llvm/ADT/SmallString.h" @@ -3849,6 +3850,133 @@ } } + // Bytes in an alloca already stored to + using StoredSet = BitVector; + // Mapping of Value modifying a pointer => Offset it's pointing to + using RedefOffsets = DenseMap; + // Set of instructions consuming a given alloca + using UserSet = SmallPtrSet; + +#define DEBUG_STORE_SEARCH(Type, Inst) \ + errs() << " Found " #Type ": "; \ + (Inst)->print(errs()); \ + errs() << "\n"; + + // 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, + UserSet &Users, RedefOffsets &Defs, + Instruction *CurrInst, StoredSet StoredBytes) { + if (TraversedSet.size() == 100) { + // Have to cut it off somewhere + errs() << " Reached max recursion.\n"; + return false; + } + BasicBlock *BB = CurrInst->getParent(); + TraversedSet.insert(BB); + + while (CurrInst) { + if (Users.count(CurrInst)) { + if (StoreInst *Store = dyn_cast(CurrInst)) { + DEBUG_STORE_SEARCH("store", CurrInst); + + const DataLayout &DL = F.getParent()->getDataLayout(); + uint64_t StoreSize = + DL.getTypeStoreSize(Store->getValueOperand()->getType()); + // The store could be a use if it's storing the alloca + // pointer somewhere. But we don't want that. + auto Def = Defs.find(Store->getPointerOperand()); + if (Def != Defs.end()) { + StoredBytes.set(Def->second, Def->second + StoreSize); + } else { + return false; + } + } else if (IntrinsicInst *Intrinsic = + dyn_cast(CurrInst)) { + DEBUG_STORE_SEARCH("intrinsic", 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 if (CallInst *Call = dyn_cast(CurrInst)) { + DEBUG_STORE_SEARCH("call", CurrInst); + + // We want to consider instrumented memory writing functions + // as stores here. + if (Call->getCalledFunction() == MS.MemcpyFn.getCallee() || + Call->getCalledFunction() == MS.MemsetFn.getCallee() || + Call->getCalledFunction() == MS.MemmoveFn.getCallee()) { + Value *Dest = Call->getArgOperand(0); + Value *StoreSizeVal = Call->getArgOperand(2); + + auto Def = Defs.find(Dest); + if (Def == Defs.end()) + return false; + + // Can only decide this covers the whole shadow if we can + // check the constant store size. + if (ConstantInt *StoreSizeConst = + dyn_cast(StoreSizeVal)) { + uint64_t StoreSize = StoreSizeConst->getLimitedValue(); + StoredBytes.set(Def->second, Def->second + StoreSize); + } else { + return false; + } + } else { + return false; + } + } else { + DEBUG_STORE_SEARCH("use", CurrInst); + return false; + } + } + + if (StoredBytes.all()) { + errs() << " Stored whole region.\n"; + return true; + } + + 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, + StoredBytes)) + return false; + } + } + return true; + }; + + void populateRelevantAllocaUsers(Instruction *I, RedefOffsets &Defs, + UserSet &Users, uint64_t StoreOffs = 0) { + Defs.insert({I, StoreOffs}); + + for (User *User : I->users()) { + if (isa(User)) { + Defs.insert({User, StoreOffs}); + continue; + } else if (GetElementPtrInst *GEP = dyn_cast(User)) { + const DataLayout &DL = F.getParent()->getDataLayout(); + APInt Offs(DL.getIndexTypeSizeInBits(GEP->getType()), 0); + if (GEP->accumulateConstantOffset(DL, Offs)) { + Defs.insert({GEP, StoreOffs + Offs.getLimitedValue()}); + continue; + } + } + Users.insert(User); + } + } + void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) { if (!InsPoint) InsPoint = &I; @@ -3856,8 +3984,22 @@ 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 { + RedefOffsets Defs; + UserSet Users; + SmallPtrSet TraversedSet; + StoredSet StoredBytes(TypeSize, false); + + errs() << "Instrumenting alloca: "; + I.print(errs()); + errs() << "\n"; + + populateRelevantAllocaUsers(&I, Defs, Users); + if (firstUsesAreStore(TraversedSet, Users, Defs, &I, StoredBytes)) + return; + } if (MS.CompileKernel) poisonAllocaKmsan(I, IRB, Len); 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,67 @@ +; 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 +} + +define i32* @with_memset(i8 %val) sanitize_memory { + ; CHECK: define {{.*}} @with_memset + %p = alloca i32 + %p8 = bitcast i32* %p to i8* + call void @llvm.memset(i8* %p8, i8 0, i64 4, i1 false) + ; CHECK-NOT: call void @__msan_poison_stack + ret i32* %p +} + +%struct.multi = type { i32, i8, i8, i8, i8 } + +define i32* @piecemeal(i8 %val) sanitize_memory { + ; CHECK: define {{.*}} @piecemeal + %p = alloca %struct.multi + %p0 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 0 + %p1 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 1 + %p2 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 2 + %p3 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 3 + %p4 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 4 + store i32 0, i32* %p0 + store i8 0, i8* %p1 + store i8 0, i8* %p2 + store i8 0, i8* %p3 + store i8 0, i8* %p4 + ; CHECK-NOT: call void @__msan_poison_stack + ret i32* %p0 +} + +declare void @llvm.memset(i8*, i8, i64, i1) +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]])