From d55ef5ce5f92fc02063e65ef328b89a7d66a3636 Mon Sep 17 00:00:00 2001
From: Evgeniy Stepanov <eugeni.stepanov@gmail.com>
Date: Wed, 3 Jul 2013 14:39:14 +0000
Subject: [PATCH] [msan] Unpoison stack allocations and undef values in
 blacklisted functions.

This changes behavior of -msan-poison-stack=0 flag from not poisoning stack
allocations to actively unpoisoning them.



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@185538 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../Instrumentation/MemorySanitizer.cpp       | 24 +++++++------
 .../MemorySanitizer/msan_basic.ll             | 35 +++++++++++++++++++
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index a3a688d7c8d..8e2baefd612 100644
--- a/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -425,6 +425,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   ValueMap<Value*, Value*> ShadowMap, OriginMap;
   bool InsertChecks;
   bool LoadShadow;
+  bool PoisonStack;
+  bool PoisonUndef;
   OwningPtr<VarArgHelper> VAHelper;
 
   struct ShadowOriginAndInsertPoint {
@@ -440,10 +442,13 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)) {
-    LoadShadow = InsertChecks =
-        !MS.BL->isIn(F) &&
-        F.getAttributes().hasAttribute(AttributeSet::FunctionIndex,
-                                       Attribute::SanitizeMemory);
+    bool SanitizeFunction = !MS.BL->isIn(F) && F.getAttributes().hasAttribute(
+                                                   AttributeSet::FunctionIndex,
+                                                   Attribute::SanitizeMemory);
+    InsertChecks = SanitizeFunction;
+    LoadShadow = SanitizeFunction;
+    PoisonStack = SanitizeFunction && ClPoisonStack;
+    PoisonUndef = SanitizeFunction && ClPoisonUndef;
 
     DEBUG(if (!InsertChecks)
           dbgs() << "MemorySanitizer is not inserting checks into '"
@@ -744,7 +749,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       return Shadow;
     }
     if (UndefValue *U = dyn_cast<UndefValue>(V)) {
-      Value *AllOnes = ClPoisonUndef ? getPoisonedShadow(V) : getCleanShadow(V);
+      Value *AllOnes = PoisonUndef ? getPoisonedShadow(V) : getCleanShadow(V);
       DEBUG(dbgs() << "Undef: " << *U << " ==> " << *AllOnes << "\n");
       (void)U;
       return AllOnes;
@@ -1704,20 +1709,19 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
   void visitAllocaInst(AllocaInst &I) {
     setShadow(&I, getCleanShadow(&I));
-    if (!ClPoisonStack) return;
     IRBuilder<> IRB(I.getNextNode());
     uint64_t Size = MS.TD->getTypeAllocSize(I.getAllocatedType());
-    if (ClPoisonStackWithCall) {
+    if (PoisonStack && ClPoisonStackWithCall) {
       IRB.CreateCall2(MS.MsanPoisonStackFn,
                       IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()),
                       ConstantInt::get(MS.IntptrTy, Size));
     } else {
       Value *ShadowBase = getShadowPtr(&I, Type::getInt8PtrTy(*MS.C), IRB);
-      IRB.CreateMemSet(ShadowBase, IRB.getInt8(ClPoisonStackPattern),
-                       Size, I.getAlignment());
+      Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0);
+      IRB.CreateMemSet(ShadowBase, PoisonValue, Size, I.getAlignment());
     }
 
-    if (MS.TrackOrigins) {
+    if (PoisonStack && MS.TrackOrigins) {
       setOrigin(&I, getCleanOrigin());
       SmallString<2048> StackDescriptionStorage;
       raw_svector_ostream StackDescription(StackDescriptionStorage);
diff --git a/test/Instrumentation/MemorySanitizer/msan_basic.ll b/test/Instrumentation/MemorySanitizer/msan_basic.ll
index 3a8092bdc46..4fa031924f5 100644
--- a/test/Instrumentation/MemorySanitizer/msan_basic.ll
+++ b/test/Instrumentation/MemorySanitizer/msan_basic.ll
@@ -638,6 +638,41 @@ declare void @bar()
 ; CHECK: ret i32
 
 
+; Test that stack allocations are unpoisoned in functions missing
+; sanitize_memory attribute
+
+define i32 @NoSanitizeMemoryAlloca() {
+entry:
+  %p = alloca i32, align 4
+  %x = call i32 @NoSanitizeMemoryAllocaHelper(i32* %p)
+  ret i32 %x
+}
+
+declare i32 @NoSanitizeMemoryAllocaHelper(i32* %p)
+
+; CHECK: @NoSanitizeMemoryAlloca
+; CHECK: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 0, i64 4, i32 4, i1 false)
+; CHECK: call i32 @NoSanitizeMemoryAllocaHelper(i32*
+; CHECK: ret i32
+
+
+; Test that undef is unpoisoned in functions missing
+; sanitize_memory attribute
+
+define i32 @NoSanitizeMemoryUndef() {
+entry:
+  %x = call i32 @NoSanitizeMemoryUndefHelper(i32 undef)
+  ret i32 %x
+}
+
+declare i32 @NoSanitizeMemoryUndefHelper(i32 %x)
+
+; CHECK: @NoSanitizeMemoryAlloca
+; CHECK: store i32 0, i32* {{.*}} @__msan_param_tls
+; CHECK: call i32 @NoSanitizeMemoryUndefHelper(i32 undef)
+; CHECK: ret i32
+
+
 ; Test argument shadow alignment
 
 define <2 x i64> @ArgumentShadowAlignment(i64 %a, <2 x i64> %b) sanitize_memory {
-- 
2.34.1