From e03579253a17afa9ea65dcbe5fae9dd4c583f99e Mon Sep 17 00:00:00 2001
From: Daniel Micay <daniel.micay@grapheneos.org>
Date: Sat, 12 Oct 2024 03:07:07 -0400
Subject: [PATCH] preserve PROT_MTE when releasing memory

---
 h_malloc.c | 13 +++++++++++--
 memory.c   | 15 +++++++++++++--
 memory.h   |  3 +++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/h_malloc.c b/h_malloc.c
index 3ceb1d22..6221d0bc 100644
--- a/h_malloc.c
+++ b/h_malloc.c
@@ -103,6 +103,15 @@ static void *memory_map_tagged(size_t size) {
     return memory_map(size);
 }
 
+static bool memory_map_fixed_tagged(void *ptr, size_t size) {
+#ifdef HAS_ARM_MTE
+    if (likely51(is_memtag_enabled())) {
+        return memory_map_fixed_mte(ptr, size);
+    }
+#endif
+    return memory_map_fixed(ptr, size);
+}
+
 #define SLAB_METADATA_COUNT
 
 struct slab_metadata {
@@ -899,7 +908,7 @@ static inline void deallocate_small(void *p, const size_t *expected_size) {
 
         if (c->empty_slabs_total + slab_size > max_empty_slabs_total) {
             int saved_errno = errno;
-            if (!memory_map_fixed(slab, slab_size)) {
+            if (!memory_map_fixed_tagged(slab, slab_size)) {
                 label_slab(slab, slab_size, class);
                 stats_slab_deallocate(c, slab_size);
                 enqueue_free_slab(c, metadata);
@@ -1896,7 +1905,7 @@ EXPORT int h_malloc_trim(UNUSED size_t pad) {
             struct slab_metadata *iterator = c->empty_slabs;
             while (iterator) {
                 void *slab = get_slab(c, slab_size, iterator);
-                if (memory_map_fixed(slab, slab_size)) {
+                if (memory_map_fixed_tagged(slab, slab_size)) {
                     break;
                 }
                 label_slab(slab, slab_size, class);
diff --git a/memory.c b/memory.c
index e1bd2ef4..2e54f6d8 100644
--- a/memory.c
+++ b/memory.c
@@ -39,8 +39,8 @@ void *memory_map_mte(size_t size) {
 }
 #endif
 
-bool memory_map_fixed(void *ptr, size_t size) {
-    void *p = mmap(ptr, size, PROT_NONE, MAP_ANONYMOUS|MAP_PRIVATE|MAP_FIXED, -1, 0);
+static bool memory_map_fixed_prot(void *ptr, size_t size, int prot) {
+    void *p = mmap(ptr, size, prot, MAP_ANONYMOUS|MAP_PRIVATE|MAP_FIXED, -1, 0);
     bool ret = p == MAP_FAILED;
     if (unlikely(ret) && errno != ENOMEM) {
         fatal_error("non-ENOMEM MAP_FIXED mmap failure");
@@ -48,6 +48,17 @@ bool memory_map_fixed(void *ptr, size_t size) {
     return ret;
 }
 
+bool memory_map_fixed(void *ptr, size_t size) {
+    return memory_map_fixed_prot(ptr, size, PROT_NONE);
+}
+
+#ifdef HAS_ARM_MTE
+// Note that PROT_MTE can't be cleared via mprotect
+bool memory_map_fixed_mte(void *ptr, size_t size) {
+    return memory_map_fixed_prot(ptr, size, PROT_MTE);
+}
+#endif
+
 bool memory_unmap(void *ptr, size_t size) {
     bool ret = munmap(ptr, size);
     if (unlikely(ret) && errno != ENOMEM) {
diff --git a/memory.h b/memory.h
index 6e4cd4da..d5e336be 100644
--- a/memory.h
+++ b/memory.h
@@ -15,6 +15,9 @@ void *memory_map(size_t size);
 void *memory_map_mte(size_t size);
 #endif
 bool memory_map_fixed(void *ptr, size_t size);
+#ifdef HAS_ARM_MTE
+bool memory_map_fixed_mte(void *ptr, size_t size);
+#endif
 bool memory_unmap(void *ptr, size_t size);
 bool memory_protect_ro(void *ptr, size_t size);
 bool memory_protect_rw(void *ptr, size_t size);