scudo: Split setRandomTag in two. NFCI.

[pcc]

Summary

Separate the IRG part from the STZG part since we will need to use
the latter on its own for some upcoming changes.

[hctim]

LGTM w/ nits.

compiler-rt/lib/scudo/standalone/memtag.h

106–122

Worth a DCHECK that Begin and End are granule-aligned (and in selectRandomTag)?

107

Probably better to be safe and make this if (Begin < End)?

[pcc]

compiler-rt/lib/scudo/standalone/memtag.h

106–122

End doesn't need to be granule-aligned (we can pass an unaligned End from the primary allocator). This is the reason why we return the address of the end of the tagged region instead of just letting the caller use End. I suppose that we could check Begin here though.

107

I'm not sure about that. I think that if Begin > End we will want to crash here (i.e. the likely source of the bug) instead of at some indeterminate point later.

compiler-rt/lib/scudo/standalone/memtag.h

 #include <sys/prctl.h>
 #if defined(ANDROID_EXPERIMENTAL_MTE)
 #include <bionic/mte_kernel.h>
 #endif
 #endif
 
 namespace scudo {
 
+void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask, uptr *TaggedBegin,
+                  uptr *TaggedEnd);
+
 #if defined(__aarch64__) || defined(SCUDO_FUZZ)
 
 inline constexpr bool archSupportsMemoryTagging() { return true; }
 inline constexpr uptr archMemoryTagGranuleSize() { return 16; }
 
 inline uptr untagPointer(uptr Ptr) { return Ptr & ((1ULL << 56) - 1); }
 
 inline uint8_t extractTag(uptr Ptr) { return (Ptr >> 56) & 0xf; }
__asm__ __volatile__(".arch_extension mte; mrs %0, tco; msr tco, #1"
                          : "=r"(PrevTCO));
   }
 
   ~ScopedDisableMemoryTagChecks() {
     __asm__ __volatile__(".arch_extension mte; msr tco, %0" : : "r"(PrevTCO));
   }
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
-                         uptr *TaggedBegin, uptr *TaggedEnd) {
-  void *End;
+inline uptr selectRandomTag(uptr Ptr, uptr ExcludeMask) {
+  uptr TaggedPtr;
+  __asm__ __volatile__(
+      ".arch_extension mte; irg %[TaggedPtr], %[Ptr], %[ExcludeMask]"
+      : [TaggedPtr] "=r"(TaggedPtr)
+      : [Ptr] "r"(Ptr), [ExcludeMask] "r"(ExcludeMask));
+  return TaggedPtr;
+}
+
+inline uptr storeTags(uptr Begin, uptr End) {
+  DCHECK(Begin % 16 == 0);
+  if (Begin != End) {
     __asm__ __volatile__(
         R"(
       .arch_extension mte
 
-    // Set a random tag for Ptr in TaggedPtr. This needs to happen even if
-    // Size = 0 so that TaggedPtr ends up pointing at a valid address.
-    irg %[TaggedPtr], %[Ptr], %[ExcludeMask]
-    mov %[Cur], %[TaggedPtr]
-
-    // Skip the loop if Size = 0. We don't want to do any tagging in this case.
-    cbz %[Size], 2f
-
-    // Set the memory tag of the region
-    // [TaggedPtr, TaggedPtr + roundUpTo(Size, 16))
-    // to the pointer tag stored in TaggedPtr.
-    add %[End], %[TaggedPtr], %[Size]
-
     1:
       stzg %[Cur], [%[Cur]], #16
       cmp %[Cur], %[End]
       b.lt 1b
-
-  2:
     )"
-      :
-      [TaggedPtr] "=&r"(*TaggedBegin), [Cur] "=&r"(*TaggedEnd), [End] "=&r"(End)
-      : [Ptr] "r"(Ptr), [Size] "r"(Size), [ExcludeMask] "r"(ExcludeMask)
+        : [Cur] "+&r"(Begin)
+        : [End] "r"(End)
         : "memory");
   }
+  return Begin;
+}
 
 inline void *prepareTaggedChunk(void *Ptr, uptr Size, uptr ExcludeMask,
                                 uptr BlockEnd) {
   // Prepare the granule before the chunk to store the chunk header by setting
   // its tag to 0. Normally its tag will already be 0, but in the case where a
   // chunk holding a low alignment allocation is reused for a higher alignment
   // allocation, the chunk may already have a non-zero tag from the previous
   // allocation.

 inline void enableMemoryTagChecksTestOnly() {
   UNREACHABLE("memory tagging not supported");
 }
 
 struct ScopedDisableMemoryTagChecks {
   ScopedDisableMemoryTagChecks() {}
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
-                         uptr *TaggedBegin, uptr *TaggedEnd) {
+inline uptr selectRandomTag(uptr Ptr, uptr ExcludeMask) {
   (void)Ptr;
-  (void)Size;
   (void)ExcludeMask;
-  (void)TaggedBegin;
-  (void)TaggedEnd;
+  UNREACHABLE("memory tagging not supported");
+}
+
+inline uptr storeTags(uptr Begin, uptr End) {
+  (void)Begin;
+  (void)End;
   UNREACHABLE("memory tagging not supported");
 }
 
 inline void *prepareTaggedChunk(void *Ptr, uptr Size, uptr ExcludeMask,
                                 uptr BlockEnd) {
   (void)Ptr;
   (void)Size;
   (void)ExcludeMask;

 
 inline uptr loadTag(uptr Ptr) {
   (void)Ptr;
   UNREACHABLE("memory tagging not supported");
 }
 
 #endif
 
+inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
+                         uptr *TaggedBegin, uptr *TaggedEnd) {
+  *TaggedBegin = selectRandomTag(reinterpret_cast<uptr>(Ptr), ExcludeMask);
+  *TaggedEnd = storeTags(*TaggedBegin, *TaggedBegin + Size);
+}
+
 } // namespace scudo
 
 #endif