Not overlapping containers

AdvenamTacet · AdvenamTacet · commit 0ab9216ab9f2 · 2024-08-26T11:21:51.000+02:00
Implemented correct container annotations copying function (non-overlapping containers only)

Refactored the initial PoC to a fully functional implementation for copying container annotations without moving/copying the memory content itself. This function currently supports cases where the source and destination containers do not overlap. Handling overlapping containers would add complexity.

The function is designed to work irrespective of whether the buffers are granule-aligned or the distance between them is granule-aligned. However, such scenarios may have an impact on performance.

A Test case has been included to verify the correctness of the implementation.

Removed unpoisoning oryginal buffer at the end, users can do it by themselves if they want to.
diff --git a/compiler-rt/lib/asan/asan_poisoning.cpp b/compiler-rt/lib/asan/asan_poisoning.cpp
@@ -576,8 +576,46 @@ void __sanitizer_annotate_double_ended_contiguous_container(
   }
 }
 
-// This function moves annotation from one buffer to another.
-// Old buffer is unpoisoned at the end.
+static bool WithinOneGranule(uptr p, uptr q) {
+  if (p == q)
+    return true;
+  return RoundDownTo(p, ASAN_SHADOW_GRANULARITY) ==
+         RoundDownTo(q - 1, ASAN_SHADOW_GRANULARITY);
+}
+
+static void PoisonContainer(uptr storage_beg, uptr storage_end) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  uptr internal_beg = RoundUpTo(storage_beg, granularity);
+  uptr external_beg = RoundDownTo(storage_beg, granularity);
+  uptr internal_end = RoundDownTo(storage_end, granularity);
+
+  if (internal_end > internal_beg)
+    PoisonShadow(internal_beg, internal_end - internal_beg,
+                 kAsanContiguousContainerOOBMagic);
+  // The new buffer may start in the middle of a granule.
+  if (internal_beg != storage_beg && internal_beg < internal_end &&
+      !AddressIsPoisoned(storage_beg)) {
+    *(u8 *)MemToShadow(external_beg) =
+        static_cast<u8>(storage_beg - external_beg);
+  }
+  // The new buffer may end in the middle of a granule.
+  if (internal_end != storage_end && AddressIsPoisoned(storage_end)) {
+    *(u8 *)MemToShadow(internal_end) =
+        static_cast<u8>(kAsanContiguousContainerOOBMagic);
+  }
+}
+
+// This function copies ASan memory annotations (poisoned/unpoisoned states)
+// from one buffer to another.
+// It's main purpose is to help with relocating trivially relocatable objects,
+// which memory may be poisoned, without calling copy constructor.
+// However, it does not move memory content itself, only annotations.
+// If the buffers aren't aligned (the distance between buffers isn't
+// granule-aligned)
+//     // old_storage_beg % granularity != new_storage_beg % granularity
+// the function handles this by going byte by byte, slowing down performance.
+// The old buffer annotations are not removed. If necessary,
+// user can unpoison old buffer with __asan_unpoison_memory_region.
 void __sanitizer_move_contiguous_container_annotations(
     const void *old_storage_beg_p, const void *old_storage_end_p,
     const void *new_storage_beg_p, const void *new_storage_end_p) {
@@ -606,6 +644,9 @@ void __sanitizer_move_contiguous_container_annotations(
         &stack);
   }
 
+  if (old_storage_beg == old_storage_end)
+    return;
+
   uptr new_internal_beg = RoundUpTo(new_storage_beg, granularity);
   uptr old_internal_beg = RoundUpTo(old_storage_beg, granularity);
   uptr new_external_beg = RoundDownTo(new_storage_beg, granularity);
