[glibc/arm/malloc-mte-v3] malloc: aarch64: Add MTE memory tagging

https://sourceware.org/git/gitweb.cgi?p=glibc.git;h=d5d5f03c8facfc0e72514dba42e3bbaeb1bf1fe6

commit d5d5f03c8facfc0e72514dba42e3bbaeb1bf1fe6
Author: Yury Khrustalev <[email protected]>
Date:   Fri May 29 13:44:25 2026 +0100

    malloc: aarch64: Add MTE memory tagging
    
    Tag memory on systems that support MTE and when memory
    tagging is enabled at runtime.
    
    Currently we implement tagging as a wrapper around core malloc
    functions. User pointers are expected to be tagged and internal
    malloc pointers are untagged (or rather tagged with 0 tag).
    
    Also add tests that check logical and allocation tags.

Diff:
---
 sysdeps/aarch64/multiarch/malloc-mte.c             | 119 +++++++++++++---
 sysdeps/unix/sysv/linux/aarch64/Makefile           |  27 ++++
 sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h   | 110 +++++++++++++++
 .../sysv/linux/aarch64/tst-mte-malloc-static.c     |   1 +
 sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc.c   | 154 +++++++++++++++++++++
 .../sysv/linux/aarch64/tst-mte-realloc-static.c    |   1 +
 sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc.c  |  72 ++++++++++
 7 files changed, 467 insertions(+), 17 deletions(-)

diff --git a/sysdeps/aarch64/multiarch/malloc-mte.c b/sysdeps/aarch64/multiarch/malloc-mte.c
index aeb5e1d3ca..8d9125ef34 100644
--- a/sysdeps/aarch64/multiarch/malloc-mte.c
+++ b/sysdeps/aarch64/multiarch/malloc-mte.c
@@ -16,76 +16,161 @@
    License along with the GNU C Library; if not, see
    <http://www.gnu.org/licenses/>.  */
 
+#include "aarch64-mte.h"
+
 #include <malloc-ifuncs.h>
+#include <errno.h>
+
+#define TAG_MEM(ptr, tagfun) __glibc_unlikely (ptr == NULL) ? NULL : ({ \
+  size_t size = __malloc_usable_size (ptr); \
+  tagfun (__mte_new_tag (ptr), size); \
+})
+
+#define UNTAG_MEM(non_null_ptr) ({ \
+  void *untagged = __mte_clear_tag (non_null_ptr); \
+  size_t size = __malloc_usable_size (untagged); \
+  __mte_tag_region (untagged, size); \
+})
 
 void *__libc_malloc_mte (size_t bytes)
 {
-  return __libc_malloc (bytes);
+  void *untagged = __libc_malloc (bytes);
+  return TAG_MEM (untagged, __mte_tag_region);
 }
 libc_hidden_def (__libc_malloc_mte)
 
 void *__libc_calloc_mte (size_t n, size_t elem_size)
 {
-  return __libc_calloc (n, elem_size);
+  /* We use core malloc instead of calloc because we can
+     take advantage of MTE to zero memory region.  */
+  void *untagged = __libc_malloc (n * elem_size);
+  return TAG_MEM (untagged, __mte_tag_region_zero);
 }
 libc_hidden_def (__libc_calloc_mte)
 
 void *__libc_memalign_mte (size_t alignment, size_t bytes)
 {
-  return __libc_memalign (alignment, bytes);
+  void *untagged = __libc_memalign (alignment, bytes);
+  return TAG_MEM (untagged, __mte_tag_region);
 }
 libc_hidden_def (__libc_memalign_mte)
 
 void *__libc_valloc_mte (size_t bytes)
 {
-  return __libc_valloc (bytes);
+  void *untagged = __libc_valloc (bytes);
+  return TAG_MEM (untagged, __mte_tag_region);
 }
 libc_hidden_def (__libc_valloc_mte)
 
 void *__libc_pvalloc_mte (size_t bytes)
 {
-  return __libc_pvalloc (bytes);
+  void *untagged = __libc_pvalloc (bytes);
+  return TAG_MEM (untagged, __mte_tag_region);
 }
 libc_hidden_def (__libc_pvalloc_mte)
 
