[glibc/azanella/aarch64-memtag] aarch64: Add memory tagging support for setjmp/longjmp routines

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

commit 52746e080c07175c6fb48dbbc4e61a07af7f0927
Author: Adhemerval Zanella <[email protected]>
Date:   Thu Apr 2 08:40:50 2026 -0300

    aarch64: Add memory tagging support for setjmp/longjmp routines
    
    When AArch64 MTE is enabled for stack variables, the compiler inserts
    instructions in the function prologue to "tag" the memory allocated on
    the stack and generates a matching tagged pointer.  In the function
    epilogue, the compiler inserts instructions to "untag" that memory
    (resetting the tag to 0) before returning.
    
    longjmp breaks this symmetry (and any non-local jumps), because it
    immediately rewinds the execution context to the setjmp location.
    It skips the function epilogues of all the intermediate functions on
    the call stack.  And if those functions used stack tagging, their
    stack frames are left behind with non-zero MTE tags.
    
    longjmp is modified to manually clear the memory tags for the portion
    of the stack that is being unwound.  To keep compatibility, if
    MTE_STACK is not set (meaning either the binary nor any dependency
    requires DT_AARCH64_MEMTAG_STACK) the untag is skipped.
    
    This change also adds a performance overhead of the stack scanning and
    untagging
    
    TODO: add mode tests.

Diff:
---
 sysdeps/aarch64/__longjmp.S                        |  45 ++++++++++
 sysdeps/aarch64/rtld-global-offsets.sym            |   1 +
 sysdeps/unix/sysv/linux/aarch64/Makefile           |  12 +++
 .../unix/sysv/linux/aarch64/tst-memtag-setjmp.c    | 100 +++++++++++++++++++++
 sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h   |  54 ++++++++++-
 5 files changed, 208 insertions(+), 4 deletions(-)

diff --git a/sysdeps/aarch64/__longjmp.S b/sysdeps/aarch64/__longjmp.S
index c4c8d72e51..4740ae8d77 100644
--- a/sysdeps/aarch64/__longjmp.S
+++ b/sysdeps/aarch64/__longjmp.S
@@ -20,9 +20,18 @@
 #include <pointer_guard.h>
 #include <jmpbuf-offsets.h>
 #include <stap-probe.h>
+#include <rtld-global-offsets.h>
 
 /* __longjmp(jmpbuf, val) */
 
+/* From sysdeps/aarch64/cpu-features.h */
+#define MTE_STACK 0x8
+
+#ifdef USE_AARCH64_MEMTAG_ABI
+	.arch armv8.5-a
+        .arch_extension memtag
+#endif
+
 ENTRY (__longjmp)
 
 #if IS_IN(libc)
@@ -136,6 +145,42 @@ L(gcs_done):
 #ifdef PTR_DEMANGLE
 	PTR_DEMANGLE (x4, x4, x3)
 #endif
