[msan] Precise origin handling in __unaligned_(load|store)*.

llvm-svn: 205412

4 files changed, 160 insertions, 35 deletions

diff --git a/compiler-rt/lib/msan/msan.cc b/compiler-rt/lib/msan/msan.cc
index 7a93b663185..4b6578ef4ce 100644
--- a/compiler-rt/lib/msan/msan.cc
+++ b/compiler-rt/lib/msan/msan.cc
@@ -509,40 +509,43 @@ u32 __msan_get_umr_origin() {
 u16 __sanitizer_unaligned_load16(const uu16 *p) {
   __msan_retval_tls[0] = *(uu16 *)MEM_TO_SHADOW((uptr)p);
   if (__msan_get_track_origins())
-    __msan_retval_origin_tls = *(uu32 *)(MEM_TO_ORIGIN((uptr)p) & ~3UL);
+    __msan_retval_origin_tls = GetOriginIfPoisoned((uptr)p, sizeof(*p));
   return *p;
 }
 u32 __sanitizer_unaligned_load32(const uu32 *p) {
   __msan_retval_tls[0] = *(uu32 *)MEM_TO_SHADOW((uptr)p);
   if (__msan_get_track_origins())
-    __msan_retval_origin_tls = *(uu32 *)(MEM_TO_ORIGIN((uptr)p) & ~3UL);
+    __msan_retval_origin_tls = GetOriginIfPoisoned((uptr)p, sizeof(*p));
   return *p;
 }
 u64 __sanitizer_unaligned_load64(const uu64 *p) {
   __msan_retval_tls[0] = *(uu64 *)MEM_TO_SHADOW((uptr)p);
   if (__msan_get_track_origins())
-    __msan_retval_origin_tls = *(uu32 *)(MEM_TO_ORIGIN((uptr)p) & ~3UL);
+    __msan_retval_origin_tls = GetOriginIfPoisoned((uptr)p, sizeof(*p));
   return *p;
 }
 void __sanitizer_unaligned_store16(uu16 *p, u16 x) {
-  *(uu16 *)MEM_TO_SHADOW((uptr)p) = __msan_param_tls[1];
-  if (__msan_get_track_origins())
+  u16 s = __msan_param_tls[1];
+  *(uu16 *)MEM_TO_SHADOW((uptr)p) = s;
+  if (s && __msan_get_track_origins())
     if (uu32 o = __msan_param_origin_tls[2])
-      __msan_set_origin(p, 2, o);
+      SetOriginIfPoisoned((uptr)p, (uptr)&s, sizeof(s), o);
   *p = x;
 }
 void __sanitizer_unaligned_store32(uu32 *p, u32 x) {
-  *(uu32 *)MEM_TO_SHADOW((uptr)p) = __msan_param_tls[1];
-  if (__msan_get_track_origins())
+  u32 s = __msan_param_tls[1];
+  *(uu32 *)MEM_TO_SHADOW((uptr)p) = s;
+  if (s && __msan_get_track_origins())
     if (uu32 o = __msan_param_origin_tls[2])
-      __msan_set_origin(p, 4, o);
+      SetOriginIfPoisoned((uptr)p, (uptr)&s, sizeof(s), o);
   *p = x;
 }
 void __sanitizer_unaligned_store64(uu64 *p, u64 x) {
-  *(uu64 *)MEM_TO_SHADOW((uptr)p) = __msan_param_tls[1];
-  if (__msan_get_track_origins())
+  u64 s = __msan_param_tls[1];
+  *(uu64 *)MEM_TO_SHADOW((uptr)p) = s;
+  if (s && __msan_get_track_origins())
     if (uu32 o = __msan_param_origin_tls[2])
-      __msan_set_origin(p, 8, o);
+      SetOriginIfPoisoned((uptr)p, (uptr)&s, sizeof(s), o);
   *p = x;
 }
 
diff --git a/compiler-rt/lib/msan/msan.h b/compiler-rt/lib/msan/msan.h
index 919415e4f04..0293b376323 100644
--- a/compiler-rt/lib/msan/msan.h
+++ b/compiler-rt/lib/msan/msan.h
@@ -88,6 +88,8 @@ void ReportAtExitStatistics();
 void UnpoisonParam(uptr n);
 void UnpoisonThreadLocalState();
 
+u32 GetOriginIfPoisoned(uptr a, uptr size);
+void SetOriginIfPoisoned(uptr addr, uptr src_shadow, uptr size, u32 src_origin);
 void CopyOrigin(void *dst, const void *src, uptr size, StackTrace *stack);
 void MovePoison(void *dst, const void *src, uptr size, StackTrace *stack);
 void CopyPoison(void *dst, const void *src, uptr size, StackTrace *stack);
diff --git a/compiler-rt/lib/msan/msan_interceptors.cc b/compiler-rt/lib/msan/msan_interceptors.cc
index 0bac746f22f..6f7f4e6f884 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cc
+++ b/compiler-rt/lib/msan/msan_interceptors.cc
@@ -1367,14 +1367,6 @@ void __msan_clear_and_unpoison(void *a, uptr size) {
   PoisonShadow((uptr)a, size, 0);
 }
 
-u32 get_origin_if_poisoned(uptr a, uptr size) {
-  unsigned char *s = (unsigned char *)MEM_TO_SHADOW(a);
-  for (uptr i = 0; i < size; ++i)
-    if (s[i])
-      return *(u32 *)SHADOW_TO_ORIGIN((s + i) & ~3UL);
-  return 0;
-}
-
 void *__msan_memcpy(void *dest, const void *src, SIZE_T n) {
   ENSURE_MSAN_INITED();
   GET_STORE_STACK_TRACE;
@@ -1405,6 +1397,24 @@ void __msan_unpoison_string(const char* s) {
 
 namespace __msan {
 
+u32 GetOriginIfPoisoned(uptr addr, uptr size) {
+  unsigned char *s = (unsigned char *)MEM_TO_SHADOW(addr);
+  for (uptr i = 0; i < size; ++i)
+    if (s[i])
+      return *(u32 *)SHADOW_TO_ORIGIN((s + i) & ~3UL);
+  return 0;
+}
+
+void SetOriginIfPoisoned(uptr addr, uptr src_shadow, uptr size,
+                         u32 src_origin) {
+  uptr dst_s = MEM_TO_SHADOW(addr);
+  uptr src_s = src_shadow;
+  uptr src_s_end = src_s + size;
+
+  for (; src_s < src_s_end; ++dst_s, ++src_s)
+    if (*(u8 *)src_s) *(u32 *)SHADOW_TO_ORIGIN(dst_s &~3UL) = src_origin;
+}
+
 void CopyOrigin(void *dst, const void *src, uptr size, StackTrace *stack) {
   if (!__msan_get_track_origins()) return;
   if (!MEM_IS_APP(dst) || !MEM_IS_APP(src)) return;
@@ -1413,7 +1423,7 @@ void CopyOrigin(void *dst, const void *src, uptr size, StackTrace *stack) {
   uptr beg = d & ~3UL;
   // Copy left unaligned origin if that memory is poisoned.
   if (beg < d) {
-    u32 o = get_origin_if_poisoned(beg, d - beg);
+    u32 o = GetOriginIfPoisoned(beg, d - beg);
     if (o) {
       if (__msan_get_track_origins() > 1) o = ChainOrigin(o, stack);
       *(u32 *)MEM_TO_ORIGIN(beg) = o;
@@ -1424,7 +1434,7 @@ void CopyOrigin(void *dst, const void *src, uptr size, StackTrace *stack) {
   uptr end = (d + size + 3) & ~3UL;
   // Copy right unaligned origin if that memory is poisoned.
   if (end > d + size) {
-    u32 o = get_origin_if_poisoned(d + size, end - d - size);
+    u32 o = GetOriginIfPoisoned(d + size, end - d - size);
     if (o) {
       if (__msan_get_track_origins() > 1) o = ChainOrigin(o, stack);
       *(u32 *)MEM_TO_ORIGIN(end - 4) = o;
diff --git a/compiler-rt/lib/msan/tests/msan_test.cc b/compiler-rt/lib/msan/tests/msan_test.cc
index a239a69becf..52719eab776 100644
--- a/compiler-rt/lib/msan/tests/msan_test.cc
+++ b/compiler-rt/lib/msan/tests/msan_test.cc
@@ -3343,34 +3343,35 @@ TEST(MemorySanitizer, VolatileBitfield) {
 TEST(MemorySanitizer, UnalignedLoad) {
   char x[32];
   U4 origin = __LINE__;
-  __msan_set_origin(&x, sizeof(x), origin);
+  for (unsigned i = 0; i < sizeof(x) / 4; ++i)
+    __msan_set_origin(x + 4 * i, 4, origin + i);
 
   memset(x + 8, 0, 16);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 6), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 7), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 6), origin + 1);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 7), origin + 1);
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load16(x + 8));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load16(x + 9));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load16(x + 22));
-  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 23), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 24), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 23), origin + 6);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load16(x + 24), origin + 6);
 
