[MTE] move stack buffer calculation to mte.h

This is so we can use them in debuggerd.

Bug: 309446520
Change-Id: Ia7d20ec12bfaee164e259d0f34b632142b43d36b
diff --git a/libc/bionic/pthread_internal.cpp b/libc/bionic/pthread_internal.cpp
index e8a8ba2..14cc7da 100644
--- a/libc/bionic/pthread_internal.cpp
+++ b/libc/bionic/pthread_internal.cpp
@@ -36,9 +36,10 @@
 #include <sys/prctl.h>
 
 #include <async_safe/log.h>
+#include <bionic/mte.h>
 #include <bionic/reserved_signals.h>
+#include <bionic/tls_defines.h>
 
-#include "bionic/tls_defines.h"
 #include "private/ErrnoRestorer.h"
 #include "private/ScopedRWLock.h"
 #include "private/bionic_futex.h"
@@ -73,16 +74,12 @@
     g_thread_list = thread->next;
   }
 }
-// N.B. that this is NOT the pagesize, but 4096. This is hardcoded in the codegen.
-// See
-// https://github.com/search?q=repo%3Allvm/llvm-project%20AArch64StackTagging%3A%3AinsertBaseTaggedPointer&type=code
-constexpr size_t kStackMteRingbufferSizeMultiplier = 4096;
 
 static void __pthread_internal_free(pthread_internal_t* thread) {
 #ifdef __aarch64__
   if (void* stack_mte_tls = thread->bionic_tcb->tls_slot(TLS_SLOT_STACK_MTE)) {
     size_t size =
-        kStackMteRingbufferSizeMultiplier * (reinterpret_cast<uintptr_t>(stack_mte_tls) >> 56ULL);
+        stack_mte_ringbuffer_size_from_pointer(reinterpret_cast<uintptr_t>(stack_mte_tls));
     void* ptr = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(stack_mte_tls) &
                                         ((1ULL << 56ULL) - 1ULL));
     munmap(ptr, size);
@@ -191,38 +188,8 @@
   async_safe_fatal("stack not found in /proc/self/maps");
 }
 
+#if defined(__aarch64__)
 __LIBC_HIDDEN__ void* __allocate_stack_mte_ringbuffer(size_t n, pthread_internal_t* thread) {
-  if (n > 7) async_safe_fatal("error: invalid mte stack ring buffer size");
-  // Allocation needs to be aligned to 2*size to make the fancy code-gen work.
-  // So we allocate 3*size - pagesz bytes, which will always contain size bytes
-  // aligned to 2*size, and unmap the unneeded part.
-  // See
-  // https://github.com/search?q=repo%3Allvm/llvm-project%20AArch64StackTagging%3A%3AinsertBaseTaggedPointer&type=code
-  //
-  // In the worst case, we get an allocation that is one page past the properly
-  // aligned address, in which case we have to unmap the previous
-  // 2*size - pagesz bytes. In that case, we still have size properly aligned
-  // bytes left.
-  size_t size = (1 << n) * kStackMteRingbufferSizeMultiplier;
-  size_t pgsize = page_size();
-
-  size_t alloc_size = __BIONIC_ALIGN(3 * size - pgsize, pgsize);
-  void* allocation_ptr =
-      mmap(nullptr, alloc_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-  if (allocation_ptr == MAP_FAILED)
-    async_safe_fatal("error: failed to allocate stack mte ring buffer");
-  uintptr_t allocation = reinterpret_cast<uintptr_t>(allocation_ptr);
-
-  size_t alignment = 2 * size;
-  uintptr_t aligned_allocation = __BIONIC_ALIGN(allocation, alignment);
-  if (allocation != aligned_allocation) {
-    munmap(reinterpret_cast<void*>(allocation), aligned_allocation - allocation);
-  }
-  if (aligned_allocation + size != allocation + alloc_size) {
-    munmap(reinterpret_cast<void*>(aligned_allocation + size),
-           (allocation + alloc_size) - (aligned_allocation + size));
-  }
-
   const char* name;
   if (thread == nullptr) {
     name = "stack_mte_ring:main";
@@ -234,11 +201,11 @@
     async_safe_format_buffer(name_buffer, arraysize(name_buffer), "stack_mte_ring:%d", thread->tid);
     name = name_buffer;
   }
-  prctl(PR_SET_VMA, PR_SET_VMA_ANON_NAME, reinterpret_cast<void*>(aligned_allocation), size, name);
-
-  // We store the size in the top byte of the pointer (which is ignored)
-  return reinterpret_cast<void*>(aligned_allocation | ((1ULL << n) << 56ULL));
+  void* ret = stack_mte_ringbuffer_allocate(n, name);
+  if (!ret) async_safe_fatal("error: failed to allocate stack mte ring buffer");
+  return ret;
 }
+#endif
 
 bool __pthread_internal_remap_stack_with_mte() {
 #if defined(__aarch64__)
diff --git a/libc/bionic/pthread_internal.h b/libc/bionic/pthread_internal.h
index b270a06..5db42ab 100644
--- a/libc/bionic/pthread_internal.h
+++ b/libc/bionic/pthread_internal.h
@@ -213,7 +213,9 @@
 __LIBC_HIDDEN__ void __pthread_internal_remove(pthread_internal_t* thread);
 __LIBC_HIDDEN__ void __pthread_internal_remove_and_free(pthread_internal_t* thread);
 __LIBC_HIDDEN__ void __find_main_stack_limits(uintptr_t* low, uintptr_t* high);
+#if defined(__aarch64__)
 __LIBC_HIDDEN__ void* __allocate_stack_mte_ringbuffer(size_t n, pthread_internal_t* thread);
+#endif
 
 static inline __always_inline bionic_tcb* __get_bionic_tcb() {
   return reinterpret_cast<bionic_tcb*>(&__get_tls()[MIN_TLS_SLOT]);
diff --git a/libc/platform/bionic/mte.h b/libc/platform/bionic/mte.h
index 73cd821..98b3d27 100644
--- a/libc/platform/bionic/mte.h
+++ b/libc/platform/bionic/mte.h
@@ -29,8 +29,11 @@
 #pragma once
 
 #include <sys/auxv.h>
+#include <sys/mman.h>
 #include <sys/prctl.h>
 
+#include "page.h"
+
 // Note: Most PR_MTE_* constants come from the upstream kernel. This tag mask
 // allows for the hardware to provision any nonzero tag. Zero tags are reserved
 // for scudo to use for the chunk headers in order to prevent linear heap
@@ -63,6 +66,65 @@
     }
   }
 };
