[sanitizer] [msan] Fix origin store of array types

[zatrazz]

Summary

This patch fixes the memory sanitizer origin store instrumentation for
array types. This can be triggered by cases where frontend lowers
function return to array type instead of aggregation.

For instance, the C code:

---

struct mypair {
int64_t x;
int y;
};

mypair my_make_pair(int64_t x, int y) {
mypair p;
p.x = x;
p.y = y;
return p;
}

int foo (int p)
{

mypair z = my_make_pair(p, 0);
return z.y + z.x;

}

It will be lowered with target set to aarch64-linux and -O0 to:

---

[...]
define i32 @_Z3fooi(i32 %p) #0 {
[...]
%call = call [2 x i64] @_Z12my_make_pairxi(i64 %conv, i32 0)
%1 = bitcast %struct.mypair* %z to [2 x i64]*
store [2 x i64] %call, [2 x i64]* %1, align 8

[...]

The origin store will emit a 'icmp' to test each store value again the
TLS origin array. However since 'icmp' does not support ArrayType the
memory instrumentation phase will bail out with an error.

This patch change it by extracting and applying the 'icmp' operation on
each array element.

It fixes the 'test/msan/insertvalue_origin.cc' for aarch64 (the -O0 case).

[eugenis]

The origin store will emit a 'icmp' to test each store value again the

TLS origin array.
Are you sure? It is supposed to test against a "zeroinitializer" constant.

Origin is 4 byte aligned, so this code would not work correctly for [i16 x 2], for example.
There is a function in compiler-rt that does this sort of thing correctly for larger arrays (in memcpy, for example), but calling it for all array stores would be too slow.
This patch does not do anything for StructType (line 695), which is a very similar case.

To do this 100% correctly, one would bitcast the shadow to a long integer, call paintOrigin for the 4-aligned part of the 4-aligned size in the middle of the store location, and conditionally (on the corresponding parts of the shadow being non-zero) store origin for the left and right partial quads.

I'd suggest not going there until we have evidence that it matters.

The simplest fix would be to bitcast the shadow to a long integer before the icmp. While you are at it, wrap the StructType case into this branch as well.

lib/Transforms/Instrumentation/MemorySanitizer.cpp

728

getCleanShadow(ConvertedShadowElem)

737

Do you need to update Addr here?
Move updateOrigin call out of the loop.

test/Instrumentation/MemorySanitizer/origin-array.ll

41

Does this test really need all this code?
Just write one function that accepts [2 x i64] %v, [2 x i64]* %p and stores one to the other.

[zatrazz]

In D15607#313344, @eugenis wrote:

The origin store will emit a 'icmp' to test each store value again the

TLS origin array.
Are you sure? It is supposed to test against a "zeroinitializer" constant.

My mistake here, you are right.

Origin is 4 byte aligned, so this code would not work correctly for [i16 x 2], for example.
There is a function in compiler-rt that does this sort of thing correctly for larger arrays (in memcpy, for example), but calling it for all array stores would be too slow.
This patch does not do anything for StructType (line 695), which is a very similar case.

To do this 100% correctly, one would bitcast the shadow to a long integer, call paintOrigin for the 4-aligned part of the 4-aligned size in the middle of the store location, and conditionally (on the corresponding parts of the shadow being non-zero) store origin for the left and right partial quads.

I'd suggest not going there until we have evidence that it matters.

Right, In fact I noted at least for aarch64 only very specific vector types are generated from functions that return aggregated type (test/msan/insertvalue_origin.cc is an example and change the struct size also changes clang IR by not using a vector type as well).

The simplest fix would be to bitcast the shadow to a long integer before the icmp. While you are at it, wrap the StructType case into this branch as well.

Well the problem is I do not think it is possible to cast the shadow to a long integer with the shadow being a array type. The problem is bitcast explicit forbid cast from aggregate types:

lib/IR/Instructions.cpp:

