[msan] Copy origin of byval arguments

[vitalybuka]

Summary

Depends on D117278

[vitalybuka]

fix condition

[eugenis]

LGTM

A compiler-rt test for origin propagation would be great.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1741

This can get weird if the byval pointer is less than 4 byte aligned, but the alternatives are not pretty, either.

Also, you can bump the alignment std::max(CopyAlign, kMinOriginAlignment).

[kda]

LGTM

[vitalybuka]

alignment

[vitalybuka]

format

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1741

Not sure if OriginPtr can be aligned higher then kMinOriginAlignment, as origin_tls base is 4 bytes aligned.
CpOriginPtr can be aligned higher, but I guess it does not matter here.

Also please notice the I changed copy size.

[vitalybuka]

Add FIXME about size

[vitalybuka]

ping

[vitalybuka]

ping again

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1741

yes, you are right

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

if (Argument *A = dyn_cast<Argument>(V)) {
         if (A == &FArg) {
           bool Overflow = ArgOffset + Size > kParamTLSSize;
           if (FArg.hasByValAttr()) {
             // ByVal pointer itself has clean shadow. We copy the actual
             // argument shadow to the underlying memory.
             // Figure out maximal valid memcpy alignment.
             const Align ArgAlign = DL.getValueOrABITypeAlignment(
                 MaybeAlign(FArg.getParamAlignment()), FArg.getParamByValType());
-            Value *CpShadowPtr =
+            Value *CpShadowPtr, *CpOriginPtr;
+            std::tie(CpShadowPtr, CpOriginPtr) =
                 getShadowOriginPtr(V, EntryIRB, EntryIRB.getInt8Ty(), ArgAlign,
-                                   /*isStore*/ true)
-                    .first;
-            // TODO(glider): need to copy origins.
+                                   /*isStore*/ true);
             if (!PropagateShadow || Overflow) {
               // ParamTLS overflow.
               EntryIRB.CreateMemSet(
                   CpShadowPtr, Constant::getNullValue(EntryIRB.getInt8Ty()),
                   Size, ArgAlign);
             } else {
               Value *Base = getShadowPtrForArgument(&FArg, EntryIRB, ArgOffset);
               const Align CopyAlign = std::min(ArgAlign, kShadowTLSAlignment);
               Value *Cpy = EntryIRB.CreateMemCpy(CpShadowPtr, CopyAlign, Base,
                                                  CopyAlign, Size);
               LLVM_DEBUG(dbgs() << "  ByValCpy: " << *Cpy << "\n");
               (void)Cpy;
+
+              if (MS.TrackOrigins) {
+                Value *OriginPtr =
+                    getOriginPtrForArgument(&FArg, EntryIRB, ArgOffset);
+                EntryIRB.CreateMemCpy(CpOriginPtr, CopyAlign, OriginPtr,
+                                      CopyAlign, Size);
+              }
             }
           }
 
           if (!PropagateShadow || Overflow || FArg.hasByValAttr() ||
               (MS.EagerChecks && FArg.hasAttribute(Attribute::NoUndef))) {
             *ShadowPtr = getCleanShadow(V);
             setOrigin(A, getCleanOrigin());
           } else {
for (auto ArgIt = CB.arg_begin(), End = CB.arg_end(); ArgIt != End;
       bool ByVal = CB.paramHasAttr(i, Attribute::ByVal);
       bool NoUndef = CB.paramHasAttr(i, Attribute::NoUndef);
       bool EagerCheck = MayCheckCall && !ByVal && NoUndef;
 
       if (EagerCheck) {
         insertShadowCheck(A, &CB);
         Size = DL.getTypeAllocSize(A->getType());
       } else {
-        bool ArgIsInitialized = false;
         Value *Store = nullptr;
         // Compute the Shadow for arg even if it is ByVal, because
         // in that case getShadow() will copy the actual arg shadow to
         // __msan_param_tls.
         Value *ArgShadow = getShadow(A);
         Value *ArgShadowBase = getShadowPtrForArgument(A, IRB, ArgOffset);
         LLVM_DEBUG(dbgs() << "  Arg#" << i << ": " << *A
                           << " Shadow: " << *ArgShadow << "\n");
         if (ByVal) {
           // ByVal requires some special handling as it's too big for a single
           // load
           assert(A->getType()->isPointerTy() &&
                  "ByVal argument is not a pointer!");
           Size = DL.getTypeAllocSize(CB.getParamByValType(i));
           if (ArgOffset + Size > kParamTLSSize)
             break;
           const MaybeAlign ParamAlignment(CB.getParamAlign(i));
           MaybeAlign Alignment = llvm::None;
           if (ParamAlignment)
             Alignment = std::min(*ParamAlignment, kShadowTLSAlignment);
-          Value *AShadowPtr =
+          Value *AShadowPtr, *AOriginPtr;
+          std::tie(AShadowPtr, AOriginPtr) =
               getShadowOriginPtr(A, IRB, IRB.getInt8Ty(), Alignment,
-                                 /*isStore*/ false)
-                  .first;
+                                 /*isStore*/ false);
 
           Store = IRB.CreateMemCpy(ArgShadowBase, Alignment, AShadowPtr,
                                    Alignment, Size);
-          // TODO(glider): need to copy origins.
+          if (MS.TrackOrigins) {
+            Value *ArgOriginBase = getOriginPtrForArgument(A, IRB, ArgOffset);
+        
+            Store = IRB.CreateMemCpy(ArgOriginBase, Alignment, AOriginPtr,
+                                    Alignment, Size);
+          }
         } else {
           // Any other parameters mean we need bit-grained tracking of uninit
           // data
           Size = DL.getTypeAllocSize(A->getType());
           if (ArgOffset + Size > kParamTLSSize)
             break;
           Store = IRB.CreateAlignedStore(ArgShadow, ArgShadowBase,
                                          kShadowTLSAlignment);
           Constant *Cst = dyn_cast<Constant>(ArgShadow);
-          if (Cst && Cst->isNullValue())
-            ArgIsInitialized = true;
-        }
-        if (MS.TrackOrigins && !ArgIsInitialized)
+          if (MS.TrackOrigins && Cst && Cst->isNullValue()) {
             IRB.CreateStore(getOrigin(A),
                             getOriginPtrForArgument(A, IRB, ArgOffset));
+          }
+        }
         (void)Store;
         assert(Store != nullptr);
         LLVM_DEBUG(dbgs() << "  Param:" << *Store << "\n");
       }
       assert(Size != 0);
       ArgOffset += alignTo(Size, kShadowTLSAlignment);
     }
     LLVM_DEBUG(dbgs() << "  done with call args\n");

