[compiler-rt] [compiler-rt][ASan] Add function copying annotations (PR #91702)

[via llvm-commits]

================
@@ -576,6 +576,235 @@ void __sanitizer_annotate_double_ended_contiguous_container(
   }
 }
 
+// Checks if a buffer [p; q) falls into a single granule.
+static bool WithinOneGranule(uptr p, uptr q) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  if (p == q)
+    return true;
+  return RoundDownTo(p, granularity) == RoundDownTo(q - 1, granularity);
+}
+
+// Copies ASan memory annotation (a shadow memory value)
+// from one granule to another.
+static void CopyGranuleAnnotation(uptr dst, uptr src) {
+  *(u8 *)MemToShadow(dst) = *(u8 *)MemToShadow(src);
+}
+
+// Marks the specified number of bytes in a granule as accessible or
+// poisones the whole granule with kAsanContiguousContainerOOBMagic value.
+static void AnnotateContainerGranuleAccessibleBytes(uptr ptr, u8 n) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  if (n == granularity) {
+    *(u8 *)MemToShadow(ptr) = 0;
+  } else if (n == 0) {
+    *(u8 *)MemToShadow(ptr) = static_cast<u8>(kAsanContiguousContainerOOBMagic);
+  } else {
+    *(u8 *)MemToShadow(ptr) = n;
+  }
+}
+
+// Performs a byte-by-byte copy of ASan annotations (shadow memory values).
+// Result may be different due to ASan limitations, but result cannot lead
+// to false positives (more memory than requested may get unpoisoned).
+static void SlowCopyContainerAnnotations(uptr src_storage_beg,
+                                         uptr src_storage_end,
+                                         uptr dst_storage_beg,
+                                         uptr dst_storage_end) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  uptr dst_internal_end = RoundDownTo(dst_storage_end, granularity);
+  uptr src_ptr = src_storage_beg;
+  uptr dst_ptr = dst_storage_beg;
+
+  while (dst_ptr < dst_storage_end) {
+    uptr next_new = RoundUpTo(dst_ptr + 1, granularity);
+    uptr granule_begin = next_new - granularity;
+    uptr unpoisoned_bytes = 0;
+
+    for (; dst_ptr != next_new && dst_ptr != dst_storage_end;
+         ++dst_ptr, ++src_ptr) {
+      if (!AddressIsPoisoned(src_ptr)) {
+        unpoisoned_bytes = dst_ptr - granule_begin + 1;
+      }
+    }
+    if (dst_ptr < dst_storage_end || dst_ptr == dst_internal_end ||
+        AddressIsPoisoned(dst_storage_end)) {
+      if (unpoisoned_bytes != 0 || granule_begin >= dst_storage_beg) {
+        AnnotateContainerGranuleAccessibleBytes(granule_begin,
+                                                unpoisoned_bytes);
+      } else if (!AddressIsPoisoned(dst_storage_beg)) {
+        AnnotateContainerGranuleAccessibleBytes(
+            granule_begin, dst_storage_beg - granule_begin);
+      }
+    }
+  }
+}
+
+// Performs a byte-by-byte copy of ASan annotations (shadow memory values),
+// going through bytes in reversed order, but not reversing annotations.
+// Result may be different due to ASan limitations, but result cannot lead
+// to false positives (more memory than requested may get unpoisoned).
+static void SlowReversedCopyContainerAnnotations(uptr src_storage_beg,
+                                                 uptr src_storage_end,
+                                                 uptr dst_storage_beg,
+                                                 uptr dst_storage_end) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  uptr dst_internal_beg = RoundDownTo(dst_storage_beg, granularity);
+  uptr dst_internal_end = RoundDownTo(dst_storage_end, granularity);
+  uptr src_ptr = src_storage_end;
+  uptr dst_ptr = dst_storage_end;
+
+  while (dst_ptr > dst_storage_beg) {
+    uptr granule_begin = RoundDownTo(dst_ptr - 1, granularity);
+    uptr unpoisoned_bytes = 0;
+
+    for (; dst_ptr != granule_begin && dst_ptr != dst_storage_beg;
+         --dst_ptr, --src_ptr) {
+      if (unpoisoned_bytes == 0 && !AddressIsPoisoned(src_ptr - 1)) {
+        unpoisoned_bytes = dst_ptr - granule_begin;
+      }
+    }
+
+    if (dst_ptr >= dst_internal_end && !AddressIsPoisoned(dst_storage_end)) {
+      continue;
+    }
+
+    if (granule_begin == dst_ptr || unpoisoned_bytes != 0) {
+      AnnotateContainerGranuleAccessibleBytes(granule_begin, unpoisoned_bytes);
+    } else if (!AddressIsPoisoned(dst_storage_beg)) {
+      AnnotateContainerGranuleAccessibleBytes(granule_begin,
+                                              dst_storage_beg - granule_begin);
+    }
+  }
+}
+
+// A helper function for __sanitizer_copy_contiguous_container_annotations,
+// has assumption about begin and end of the container.
+// Should not be used stand alone.
+static void CopyContainerFirstGranuleAnnotation(uptr src_storage_begin,
+                                                uptr dst_storage_begin) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  // First granule
+  uptr dst_external_begin = RoundDownTo(dst_storage_begin, granularity);
+  uptr src_external_begin = RoundDownTo(src_storage_begin, granularity);
+  if (!AddressIsPoisoned(src_storage_begin)) {
+    CopyGranuleAnnotation(dst_external_begin, src_external_begin);
+  } else if (!AddressIsPoisoned(dst_storage_begin)) {
+    AnnotateContainerGranuleAccessibleBytes(
+        dst_external_begin, dst_storage_begin - dst_external_begin);
+  }
+}
+
+// A helper function for __sanitizer_copy_contiguous_container_annotations,
+// has assumption about begin and end of the container.
+// Should not be used stand alone.
+static void CopyContainerLastGranuleAnnotation(uptr src_storage_end,
+                                               uptr dst_internal_end) {
+  constexpr uptr granularity = ASAN_SHADOW_GRANULARITY;
+  // Last granule
+  uptr src_internal_end = RoundDownTo(src_storage_end, granularity);
+  if (AddressIsPoisoned(src_storage_end)) {
+    CopyGranuleAnnotation(dst_internal_end, src_internal_end);
+  } else {
+    AnnotateContainerGranuleAccessibleBytes(dst_internal_end,
+                                            src_storage_end - src_internal_end);
+  }
+}
+
+// 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)
+//     // src_storage_beg % granularity != dst_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_copy_contiguous_container_annotations(
+    const void *src_begin_p, const void *src_end_p, const void *dst_begin_p) {
+  if (!flags()->detect_container_overflow)
+    return;
+
+  VPrintf(2, "contiguous_container_src: %p %p\n", src_begin_p, src_end_p);
+  VPrintf(2, "contiguous_container_dst: %p\n", dst_begin_p);
+
+  uptr src_storage_begin = reinterpret_cast<uptr>(src_begin_p);
+  uptr src_storage_end = reinterpret_cast<uptr>(src_end_p);
+  uptr dst_storage_begin = reinterpret_cast<uptr>(dst_begin_p);
+  uptr dst_storage_end =
----------------
AdvenamTacet wrote:

That was my first implementation and original reasoning, but later I changed it. My motivation behind removing that argument was intended use of that function when (without ASan) memmove or memcpy is used and those functions assume same size of both containers. But I'm happy with either, reverted that change for now.

Also, should we keep this order of arguments (src, dst) or should we change the order to dst be before src? I based order on old/new container annotation functions, like:
```
void __sanitizer_annotate_double_ended_contiguous_container(
    const void *storage_beg_p, const void *storage_end_p,
    const void *old_container_beg_p, const void *old_container_end_p,
    const void *new_container_beg_p, const void *new_container_end_p) {
```
But maybe it's better to mimic memmove/memcpy with destination being first?

https://github.com/llvm/llvm-project/pull/91702