[MSan] store origins for variadic function parameters in __msan_va_arg_origin_tls

Add the __msan_va_arg_origin_tls TLS array to keep the origins for
variadic function parameters.
Change the instrumentation pass to store parameter origins in this array.

llvm-svn: 341528

4 files changed, 232 insertions, 4 deletions

diff --git a/compiler-rt/lib/msan/msan.cc b/compiler-rt/lib/msan/msan.cc
index 06bcbdf8869..ba2d5d59330 100644
--- a/compiler-rt/lib/msan/msan.cc
+++ b/compiler-rt/lib/msan/msan.cc
@@ -59,6 +59,10 @@ SANITIZER_INTERFACE_ATTRIBUTE
 ALIGNED(16) THREADLOCAL u64 __msan_va_arg_tls[kMsanParamTlsSize / sizeof(u64)];
 
 SANITIZER_INTERFACE_ATTRIBUTE
+ALIGNED(16)
+THREADLOCAL u32 __msan_va_arg_origin_tls[kMsanParamTlsSize / sizeof(u32)];
+
+SANITIZER_INTERFACE_ATTRIBUTE
 THREADLOCAL u64 __msan_va_arg_overflow_size_tls;
 
 SANITIZER_INTERFACE_ATTRIBUTE
@@ -277,6 +281,8 @@ void ScopedThreadLocalStateBackup::Restore() {
   internal_memset(__msan_param_tls, 0, sizeof(__msan_param_tls));
   internal_memset(__msan_retval_tls, 0, sizeof(__msan_retval_tls));
   internal_memset(__msan_va_arg_tls, 0, sizeof(__msan_va_arg_tls));
+  internal_memset(__msan_va_arg_origin_tls, 0,
+                  sizeof(__msan_va_arg_origin_tls));
 
   if (__msan_get_track_origins()) {
     internal_memset(&__msan_retval_origin_tls, 0,
diff --git a/compiler-rt/test/msan/vararg.cc b/compiler-rt/test/msan/vararg.cc
new file mode 100644
index 00000000000..efdac40a9c2
--- /dev/null
+++ b/compiler-rt/test/msan/vararg.cc
@@ -0,0 +1,55 @@
+// RUN: %clangxx_msan -fsanitize-memory-track-origins=0 -O3 %s -o %t && \
+// RUN:     not %run %t va_arg_tls >%t.out 2>&1
+// RUN: FileCheck %s --check-prefix=CHECK < %t.out
+
+// RUN: %clangxx_msan -fsanitize-memory-track-origins=0 -O3 %s -o %t && \
+// RUN:     not %run %t overflow >%t.out 2>&1
+// RUN: FileCheck %s --check-prefix=CHECK < %t.out
+
+// RUN: %clangxx_msan -fsanitize-memory-track-origins=2 -O3 %s -o %t && \
+// RUN:     not %run %t va_arg_tls >%t.out 2>&1
+// RUN: FileCheck %s --check-prefixes=CHECK,CHECK-ORIGIN < %t.out
+
+// RUN: %clangxx_msan -fsanitize-memory-track-origins=2 -O3 %s -o %t && \
+// RUN:     not %run %t overflow >%t.out 2>&1
+// RUN: FileCheck %s --check-prefixes=CHECK,CHECK-ORIGIN < %t.out
+
+// Check that shadow and origin are passed through va_args.
+
+#include <stdarg.h>
+#include <string.h>
+
+__attribute__((noinline))
+int sum(int n, ...) {
+  va_list args;
+  int i, sum = 0, arg;
+  volatile int temp;
+  va_start(args, n);
+  for (i = 0; i < n; i++) {
+    arg = va_arg(args, int);
+    sum += arg;
+  }
+  va_end(args);
+  return sum;
+}
+
+int main(int argc, char *argv[]) {
+  volatile int uninit;
+  volatile int a = 1, b = 2;
+  if (argc == 2) {
+    // Shadow/origin will be passed via va_arg_tls/va_arg_origin_tls.
+    if (strcmp(argv[1], "va_arg_tls") == 0) {
+      return sum(3, uninit, a, b);
+    }
+    // Shadow/origin of |uninit| will be passed via overflow area.
+    if (strcmp(argv[1], "overflow") == 0) {
+      return sum(7,
+        a, a, a, a, a, a, uninit
+      );
+    }
+  }
+  return 0;
+}
+
+// CHECK: WARNING: MemorySanitizer: use-of-uninitialized-value
+// CHECK-ORIGIN: Uninitialized value was created by an allocation of 'uninit' in the stack frame of function 'main'
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index b3c31f5d53b..80272f0da71 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -448,6 +448,10 @@ private:
   /// parameters (x86_64-specific).
   GlobalVariable *VAArgTLS;
 
+  /// Thread-local shadow storage for in-register va_arg function
+  /// parameters (x86_64-specific).
+  GlobalVariable *VAArgOriginTLS;
+
   /// Thread-local shadow storage for va_arg overflow area
   /// (x86_64-specific).
   GlobalVariable *VAArgOverflowSizeTLS;
@@ -560,6 +564,12 @@ void MemorySanitizer::createUserspaceApi(Module &M) {
       M, ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), false,
       GlobalVariable::ExternalLinkage, nullptr, "__msan_va_arg_tls", nullptr,
       GlobalVariable::InitialExecTLSModel);
+
+  VAArgOriginTLS = new GlobalVariable(
+      M, ArrayType::get(OriginTy, kParamTLSSize / 4), false,
+      GlobalVariable::ExternalLinkage, nullptr, "__msan_va_arg_origin_tls",
+      nullptr, GlobalVariable::InitialExecTLSModel);
+
   VAArgOverflowSizeTLS = new GlobalVariable(
       M, IRB.getInt64Ty(), false, GlobalVariable::ExternalLinkage, nullptr,
       "__msan_va_arg_overflow_size_tls", nullptr,
@@ -3258,6 +3268,7 @@ struct VarArgAMD64Helper : public VarArgHelper {
   MemorySanitizer &MS;
   MemorySanitizerVisitor &MSV;
   Value *VAArgTLSCopy = nullptr;
+  Value *VAArgTLSOriginCopy = nullptr;
   Value *VAArgOverflowSize = nullptr;
 
   SmallVector<CallInst*, 16> VAStartInstrumentationList;
@@ -3320,6 +3331,9 @@ struct VarArgAMD64Helper : public VarArgHelper {
         uint64_t ArgSize = DL.getTypeAllocSize(RealTy);
         Value *ShadowBase = getShadowPtrForVAArgument(
             RealTy, IRB, OverflowOffset, alignTo(ArgSize, 8));
+        Value *OriginBase = nullptr;
+        if (MS.TrackOrigins)
+          OriginBase = getOriginPtrForVAArgument(RealTy, IRB, OverflowOffset);
         OverflowOffset += alignTo(ArgSize, 8);
         if (!ShadowBase)
           continue;
@@ -3330,22 +3344,31 @@ struct VarArgAMD64Helper : public VarArgHelper {
 
         IRB.CreateMemCpy(ShadowBase, kShadowTLSAlignment, ShadowPtr,
                          kShadowTLSAlignment, ArgSize);
+        if (MS.TrackOrigins)
+          IRB.CreateMemCpy(OriginBase, kShadowTLSAlignment, OriginPtr,
+                           kShadowTLSAlignment, ArgSize);
       } else {
         ArgKind AK = classifyArgument(A);
         if (AK == AK_GeneralPurpose && GpOffset >= AMD64GpEndOffset)
           AK = AK_Memory;
         if (AK == AK_FloatingPoint && FpOffset >= AMD64FpEndOffset)
           AK = AK_Memory;
-        Value *ShadowBase;
+        Value *ShadowBase, *OriginBase = nullptr;
         switch (AK) {
           case AK_GeneralPurpose:
             ShadowBase =
                 getShadowPtrForVAArgument(A->getType(), IRB, GpOffset, 8);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, GpOffset);
             GpOffset += 8;
             break;
           case AK_FloatingPoint:
             ShadowBase =
                 getShadowPtrForVAArgument(A->getType(), IRB, FpOffset, 16);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, FpOffset);
             FpOffset += 16;
             break;
           case AK_Memory:
@@ -3354,16 +3377,26 @@ struct VarArgAMD64Helper : public VarArgHelper {
             uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
             ShadowBase =
                 getShadowPtrForVAArgument(A->getType(), IRB, OverflowOffset, 8);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, OverflowOffset);
             OverflowOffset += alignTo(ArgSize, 8);
         }
         // Take fixed arguments into account for GpOffset and FpOffset,
         // but don't actually store shadows for them.
+        // TODO(glider): don't call get*PtrForVAArgument() for them.
         if (IsFixed)
           continue;
         if (!ShadowBase)
           continue;
-        IRB.CreateAlignedStore(MSV.getShadow(A), ShadowBase,
-                               kShadowTLSAlignment);
+        Value *Shadow = MSV.getShadow(A);
+        IRB.CreateAlignedStore(Shadow, ShadowBase, kShadowTLSAlignment);
+        if (MS.TrackOrigins) {
+          Value *Origin = MSV.getOrigin(A);
+          unsigned StoreSize = DL.getTypeStoreSize(Shadow->getType());
+          MSV.paintOrigin(IRB, Origin, OriginBase, StoreSize,
+                          std::max(kShadowTLSAlignment, kMinOriginAlignment));
+        }
       }
     }
     Constant *OverflowSize =
@@ -3380,7 +3413,18 @@ struct VarArgAMD64Helper : public VarArgHelper {
     Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
     Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
     return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
-                              "_msarg");
+                              "_msarg_va_s");
+  }
+
+  /// Compute the origin address for a given va_arg.
+  Value *getOriginPtrForVAArgument(Type *Ty, IRBuilder<> &IRB, int ArgOffset) {
+    Value *Base = IRB.CreatePointerCast(MS.VAArgOriginTLS, MS.IntptrTy);
+    // getOriginPtrForVAArgument() is always called after
+    // getShadowPtrForVAArgument(), so __msan_va_arg_origin_tls can never
+    // overflow.
+    Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
+    return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
+                              "_msarg_va_o");
   }
 
   void unpoisonVAListTagForInst(IntrinsicInst &I) {
@@ -3425,6 +3469,10 @@ struct VarArgAMD64Helper : public VarArgHelper {
                       VAArgOverflowSize);
       VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
       IRB.CreateMemCpy(VAArgTLSCopy, 8, MS.VAArgTLS, 8, CopySize);
+      if (MS.TrackOrigins) {
+        VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+        IRB.CreateMemCpy(VAArgTLSOriginCopy, 8, MS.VAArgOriginTLS, 8, CopySize);
+      }
     }
 
     // Instrument va_start.