llvm/test/Instrumentation/MemorySanitizer/byval.ll

 ; RUN: opt < %s -S -passes="msan<track-origins=1>" 2>&1 | FileCheck %s --implicit-check-not "call void @llvm.mem" --implicit-check-not " load" --implicit-check-not " store"
 ; RUN: opt < %s -S -msan -msan-track-origins=1 | FileCheck %s --implicit-check-not "call void @llvm.mem" --implicit-check-not " load" --implicit-check-not " store"
 
 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"
 
 declare void @FnByVal(i128* byval(i128) %p);
 declare void @Fn(i128* %p);
 
 define i128 @ByValArgument(i32, i128* byval(i128) %p) sanitize_memory {
 ; CHECK-LABEL: @ByValArgument(
 ; CHECK-NEXT:  entry:
 ; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_tls to i64), i64 8) to i8*), i64 16, i1 false)
+; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_origin_tls to i64), i64 8) to i8*), i64 16, i1 false)
 ; CHECK:         %x = load i128, i128* %p
 ; CHECK:         load i128
 ; CHECK:         load i32
 ; CHECK:         store i128 {{.*}}, i128* bitcast ([100 x i64]* @__msan_retval_tls to i128*)
 ; CHECK-NEXT:    store i32 {{.*}}, i32* @__msan_retval_origin_tls
 ; CHECK-NEXT:    ret i128
 ;
 entry:

 ; CHECK-NEXT:    store i32 0, i32* @__msan_retval_origin_tls
 ; CHECK-NEXT:    ret i128
 ;
 entry:
   %x = load i128, i128* %p
   ret i128 %x
 }
 