3019 bool
3020 CastInst::castIsValid(Instruction::CastOps op, Value *S, Type *DstTy) {
3021 
3022   // Check for type sanity on the arguments
3023   Type *SrcTy = S->getType();
3024 
3025   if (!SrcTy->isFirstClassType() || !DstTy->isFirstClassType() ||
3026       SrcTy->isAggregateType() || DstTy->isAggregateType())
3027     return false;

I tried to simple cast the converted shadow to a long integer and it bails with an invalid cast due 'SrcTy->isAggregateType()' being true (and it will be true as well for struct types).

I will update the patch with the comments you suggested.

lib/Transforms/Instrumentation/MemorySanitizer.cpp

728

I will change that.

737

I will change that.

test/Instrumentation/MemorySanitizer/origin-array.ll

41

I will simplify the testcase. I created this based on the original bug trigger (test/msan/insertvalue_origin.cc).

[zatrazz]

Changes from previous versions:

• Use getCleanShadow on the extracted elements;
• Moved updateOrigin call out of the loop;
• Simplified the testcase.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

733

This code seems to build nested if() blocks, effectively testing that all shadow elements are non-null. This is wrong, it should test that any of those is non-null. And even that is a rough approximation.

Would it be enough to change the isa<StructType> test above to isAggregateType() ?

[zatrazz]

Changes from previous version:

• Drop the extract array logic in favor or handling as the struct type.

[eugenis]

LGTM

lib/Transforms/Instrumentation/MemorySanitizer.cpp

if (isa<StructType>(Shadow->getType())) {
       if (AsCall && SizeIndex < kNumberOfAccessSizes) {
         Value *Fn = MS.MaybeStoreOriginFn[SizeIndex];
         Value *ConvertedShadow2 = IRB.CreateZExt(
             ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
         IRB.CreateCall(Fn, {ConvertedShadow2,
                             IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()),
                             Origin});
       } else {
+        if (ArrayType *AT = dyn_cast<ArrayType>(ConvertedShadow->getType())) {
+          // For ArrayType we need to apply the icmp to each element (since the
+          // the instruction does not support array types).
+          Value *CleanShadow = getCleanShadow(ConvertedShadow);
+          for (int i = 0, s = AT->getNumElements(); i < s; i++) {
+            Value *ConvertedShadowElem = IRB.CreateExtractValue(
+                ConvertedShadow, i);
+            Value *CleanShadowElem = IRB.CreateExtractValue(CleanShadow, i);
+            Value *Cmp = IRB.CreateICmpNE(
+                ConvertedShadowElem, CleanShadowElem, "_mscmp");
+
+            Instruction *CheckTerm = SplitBlockAndInsertIfThen(
+                Cmp, &*IRB.GetInsertPoint(), false, MS.OriginStoreWeights);
+            IRB.SetInsertPoint(CheckTerm);
+
+            IRBuilder<> IRBNew(CheckTerm);
+            paintOrigin(IRBNew, updateOrigin(Origin, IRBNew),
+                        getOriginPtr(Addr, IRBNew, Alignment), StoreSize,
+                        OriginAlignment);
+          }
+        } else {
           Value *Cmp = IRB.CreateICmpNE(
               ConvertedShadow, getCleanShadow(ConvertedShadow), "_mscmp");
           Instruction *CheckTerm = SplitBlockAndInsertIfThen(
               Cmp, &*IRB.GetInsertPoint(), false, MS.OriginStoreWeights);
           IRBuilder<> IRBNew(CheckTerm);
           paintOrigin(IRBNew, updateOrigin(Origin, IRBNew),
                       getOriginPtr(Addr, IRBNew, Alignment), StoreSize,
                       OriginAlignment);
         }
       }
     }
+  }
 
   void materializeStores(bool InstrumentWithCalls) {
     for (auto Inst : StoreList) {
       StoreInst &SI = *dyn_cast<StoreInst>(Inst);
 
       IRBuilder<> IRB(&SI);
       Value *Val = SI.getValueOperand();
       Value *Addr = SI.getPointerOperand();

test/Instrumentation/MemorySanitizer/origin-array.ll

+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+; Check origin handling of array types.  Since the instrumentation will
+; potentially creates icmp instruction to check origin values, it requires
+; the intrumentation pass to create icmp for each element.
+
+%struct.mypair = type { i64, i32 }
+
+; Function Attrs: nounwind
+define [2 x i64] @_Z12my_make_pairli(i64 %x, i32 %y) {
+entry:
+  %retval = alloca %struct.mypair, align 8
+  %x.addr = alloca i64, align 8
+  %y.addr = alloca i32, align 4
+  store i64 %x, i64* %x.addr, align 8
+  store i32 %y, i32* %y.addr, align 4
+  %0 = load i64, i64* %x.addr, align 8
+  %x1 = getelementptr inbounds %struct.mypair, %struct.mypair* %retval, i32 0, i32 0
+  store i64 %0, i64* %x1, align 8
+  %1 = load i32, i32* %y.addr, align 4
+  %y2 = getelementptr inbounds %struct.mypair, %struct.mypair* %retval, i32 0, i32 1
+  store i32 %1, i32* %y2, align 8
+  %2 = bitcast %struct.mypair* %retval to [2 x i64]*
+  %3 = load [2 x i64], [2 x i64]* %2, align 8
+  ret [2 x i64] %3
+}
+
+; Function Attrs: nounwind
+define i32 @_Z3fooi(i32 %p) #0 {
+entry:
+  %p.addr = alloca i32, align 4
+  %z = alloca %struct.mypair, align 8
+  store i32 %p, i32* %p.addr, align 4
+  %0 = load i32, i32* %p.addr, align 4
+  %conv = sext i32 %0 to i64
+  %call = call [2 x i64] @_Z12my_make_pairli(i64 %conv, i32 0)
+  %1 = bitcast %struct.mypair* %z to [2 x i64]*
+  store [2 x i64] %call, [2 x i64]* %1, align 8
+  %x = getelementptr inbounds %struct.mypair, %struct.mypair* %z, i32 0, i32 0
+  %2 = load i64, i64* %x, align 8
+  %y = getelementptr inbounds %struct.mypair, %struct.mypair* %z, i32 0, i32 1
+  %3 = load i32, i32* %y, align 8
+  %conv1 = sext i32 %3 to i64
+  %add = add nsw i64 %2, %conv1
+  %conv2 = trunc i64 %add to i32
+  ret i32 %conv2
+}
+
+; CHECK-LABEL: @_Z3fooi
+; CHECK: load {{.*}} @__msan_param_tls
+; CHECK: [[ORIGIN:%[01-9a-z]+]] = load {{.*}} @__msan_retval_origin_tls
+; CHECK: store {{.*}}
+
+; Extract and compare the first element.
+; CHECK: {{.*}} extractvalue {{.*}} 0
+; CHECK: icmp
+; CHECK: br i1 {{.*}}
+; CHECK: ; <label>{{.*}}
+; CHECK: [[ORIGIN2:%[01-9a-z]+]] = call i32 @__msan_chain_origin(i32 {{.*}}[[ORIGIN]])
+
+; And then the second element.
+; CHECK: store {{.*}}
+; CHECK: {{.*}} extractvalue {{.*}} 1
+; CHECK: icmp
+; CHECK: br i1 {{.*}}
+; CHECK: ; <label>{{.*}}
+; CHECK: [[ORIGIN2:%[01-9a-z]+]] = call i32 @__msan_chain_origin(i32 {{.*}}[[ORIGIN]])
+
+attributes #0 = { nounwind sanitize_memory "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" "unsafe-fp-math"="false" "use-soft-float"="false" }