[sanitizer] Add MemCpyAccessible

[vitalybuka]

A layer over TryMemCpy to copy only availible
pages.

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Vitaly Buka (vitalybuka)

Changes

A layer over TryMemCpy to copy only availible
pages.

---

Full diff: https://github.com/llvm/llvm-project/pull/112794.diff

3 Files Affected:

• (modified) compiler-rt/lib/sanitizer_common/sanitizer_common.h (+2)
• (modified) compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cpp (+26)
• (modified) compiler-rt/lib/sanitizer_common/tests/sanitizer_posix_test.cpp (+28-6)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common.h b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
index 9b1e58f5e7a61d..a7b6a4033a5276 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
@@ -273,6 +273,8 @@ bool IsAccessibleMemoryRange(uptr beg, uptr size);
 // Returns true if we can read a memory range starting at `src`, and copies
 // content into `dest`.
 bool TryMemCpy(void *dest, const void *src, uptr n);
+// Copies accesible memory, and zero fill the rest.
+void MemCpyAccessible(void *dest, const void *src, uptr n);
 
 // Error report formatting.
 const char *StripPathPrefix(const char *filepath,
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cpp
index 684720963a8dcb..3ad5aaa0a7d4ab 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cpp
@@ -219,6 +219,32 @@ static void StopStackDepotBackgroundThread() {
 static void StopStackDepotBackgroundThread() {}
 #endif
 
+void MemCpyAccessible(void *dest, const void *src, uptr n) {
+  if (TryMemCpy(dest, src, n))
+    return;
+
+  const uptr page_size = GetPageSize();
+  uptr b = reinterpret_cast<uptr>(src);
+  uptr b_up = RoundUpTo(b, page_size);
+
+  uptr e = reinterpret_cast<uptr>(src) + n;
+  uptr e_down = RoundDownTo(e, page_size);
+
+  auto copy_or_zero = [dest, src](uptr b, uptr e) {
+    uptr d = reinterpret_cast<uptr>(dest) + (b - reinterpret_cast<uptr>(src));
+    if (!TryMemCpy(reinterpret_cast<void *>(d), reinterpret_cast<void *>(b),
+                   e - b))
+      internal_memset(reinterpret_cast<void *>(d), 0, e - b);
+  };
+
+  copy_or_zero(b, b_up);
+
+  for (uptr p = b_up; p < e_down; p += page_size)
+    copy_or_zero(p, p + page_size);
+
+  copy_or_zero(e_down, e);
+}
+
 }  // namespace __sanitizer
 
 SANITIZER_INTERFACE_WEAK_DEF(void, __sanitizer_sandbox_on_notify,
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_posix_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_posix_test.cpp
index 03a841a6f438cb..5016b09c15307f 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_posix_test.cpp
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_posix_test.cpp
@@ -127,19 +127,41 @@ TEST(SanitizerCommon, TryMemCpyNull) {
   EXPECT_FALSE(TryMemCpy(dst.data(), nullptr, dst.size()));
 }
 
-TEST(SanitizerCommon, TryMemCpyProtected) {
+TEST(SanitizerCommon, MemCpyAccessible) {
+  const int page_num = 1000;
   const int page_size = GetPageSize();
-  InternalMmapVector<char> src(3 * page_size);
+  InternalMmapVector<char> src(page_num * page_size);
   std::iota(src.begin(), src.end(), 123);
   std::vector<char> dst;
-  // Protect the middle page.
-  mprotect(src.data() + page_size, page_size, PROT_NONE);
+  std::vector<char> exp = {src.begin(), src.end()};
+
+  // Protect some pages.
+  for (int i = 7; i < page_num; i *= 2) {
+    mprotect(src.data() + i * page_size, page_size, PROT_NONE);
+    std::fill(exp.data() + i * page_size, exp.data() + (i + 1) * page_size, 0);
+  }
 
   dst.assign(src.size(), 0);
   EXPECT_FALSE(TryMemCpy(dst.data(), src.data(), dst.size()));
 
-  mprotect(src.data() + page_size, page_size, PROT_READ | PROT_WRITE);
-  EXPECT_TRUE(std::equal(dst.begin(), dst.end(), src.begin()));
+  // Full page aligned range with mprotect pages.
+  dst.assign(src.size(), 0);
+  MemCpyAccessible(dst.data(), src.data(), dst.size());
+  EXPECT_TRUE(std::equal(dst.begin(), dst.end(), exp.begin()));
+
+  // Misaligned range with mprotect pages.
+  size_t offb = 3;
+  size_t offe = 7;
+  dst.assign(src.size() - offb - offe, 0);
+  MemCpyAccessible(dst.data(), src.data() + offb, dst.size());
+  EXPECT_TRUE(std::equal(dst.begin(), dst.end(), exp.begin() + offb));
+
+  // Misaligned range with ends in mprotect pages.
+  offb = 3 + 7 * page_size;
+  offe = 7 + 14 * page_size;
+  dst.assign(src.size() - offb - offe, 0);
+  MemCpyAccessible(dst.data(), src.data() + offb, dst.size());
+  EXPECT_TRUE(std::equal(dst.begin(), dst.end(), exp.begin() + offb));
 }
 
 }  // namespace __sanitizer

[fmayer]

I am not fully sure about this optimization. In the case where there are holes this makes it more expensive. Does the caller generally know that?

[vitalybuka]

Not sure as well. On some examples fast pass 4x faster than page by page.

[fmayer]

nit '\0'