[msan] Fix instrumentation of array allocas.

Before this, MSan poisoned exactly one element of any array alloca,
even if the number of elements was zero.

llvm-svn: 296050

2 files changed, 66 insertions, 6 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index fafb0fcbd01..69b1dce63d5 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -2717,15 +2717,17 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     setOrigin(&I, getCleanOrigin());
     IRBuilder<> IRB(I.getNextNode());
     const DataLayout &DL = F.getParent()->getDataLayout();
-    uint64_t Size = DL.getTypeAllocSize(I.getAllocatedType());
+    uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
+    Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
+    if (I.isArrayAllocation())
+      Len = IRB.CreateMul(Len, I.getArraySize());
     if (PoisonStack && ClPoisonStackWithCall) {
       IRB.CreateCall(MS.MsanPoisonStackFn,
-                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()),
-                      ConstantInt::get(MS.IntptrTy, Size)});
+                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
     } else {
       Value *ShadowBase = getShadowPtr(&I, Type::getInt8PtrTy(*MS.C), IRB);
       Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0);
-      IRB.CreateMemSet(ShadowBase, PoisonValue, Size, I.getAlignment());
+      IRB.CreateMemSet(ShadowBase, PoisonValue, Len, I.getAlignment());
     }
 
     if (PoisonStack && MS.TrackOrigins) {
@@ -2742,8 +2744,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
                                                StackDescription.str());
 
       IRB.CreateCall(MS.MsanSetAllocaOrigin4Fn,
-                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()),
-                      ConstantInt::get(MS.IntptrTy, Size),
+                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
                       IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy()),
                       IRB.CreatePointerCast(&F, MS.IntptrTy)});
     }
diff --git a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll
new file mode 100644
index 00000000000..57ee9120ae8
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll
@@ -0,0 +1,59 @@
+; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck %s --check-prefixes=CHECK,INLINE
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-poison-stack-with-call=1 -S | FileCheck %s --check-prefixes=CHECK,CALL
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck %s --check-prefixes=CHECK,ORIGIN
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck %s --check-prefixes=CHECK,ORIGIN
+
+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 void @static() sanitize_memory {
+entry:
+  %x = alloca i32, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @static(
+; INLINE: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 -1, i64 4, i32 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; CHECK: ret void
+
+
+define void @dynamic() sanitize_memory {
+entry:
+  br label %l
+l:
+  %x = alloca i32, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @dynamic(
+; INLINE: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 -1, i64 4, i32 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; CHECK: ret void
+
+define void @array() sanitize_memory {
+entry:
+  %x = alloca i32, i64 5, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @array(
+; INLINE: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 -1, i64 20, i32 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 20)
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
+; CHECK: ret void
+
+define void @array_non_const(i64 %cnt) sanitize_memory {
+entry:
+  %x = alloca i32, i64 %cnt, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @array_non_const(
+; CHECK: %[[A:.*]] = mul i64 4, %cnt
+; INLINE: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 -1, i64 %[[A]], i32 4, i1 false)
+; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 %[[A]])
+; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 %[[A]],
+; CHECK: ret void