-void *__libc_realloc_mte (void *oldmem, size_t bytes)
+/* See malloc.c for details.  */
+#ifndef REALLOC_ZERO_BYTES_FREES
+#define REALLOC_ZERO_BYTES_FREES 1
+#endif
+
+void *__libc_realloc_mte (void *tagged_oldmem, size_t bytes)
 {
-  return __libc_realloc (oldmem, bytes);
+  /* Quick check: realloc of null is supposed to be same as malloc.  */
+  if (tagged_oldmem == NULL)
+    return __libc_malloc_mte (bytes);
+
+#if REALLOC_ZERO_BYTES_FREES
+  /* Quick check: realloc with 0 size is supposed to be same as free.  */
+  if (bytes == 0)
+    {
+      __libc_free_mte (tagged_oldmem);
+      return NULL;
+    }
+#endif
+
+  /* Bad size, old memory remains unchanged.  */
+  if (bytes > PTRDIFF_MAX)
+    {
+      __set_errno (ENOMEM);
+      return NULL;
+    }
+
+  /* At this point we untag oldmem allocation.  */
+  void *untagged_oldmem = __mte_clear_tag (tagged_oldmem);
+
+  /* Mark the chunk as belonging to the library again.  */
+  size_t size_old = __malloc_usable_size (untagged_oldmem);
+  untagged_oldmem = __mte_tag_region (untagged_oldmem, size_old);
+
+  /* Call realloc core.  */
+  void *untagged_newmem = __libc_realloc (untagged_oldmem, bytes);
+  if (untagged_newmem == NULL)
+    return NULL;
+  size_t size_new = __malloc_usable_size (untagged_newmem);
+
+  /* If realloc core returns old pointer, we need re-tag it.  */
+  if (size_new == size_old && untagged_newmem == untagged_oldmem)
+    return __mte_tag_region (tagged_oldmem, size_new);
+
+  /* Otherwise, assign new tag.  */
+  void *tagged_newmem = __mte_new_tag (untagged_newmem);
+  return __mte_tag_region (tagged_newmem, size_new);
 }
 libc_hidden_def (__libc_realloc_mte)
 
-void __libc_free_mte (void *mem)
+void __libc_free_mte (void *tagged)
 {
-  __libc_free (mem);
+  if (__glibc_unlikely (tagged == NULL))
+    return;
+  /* Mark the chunk as belonging to the library again.  */
+  void *untagged = UNTAG_MEM (tagged);
+  /* Call free core.  */
+  __libc_free (untagged);
 }
 libc_hidden_def (__libc_free_mte)
 
-size_t __malloc_usable_size_mte (void *m)
+size_t __malloc_usable_size_mte (void *tagged)
 {
-  return __malloc_usable_size (m);
+  /* Clear only logical tag to allow accessing internal malloc
+     structures via offset from this pointer.  */
+  void *untagged = __mte_clear_tag (tagged);
+  return __malloc_usable_size (untagged);
 }
 libc_hidden_def (__malloc_usable_size_mte)
 
 int __posix_memalign_mte (void **memptr, size_t alignment, size_t size)
 {
-  return __posix_memalign (memptr, alignment, size);
+  /* Call core function.  */
+  int err = __posix_memalign (memptr, alignment, size);
+  if (err != 0)
+    return err;
+  *memptr = TAG_MEM (*memptr, __mte_tag_region);
+  return err;
 }
 libc_hidden_def (__posix_memalign_mte)
 
 void *__aligned_alloc_mte (size_t alignment, size_t bytes)
 {
-  return __aligned_alloc (alignment, bytes);
+  void *untagged = __aligned_alloc (alignment, bytes);
+  return TAG_MEM (untagged, __mte_tag_region);
 }
 libc_hidden_def (__aligned_alloc_mte)
 
-void __free_sized_mte (void *ptr, size_t size)
+void __free_sized_mte (void *tagged, size_t size)
 {
-  __free_sized (ptr, size);
+  /* Mark the chunk as belonging to the library again.  */
+  void *untagged = UNTAG_MEM (tagged);
+  /* Call core function.  */
+  __free_sized (untagged, size);
 }
 libc_hidden_def (__free_sized_mte)
 
-void __free_aligned_sized_mte (void *ptr, size_t alignment, size_t size)
+void __free_aligned_sized_mte (void *tagged, size_t alignment, size_t size)
 {
-  __free_aligned_sized (ptr, alignment, size);
+  /* Mark the chunk as belonging to the library again.  */
+  void *untagged = UNTAG_MEM (tagged);
+  /* Call core function.  */
+  __free_aligned_sized (untagged, alignment, size);
 }
 libc_hidden_def (__free_aligned_sized_mte)