+
+// N.B. that this is NOT the pagesize, but 4096. This is hardcoded in the codegen.
+// See
+// https://github.com/search?q=repo%3Allvm/llvm-project%20AArch64StackTagging%3A%3AinsertBaseTaggedPointer&type=code
+constexpr size_t kStackMteRingbufferSizeMultiplier = 4096;
+
+inline size_t stack_mte_ringbuffer_size(uintptr_t size_cls) {
+  return kStackMteRingbufferSizeMultiplier * (1 << size_cls);
+}
+
+inline size_t stack_mte_ringbuffer_size_from_pointer(uintptr_t ptr) {
+  // The size in the top byte is not the size_cls, but the number of "pages" (not OS pages, but
+  // kStackMteRingbufferSizeMultiplier).
+  return kStackMteRingbufferSizeMultiplier * (ptr >> 56ULL);
+}
+
+inline uintptr_t stack_mte_ringbuffer_size_add_to_pointer(uintptr_t ptr, uintptr_t size_cls) {
+  return ptr | ((1ULL << size_cls) << 56ULL);
+}
+
+inline void* stack_mte_ringbuffer_allocate(size_t n, const char* name) {
+  if (n > 7) return nullptr;
+  // Allocation needs to be aligned to 2*size to make the fancy code-gen work.
+  // So we allocate 3*size - pagesz bytes, which will always contain size bytes
+  // aligned to 2*size, and unmap the unneeded part.
+  // See
+  // https://github.com/search?q=repo%3Allvm/llvm-project%20AArch64StackTagging%3A%3AinsertBaseTaggedPointer&type=code
+  //
+  // In the worst case, we get an allocation that is one page past the properly
+  // aligned address, in which case we have to unmap the previous
+  // 2*size - pagesz bytes. In that case, we still have size properly aligned
+  // bytes left.
+  size_t size = stack_mte_ringbuffer_size(n);
+  size_t pgsize = page_size();
+
+  size_t alloc_size = __BIONIC_ALIGN(3 * size - pgsize, pgsize);
+  void* allocation_ptr =
+      mmap(nullptr, alloc_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  if (allocation_ptr == MAP_FAILED)
+    return nullptr;
+  uintptr_t allocation = reinterpret_cast<uintptr_t>(allocation_ptr);
+
+  size_t alignment = 2 * size;
+  uintptr_t aligned_allocation = __BIONIC_ALIGN(allocation, alignment);
+  if (allocation != aligned_allocation) {
+    munmap(reinterpret_cast<void*>(allocation), aligned_allocation - allocation);
+  }
+  if (aligned_allocation + size != allocation + alloc_size) {
+    munmap(reinterpret_cast<void*>(aligned_allocation + size),
+           (allocation + alloc_size) - (aligned_allocation + size));
+  }
+
+  if (name) {
+    prctl(PR_SET_VMA, PR_SET_VMA_ANON_NAME, reinterpret_cast<void*>(aligned_allocation), size, name);
+  }
+
+  // We store the size in the top byte of the pointer (which is ignored)
+  return reinterpret_cast<void*>(stack_mte_ringbuffer_size_add_to_pointer(aligned_allocation, n));
+}
 #else
 struct ScopedDisableMTE {
   // Silence unused variable warnings in non-aarch64 builds.