-  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 4), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 7), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 4), origin + 1);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 7), origin + 1);
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load32(x + 8));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load32(x + 9));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load32(x + 20));
-  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 21), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 24), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 21), origin + 6);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load32(x + 24), origin + 6);
 
   EXPECT_POISONED_O(__sanitizer_unaligned_load64(x), origin);
   EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 1), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 7), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 7), origin + 1);
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load64(x + 8));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load64(x + 9));
   EXPECT_NOT_POISONED(__sanitizer_unaligned_load64(x + 16));
-  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 17), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 21), origin);
-  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 24), origin);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 17), origin + 6);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 21), origin + 6);
+  EXPECT_POISONED_O(__sanitizer_unaligned_load64(x + 24), origin + 6);
 }
 
 TEST(MemorySanitizer, UnalignedStore16) {
@@ -3429,6 +3430,115 @@ TEST(MemorySanitizer, UnalignedStore64) {
   EXPECT_POISONED_O(x[11], origin);
 }
 
+TEST(MemorySanitizer, UnalignedStore16_precise) {
+  char x[8];
+  U2 y = 0;
+  U4 originx1 = __LINE__;
+  U4 originx2 = __LINE__;
+  U4 originy = __LINE__;
+  __msan_poison(x, sizeof(x));
+  __msan_set_origin(x, 4, originx1);
+  __msan_set_origin(x + 4, 4, originx2);
+  __msan_poison(((char *)&y) + 1, 1);
+  __msan_set_origin(&y, sizeof(y), originy);
+
+  __sanitizer_unaligned_store16(x + 3, y);
+  EXPECT_POISONED_O(x[0], originx1);
+  EXPECT_POISONED_O(x[1], originx1);
+  EXPECT_POISONED_O(x[2], originx1);
+  EXPECT_NOT_POISONED(x[3]);
+  EXPECT_POISONED_O(x[4], originy);
+  EXPECT_POISONED_O(x[5], originy);
+  EXPECT_POISONED_O(x[6], originy);
+  EXPECT_POISONED_O(x[7], originy);
+}
+
+TEST(MemorySanitizer, UnalignedStore16_precise2) {
+  char x[8];
+  U2 y = 0;
+  U4 originx1 = __LINE__;
+  U4 originx2 = __LINE__;
+  U4 originy = __LINE__;
+  __msan_poison(x, sizeof(x));
+  __msan_set_origin(x, 4, originx1);
+  __msan_set_origin(x + 4, 4, originx2);
+  __msan_poison(((char *)&y), 1);
+  __msan_set_origin(&y, sizeof(y), originy);
+
+  __sanitizer_unaligned_store16(x + 3, y);
+  EXPECT_POISONED_O(x[0], originy);
+  EXPECT_POISONED_O(x[1], originy);
+  EXPECT_POISONED_O(x[2], originy);
+  EXPECT_POISONED_O(x[3], originy);
+  EXPECT_NOT_POISONED(x[4]);
+  EXPECT_POISONED_O(x[5], originx2);
+  EXPECT_POISONED_O(x[6], originx2);
+  EXPECT_POISONED_O(x[7], originx2);
+}
+
+TEST(MemorySanitizer, UnalignedStore64_precise) {
+  char x[12];
+  U8 y = 0;
+  U4 originx1 = __LINE__;
+  U4 originx2 = __LINE__;
+  U4 originx3 = __LINE__;
+  U4 originy = __LINE__;
+  __msan_poison(x, sizeof(x));
+  __msan_set_origin(x, 4, originx1);
+  __msan_set_origin(x + 4, 4, originx2);
+  __msan_set_origin(x + 8, 4, originx3);
+  __msan_poison(((char *)&y) + 1, 1);
+  __msan_poison(((char *)&y) + 7, 1);
+  __msan_set_origin(&y, sizeof(y), originy);
+
+  __sanitizer_unaligned_store64(x + 2, y);
+  EXPECT_POISONED_O(x[0], originy);
+  EXPECT_POISONED_O(x[1], originy);
+  EXPECT_NOT_POISONED(x[2]);
+  EXPECT_POISONED_O(x[3], originy);
+
+  EXPECT_NOT_POISONED(x[4]);
+  EXPECT_NOT_POISONED(x[5]);
+  EXPECT_NOT_POISONED(x[6]);
+  EXPECT_NOT_POISONED(x[7]);
+
+  EXPECT_NOT_POISONED(x[8]);
+  EXPECT_POISONED_O(x[9], originy);
+  EXPECT_POISONED_O(x[10], originy);
+  EXPECT_POISONED_O(x[11], originy);
+}
+
+TEST(MemorySanitizer, UnalignedStore64_precise2) {
+  char x[12];
+  U8 y = 0;
+  U4 originx1 = __LINE__;
+  U4 originx2 = __LINE__;
+  U4 originx3 = __LINE__;
+  U4 originy = __LINE__;
+  __msan_poison(x, sizeof(x));
+  __msan_set_origin(x, 4, originx1);
+  __msan_set_origin(x + 4, 4, originx2);
+  __msan_set_origin(x + 8, 4, originx3);
+  __msan_poison(((char *)&y) + 3, 3);
+  __msan_set_origin(&y, sizeof(y), originy);
+
+  __sanitizer_unaligned_store64(x + 2, y);
+  EXPECT_POISONED_O(x[0], originx1);
+  EXPECT_POISONED_O(x[1], originx1);
+  EXPECT_NOT_POISONED(x[2]);
+  EXPECT_NOT_POISONED(x[3]);
+
+  EXPECT_NOT_POISONED(x[4]);
+  EXPECT_POISONED_O(x[5], originy);
+  EXPECT_POISONED_O(x[6], originy);
+  EXPECT_POISONED_O(x[7], originy);
+
+  EXPECT_NOT_POISONED(x[8]);
+  EXPECT_NOT_POISONED(x[9]);
+  EXPECT_POISONED_O(x[10], originx3);
+  EXPECT_POISONED_O(x[11], originx3);
+}
+
 namespace {
 typedef U2 V8x16 __attribute__((__vector_size__(16)));
 typedef U4 V4x32 __attribute__((__vector_size__(16)));