arm64: Manage the MTE state like pointer auth

Add the same group of functions we use to manage pointer authentication
in userspace threads.

Sponsored by:	Arm Ltd
Differential Revision:	https://reviews.freebsd.org/D55951

sys/arm64/arm64/exec_machdep.c

@@ -471,6 +471,7 @@ exec_setregs(struct thread *td, struct image_params *imgp, uintptr_t stack)
 
 	/* Generate new pointer authentication keys */
 	ptrauth_exec(td);
+	mte_exec(td);
 }
 
 /* Sanity check these are the same size, they will be memcpy'd to and from */

sys/arm64/arm64/machdep.c

@@ -451,6 +451,7 @@ init_proc0(void *kstack)
 	thread0.td_pcb->pcb_vfpcpu = UINT_MAX;
 	thread0.td_frame = &proc0_tf;
 	ptrauth_thread0(&thread0);
+	mte_thread0(&thread0);
 	pcpup->pc_curpcb = thread0.td_pcb;
 
 	/*

sys/arm64/arm64/mte.c

@@ -0,0 +1,112 @@
+/*-
+ * SPDX-License-Identifier: BSD-2-Clause
+ *
+ * Copyright (c) 2024-2026 Arm Ltd
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+ * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+ * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+ * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include <sys/param.h>
+#include <sys/kernel.h>
+#include <sys/libkern.h>
+#include <sys/proc.h>
+
+#include <machine/cpu_feat.h>
+#include <machine/pcb.h>
+#include <machine/pte.h>
+#include <machine/sysarch.h>
+#include <vm/vm.h>
+#include <vm/vm_page.h>
+
+/* Version of MTE implemented. 0 == unimplemented */
+static u_int __read_mostly mte_version = 0;
+
+/*
+ * FEAT_MTE (mte_version == 1) has userspace instructions, but no tag
+ * checking. May of the registers/fields need FEAT_MTE2 to be implemented
+ * before we can access them.
+ */
+#define	MTE_HAS_TAG_CHECK	(mte_version >= 2)
+
+struct thread *mte_switch(struct thread *);
+
+static void
+mte_update_sctlr(struct thread *td, uint64_t sctlr)
+{
+	MPASS((sctlr & ~(SCTLR_ATA0 | SCTLR_TCF0_MASK)) == 0);
+	td->td_md.md_sctlr &= ~(SCTLR_ATA0 | SCTLR_TCF0_MASK);
+	td->td_md.md_sctlr |= sctlr;
+}
+
+void
+mte_fork(struct thread *new_td, struct thread *orig_td)
+{
+	if (!MTE_HAS_TAG_CHECK)
+		return;
+
+	mte_update_sctlr(new_td,
+	    orig_td->td_md.md_sctlr & SCTLR_TCF0_MASK);
+	new_td->td_md.md_gcr = orig_td->td_md.md_gcr;
+}
+
+void
+mte_exec(struct thread *td)
+{
+	if (!MTE_HAS_TAG_CHECK)
+		return;
+
+	mte_update_sctlr(td, SCTLR_TCF0_NONE);
+	td->td_md.md_gcr = GCR_RRND;
+}
+
+void
+mte_copy_thread(struct thread *new_td, struct thread *orig_td)
+{
+	if (!MTE_HAS_TAG_CHECK)
+		return;
+
+	mte_update_sctlr(new_td,
+	    orig_td->td_md.md_sctlr & SCTLR_TCF0_MASK);
+	new_td->td_md.md_gcr = orig_td->td_md.md_gcr;
+}
+
+/* Only for kernel threads */
+void
+mte_thread_alloc(struct thread *td)
+{
+}
+
+/* Only for a kernel thread */
+void
+mte_thread0(struct thread *td)
+{
+}
+
+
+struct thread *
+mte_switch(struct thread *td)
+{
+	if (MTE_HAS_TAG_CHECK) {
+		WRITE_SPECIALREG(GCR_EL1_REG, td->td_md.md_gcr);
+	}
+	return (td);
+}

sys/arm64/arm64/swtch.S

@@ -96,8 +96,9 @@ ENTRY(cpu_throw)
 	mov	x0, x1
 #endif
 
-	/* This returns the thread pointer so no need to save it */
+	/* These return the thread pointer so no need to save it */
 	bl	ptrauth_switch
+	bl	mte_switch
 #ifdef PERTHREAD_SSP
 	mov	x19, x0
 #endif
@@ -176,8 +177,9 @@ ENTRY(cpu_switch)
 	mov	x0, x1
 #endif
 
-	/* This returns the thread pointer so no need to save it */
+	/* These return the thread pointer so no need to save it */
 	bl	ptrauth_switch
+	bl	mte_switch
 	/* This returns the thread pcb */
 	bl	pmap_switch
 	/* Move the new pcb out of the way */

sys/arm64/arm64/vm_machdep.c

@@ -98,6 +98,7 @@ cpu_fork(struct thread *td1, struct proc *p2, struct thread *td2, int flags)
 	bzero(&pcb2->pcb_dbg_regs, sizeof(pcb2->pcb_dbg_regs));
 
 	ptrauth_fork(td2, td1);
+	mte_fork(td2, td1);
 
 	tf = td2->td_frame;
 	bcopy(td1->td_frame, tf, sizeof(*tf));
@@ -200,6 +201,7 @@ cpu_copy_thread(struct thread *td, struct thread *td0)
 
 	/* Generate new pointer authentication keys. */
 	ptrauth_copy_thread(td, td0);
+	mte_copy_thread(td, td0);
 }
 
 /*
@@ -264,6 +266,7 @@ cpu_thread_alloc(struct thread *td)
 {
 	td->td_pcb = uma_zalloc(pcb_zone, M_WAITOK);
 	ptrauth_thread_alloc(td);
+	mte_thread_alloc(td);
 }
 
 void

sys/arm64/include/cpu.h

@@ -277,6 +277,13 @@ void	ptrauth_thread0(struct thread *);
 void	ptrauth_mp_start(uint64_t);
 #endif
 
+/* Memory Tagging Extension (MTE) support */
+void	mte_fork(struct thread *, struct thread *);
+void	mte_exec(struct thread *);
+void	mte_copy_thread(struct thread *, struct thread *);
+void	mte_thread_alloc(struct thread *);
+void	mte_thread0(struct thread *);
+
 /* Functions to read the sanitised view of the special registers */
 void	update_special_regs(u_int);
 void	update_special_reg_iss(u_int, uint64_t, uint64_t);

sys/arm64/include/proc.h

@@ -71,7 +71,7 @@ struct mdthread {
 
 	int md_reserved0;
 	uint64_t md_sctlr;
-	uint64_t md_reserved1;
+	uint64_t md_gcr;		/* FEAT_MTE: Tag Control Register */
 };
 
 struct mdproc {

sys/conf/files.arm64

@@ -76,6 +76,7 @@ arm64/arm64/ofw_machdep.c			optional fdt
 arm64/arm64/pl031_rtc.c				optional fdt pl031
 arm64/arm64/ptrauth.c				standard \
 	compile-with	"${NORMAL_C:N-mbranch-protection*} -mbranch-protection=bti"
+arm64/arm64/mte.c				standard
 arm64/arm64/pmap.c				standard
 arm64/arm64/ptrace_machdep.c			standard
 arm64/arm64/rsi.c				standard