diff --git a/sysdeps/unix/sysv/linux/aarch64/Makefile b/sysdeps/unix/sysv/linux/aarch64/Makefile
index e7baa47458..0c15ae4ae2 100644
--- a/sysdeps/unix/sysv/linux/aarch64/Makefile
+++ b/sysdeps/unix/sysv/linux/aarch64/Makefile
@@ -372,6 +372,33 @@ gen-as-const-headers += ucontext_i.sym
 endif
 
 ifeq ($(subdir),malloc)
+ifneq ($(multi-arch),no)
+# Special test that check MTE memory tagging
+tests-for-mte += \
+  tst-mte-malloc \
+  tst-mte-malloc-static \
+  tst-mte-realloc \
+  tst-mte-realloc-static \
+  # tests-for-mte
+
+tests += $(tests-for-mte)
+
+tests-static += \
+  tst-mte-malloc-static \
+  tst-mte-realloc-static \
+  # tests-static
+
+CFLAGS-tst-mte-malloc.o += -march=armv9-a+memtag
+CFLAGS-tst-mte-malloc-static.o += -march=armv9-a+memtag
+CFLAGS-tst-mte-realloc.o += -march=armv9-a+memtag
+CFLAGS-tst-mte-realloc-static.o += -march=armv9-a+memtag
+
+tests-exclude-malloc-check += $(tests-for-mte)
+tests-exclude-mcheck += $(tests-for-mte)
+tests-exclude-hugetlb1 += $(tests-for-mte)
+tests-exclude-hugetlb2 += $(tests-for-mte)
+endif # ifneq ($(multi-arch),no)
+
 # Add MTE tunable to all malloc tests except malloc-check and mcheck variants
 tests-malloc-mte = \
   $(tests) \
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h b/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h
new file mode 100644
index 0000000000..acfb189530
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h
@@ -0,0 +1,110 @@
+/* AArch64 test helper functions for MTE.
+   Copyright (C) 2026 Free Software Foundation, Inc.
+   This file is part of the GNU C Library.
+
+   The GNU C Library is free software; you can redistribute it and/or
+   modify it under the terms of the GNU Lesser General Public
+   License as published by the Free Software Foundation; either
+   version 2.1 of the License, or (at your option) any later version.
+
+   The GNU C Library is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+   Lesser General Public License for more details.
+
+   You should have received a copy of the GNU Lesser General Public
+   License along with the GNU C Library; if not, see
+   <https://www.gnu.org/licenses/>.  */
+
+#ifndef TST_MTE_HELPER_H
+#define TST_MTE_HELPER_H
+
+#include <support/check.h>
+
+#include <stdio.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <malloc.h>
+#include <sys/auxv.h>
+#include <sys/prctl.h>
+
+#define GRANULE_SIZE 16
+
+/* Extract logical tag from pointer PTR.  */
+static __always_inline
+uint64_t get_logical_tag (const void *ptr)
+{
+  uint64_t t = (uint64_t)ptr;
+  return t >> 56ul & 0xf;
+}
+
+/* Load allocation tag from memory pointed-to by the PTR pointer.  */
+static __always_inline
+uint64_t get_allocation_tag (const void *ptr)
+{
+  uint64_t t;
+  asm volatile ("ldg %0, [%1]" : "=r" (t) : "r" (ptr));
+  return t >> 56ul & 0xf;
+}
+
+/* Read the Tag Check Override bit.  */
+static __always_inline
+uint64_t get_pstate_tco (void) {
+  uint64_t t;
+  asm volatile ("mrs %0, tco" : "=r" (t));
+  return t;
+}
+
+static __always_inline
+bool check_tags (void *tm)
+{
+  size_t len = malloc_usable_size (tm);
+  TEST_VERIFY (len % GRANULE_SIZE == 0);
+
+  uint64_t ltag = get_logical_tag (tm);
+  TEST_VERIFY (ltag != 0);
+
+  for (size_t offset = 0; offset < len; offset += GRANULE_SIZE)
+    {
+      const char *g = (char *)tm + offset;
+      uint64_t atag = get_allocation_tag (g);
+      TEST_COMPARE (ltag, atag);
+      if (ltag != atag)
+	{
+          printf ("tagged ptr: %p usable size: %zu\n", tm, len);
+	  printf ("tags mismatch at offset %zu: logical=%lu, allocation=%lu\n",
+		  offset, ltag, atag);
+          return false;
+	}
+    }
+  return ltag != 0;
+}
+
+static __always_inline
+void check_mte_enabled (void)
+{
+  /* Check if MTE is supported.  */
+  if (!(getauxval (AT_HWCAP2) & HWCAP2_MTE))
+    FAIL_UNSUPPORTED ("kernel or CPU does not support HWCAP2_MTE");
+
+  /* Check if Tag Check Override bit is set.  */
+  if (get_pstate_tco () != 0)
+    FAIL_UNSUPPORTED ("MTE tag check override is enabled");
+
+  /* Check applied MTE params.  */
+  uint64_t x = (uint64_t) prctl (PR_GET_TAGGED_ADDR_CTRL, 0, 0, 0, 0);
+  uint64_t status = (x & 1ul);
+  uint64_t mode = (x & PR_MTE_TCF_MASK) >> PR_MTE_TCF_SHIFT;
+  uint64_t tags = (x & PR_MTE_TAG_MASK) >> PR_MTE_TAG_SHIFT;
+
+  printf ("MTE status: %4lx\n", status);
+  printf ("MTE mode:   %4lx\n", mode);
+  printf ("MTE tags:   %4lx\n", tags);
+
+  /* This test should be run in sync mode for tag checks.  */
+  TEST_VERIFY (status == 1);
+  TEST_VERIFY (mode == PR_MTE_TCF_SYNC >> PR_MTE_TCF_SHIFT);
+  TEST_VERIFY (tags == 0xfffe);
+}
+
+#endif // TST_MTE_HELPER_H
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc-static.c b/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc-static.c
new file mode 100644
index 0000000000..43f7f54008
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc-static.c
@@ -0,0 +1 @@
+#include "tst-mte-malloc.c"
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc.c b/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc.c
new file mode 100644
index 0000000000..1747224e47
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-malloc.c
@@ -0,0 +1,154 @@
+/* AArch64 tests for heap memory tagging.
+   Copyright (C) 2026 Free Software Foundation, Inc.
+   This file is part of the GNU C Library.
+
+   The GNU C Library is free software; you can redistribute it and/or
+   modify it under the terms of the GNU Lesser General Public
+   License as published by the Free Software Foundation; either
+   version 2.1 of the License, or (at your option) any later version.
+
+   The GNU C Library is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+   Lesser General Public License for more details.
+
+   You should have received a copy of the GNU Lesser General Public
+   License along with the GNU C Library; if not, see
+   <https://www.gnu.org/licenses/>.  */
+
+#include <support/check.h>
+#include <support/support.h>
+#include <support/xsignal.h>
+#include <support/test-driver.h>
+#include <array_length.h>
+
+#include "tst-mte-helper.h"
+
+/* Characteristic malloc sizes to cover various allocation methods.  */
+size_t sizes[] = {
+  1,
+  16, 40, 64, 120,
+  128, 500, 1000,
+  1050, 4096, 5000, 65000,
+  131072, 2000000
+};
+
+static void check_malloc (size_t len)
+{
+  printf ("testing malloc with req size %zu\n", len);
+  void *tm = malloc (len);
+  if (!check_tags (tm))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free (tm);
+}
+
+static void check_calloc (size_t len)
+{
+  size_t num = len / sizeof (uint64_t) + 1;
+  printf ("testing calloc with req size %zu\n", num * sizeof (uint64_t));
+  uint64_t *tm = calloc (num, sizeof (uint64_t));
+  if (check_tags (tm))
+    for (int n = 0; n < num; n ++)
+      TEST_VERIFY_EXIT (tm[n] == 0);
+  else
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free (tm);
+}
+
+static void check_memalign (size_t len, size_t alignment)
+{
+  printf ("testing memalign(%zu) with req size %zu\n", alignment, len);
+  void *tm = memalign (alignment, len);
+  if (!check_tags (tm))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free (tm);
+}
+
+static void check_valloc (size_t len)
+{
+  printf ("testing valloc with req size %zu\n", len);
+  void *tm = valloc (len);
+  if (!check_tags (tm))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free_sized (tm, len);
+}
+
+static void check_pvalloc (size_t len)
+{
+  printf ("testing pvalloc with req size %zu\n", len);
+  void *tm = pvalloc (len);
+  if (!check_tags (tm))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free_sized (tm, len);
+}
+
+static void check_posix_memalign (size_t len, size_t alignment)
+{
+  printf ("testing posix_memalign(%zu) with req size %zu\n", alignment, len);
+  void *p = NULL;
+  int err = posix_memalign (&p, alignment, len);
+  if (err)
+    perror ("posix_memalign");
+  TEST_VERIFY (p != NULL);
+  TEST_VERIFY (err == 0);
+  if (!check_tags (p))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)p);
+  free_aligned_sized (p, alignment, len);
+}
+
+static void check_aligned_alloc (size_t len, size_t alignment)
+{
+  printf ("testing aligned_alloc(%zu) with req size %zu\n", alignment, len);
+  void *tm = aligned_alloc (alignment, len);
+  if (!check_tags (tm))
+    printf ("tagged pointer?: %016lx\n", (uintptr_t)tm);
+  free_aligned_sized (tm, alignment, len);
+}
+
+static int
+do_test (void)
+{
+  /* Check if MTE is supported, configured and enabled.  */
+  check_mte_enabled ();
+
+  array_foreach_const (plen, sizes)
+    check_malloc (*plen);
+
+  array_foreach_const (plen, sizes)
+    check_calloc (*plen);
+
+  array_foreach_const (plen, sizes)
+    {
+      check_memalign (*plen, 2);
+      check_memalign (*plen, 4);
+      check_memalign (*plen, 8);
+      check_memalign (*plen, 16);
+      check_memalign (*plen, 32);
+    }
+
+  array_foreach_const (plen, sizes)
+    check_valloc (*plen);
+
+  array_foreach_const (plen, sizes)
+    check_pvalloc (*plen);
+
+  array_foreach_const (plen, sizes)
+    {
+      check_posix_memalign (*plen, sizeof (void *) * 1);
+      check_posix_memalign (*plen, sizeof (void *) * 2);
+      check_posix_memalign (*plen, sizeof (void *) * 4);
+    }
+
+  array_foreach_const (plen, sizes)
+    {
+      check_aligned_alloc (*plen, 2);
+      check_aligned_alloc (*plen, 4);
+      check_aligned_alloc (*plen, 8);
+      check_aligned_alloc (*plen, 16);
+      check_aligned_alloc (*plen, 32);
+    }
+
+  return 0;
+}
+
+#include <support/test-driver.c>
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc-static.c b/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc-static.c
new file mode 100644
index 0000000000..9ee758128c
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc-static.c
@@ -0,0 +1 @@
+#include "tst-mte-realloc.c"
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc.c b/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc.c
new file mode 100644
index 0000000000..e7197b58fe
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-realloc.c
@@ -0,0 +1,72 @@
+/* AArch64 tests for heap memory tagging.
+   Copyright (C) 2026 Free Software Foundation, Inc.
+   This file is part of the GNU C Library.
+
+   The GNU C Library is free software; you can redistribute it and/or
+   modify it under the terms of the GNU Lesser General Public
+   License as published by the Free Software Foundation; either
+   version 2.1 of the License, or (at your option) any later version.
+
+   The GNU C Library is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+   Lesser General Public License for more details.
+
+   You should have received a copy of the GNU Lesser General Public
+   License along with the GNU C Library; if not, see
+   <https://www.gnu.org/licenses/>.  */
+
+#include <support/check.h>
+#include <support/support.h>
+#include <support/xsignal.h>
+#include <support/test-driver.h>
+#include <array_length.h>
+
+#include "tst-mte-helper.h"
+
+/* Characteristic malloc sizes to cover various allocation methods.  */
+size_t sizes[] = {
+  1,
+  16, 40, 64, 120,
+  128, 500, 1000,
+  1050, 4096, 5000, 65000,
+  131072, 2000000
+};
+
+static void check_realloc (size_t len)
+{
+  /* Tagged pointers.  */
+  void *tm, *new_tm;
+
+  printf ("testing realloc (NULL) for req size %zu\n", len);
+  tm = realloc (NULL, len);
+  check_tags (tm);
+
+  /* Reduce size.  */
+  printf ("testing realloc (decreased size) for req size %zu\n", len);
+  new_tm = realloc (tm, len / 2 + 1);
+  check_tags (new_tm);
+
+  /* Increase size.  */
+  printf ("testing realloc (increased size) for req size %zu\n", len);
+  new_tm = realloc (new_tm, len + 2);
+  check_tags (new_tm);
+
+  free (new_tm);
+}
+
+static int
+do_test (void)
+{
+
+  /* Check if MTE is supported, configured and enabled.  */
+  check_mte_enabled ();
+
+  array_foreach_const (plen, sizes)
+    check_realloc (*plen);
+
+  return 0;
+}
+
+#include <support/test-driver.c>
+