@@ -3446,6 +3494,9 @@ struct VarArgAMD64Helper : public VarArgHelper {
                                  Alignment, /*isStore*/ true);
       IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
                        AMD64FpEndOffset);
+      if (MS.TrackOrigins)
+        IRB.CreateMemCpy(RegSaveAreaOriginPtr, Alignment, VAArgTLSOriginCopy,
+                         Alignment, AMD64FpEndOffset);
       Value *OverflowArgAreaPtrPtr = IRB.CreateIntToPtr(
           IRB.CreateAdd(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
                         ConstantInt::get(MS.IntptrTy, 8)),
@@ -3459,6 +3510,12 @@ struct VarArgAMD64Helper : public VarArgHelper {
                                              AMD64FpEndOffset);
       IRB.CreateMemCpy(OverflowArgAreaShadowPtr, Alignment, SrcPtr, Alignment,
                        VAArgOverflowSize);
+      if (MS.TrackOrigins) {
+        SrcPtr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), VAArgTLSOriginCopy,
+                                        AMD64FpEndOffset);
+        IRB.CreateMemCpy(OverflowArgAreaOriginPtr, Alignment, SrcPtr, Alignment,
+                         VAArgOverflowSize);
+      }
     }
   }
 };
diff --git a/llvm/test/Instrumentation/MemorySanitizer/X86/vararg_call.ll b/llvm/test/Instrumentation/MemorySanitizer/X86/vararg_call.ll
new file mode 100644
index 00000000000..79817a889ca
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/X86/vararg_call.ll
@@ -0,0 +1,110 @@
+; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck %s
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck %s --check-prefixes=CHECK,CHECK-ORIGIN
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck %s --check-prefixes=CHECK,CHECK-ORIGIN
+
+; Test that shadow and origin are stored for variadic function params.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.__va_list_tag = type { i32, i32, i8*, i8* }
+
+define dso_local i32 @test(i32 %a, i32 %b, i32 %c) local_unnamed_addr {
+entry:
+  %call = tail call i32 (i32, ...) @sum(i32 3, i32 %a, i32 %b, i32 %c)
+  ret i32 %call
+}
+
+; CHECK: store i32 0, {{.*}} @__msan_param_tls {{.*}} i64 8
+; CHECK: store i32 0, {{.*}} @__msan_param_tls {{.*}} i64 16
+; CHECK: store i32 0, {{.*}} @__msan_param_tls {{.*}} i64 24
+; CHECK: store i32 0, {{.*}} @__msan_va_arg_tls {{.*}} i64 8
+; CHECK-ORIGIN: store i32 0, {{.*}} @__msan_va_arg_origin_tls {{.*}} i64 8
+; CHECK: store i32 0, {{.*}} @__msan_va_arg_tls {{.*}} i64 16
+; CHECK-ORIGIN: store i32 0, {{.*}} @__msan_va_arg_origin_tls {{.*}} i64 16
+; CHECK: store i32 0, {{.*}} @__msan_va_arg_tls {{.*}} i64 24
+; CHECK-ORIGIN: store i32 0, {{.*}} @__msan_va_arg_origin_tls {{.*}} i64 24
+
+define dso_local i32 @sum(i32 %n, ...) local_unnamed_addr #0 {
+entry:
+  %args = alloca [1 x %struct.__va_list_tag], align 16
+  %0 = bitcast [1 x %struct.__va_list_tag]* %args to i8*
+  call void @llvm.lifetime.start.p0i8(i64 24, i8* nonnull %0) #2
+  call void @llvm.va_start(i8* nonnull %0)
+  %cmp9 = icmp sgt i32 %n, 0
+  br i1 %cmp9, label %for.body.lr.ph, label %for.end
+
+; CHECK: [[VA_ARG_TLS:%[_0-9a-z]+]] = bitcast {{.*}} @__msan_va_arg_tls
+; CHECK: call void @llvm.memcpy.{{.*}} [[SHADOW_COPY:%[_0-9a-z]+]], {{.*}} [[VA_ARG_TLS]]
+; CHECK-ORIGIN: [[VA_ARG_ORIGIN_TLS:%[_0-9a-z]+]] = bitcast {{.*}} @__msan_va_arg_origin_tls
+; CHECK-ORIGIN: call void @llvm.memcpy{{.*}} [[ORIGIN_COPY:%[_0-9a-z]+]], {{.*}} [[VA_ARG_ORIGIN_TLS]]
+
+; CHECK: call void @llvm.va_start
+; CHECK: call void @llvm.memcpy.{{.*}}, {{.*}} [[SHADOW_COPY]], i{{.*}} [[REGSAVE:[0-9]+]]
+; CHECK-ORIGIN: call void @llvm.memcpy.{{.*}}, {{.*}} [[ORIGIN_COPY]], i{{.*}} [[REGSAVE]]
+
+; CHECK: [[OVERFLOW_SHADOW:%[_0-9a-z]+]] = getelementptr i8, i8* [[SHADOW_COPY]], i{{.*}} [[REGSAVE]]
+; CHECK: call void @llvm.memcpy.{{.*}}[[OVERFLOW_SHADOW]]
+; CHECK-ORIGIN: [[OVERFLOW_ORIGIN:%[_0-9a-z]+]] = getelementptr i8, i8* [[ORIGIN_COPY]], i{{.*}} [[REGSAVE]]
+; CHECK-ORIGIN: call void @llvm.memcpy.{{.*}}[[OVERFLOW_ORIGIN]]
+
+for.body.lr.ph:                                   ; preds = %entry
+  %gp_offset_p = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %args, i64 0, i64 0, i32 0
+  %1 = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %args, i64 0, i64 0, i32 3
+  %overflow_arg_area_p = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %args, i64 0, i64 0, i32 2
+  %gp_offset.pre = load i32, i32* %gp_offset_p, align 16
+  br label %for.body
+
+for.body:                                         ; preds = %vaarg.end, %for.body.lr.ph
+  %gp_offset = phi i32 [ %gp_offset.pre, %for.body.lr.ph ], [ %gp_offset12, %vaarg.end ]
+  %sum.011 = phi i32 [ 0, %for.body.lr.ph ], [ %add, %vaarg.end ]
+  %i.010 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %vaarg.end ]
+  %fits_in_gp = icmp ult i32 %gp_offset, 41
+  br i1 %fits_in_gp, label %vaarg.in_reg, label %vaarg.in_mem
+
+vaarg.in_reg:                                     ; preds = %for.body
+  %reg_save_area = load i8*, i8** %1, align 16
+  %2 = sext i32 %gp_offset to i64
+  %3 = getelementptr i8, i8* %reg_save_area, i64 %2
+  %4 = add i32 %gp_offset, 8
+  store i32 %4, i32* %gp_offset_p, align 16
+  br label %vaarg.end
+
+vaarg.in_mem:                                     ; preds = %for.body
+  %overflow_arg_area = load i8*, i8** %overflow_arg_area_p, align 8
+  %overflow_arg_area.next = getelementptr i8, i8* %overflow_arg_area, i64 8
+  store i8* %overflow_arg_area.next, i8** %overflow_arg_area_p, align 8
+  br label %vaarg.end
+
+vaarg.end:                                        ; preds = %vaarg.in_mem, %vaarg.in_reg
+  %gp_offset12 = phi i32 [ %4, %vaarg.in_reg ], [ %gp_offset, %vaarg.in_mem ]
+  %vaarg.addr.in = phi i8* [ %3, %vaarg.in_reg ], [ %overflow_arg_area, %vaarg.in_mem ]
+  %vaarg.addr = bitcast i8* %vaarg.addr.in to i32*
+  %5 = load i32, i32* %vaarg.addr, align 4
+  %add = add nsw i32 %5, %sum.011
+  %inc = add nuw nsw i32 %i.010, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %vaarg.end, %entry
+  %sum.0.lcssa = phi i32 [ 0, %entry ], [ %add, %vaarg.end ]
+  call void @llvm.va_end(i8* nonnull %0)
+  call void @llvm.lifetime.end.p0i8(i64 24, i8* nonnull %0) #2
+  ret i32 %sum.0.lcssa
+}
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+; Function Attrs: nounwind
+declare void @llvm.va_start(i8*) #2
+
+; Function Attrs: nounwind
+declare void @llvm.va_end(i8*) #2
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+
+;declare dso_local i32 @sum(i32, ...) local_unnamed_addr