+
+#ifdef USE_AARCH64_MEMTAG_ABI
+	/* Skip if MTE stack tagging is not enabled.  */
+# ifdef SHARED
+        /* In libc.so.  */
+        adrp    x2, :got:_rtld_global
+	ldr	x2, [x2, :got_lo12:_rtld_global]
+	ldr	x2, [x2, GL_DL_AARCH64_MTE_OFFSET]
+# else
+	/* In libc.a (PIC or !PIC).  */
+        adrp    x2, :got:_dl_aarch64_mte
+        ldr     x2, [x2, :got_lo12:_dl_aarch64_mte]
+	ldr	x2, [x2]
+# endif
+	tbnz	x2, MTE_STACK, L(tag_done)
+
+	subps	x2, x4, sp
+	beq	L(tag_done)
+	tbz	x2, 4, L(tag_clean32)
+	stg	sp, [sp], 16
+L(tag_clean32):
+	tbz	x2, 5, L(tag_clean64)
+	st2g	sp, [sp], 32
+L(tag_clean64):
+	lsr	x2, x2, 6
+	cbz	x2, L(tag_done)
+
+	mov	x2, sp
+L(tag_loop):
+	st2g	sp, [x2], 32
+	st2g	sp, [x2], 32
+	cmp	x4, x2
+	bhi	L(tag_loop)
+L(tag_done):
+#endif
+
 	mov	sp, x4
 
 	/* longjmp_target probe takes 3 arguments, address of jump buffer
diff --git a/sysdeps/aarch64/rtld-global-offsets.sym b/sysdeps/aarch64/rtld-global-offsets.sym
index 6c0690bb95..61e1032a28 100644
--- a/sysdeps/aarch64/rtld-global-offsets.sym
+++ b/sysdeps/aarch64/rtld-global-offsets.sym
@@ -13,3 +13,4 @@ GLRO_DL_HWCAP2_OFFSET	GLRO_offsetof (dl_hwcap2)
 -- Offsets of _rtld_global in libc.so
 
 GL_DL_AARCH64_GCS_OFFSET	GL_offsetof (dl_aarch64_gcs)
+GL_DL_AARCH64_MTE_OFFSET	GL_offsetof (dl_aarch64_mte)
diff --git a/sysdeps/unix/sysv/linux/aarch64/Makefile b/sysdeps/unix/sysv/linux/aarch64/Makefile
index fe776c52e2..1a686e3994 100644
--- a/sysdeps/unix/sysv/linux/aarch64/Makefile
+++ b/sysdeps/unix/sysv/linux/aarch64/Makefile
@@ -423,6 +423,18 @@ endif # ifeq ($(aarch64-memtag-abi)$(have-test-memtag-stack),yesyes)
 
 endif # ifeq ($(subdir),misc)
 
+ifeq ($(subdir),setjmp)
+ifeq ($(aarch64-memtag-abi)$(have-test-memtag-stack),yesyes)
+tests += \
+  tst-memtag-setjmp \
+  # tests
+
+CFLAGS-tst-memtag-setjmp.c += -march=armv8.5-a+memtag -fsanitize=memtag-stack \
+	$(no-stack-protector)
+LDFLAGS-tst-memtag-setjmp += -Wl,-z,memtag-stack -Wl,-z,memtag-mode=sync
+endif # ifeq ($(aarch64-memtag-abi)$(have-test-memtag-stack),yesyes)
+endif # $subdir = setjmp
+
 ifeq ($(subdir),stdlib)
 gen-as-const-headers += ucontext_i.sym
 endif
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-memtag-setjmp.c b/sysdeps/unix/sysv/linux/aarch64/tst-memtag-setjmp.c
new file mode 100644
index 0000000000..aa22078365
--- /dev/null
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-memtag-setjmp.c
@@ -0,0 +1,100 @@
+/* Test for longjmp when we use MTE stack sanitation.
+   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 <setjmp.h>
+#include <sys/auxv.h>
+#include <stdlib.h>
+
+#include <support/check.h>
+#include <support/support.h>
+#include <support/xsignal.h>
+#include <support/test-driver.h>
+#include <support/xthread.h>
+
+#include "tst-mte-helper.h"
+
+static jmp_buf jump_buffer;
+
+/* The longjmp call is an abnormal exit, thus it lets the stack memory
+   mte dirty.  No exception is thrown here.  */
+static void
+__attribute_noinline__
+__attribute_optimization_barrier__
+foo (void)
+{
+  volatile unsigned char x[15];
+  volatile unsigned char *ptr = &x[0];
+
+  for (int i = 0; i < sizeof (x); ++i)
+    ptr[i] = 0xcc;
+
+  longjmp (jump_buffer, 1);  // simulate throw
+}
+
+/* Reads data, from a pointer of a given size.  No MTE sanitization
+   for this one, any dirty memory should trap here.  */
+static void
+__attribute_noinline__
+hwasan_read (char *p, int size)
+{
+  volatile char __attribute__ ((unused)) sink;
+  for (int i = 0; i < size; ++i)
+    sink = p[i];
+}
+
+/* Creates a local array, large enough to include any previously tag
+   dirty stack memory.  */
+static void
+__attribute_noinline__
+after_catch (void)
+{
+  char x[10000];
+  hwasan_read (&x[0], sizeof(x));
+}
+
+static void *
+__attribute_noinline__
+f (void *closure)
+{
+  if (setjmp (jump_buffer) == 0)
+    foo ();
+  else
+    after_catch ();
+
+  return NULL;
+}
+
+static int
+do_test (void)
+{
+  if (!(getauxval (AT_HWCAP2) & HWCAP2_MTE))
+    FAIL_UNSUPPORTED ("kernel or CPU does not support or enable MTE");
+
+  TEST_VERIFY_EXIT (mte_enable ());
+  TEST_VERIFY_EXIT (mte_mode () == PR_MTE_TCF_SYNC);
+
+  install_sigsegv_handler_failure ();
+
+  f (NULL);
+
+  xpthread_join (xpthread_create (NULL, f, NULL));
+
+  return 0;
+}
+
+#include <support/test-driver.c>
diff --git a/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h b/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h
index bcf324a36f..016bbe94b0 100644
--- a/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h
+++ b/sysdeps/unix/sysv/linux/aarch64/tst-mte-helper.h
@@ -19,12 +19,11 @@
 #ifndef TST_MTE_HELPER_H
 #define TST_MTE_HELPER_H
 
-#include <fcntl.h>
-#include <limits.h>
-#include <string.h>
+#include <stdlib.h>
 #include <stdbool.h>
 #include <sys/prctl.h>
-#include <stdio.h>
+#include <support/xsignal.h>
+#include <support/check.h>
 
 typedef enum
 {
@@ -54,4 +53,51 @@ mte_mode (void)
 #define __attribute_disable_mte_stack__ \
  __attribute__((no_sanitize("memtag-stack")))
 
+#define MTE_GRANULE_SIZE 16
+
+static void
+__attribute_maybe_unused__
+sigsegv_handler_expected (int signum, siginfo_t *si, void *context)
+{
+  if (si->si_signo == SIGSEGV
+      && (si->si_code == SEGV_MTESERR || si->si_code == SEGV_MTEAERR))
+    _exit (EXIT_MTESERR);
+  else
+    _exit (EXIT_FAILURE);
+}
+
+static void
+__attribute_maybe_unused__
+sigsegv_handler_failure (int signum, siginfo_t *si, void *context)
+{
+  support_record_failure ();
+  _exit (EXIT_FAILURE);
+}
+
+static inline void
+install_sigsegv_handler (void)
+{
+  {
+    struct sigaction sa = {
+      .sa_sigaction = sigsegv_handler_expected,
+      .sa_flags = SA_NODEFER | SA_SIGINFO,
+    };
+    sigemptyset (&sa.sa_mask);
+    xsigaction (SIGSEGV, &sa, NULL);
+  }
+}
+
+static inline void
+install_sigsegv_handler_failure (void)
+{
+  {
+    struct sigaction sa = {
+      .sa_sigaction = sigsegv_handler_failure,
+      .sa_flags = SA_NODEFER | SA_SIGINFO,
+    };
+    sigemptyset (&sa.sa_mask);
+    xsigaction (SIGSEGV, &sa, NULL);
+  }
+}
+
 #endif