[PATCH 02/24] kselftest/arm64: Test lazy PSTATE.TCO handling

["Pierre-Clément Tosi" <ptosi@google.com>]

Reproduce a bug introduced by commit 38ddf7dafaea ("arm64: mte: avoid
clearing PSTATE.TCO on entry unless necessary") where TCO wasn't cleared
when it should have, by detecting that the kernel misses MTE faults when
accessing (tagged) memory of a process running sync MTE.

Cc: Peter Collingbourne <peter@pcc.me.uk>
Cc: Shuah Khan <shuah@kernel.org>
Cc: linux-kselftest@vger.kernel.org
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 tools/testing/selftests/arm64/mte/.gitignore       |   1 +
 tools/testing/selftests/arm64/mte/check_lazy_tco.c | 152 +++++++++++++++++++++
 2 files changed, 153 insertions(+)

diff --git a/tools/testing/selftests/arm64/mte/.gitignore b/tools/testing/selftests/arm64/mte/.gitignore
index f6937f890039..70ebc3883a26 100644
--- a/tools/testing/selftests/arm64/mte/.gitignore
+++ b/tools/testing/selftests/arm64/mte/.gitignore
@@ -4,6 +4,7 @@ check_tags_inclusion
 check_child_memory
 check_mmap_options
 check_prctl
+check_lazy_tco
 check_ksm_options
 check_user_mem
 check_hugetlb_options
diff --git a/tools/testing/selftests/arm64/mte/check_lazy_tco.c b/tools/testing/selftests/arm64/mte/check_lazy_tco.c
new file mode 100644
index 000000000000..28174495056a
--- /dev/null
+++ b/tools/testing/selftests/arm64/mte/check_lazy_tco.c
@@ -0,0 +1,152 @@
+// SPDX-License-Identifier: GPL-2.0
+// Copyright (C) 2026 Google LLC
+/*
+ * Validate that the kernel does not leave PSTATE.TCO set on kernel entry and/or
+ * thread switching, when it should be clear i.e. when kernel accesses (such as
+ * uaccess) are expected to be tag checked.
+ */
+
+#include <sys/wait.h>
+
+#include "kselftest.h"
+#include "mte_common_util.h"
+#include "mte_def.h"
+
+#define BUG_EXIT_CODE	99
+#define MAX_TRIALS	32
+
+static void handle_signal(int sig);
+
+/*
+ * We validate that uaccesses are tag checked by tricking the kernel into
+ * triggering a tag check fault while accessing user (tagged) memory. If the
+ * kernel does not report an error, we assume that PSTATE.TCO was set, silencing
+ * the hardware exception.
+ *
+ * To achieve this with robust userspace-kernel synchronization, we allocate a
+ * signal handler stack in PROT_MTE memory, tagged with 0x0 so that our signal
+ * handler can safely run without triggering tag check violations from SP
+ * (untagged) accesses. We only re-tag (with a non-zero value) the uppermost
+ * granule of the stack, where the kernel must write user->next_frame.
+ *
+ * When the parent process sends SIGUSR1 to its child, the following happens:
+ * - TCO should be cleared before the uaccess
+ * - the kernel writes user->next_frame to the tagged stack granule
+ * - (if MTE is properly enabled, a synchronous tag fault is raised)
+ * - if the fault was detected and handled, SIGSEGV should be delivered
+ * - if it was not detected or raised, the initial SIGUSR1 is received
+ */
+static void setup_sigaltstack_tripwire(int sig)
+{
+	size_t sz = MT_ALIGN_UP((size_t)MINSIGSTKSZ + 2 * (size_t)getpagesize());
+	void *altstack, *top_granule;
+	struct sigaction sa = {};
+	stack_t ss = {};
+
+	/* Enable non-zero IRG, needed for tagging the sigaltstack. */
+	if (mte_switch_mode(MTE_NONE_ERR, MTE_ALLOW_NON_ZERO_TAG, false) != 0)
+		ksft_exit_fail_msg("Failed to enable MTE tag generation\n");
+
+	altstack = mte_allocate_memory(sz, USE_MMAP, 0, false);
+	if (check_allocated_memory(altstack, sz, USE_MMAP, false) != KSFT_PASS)
+		ksft_exit_fail_msg("Failed to allocate MTE sigaltstack\n");
+
+	top_granule = (char *)altstack + sz - MT_GRANULE_SIZE;
+	if (!mte_insert_tags(top_granule, MT_GRANULE_SIZE))
+		ksft_exit_fail_msg("Failed to tag sigaltstack top granule\n");
+
+	ss.ss_sp = altstack;
+	ss.ss_size = sz;
+	if (sigaltstack(&ss, NULL) != 0)
+		ksft_exit_fail_perror("sigaltstack()");
+
+	sa.sa_handler = handle_signal;
+	sa.sa_flags = SA_ONSTACK;
+	sigemptyset(&sa.sa_mask);
+	if (sigaction(sig, &sa, NULL) != 0)
+		ksft_exit_fail_perror("sigaction()");
+}
+
+static __noreturn void stress_test_lazy_tco(void)
+{
+	const unsigned long sync_ctrl =
+		PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC |
+		(MTE_ALLOW_NON_ZERO_TAG << PR_MTE_TAG_SHIFT);
+
+	if (prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0) != 0)
+		_exit(KSFT_FAIL);
+
+	for (;;) {
+		if (prctl(PR_SET_TAGGED_ADDR_CTRL, sync_ctrl, 0, 0, 0) != 0)
+			_exit(KSFT_FAIL);
+		if (prctl(PR_SET_TAGGED_ADDR_CTRL, 0, 0, 0, 0) != 0)
+			_exit(KSFT_FAIL);
+	}
+}
+
+static void handle_signal(int sig __always_unused)
+{
+	if (prctl(PR_GET_TAGGED_ADDR_CTRL, 0, 0, 0, 0) & PR_MTE_TCF_SYNC)
+		_exit(BUG_EXIT_CODE);
+}
+
+static int spam_signal_until_child_exit(pid_t pid, int sig)
+{
+	int status = 0;
+	pid_t ret;
+
+	do {
+		if (kill(pid, sig) != 0)
+			ksft_exit_fail_perror("kill()");
+		mte_wait_after_trig();
+		ret = waitpid(pid, &status, WNOHANG);
+	} while (ret == 0);
+	if (ret != pid)
+		ksft_exit_fail_perror("waitpid()");
+
+	if (WIFSIGNALED(status) && WTERMSIG(status) == SIGSEGV)
+		return KSFT_PASS;
+
+	if (WIFEXITED(status) && WEXITSTATUS(status) == BUG_EXIT_CODE)
+		return KSFT_FAIL;
+
+	ksft_exit_fail_msg("Unexpected child wait status %#x\n", status);
+}
+
+static int check_lazy_tco(void)
+{
+	for (int trial = 0; trial < MAX_TRIALS; trial++) {
+		pid_t child = fork();
+
+		if (child < 0)
+			ksft_exit_fail_perror("fork()");
+
+		if (!child)
+			stress_test_lazy_tco();
+
+		if (spam_signal_until_child_exit(child, SIGUSR1) == KSFT_PASS)
+			continue;
+
+		ksft_print_msg("Detected invalid TCO state\n");
+		return KSFT_FAIL;
+	}
+
+	return KSFT_PASS;
+}
+
+int main(void)
+{
+	int err;
+
+	ksft_print_header();
+
+	err = mte_default_setup();
+	if (err)
+		return err;
+
+	ksft_set_plan(1);
+	setup_sigaltstack_tripwire(SIGUSR1);
+	ksft_test_result_report(check_lazy_tco(), "Check lazy TCO clearing\n");
+	mte_restore_setup();
+	ksft_finished();
+}

-- 
2.56.0.455.g6c5a503b38-goog