3 files changed, 67 insertions, 11 deletions

diff --git a/lib/msan/msan.h b/lib/msan/msan.h
index fe7f20ae..51fa2ebb 100644
--- a/lib/msan/msan.h
+++ b/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/lib/msan/msan_interceptors.cc b/lib/msan/msan_interceptors.cc
index c200e6c5..849949d3 100644
--- a/lib/msan/msan_interceptors.cc
+++ b/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/lib/msan/tests/msan_test.cc b/lib/msan/tests/msan_test.cc
index fcb5fc81..f185a31d 100644
--- a/lib/msan/tests/msan_test.cc
+++ b/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];