-; FIXME: Origin of byval pointee is not propagated.
 define void @ByValForward(i32, i128* byval(i128) %p) sanitize_memory {
 ; CHECK-LABEL: @ByValForward(
 ; CHECK-NEXT:  entry:
 ; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_tls to i64), i64 8) to i8*), i64 16, i1 false)
+; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_origin_tls to i64), i64 8) to i8*), i64 16, i1 false)
 ; CHECK:         store i64 0, i64* getelementptr inbounds ([100 x i64], [100 x i64]* @__msan_param_tls, i32 0, i32 0)
+; CHECK:         store i32 0, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0)
 ; CHECK-NEXT:    call void @Fn(
 ; CHECK-NEXT:    ret void
 ;
 entry:
   call void @Fn(i128* %p)
   ret void
 }
 
 define void @ByValForwardNoSanitize(i32, i128* byval(i128) %p) {
 ; CHECK-LABEL: @ByValForwardNoSanitize(
 ; CHECK-NEXT:  entry:
 ; CHECK:         call void @llvm.memset.p0i8.i64(i8* align 8 {{.*}}, i8 0, i64 16, i1 false)
 ; CHECK:         store i64 0, i64* getelementptr inbounds ([100 x i64], [100 x i64]* @__msan_param_tls, i32 0, i32 0)
+; CHECK:         store i32 0, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0)
 ; CHECK-NEXT:    call void @Fn(
 ; CHECK-NEXT:    ret void
 ;
 entry:
   call void @Fn(i128* %p)
   ret void
 }
 
-; FIXME: Origin of %p byval pointee is not propagated.
 define void @ByValForwardByVal(i32, i128* byval(i128) %p) sanitize_memory {
 ; CHECK-LABEL: @ByValForwardByVal(
 ; CHECK-NEXT:  entry:
 ; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_tls to i64), i64 8) to i8*), i64 16, i1 false)
+; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64({{.*}} @__msan_param_origin_tls to i64), i64 8) to i8*), i64 16, i1 false)
 ; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64(i8* bitcast ([100 x i64]* @__msan_param_tls to i8*), i8* {{.*}}, i64 16, i1 false)
-; CHECK:         store i32 {{.*}}, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0)
+; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64(i8* bitcast ([200 x i32]* @__msan_param_origin_tls to i8*), i8* {{.*}}, i64 16, i1 false)
 ; CHECK-NEXT:    call void @FnByVal(
 ; CHECK-NEXT:    ret void
 ;
 entry:
   call void @FnByVal(i128* byval(i128) %p)
   ret void
 }
 
 ; FIXME: Shadow for byval should be reset not copied before the call.
 define void @ByValForwardByValNoSanitize(i32, i128* byval(i128) %p) {
 ; CHECK-LABEL: @ByValForwardByValNoSanitize(
 ; CHECK-NEXT:  entry:
 ; CHECK:         call void @llvm.memset.p0i8.i64(i8* align 8 {{.*}}, i8 0, i64 16, i1 false)
 ; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64(i8* bitcast ([100 x i64]* @__msan_param_tls to i8*), i8* {{.*}}, i64 16, i1 false) 
-; CHECK:         store i32 0, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0)
+; CHECK:         call void @llvm.memcpy.p0i8.p0i8.i64(i8* bitcast ([200 x i32]* @__msan_param_origin_tls to i8*), i8*  {{.*}}, i64 16, i1 false)
 ; CHECK-NEXT:    call void @FnByVal(
 ; CHECK-NEXT:    ret void
 ;
 entry:
   call void @FnByVal(i128* byval(i128) %p)
   ret void
 }