@@ -615,93 +656,94 @@ void __sanitizer_move_contiguous_container_annotations(
 
   // At the very beginning we poison the whole buffer.
   // Later we unpoison what is necessary.
-  PoisonShadow(new_internal_beg, new_internal_end - new_internal_beg,
-               kAsanContiguousContainerOOBMagic);
-  if (new_internal_beg != new_storage_beg) {
-    uptr new_unpoisoned = *(u8 *)MemToShadow(new_external_beg);
-    if (new_unpoisoned > (new_storage_beg - new_external_beg)) {
-      *(u8 *)MemToShadow(new_external_beg) =
-          static_cast<u8>(new_storage_beg - new_external_beg);
-    }
-  }
-  if (new_internal_end != new_storage_end) {
-    uptr new_unpoisoned = *(u8 *)MemToShadow(new_internal_end);
-    if (new_unpoisoned <= (new_storage_end - new_internal_end)) {
-      *(u8 *)MemToShadow(new_external_beg) =
-          static_cast<u8>(kAsanContiguousContainerOOBMagic);
-    }
-  }
+  PoisonContainer(new_storage_beg, new_storage_end);
 
   // There are two cases.
   // 1) Distance between buffers is granule-aligned.
-  // 2) It's not aligned, that case is slower.
+  // 2) It's not aligned, and therefore requires going byte by byte.
   if (old_storage_beg % granularity == new_storage_beg % granularity) {
     // When buffers are aligned in the same way, we can just copy shadow memory,
-    // except first and last granule.
-    __builtin_memcpy((u8 *)MemToShadow(new_internal_beg),
-                     (u8 *)MemToShadow(old_internal_beg),
-                     (new_internal_end - new_internal_beg) / granularity);
-    // In first granule we cannot poison anything before beginning of the
-    // container.
-    if (new_internal_beg != new_storage_beg) {
-      uptr old_unpoisoned = *(u8 *)MemToShadow(old_external_beg);
-      uptr new_unpoisoned = *(u8 *)MemToShadow(new_external_beg);
-
-      if (old_unpoisoned > old_storage_beg - old_external_beg) {
-        *(u8 *)MemToShadow(new_external_beg) = old_unpoisoned;
-      } else if (new_unpoisoned > new_storage_beg - new_external_beg) {
-        *(u8 *)MemToShadow(new_external_beg) =
-            new_storage_beg - new_external_beg;
-      }
-    }
-    // In last granule we cannot poison anything after the end of the container.
-    if (new_internal_end != new_storage_end) {
-      uptr old_unpoisoned = *(u8 *)MemToShadow(old_internal_end);
-      uptr new_unpoisoned = *(u8 *)MemToShadow(new_internal_end);
-      if (new_unpoisoned <= new_storage_end - new_internal_end &&
-          old_unpoisoned < new_unpoisoned) {
-        *(u8 *)MemToShadow(new_internal_end) = old_unpoisoned;
+    // except the first and the last granule.
+    if (new_internal_end > new_internal_beg)
+      __builtin_memcpy((u8 *)MemToShadow(new_internal_beg),
+                       (u8 *)MemToShadow(old_internal_beg),
+                       (new_internal_end - new_internal_beg) / granularity);
+    // If the beginning and the end of the storage are aligned, we are done.
+    // Otherwise, we have to handle remaining granules.
+    if (new_internal_beg != new_storage_beg ||
+        new_internal_end != new_storage_end) {
+      if (WithinOneGranule(new_storage_beg, new_storage_end)) {
+        if (new_internal_end == new_storage_end) {
+          if (!AddressIsPoisoned(old_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                *(u8 *)MemToShadow(old_external_beg);
+          } else if (!AddressIsPoisoned(new_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                new_storage_beg - new_external_beg;
+          }
+        } else if (AddressIsPoisoned(new_storage_end)) {
+          if (!AddressIsPoisoned(old_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                AddressIsPoisoned(old_storage_end)
+                    ? *(u8 *)MemToShadow(old_internal_end)
+                    : new_storage_end - new_external_beg;
+          } else if (!AddressIsPoisoned(new_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                (new_storage_beg == new_external_beg)
+                    ? static_cast<u8>(kAsanContiguousContainerOOBMagic)
+                    : new_storage_beg - new_external_beg;
+          }
+        }
+      } else {
+        // Buffer is not within one granule!
+        if (new_internal_beg != new_storage_beg) {
+          if (!AddressIsPoisoned(old_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                *(u8 *)MemToShadow(old_external_beg);
+          } else if (!AddressIsPoisoned(new_storage_beg)) {
+            *(u8 *)MemToShadow(new_external_beg) =
+                new_storage_beg - new_external_beg;
+          }
+        }
+        if (new_internal_end != new_storage_end &&
+            AddressIsPoisoned(new_storage_end)) {
+          *(u8 *)MemToShadow(new_internal_end) =
+              AddressIsPoisoned(old_storage_end)
+                  ? *(u8 *)MemToShadow(old_internal_end)
+                  : old_storage_end - old_internal_end;
+        }
       }
     }
   } else {
     // If buffers are not aligned, we have to go byte by byte.
     uptr old_ptr = old_storage_beg;
     uptr new_ptr = new_storage_beg;
     uptr next_new;
-    for (; new_ptr + granularity <= new_storage_end;) {
+    for (; new_ptr < new_storage_end;) {
       next_new = RoundUpTo(new_ptr + 1, granularity);
       uptr unpoison_to = 0;
-      for (; new_ptr != next_new; ++new_ptr, ++old_ptr) {
+      for (; new_ptr != next_new && new_ptr != new_storage_end;
+           ++new_ptr, ++old_ptr) {
         if (!AddressIsPoisoned(old_ptr)) {
           unpoison_to = new_ptr + 1;
         }
       }
-      if (unpoison_to != 0) {
-        uptr granule_beg = new_ptr - granularity;
-        uptr value = unpoison_to - granule_beg;
-        *(u8 *)MemToShadow(granule_beg) = static_cast<u8>(value);
-      }
-    }
-    // Only case left is the end of the container in the middle of a granule.
-    // If memory after the end is unpoisoned, we cannot change anything.
-    // But if it's poisoned, we should unpoison as little as possible.
-    if (new_ptr != new_storage_end && AddressIsPoisoned(new_storage_end)) {
-      uptr unpoison_to = 0;
-      for (; new_ptr != new_storage_end; ++new_ptr, ++old_ptr) {
-        if (!AddressIsPoisoned(old_ptr)) {
-          unpoison_to = new_ptr + 1;
+      if (new_ptr < new_storage_end || new_ptr == new_internal_end ||
+          AddressIsPoisoned(new_storage_end)) {
+        uptr granule_beg = RoundDownTo(new_ptr - 1, granularity);
+        if (unpoison_to != 0) {
+          uptr value =
+              (unpoison_to == next_new) ? 0 : unpoison_to - granule_beg;
+          *(u8 *)MemToShadow(granule_beg) = static_cast<u8>(value);
+        } else {
+          *(u8 *)MemToShadow(granule_beg) =
+              (granule_beg >= new_storage_beg)
+                  ? static_cast<u8>(kAsanContiguousContainerOOBMagic)
+                  : new_storage_beg - granule_beg;
         }
       }
-      if (unpoison_to != 0) {
-        uptr granule_beg = RoundDownTo(new_storage_end, granularity);
-        uptr value = unpoison_to - granule_beg;
-        *(u8 *)MemToShadow(granule_beg) = static_cast<u8>(value);
-      }
     }
   }
-
-  __asan_unpoison_memory_region((void *)old_storage_beg,
-                                old_storage_end - old_storage_beg);
 }
 
 static const void *FindBadAddress(uptr begin, uptr end, bool poisoned) {
diff --git a/compiler-rt/test/asan/TestCases/move_container_annotations.cpp b/compiler-rt/test/asan/TestCases/move_container_annotations.cpp
@@ -27,14 +27,15 @@ template <class T> static constexpr T RoundUp(T x) {
 
 static std::deque<int> GetPoisonedState(char *begin, char *end) {
   std::deque<int> result;
-  for (; begin != end; ++begin) {
-    result.push_back(__asan_address_is_poisoned(begin));
+  for (char *ptr = begin; ptr != end; ++ptr) {
+    result.push_back(__asan_address_is_poisoned(ptr));
   }
   return result;
 }
 
 static void RandomPoison(char *beg, char *end) {
-  if (beg != RoundDown(beg) && (rand() % 2 == 1)) {
+  if (beg != RoundDown(beg) && RoundDown(beg) != RoundDown(end) &&
+      rand() % 2 == 1) {
     __asan_poison_memory_region(beg, RoundUp(beg) - beg);
     __asan_unpoison_memory_region(beg, rand() % (RoundUp(beg) - beg + 1));
   }
@@ -70,31 +71,36 @@ void TestMove(size_t capacity, size_t off_old, size_t off_new,
   char *new_buffer_end = new_buffer + new_buffer_size;
   bool poison_old = poison_buffers % 2 == 1;
   bool poison_new = poison_buffers / 2 == 1;
-  if (poison_old)
-    __asan_poison_memory_region(old_buffer, old_buffer_size);
-  if (poison_new)
-    __asan_poison_memory_region(new_buffer, new_buffer_size);
   char *old_beg = old_buffer + off_old;
   char *new_beg = new_buffer + off_new;
   char *old_end = old_beg + capacity;
   char *new_end = new_beg + capacity;
 
-  for (int i = 0; i < 1000; i++) {
+  for (int i = 0; i < 75; i++) {
+    if (poison_old)
+      __asan_poison_memory_region(old_buffer, old_buffer_size);
+    if (poison_new)
+      __asan_poison_memory_region(new_buffer, new_buffer_size);
+
     RandomPoison(old_beg, old_end);
-    std::deque<int> poison_states(old_beg, old_end);
+    std::deque<int> poison_states = GetPoisonedState(old_beg, old_end);
     __sanitizer_move_contiguous_container_annotations(old_beg, old_end, new_beg,
                                                       new_end);
 
     // If old_buffer were poisoned, expected state of memory before old_beg
     // is undetermined.
     // If old buffer were not poisoned, that memory should still be unpoisoned.
-    // Area between old_beg and old_end should never be poisoned.
-    char *cur = poison_old ? old_beg : old_buffer;
-    for (; cur < old_end; ++cur) {
-      assert(!__asan_address_is_poisoned(cur));
+    char *cur;
+    if (!poison_old) {
+      for (cur = old_buffer; cur < old_beg; ++cur) {
+        assert(!__asan_address_is_poisoned(cur));
+      }
+    }
+    for (size_t i = 0; i < poison_states.size(); ++i) {
+      assert(__asan_address_is_poisoned(&old_beg[i]) == poison_states[i]);
     }
-    // Memory after old_beg should be the same as at the beginning.
-    for (; cur < old_buffer_end; ++cur) {
+    // Memory after old_end should be the same as at the beginning.
+    for (cur = old_end; cur < old_buffer_end; ++cur) {
       assert(__asan_address_is_poisoned(cur) == poison_old);
     }
 
@@ -118,7 +124,8 @@ void TestMove(size_t capacity, size_t off_old, size_t off_new,
       }
     }
     // [cur; new_end) is not checked yet.
-    // If new_buffer were not poisoned, it cannot be poisoned and we can ignore check.
+    // If new_buffer were not poisoned, it cannot be poisoned and we can ignore
+    // a separate check.
     // If new_buffer were poisoned, it should be same as earlier.
     if (cur < new_end && poison_new) {
       size_t unpoisoned = count_unpoisoned(poison_states, new_end - cur);
@@ -143,9 +150,9 @@ void TestMove(size_t capacity, size_t off_old, size_t off_new,
 
 int main(int argc, char **argv) {
   int n = argc == 1 ? 64 : atoi(argv[1]);
-  for (int i = 0; i <= n; i++) {
-    for (int j = 0; j < kGranularity * 2; j++) {
-      for (int k = 0; k < kGranularity * 2; k++) {
+  for (size_t j = 0; j < kGranularity * 2; j++) {
+    for (size_t k = 0; k < kGranularity * 2; k++) {
+      for (int i = 0; i <= n; i++) {
         for (int poison = 0; poison < 4; ++poison) {
           TestMove(i, j, k, poison);
         }
