[msan] Fix invalid origin copying.

Origin copying may destroy valid origin info. This is caused by
__msan_copy_origin widening the address range to the nearest 4-byte aligned
addresses both on the left and on the right. If the target buffer is
uninitialized and the source is fully initialized, this will result in
overriding valid origin of target buffer with stale (possibly 0) origin of the
source buffer.

With this change the widened origin is copied only if corresponding shadow
values are non zero.

llvm-svn: 193338

3 files changed, 67 insertions, 11 deletions

diff --git a/compiler-rt/lib/msan/msan.h b/compiler-rt/lib/msan/msan.h
index fe7f20ae272..51fa2ebbcf3 100644
--- a/compiler-rt/lib/msan/msan.h
+++ b/compiler-rt/lib/msan/msan.h
@@ -25,11 +25,12 @@
 # define MSAN_REPLACE_OPERATORS_NEW_AND_DELETE 1
 #endif
 
-#define MEM_TO_SHADOW(mem) (((uptr)mem)       & ~0x400000000000ULL)
-#define MEM_TO_ORIGIN(mem) (MEM_TO_SHADOW(mem) + 0x200000000000ULL)
-#define MEM_IS_APP(mem)    ((uptr)mem >=         0x600000000000ULL)
-#define MEM_IS_SHADOW(mem) ((uptr)mem >=         0x200000000000ULL && \
-                            (uptr)mem <=         0x400000000000ULL)
+#define MEM_TO_SHADOW(mem)       (((uptr)mem) & ~0x400000000000ULL)
+#define SHADOW_TO_ORIGIN(shadow) (((uptr)shadow) + 0x200000000000ULL)
+#define MEM_TO_ORIGIN(mem)       (SHADOW_TO_ORIGIN(MEM_TO_SHADOW(mem)))
+#define MEM_IS_APP(mem)          ((uptr)mem >= 0x600000000000ULL)
+#define MEM_IS_SHADOW(mem) \
+  ((uptr)mem >= 0x200000000000ULL && (uptr)mem <= 0x400000000000ULL)
 
 const int kMsanParamTlsSizeInWords = 100;
 const int kMsanRetvalTlsSizeInWords = 100;
diff --git a/compiler-rt/lib/msan/msan_interceptors.cc b/compiler-rt/lib/msan/msan_interceptors.cc
index c200e6c5dec..849949d3c93 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cc
+++ b/compiler-rt/lib/msan/msan_interceptors.cc
@@ -1278,15 +1278,41 @@ void __msan_clear_and_unpoison(void *a, uptr size) {
   fast_memset((void*)MEM_TO_SHADOW((uptr)a), 0, size);
 }
 
+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 *(uptr *)SHADOW_TO_ORIGIN((s + i) & ~3UL);
+  return 0;
+}
+
 void __msan_copy_origin(void *dst, const void *src, uptr size) {
   if (!__msan_get_track_origins()) return;
   if (!MEM_IS_APP(dst) || !MEM_IS_APP(src)) return;
-  uptr d = MEM_TO_ORIGIN(dst);
-  uptr s = MEM_TO_ORIGIN(src);
-  uptr beg = d & ~3UL;  // align down.
-  uptr end = (d + size + 3) & ~3UL;  // align up.
-  s = s & ~3UL;  // align down.
-  fast_memcpy((void*)beg, (void*)s, end - beg);
+  uptr d = (uptr)dst;
+  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);
+    if (o)
+      *(uptr *)MEM_TO_ORIGIN(beg) = o;
+    beg += 4;
+  }
+
+  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);
+    if (o)
+      *(uptr *)MEM_TO_ORIGIN(end - 4) = o;
+    end -= 4;
+  }
+
+  if (beg < end) {
+    // Align src up.
+    uptr s = ((uptr)src + 3) & ~3UL;
+    fast_memcpy((void*)MEM_TO_ORIGIN(beg), (void*)MEM_TO_ORIGIN(s), end - beg);
+  }
 }
 
 void __msan_copy_poison(void *dst, const void *src, uptr size) {
diff --git a/compiler-rt/lib/msan/tests/msan_test.cc b/compiler-rt/lib/msan/tests/msan_test.cc
index fcb5fc81ac9..f185a31de30 100644
--- a/compiler-rt/lib/msan/tests/msan_test.cc
+++ b/compiler-rt/lib/msan/tests/msan_test.cc
@@ -1210,6 +1210,35 @@ TEST(MemorySanitizer, memcpy) {
   EXPECT_POISONED(y[1]);
 }
 
+void TestUnalignedMemcpy(int left, int right, bool src_is_aligned) {
+  const int sz = 20;
+  char *dst = (char *)malloc(sz);
+  U4 origin = __msan_get_origin(dst);
+
+  char *src = (char *)malloc(sz);
+  memset(src, 0, sz);
+
+  memcpy(dst + left, src_is_aligned ? src + left : src, sz - left - right);
+  for (int i = 0; i < left; ++i)
+    EXPECT_POISONED_O(dst[i], origin);
+  for (int i = 0; i < right; ++i)
+    EXPECT_POISONED_O(dst[sz - i - 1], origin);
+  EXPECT_NOT_POISONED(dst[left]);
+  EXPECT_NOT_POISONED(dst[sz - right - 1]);
+
+  free(dst);
+  free(src);
+}
+
+TEST(MemorySanitizer, memcpy_unaligned) {
+  for (int i = 0; i < 4; ++i) {
+    for (int j = 0; j < 4; ++j) {
+      TestUnalignedMemcpy(i, j, true);
+      TestUnalignedMemcpy(i, j, false);
+    }
+  }
+}
+
 TEST(MemorySanitizer, memmove) {
   char* x = new char[2];
   char* y = new char[2];