[PATCH 00/24] arm64: Decouple MTE code from KASAN support

[PATCH 00/24] arm64: Decouple MTE code from KASAN support

[Pierre-Clément Tosi]

MTE is a hardware feature, KASAN is software infrastructure.

The direction of the dependency should be clear: KASAN uses MTE.

However, the code does not reflect this split: many mte_*() functions
are actually only intended for HW_TAGS such that there is no way to make
use of in-kernel MTE functionality without enabling KASAN. Furthermore,
there is a lot of unnecessary indirection in those functions, causing
duplication, code that is harder to follow than it should be, and making
it easier for subtle inconsistencies to hide.

Instead, this series aims to untangle the two into a set of reusable,
arch-level MTE functions and a thin wrapper around them that KASAN
needs. For non-KASAN MTE-related code, instead of attempting to reuse
KASAN functions, break out individual MTE primitives into KASAN-agnostic
functions and use them where needed by aligning in style with how other
features are supported e.g. during kernel entry/exit.

The result of this rework made some bugs and inefficiencies obvious so
I ordered the patches into 3 sets: first, bug fixes, then, functional
improvements, and finally the refactor. Specifically:

- patch 1 fixes a bug in the support for userspace MTE, with a kselftest
  reproducer in patch 2
- patch 3 fixes an unsound inline asm statement
- patches 4-7 address inefficiencies in PSTATE.TCO-guarded sections, TCO
  lazy clearing, and harden TTBR1 address categorization
- patches 9-12 rework MTE-specific code in arch-agnostic KASAN, with an
  arch-level early refactor in patch 8 to keep those minimal
- patches 13-23 then proceed with the methodical untangling of MTE
  from KASAN, leaving only a thin KASAN_HW_TAGS wrapper in <asm/kasan.h>
- patch 24 introduces a hidden Kconfig flag ARM64_MTE_KERNEL, to better
  document in the code when logic is required for EL1 MTE, as opposed to
  KASAN

This was tested on v7.4-rc1 and applies cleanly to arm64 for-next/core
at commit f46be4da2cbb ("selftests/arm64: GCS must not leak on a failed
clone3").

The following checkpatch.pl errors are consciously ignored:

- "WARNING: ... does MAINTAINERS need updating?" on patch 2 for the
  new kselftest but selftests/arm64/ is already properly maintained
- "WARNING: Macros with flow control statements should be avoided" on
  patch 5 but the new __raw_{get,put}_mem_tco() macros are intentionally
  consistent in style with the pre-existing __raw_{get,put}_mem()

while the other warnings/error are false-positives I fix in [1] and [2].

[1]: https://lore.kernel.org/r/20261107-checkpatch-asm-complex_macro-v1-1-9a7b3c69fc28@google.com
[2]: https://lore.kernel.org/r/20261107-checkpatch-file_path_changes-v1-1-8b11814560d8@google.com

---
Pierre-Clément Tosi (24):
      arm64: Fix lazy TCO clear when enabling user MTE
      kselftest/arm64: Test lazy PSTATE.TCO handling
      arm64: Fix volatility of IRG inline assembly
      arm64: Skip TCO on privileged uaccess without kMTE
      arm64: Check MTE once per fault-handling section
      arm64: Harden is_ttbr1_addr() on tagged addresses
      arm64: Skip redundant TCO clear post-kernel entry
      arm64: Introduce mte_get_reset_tfsr() helpers
      kasan: rework hw_{force=>flush}_async_tag_fault()
      kasan: rework hw_{get=>set}_random_tag()
      kasan: provide empty stub for kasan_report_async()
      kasan: drop mte-kasan.h from arch-agnostic header
      arm64: Introduce set_pstate_tco() helper macro
      arm64: Add mte_{enter,exit}_*_{kernel,user}_mode()
      arm64: Refactor and remove mte_disable_tco_entry()
      arm64: Refactor KERNEL_GCR_EL1 out of <sysreg.h>
      arm64: Introduce MTE & KASAN tag handling helpers
      arm64: Inline and remove mte_check_tfsr_*()
      arm64: Rework static branches for kernel MTE
      arm64: Move MTE insn wrappers to <asm/mte-def.h>
      arm64: Introduce mte_tag[_and_zero]_range()
      arm64: Consolidate MTE-based KASAN in mte-kasan.h
      arm64: Fold <asm/mte-kasan.h> into <asm/kasan.h>
      arm64: Introduce CONFIG_ARM64_MTE_KERNEL

 MAINTAINERS                                        |   2 +-
 arch/arm64/Kconfig                                 |   8 +
 arch/arm64/include/asm/kasan.h                     | 103 +++++++-
 arch/arm64/include/asm/memory.h                    |  14 --
 arch/arm64/include/asm/mte-def.h                   |  65 +++++
 arch/arm64/include/asm/mte-kasan.h                 | 266 ---------------------
 arch/arm64/include/asm/mte.h                       | 189 ++++++++++-----
 arch/arm64/include/asm/processor.h                 |   6 +-
 arch/arm64/include/asm/sysreg.h                    |  17 +-
 arch/arm64/include/asm/uaccess.h                   | 106 +++++---
 arch/arm64/include/asm/word-at-a-time.h            |  22 +-
 arch/arm64/kernel/entry-common.c                   |  58 ++++-
 arch/arm64/kernel/entry.S                          |   3 +-
 arch/arm64/kernel/mte.c                            |  72 +++---
 arch/arm64/lib/strlen.S                            |   4 +-
 arch/arm64/mm/copypage.c                           |   1 +
 arch/arm64/mm/fault.c                              |   2 +-
 include/linux/kasan.h                              |   8 +-
 mm/kasan/hw_tags.c                                 |   3 +-
 mm/kasan/kasan.h                                   |  14 +-
 tools/testing/selftests/arm64/mte/.gitignore       |   1 +
 tools/testing/selftests/arm64/mte/check_lazy_tco.c | 152 ++++++++++++
 22 files changed, 649 insertions(+), 467 deletions(-)
---
base-commit: 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4
change-id: 20261107-mte-kasan-split-3632e818f997

Best regards,
-- 
Pierre-Clément Tosi <ptosi@google.com>

[PATCH 01/24] arm64: Fix lazy TCO clear when enabling user MTE

[Pierre-Clément Tosi]

On kernel entry from EL0 with HW KASAN inactive and user MTE disabled or
in async mode, mte_disable_tco_entry() leaves PSTATE.TCO set to avoid an
expensive write, as it assumes that EL1 tag checks are not needed.

However, if the thread is in the process of enabling sync or asymm MTE,
set_mte_ctrl() updates SCTLR_EL1.TCF0 without clearing PSTATE.TCO! As a
result, kernel uaccesses miss tag faults until the next context switch.

Clear TCO in set_mte_ctrl() immediately after setting the task's SCTLR.

Cc: Peter Collingbourne <peter@pcc.me.uk>
Cc: stable@vger.kernel.org
Fixes: 38ddf7dafaea ("arm64: mte: avoid clearing PSTATE.TCO on entry unless necessary")
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/kernel/mte.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index 31f5c6b0510e..d1dd9b80173a 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -409,6 +409,7 @@ long set_mte_ctrl(struct task_struct *task, unsigned long arg)
 		mte_update_sctlr_user(task);
 		mte_update_gcr_excl(task);
 		update_sctlr_el1(task->thread.sctlr_user);
+		mte_disable_tco_entry(task);
 		preempt_enable();
 	}
 

-- 
2.56.0.455.g6c5a503b38-goog

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

[Pierre-Clément Tosi]

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

[PATCH 03/24] arm64: Fix volatility of IRG inline assembly

[Pierre-Clément Tosi]

Add the 'volatile' qualifier, without which both GCC and Clang treat the
function as pure, potentially collapsing multiple inlined calls into a
single one (returning the same not-so-random tag) through CSE or moving
an inlined call within a loop out of it. This is not compatible with
IRG's architecturally-defined RNG side effects.

As existing KASAN callers are either not inlined or guarded by runtime
conditions, the compiler isn't allowed to hoist the non-volatile asm.
Note that commit 85f49cae4dfc ("arm64: mte: add in-kernel MTE helpers")
originally introduced the helper as an out-of-line function (in mte.c),
so the risk then was even smaller than when defined in a header.

I'll mark this 'Fixes' since it addresses UB but won't request stable
backports, as none of the issues mentioned above were actually hit.

Fixes: 2cb34276427a ("arm64: kasan: simplify and inline MTE functions")
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index de910f1c19e1..43510341b965 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -111,7 +111,7 @@ static inline u8 mte_get_random_tag(void)
 {
 	void *addr;
 
-	asm(__MTE_PREAMBLE "irg %0, %0"
+	asm volatile(__MTE_PREAMBLE "irg %0, %0"
 		: "=r" (addr));
 
 	return mte_get_ptr_tag(addr);

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 04/24] arm64: Skip TCO on privileged uaccess without kMTE

[Pierre-Clément Tosi]

PSTATE.TCO toggling can be expensive but uaccess_*_privileged() (used in
futex()) were doing it unconditionally when built with KASAN_HW_TAGS and
running on MTE-capable CPUs, due to mte_{enable,disable}_tco(). This was
unnecessary when 'kasan=off', adding overhead despite the care taken in
mte_disable_tco_entry() to keep TCO set.

Instead, only toggle TCO if EL1 accesses risk triggering tag faults.

Name the new helpers by following the example of __uaccess_*_hw_pan().
Specifically, this favors a simpler terminology ("re-enabling MTE")
over confusing double-negations ("disabling the override").

Move the comment intended to document this out of a shared KASAN header,
to be next to the code it explains.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h | 17 -----------------
 arch/arm64/include/asm/uaccess.h   | 21 +++++++++++++++++++--
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 43510341b965..467d1d41a392 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -34,23 +34,6 @@ static inline bool system_uses_mte_async_or_asymm_mode(void)
 
 #ifdef CONFIG_ARM64_MTE
 
-/*
- * The Tag Check Flag (TCF) mode for MTE is per EL, hence TCF0
- * affects EL0 and TCF affects EL1 irrespective of which TTBR is
- * used.
- * The kernel accesses TTBR0 usually with LDTR/STTR instructions
- * when UAO is available, so these would act as EL0 accesses using
- * TCF0.
- * However futex.h code uses exclusives which would be executed as
- * EL1, this can potentially cause a tag check fault even if the
- * user disables TCF0.
- *
- * To address the problem we set the PSTATE.TCO bit in uaccess_enable()
- * and reset it in uaccess_disable().
- *
- * The Tag check override (TCO) bit disables temporarily the tag checking
- * preventing the issue.
- */
 static inline void mte_disable_tco(void)
 {
 	if (!IS_ENABLED(CONFIG_KASAN_HW_TAGS))
diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index 9f5bd9c69c24..f0924dc0bd64 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -16,6 +16,7 @@
  */
 #include <linux/bitops.h>
 #include <linux/kasan-checks.h>
+#include <linux/kasan-enabled.h>
 #include <linux/string.h>
 
 #include <asm/asm-extable.h>
@@ -132,9 +133,25 @@ static inline void __uaccess_enable_hw_pan(void)
 	asm(ALTERNATIVE("nop", SET_PSTATE_PAN(1), ARM64_HAS_PAN));
 }
 
+static inline void __uaccess_disable_mte(void)
+{
+	/*
+	 * Unlike get_user(), privileged uaccess (TTBR0) uses TCF (not TCF0)
+	 * such that a bad user address tag could trigger a kernel MTE fault.
+	 */
+	if (kasan_hw_tags_enabled())
+		asm volatile(SET_PSTATE_TCO(1));
+}
+
+static inline void __uaccess_enable_mte(void)
+{
+	if (kasan_hw_tags_enabled())
+		asm volatile(SET_PSTATE_TCO(0));
+}
+
 static inline void uaccess_disable_privileged(void)
 {
-	mte_disable_tco();
+	__uaccess_enable_mte();
 
 	if (uaccess_ttbr0_disable())
 		return;
@@ -144,7 +161,7 @@ static inline void uaccess_disable_privileged(void)
 
 static inline void uaccess_enable_privileged(void)
 {
-	mte_enable_tco();
+	__uaccess_disable_mte();
 
 	if (uaccess_ttbr0_enable())
 		return;

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 05/24] arm64: Check MTE once per fault-handling section

[Pierre-Clément Tosi]

Call system_uses_mte_async_or_asymm_mode() only once per section. This
improves performance (fewer instructions at runtime and fewer static
branches needing to be tracked and patched), correctness (no theoretical
race between the execution of the section and the static branch being
toggled, which could leave TCO in an unexpected state), and code clarity
(no 3-level indirection with redundant checks to flip one PSTATE bit).

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h      | 25 -----------
 arch/arm64/include/asm/uaccess.h        | 80 ++++++++++++++++++++-------------
 arch/arm64/include/asm/word-at-a-time.h | 22 +++++++--
 3 files changed, 67 insertions(+), 60 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 467d1d41a392..b474fca3ea56 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -48,23 +48,6 @@ static inline void mte_enable_tco(void)
 	asm volatile(ALTERNATIVE("nop", SET_PSTATE_TCO(1), ARM64_MTE));
 }
 
-/*
- * These functions disable tag checking only if in MTE async mode
- * since the sync mode generates exceptions synchronously and the
- * nofault or load_unaligned_zeropad can handle them.
- */
-static inline void __mte_disable_tco_async(void)
-{
-	if (system_uses_mte_async_or_asymm_mode())
-		mte_disable_tco();
-}
-
-static inline void __mte_enable_tco_async(void)
-{
-	if (system_uses_mte_async_or_asymm_mode())
-		mte_enable_tco();
-}
-
 /*
  * These functions are meant to be only used from KASAN runtime through
  * the arch_*() interface defined in asm/memory.h.
@@ -197,14 +180,6 @@ static inline void mte_enable_tco(void)
 {
 }
 
-static inline void __mte_disable_tco_async(void)
-{
-}
-
-static inline void __mte_enable_tco_async(void)
-{
-}
-
 static inline u8 mte_get_ptr_tag(void *ptr)
 {
 	return 0xFF;
diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index f0924dc0bd64..9e3957ee3642 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -240,6 +240,24 @@ do {										\
 	(x) = (__force __typeof__(*(ptr)))__gu_val;				\
 } while (0)
 
+/*
+ * We must not call into the scheduler between setting and clearing TCO. As `x`
+ * and `ptr` may contain blocking functions, we must evaluate these outside of
+ * the critical section.
+ */
+#define __raw_get_mem_tco(ldr, x, ptr, label, type)			\
+do {									\
+	__label__ __rgm_err;						\
+									\
+	asm volatile(SET_PSTATE_TCO(1));				\
+	__raw_get_mem(ldr, x, ptr, __rgm_err, type);			\
+	asm volatile(SET_PSTATE_TCO(0));				\
+	break;								\
+__rgm_err:								\
+	asm volatile(SET_PSTATE_TCO(0));				\
+	goto label;							\
+} while (0)
+
 /*
  * We must not call into the scheduler between uaccess_ttbr0_enable() and
  * uaccess_ttbr0_disable(). As `x` and `ptr` could contain blocking functions,
@@ -286,27 +304,18 @@ do {									\
 
 #define get_user	__get_user
 
-/*
- * We must not call into the scheduler between __mte_enable_tco_async() and
- * __mte_disable_tco_async(). As `dst` and `src` may contain blocking
- * functions, we must evaluate these outside of the critical section.
- */
+/* We handle synchronous exceptions so only use PSTATE.TCO in non-sync MTE. */
 #define __get_kernel_nofault(dst, src, type, err_label)			\
 do {									\
 	__typeof__(dst) __gkn_dst = (dst);				\
 	__typeof__(src) __gkn_src = (src);				\
-	do { 								\
-		__label__ __gkn_label;					\
 									\
-		__mte_enable_tco_async();				\
+	if (system_uses_mte_async_or_asymm_mode())			\
+		__raw_get_mem_tco("ldr", *((type *)(__gkn_dst)),	\
+		      (__force type *)(__gkn_src), err_label, K);	\
+	else								\
 		__raw_get_mem("ldr", *((type *)(__gkn_dst)),		\
-		      (__force type *)(__gkn_src), __gkn_label, K);	\
-		__mte_disable_tco_async();				\
-		break;							\
-	__gkn_label:							\
-		__mte_disable_tco_async();				\
-		goto err_label;						\
-	} while (0);							\
+		      (__force type *)(__gkn_src), err_label, K);	\
 } while (0)
 
 #define __put_mem_asm(store, reg, x, addr, label, type)			\
@@ -337,6 +346,24 @@ do {										\
 	}									\
 } while (0)
 
+/*
+ * We must not call into the scheduler between setting and clearing TCO. As `x`
+ * and `ptr` may contain blocking functions, we must evaluate these outside of
+ * the critical section.
+ */
+#define __raw_put_mem_tco(str, x, ptr, label, type)			\
+do {									\
+	__label__ __rpm_err;						\
+									\
+	asm volatile(SET_PSTATE_TCO(1));				\
+	__raw_put_mem(str, x, ptr, __rpm_err, type);			\
+	asm volatile(SET_PSTATE_TCO(0));				\
+	break;								\
+__rpm_err:								\
+	asm volatile(SET_PSTATE_TCO(0));				\
+	goto label;							\
+} while (0)
+
 /*
  * We must not call into the scheduler between uaccess_ttbr0_enable() and
  * uaccess_ttbr0_disable(). As `x` and `ptr` could contain blocking functions,
@@ -383,28 +410,19 @@ do {									\
 
 #define put_user	__put_user
 
-/*
- * We must not call into the scheduler between __mte_enable_tco_async() and
- * __mte_disable_tco_async(). As `dst` and `src` may contain blocking
- * functions, we must evaluate these outside of the critical section.
- */
+/* We handle synchronous exceptions so only use PSTATE.TCO in non-sync MTE. */
 #define __put_kernel_nofault(dst, src, type, err_label)			\
 do {									\
 	__typeof__(dst) __pkn_dst = (dst);				\
 	__typeof__(src) __pkn_src = (src);				\
 									\
-	do {								\
-		__label__ __pkn_err;					\
-		__mte_enable_tco_async();				\
+	if (system_uses_mte_async_or_asymm_mode())			\
+		__raw_put_mem_tco("str", *((type *)(__pkn_src)),	\
+			      (__force type *)(__pkn_dst), err_label, K); \
+	else								\
 		__raw_put_mem("str", *((type *)(__pkn_src)),		\
-			      (__force type *)(__pkn_dst), __pkn_err, K);	\
-		__mte_disable_tco_async();				\
-		break;							\
-	__pkn_err:							\
-		__mte_disable_tco_async();				\
-		goto err_label;						\
-	} while (0);							\
-} while(0)
+			      (__force type *)(__pkn_dst), err_label, K); \
+} while (0)
 
 extern unsigned long __must_check __arch_copy_from_user(void *to, const void __user *from, unsigned long n);
 #define raw_copy_from_user(to, from, n)					\
diff --git a/arch/arm64/include/asm/word-at-a-time.h b/arch/arm64/include/asm/word-at-a-time.h
index d7a3b69dc8c6..ec72126210d1 100644
--- a/arch/arm64/include/asm/word-at-a-time.h
+++ b/arch/arm64/include/asm/word-at-a-time.h
@@ -10,6 +10,9 @@
 #include <linux/bitops.h>
 #include <linux/wordpart.h>
 
+#include <asm/mte-kasan.h>
+#include <asm/sysreg.h>
+
 struct word_at_a_time {
 	const unsigned long one_bits, high_bits;
 };
@@ -41,12 +44,10 @@ static inline unsigned long zero_bytemask(unsigned long bits)
  * and the next page not being mapped, take the exception and
  * return zeroes in the non-existing part.
  */
-static inline unsigned long load_unaligned_zeropad(const void *addr)
+static __always_inline unsigned long __load_unaligned_zeropad(const void *addr)
 {
 	unsigned long ret;
 
-	__mte_enable_tco_async();
-
 	/* Load word from unaligned pointer addr */
 	asm(
 	"1:	ldr	%0, %2\n"
@@ -55,7 +56,20 @@ static inline unsigned long load_unaligned_zeropad(const void *addr)
 	: "=&r" (ret)
 	: "r" (addr), "Q" (*(unsigned long *)addr));
 
-	__mte_disable_tco_async();
+	return ret;
+}
+
+static inline unsigned long load_unaligned_zeropad(const void *addr)
+{
+	unsigned long ret;
+
+	/* We handle synchronous exceptions so only use TCO in non-sync MTE. */
+	if (!system_uses_mte_async_or_asymm_mode())
+		return __load_unaligned_zeropad(addr);
+
+	asm volatile(SET_PSTATE_TCO(1));
+	ret = __load_unaligned_zeropad(addr);
+	asm volatile(SET_PSTATE_TCO(0));
 
 	return ret;
 }

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 06/24] arm64: Harden is_ttbr1_addr() on tagged addresses

[Pierre-Clément Tosi]

When neither CONFIG_KASAN_{SW,HW}_TAGS is enabled, the KASAN function is
a no-op, leading to the check "addr >= PAGE_OFFSET". This is inaccurate
as CONFIG_ARM64_MTE=y sets TCR_EL1.TBI1, meaning that addresses with a
non-canonical top byte would get translated through TTBR1 while failing
the check above (PAGE_OFFSET's MSB is 0xff), making is_ttbr0_addr()
wrong on those corner cases.

Remove the KASAN dependency from <asm/processor.h>, as is_ttbr1_addr()
was its sole user and include <asm/mte-kasan.h> from <asm/mte.h> to
satisfy #includes which relied on the transitive include.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte.h       | 1 +
 arch/arm64/include/asm/processor.h | 6 +++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 83f3b05fc784..1d6113ae786f 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -7,6 +7,7 @@
 
 #include <asm/compiler.h>
 #include <asm/mte-def.h>
+#include <asm/mte-kasan.h>
 
 #ifndef __ASSEMBLER__
 
diff --git a/arch/arm64/include/asm/processor.h b/arch/arm64/include/asm/processor.h
index 481f1999f7c1..510e3a3046bf 100644
--- a/arch/arm64/include/asm/processor.h
+++ b/arch/arm64/include/asm/processor.h
@@ -39,8 +39,8 @@
 #include <asm/alternative.h>
 #include <asm/cpufeature.h>
 #include <asm/hw_breakpoint.h>
-#include <asm/kasan.h>
 #include <asm/lse.h>
+#include <asm/memory.h>
 #include <asm/pgtable-hwdef.h>
 #include <asm/pointer_auth.h>
 #include <asm/ptrace.h>
@@ -366,8 +366,8 @@ static __always_inline bool is_ttbr0_addr(unsigned long addr)
 
 static __always_inline bool is_ttbr1_addr(unsigned long addr)
 {
-	/* TTBR1 addresses may have a tag if KASAN_SW_TAGS is in use */
-	return arch_kasan_reset_tag(addr) >= PAGE_OFFSET;
+	/* TTBR1 addresses may have a tag */
+	return __untagged_addr(addr) >= PAGE_OFFSET;
 }
 
 /* Forward declaration, a strange C thing */

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 07/24] arm64: Skip redundant TCO clear post-kernel entry

[Pierre-Clément Tosi]

For performance reasons, kernel entry might skip clearing TCO. Still,
tasks expecting synchronous uaccess tag faults to be reported (and
therefore raised) require it to be clear; this is why thread switching
and MTE enabling functions were calling into mte_disable_tco_entry().
However, when the kernel itself uses MTE (for KASAN), it must have
cleared TCO on kernel entry.

Use that invariant to skip redundantly clearing TCO and/or checking the
task's SCTLR. As kasan_hw_tags_enabled() uses a static branch, TCO
maintenance is completely optimized out the moment HW KASAN is enabled.

Cc: Peter Collingbourne <peter@pcc.me.uk>
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/kernel/mte.c | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index d1dd9b80173a..8025ee1f7fdb 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -5,6 +5,7 @@
 
 #include <linux/bitops.h>
 #include <linux/cpu.h>
+#include <linux/kasan-enabled.h>
 #include <linux/kernel.h>
 #include <linux/mm.h>
 #include <linux/prctl.h>
@@ -290,7 +291,8 @@ void mte_thread_switch(struct task_struct *next)
 	mte_update_gcr_excl(next);
 
 	/* TCO may not have been disabled on exception entry for the current task. */
-	mte_disable_tco_entry(next);
+	if (!kasan_hw_tags_enabled())
+		mte_disable_tco_entry(next);
 
 	if (!system_uses_mte_async_or_asymm_mode())
 		return;
@@ -409,7 +411,8 @@ long set_mte_ctrl(struct task_struct *task, unsigned long arg)
 		mte_update_sctlr_user(task);
 		mte_update_gcr_excl(task);
 		update_sctlr_el1(task->thread.sctlr_user);
-		mte_disable_tco_entry(task);
+		if (!kasan_hw_tags_enabled())
+			mte_disable_tco_entry(task);
 		preempt_enable();
 	}
 

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 08/24] arm64: Introduce mte_get_reset_tfsr() helpers

[Pierre-Clément Tosi]

In an effort to untangle MTE from KASAN, split a reusable MTE helper out
of mte_check_tfsr_el1(), providing two variants, to accommodate the
contexts under which it gets called: one synchronising TFx before
reading it and another delegating that responsibility to the caller.

Name the helpers similarly to __dsu_pmu_get_reset_overflow().

This prepares for a rework of the KASAN API to simplify circular header
dependencies, as KASAN calls into mte_check_tfsr_el1(), which calls back
into a kasan_*() function.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte.h | 31 +++++++++++++++++++++++++++++++
 arch/arm64/kernel/mte.c      | 12 +-----------
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 1d6113ae786f..614fff7dc35f 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -17,7 +17,9 @@
 #include <linux/sched.h>
 #include <linux/types.h>
 
+#include <asm/barrier.h>
 #include <asm/pgtable-types.h>
+#include <asm/sysreg.h>
 
 void mte_clear_page_tags(void *addr);
 unsigned long mte_copy_tags_from_user(void *to, const void __user *from,
@@ -31,6 +33,35 @@ void mte_invalidate_tags_area(int type);
 void *mte_allocate_tag_storage(void);
 void mte_free_tag_storage(char *storage);
 
+static __always_inline bool __mte_get_reset_tfsr(void)
+{
+	u64 tfsr_el1 = read_sysreg_s(SYS_TFSR_EL1);
+
+	if (likely(!(tfsr_el1 & SYS_TFSR_EL1_TF1)))
+		return false;
+
+	/*
+	 * Note: isb() is not required after this direct write
+	 * because there is no indirect read subsequent to it
+	 * (per ARM DDI 0487F.c table D13-1).
+	 */
+	write_sysreg_s(0, SYS_TFSR_EL1);
+
+	return true;
+}
+
+static __always_inline bool mte_get_reset_tfsr(void)
+{
+	/*
+	 * The barriers are required to guarantee that the indirect writes
+	 * to TFSR_EL1 are synchronized before we report the state.
+	 */
+	dsb(nsh);
+	isb();
+
+	return __mte_get_reset_tfsr();
+}
+
 #ifdef CONFIG_ARM64_MTE
 
 /* track which pages have valid allocation tags */
diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index 8025ee1f7fdb..47432c9c4de0 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -182,18 +182,8 @@ int mte_enable_kernel_store_only(void)
 #ifdef CONFIG_KASAN_HW_TAGS
 void mte_check_tfsr_el1(void)
 {
-	u64 tfsr_el1 = read_sysreg_s(SYS_TFSR_EL1);
-
-	if (unlikely(tfsr_el1 & SYS_TFSR_EL1_TF1)) {
-		/*
-		 * Note: isb() is not required after this direct write
-		 * because there is no indirect read subsequent to it
-		 * (per ARM DDI 0487F.c table D13-1).
-		 */
-		write_sysreg_s(0, SYS_TFSR_EL1);
-
+	if (__mte_get_reset_tfsr())
 		kasan_report_async();
-	}
 }
 #endif
 

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 09/24] kasan: rework hw_{force=>flush}_async_tag_fault()

[Pierre-Clément Tosi]

Separate common arch-agnostic KASAN policy (fault reporting) from arch
HW details (fault flushing) by reworking the function's signature.

Update its only definition (in arm64) accordingly.

This resolves #include challenges arising from the circular dependency:

  mte_check_tfsr_exit() ==> kasan_report_async()
    kasan_report_async() ==> <linux/kasan.h>
      <linux/kasan.h> ==> hw_force_async_tag_fault()
        hw_force_async_tag_fault() ==> arch_force_async_tag_fault()
          arch_force_async_tag_fault() ==> mte_check_tfsr_exit()

Cc: Alexander Potapenko <glider@google.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Andrey Konovalov <andreyknvl@gmail.com>
Cc: Andrey Ryabinin <ryabinin.a.a@gmail.com>
Cc: Dmitry Vyukov <dvyukov@google.com>
Cc: kasan-dev@googlegroups.com
Cc: linux-mm@kvack.org
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/memory.h | 2 +-
 mm/kasan/hw_tags.c              | 3 ++-
 mm/kasan/kasan.h                | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index a2b7a33966ff..8e751b56d20c 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -316,7 +316,7 @@ static inline const void *__tag_set(const void *addr, u8 tag)
 #define arch_enable_tag_checks_write_only()	mte_enable_kernel_store_only()
 #define arch_suppress_tag_checks_start()	mte_enable_tco()
 #define arch_suppress_tag_checks_stop()		mte_disable_tco()
-#define arch_force_async_tag_fault()		mte_check_tfsr_exit()
+#define arch_flush_async_tag_fault()		mte_get_reset_tfsr()
 #define arch_get_random_tag()			mte_get_random_tag()
 #define arch_get_mem_tag(addr)			mte_get_mem_tag(addr)
 #define arch_set_mem_tag_range(addr, size, tag, init)	\
diff --git a/mm/kasan/hw_tags.c b/mm/kasan/hw_tags.c
index a848eb2f9910..c4772eebb80d 100644
--- a/mm/kasan/hw_tags.c
+++ b/mm/kasan/hw_tags.c
@@ -428,7 +428,8 @@ EXPORT_SYMBOL_IF_KUNIT(kasan_enable_hw_tags);
 
 VISIBLE_IF_KUNIT void kasan_force_async_fault(void)
 {
-	hw_force_async_tag_fault();
+	if (hw_flush_async_tag_fault())
+		kasan_report_async();
 }
 EXPORT_SYMBOL_IF_KUNIT(kasan_force_async_fault);
 
diff --git a/mm/kasan/kasan.h b/mm/kasan/kasan.h
index fc9169a54766..be5dcf6b052c 100644
--- a/mm/kasan/kasan.h
+++ b/mm/kasan/kasan.h
@@ -431,7 +431,7 @@ static inline const void *arch_kasan_set_tag(const void *addr, u8 tag)
 #define hw_enable_tag_checks_asymm()		arch_enable_tag_checks_asymm()
 #define hw_suppress_tag_checks_start()		arch_suppress_tag_checks_start()
 #define hw_suppress_tag_checks_stop()		arch_suppress_tag_checks_stop()
-#define hw_force_async_tag_fault()		arch_force_async_tag_fault()
+#define hw_flush_async_tag_fault()		arch_flush_async_tag_fault()
 #define hw_enable_tag_checks_write_only()	arch_enable_tag_checks_write_only()
 #define hw_get_random_tag()			arch_get_random_tag()
 #define hw_get_mem_tag(addr)			arch_get_mem_tag(addr)

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 10/24] kasan: rework hw_{get=>set}_random_tag()

[Pierre-Clément Tosi]

With MTE, the IRG instruction tags a pointer in-place, so that

    set_tag(ptr, hw_get_random_tag())

can conceptually be synthesized as a single "IRG %[ptr], %[ptr], 0".

However, with the way the API is currently defined, we confuse the
compiler into generating a random tag with invalid address bits, to then
bitfield-extract the tag and shift-and-mask it into another pointer.

Furthermore, every single (indirect or direct) use-case for a random tag
in arch-agnostic KASAN ends up passing it to set_tag(), which likely is
an indicator of a suboptimal API design.

Therefore, make the function clearer and more optimization-friendly by
directly taking the pointer to be randomly tagged.

Update the only definition (in arm64) accordingly.

Intentionally keep this patch minimal (and limited to KASAN_HW_TAGS) to
reduce its associated risks and make it as uncontentious as possible. A
future patch will refactor the arch-agnostic KASAN code to enable the
optimization mentioned above.

Cc: Alexander Potapenko <glider@google.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Andrey Konovalov <andreyknvl@gmail.com>
Cc: Andrey Ryabinin <ryabinin.a.a@gmail.com>
Cc: Dmitry Vyukov <dvyukov@google.com>
Cc: kasan-dev@googlegroups.com
Cc: linux-mm@kvack.org
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/memory.h    |  2 +-
 arch/arm64/include/asm/mte-kasan.h | 14 +++++++-------
 mm/kasan/kasan.h                   |  7 +++++--
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index 8e751b56d20c..29e7537f5e58 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -317,7 +317,7 @@ static inline const void *__tag_set(const void *addr, u8 tag)
 #define arch_suppress_tag_checks_start()	mte_enable_tco()
 #define arch_suppress_tag_checks_stop()		mte_disable_tco()
 #define arch_flush_async_tag_fault()		mte_get_reset_tfsr()
-#define arch_get_random_tag()			mte_get_random_tag()
+#define arch_set_random_tag(addr)		mte_set_random_tag(addr)
 #define arch_get_mem_tag(addr)			mte_get_mem_tag(addr)
 #define arch_set_mem_tag_range(addr, size, tag, init)	\
 			mte_set_mem_tag_range((addr), (size), (tag), (init))
diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index b474fca3ea56..17e4155781fc 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -73,14 +73,14 @@ static inline u8 mte_get_mem_tag(void *addr)
 }
 
 /* Generate a random tag. */
-static inline u8 mte_get_random_tag(void)
+static inline void *mte_set_random_tag(const void *addr)
 {
-	void *addr;
-
+	/* Writes the random MTE tag in bits [59:56]. */
 	asm volatile(__MTE_PREAMBLE "irg %0, %0"
-		: "=r" (addr));
+		: "+r" (addr));
 
-	return mte_get_ptr_tag(addr);
+	/* Sets bits [63:60] as KASAN expects. */
+	return (void *)__tag_set(addr, mte_get_ptr_tag((void *)addr));
 }
 
 static inline u64 __stg_post(u64 p)
@@ -190,9 +190,9 @@ static inline u8 mte_get_mem_tag(void *addr)
 	return 0xFF;
 }
 
-static inline u8 mte_get_random_tag(void)
+static inline void *mte_set_random_tag(const void *addr)
 {
-	return 0xFF;
+	return (void *)__tag_set(addr, 0xFF);
 }
 
 static inline void mte_set_mem_tag_range(void *addr, size_t size,
diff --git a/mm/kasan/kasan.h b/mm/kasan/kasan.h
index be5dcf6b052c..0df7daee1864 100644
--- a/mm/kasan/kasan.h
+++ b/mm/kasan/kasan.h
@@ -433,7 +433,7 @@ static inline const void *arch_kasan_set_tag(const void *addr, u8 tag)
 #define hw_suppress_tag_checks_stop()		arch_suppress_tag_checks_stop()
 #define hw_flush_async_tag_fault()		arch_flush_async_tag_fault()
 #define hw_enable_tag_checks_write_only()	arch_enable_tag_checks_write_only()
-#define hw_get_random_tag()			arch_get_random_tag()
+#define hw_set_random_tag(addr)			arch_set_random_tag(addr)
 #define hw_get_mem_tag(addr)			arch_get_mem_tag(addr)
 #define hw_set_mem_tag_range(addr, size, tag, init) \
 			arch_set_mem_tag_range((addr), (size), (tag), (init))
@@ -469,7 +469,10 @@ static inline bool kasan_write_only_enabled(void)
 #ifdef CONFIG_KASAN_SW_TAGS
 u8 kasan_random_tag(void);
 #elif defined(CONFIG_KASAN_HW_TAGS)
-static inline u8 kasan_random_tag(void) { return hw_get_random_tag(); }
+static inline u8 kasan_random_tag(void)
+{
+	return get_tag(hw_set_random_tag(NULL));
+}
 #else
 static inline u8 kasan_random_tag(void) { return 0; }
 #endif

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 11/24] kasan: provide empty stub for kasan_report_async()

[Pierre-Clément Tosi]

Allow callers to compile without #ifdef guards.

Cc: Alexander Potapenko <glider@google.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Andrey Konovalov <andreyknvl@gmail.com>
Cc: Andrey Ryabinin <ryabinin.a.a@gmail.com>
Cc: Dmitry Vyukov <dvyukov@google.com>
Cc: kasan-dev@googlegroups.com
Cc: linux-mm@kvack.org
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 include/linux/kasan.h | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/include/linux/kasan.h b/include/linux/kasan.h
index bf233bde68c7..e429fb483176 100644
--- a/include/linux/kasan.h
+++ b/include/linux/kasan.h
@@ -541,12 +541,6 @@ static inline void *kasan_reset_tag(const void *addr)
 
 #endif /* CONFIG_KASAN_SW_TAGS || CONFIG_KASAN_HW_TAGS*/
 
-#ifdef CONFIG_KASAN_HW_TAGS
-
-void kasan_report_async(void);
-
-#endif /* CONFIG_KASAN_HW_TAGS */
-
 #ifdef CONFIG_KASAN_GENERIC
 void __init kasan_init_generic(void);
 #else
@@ -562,9 +556,11 @@ static inline void kasan_init_sw_tags(void) { }
 #ifdef CONFIG_KASAN_HW_TAGS
 void kasan_init_hw_tags_cpu(void);
 void __init kasan_init_hw_tags(void);
+void kasan_report_async(void);
 #else
 static inline void kasan_init_hw_tags_cpu(void) { }
 static inline void kasan_init_hw_tags(void) { }
+static inline void kasan_report_async(void) { }
 #endif
 
 #ifdef CONFIG_KASAN_VMALLOC

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 12/24] kasan: drop mte-kasan.h from arch-agnostic header

[Pierre-Clément Tosi]

The inclusion of <asm/mte-kasan.h> (only defined by arm64) is based on
the brittle assumption that KASAN_HW_TAGS is equivalent to MTE, which
will break as soon as another architecture supports it. Furthermore,
exposing the header's contents to common KASAN_HW_TAGS code risks
spreading this wrong assumption further.

Allow KASAN_GRANULE_SIZE to be arch-defined, with a fallback.

Define KASAN_GRANULE_SIZE in arm64's KASAN header when KASAN_HW_TAGS.

Remove the #include, which is now unused as a result.

Cc: Alexander Potapenko <glider@google.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Andrey Konovalov <andreyknvl@gmail.com>
Cc: Andrey Ryabinin <ryabinin.a.a@gmail.com>
Cc: Dmitry Vyukov <dvyukov@google.com>
Cc: kasan-dev@googlegroups.com
Cc: linux-mm@kvack.org
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/kasan.h | 5 +++++
 mm/kasan/kasan.h               | 5 +----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/arch/arm64/include/asm/kasan.h b/arch/arm64/include/asm/kasan.h
index b167e9d3da91..3106a2271781 100644
--- a/arch/arm64/include/asm/kasan.h
+++ b/arch/arm64/include/asm/kasan.h
@@ -6,6 +6,7 @@
 
 #include <linux/linkage.h>
 #include <asm/memory.h>
+#include <asm/mte-def.h>
 #include <asm/mte-kasan.h>
 #include <asm/pgtable-types.h>
 
@@ -22,5 +23,9 @@ void kasan_init(void);
 static inline void kasan_init(void) { }
 #endif
 
+#ifdef CONFIG_KASAN_HW_TAGS
+#define KASAN_GRANULE_SIZE	MTE_GRANULE_SIZE
+#endif /* CONFIG_KASAN_HW_TAGS */
+
 #endif
 #endif
diff --git a/mm/kasan/kasan.h b/mm/kasan/kasan.h
index 0df7daee1864..52c2e7dca739 100644
--- a/mm/kasan/kasan.h
+++ b/mm/kasan/kasan.h
@@ -128,11 +128,8 @@ static inline bool kasan_requires_meta(void)
 
 #endif /* CONFIG_KASAN_GENERIC */
 
-#if defined(CONFIG_KASAN_GENERIC) || defined(CONFIG_KASAN_SW_TAGS)
+#ifndef KASAN_GRANULE_SIZE
 #define KASAN_GRANULE_SIZE	(1UL << KASAN_SHADOW_SCALE_SHIFT)
-#else
-#include <asm/mte-kasan.h>
-#define KASAN_GRANULE_SIZE	MTE_GRANULE_SIZE
 #endif
 
 #define KASAN_GRANULE_MASK	(KASAN_GRANULE_SIZE - 1)

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 13/24] arm64: Introduce set_pstate_tco() helper macro

[Pierre-Clément Tosi]

Like other PSTATE bits, TCO deserves its own accessor.

Note that only HW KASAN common code now calls mte_{enable,disable}_tco()
so that their KASAN_HW_TAGS and ALTERNATIVE(ARM64_MTE) are superfluous.

Move the comment, as it now applies to them.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h      | 24 +++++++++++-------------
 arch/arm64/include/asm/mte.h            |  2 +-
 arch/arm64/include/asm/sysreg.h         |  1 +
 arch/arm64/include/asm/uaccess.h        | 16 ++++++++--------
 arch/arm64/include/asm/word-at-a-time.h |  4 ++--
 5 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 17e4155781fc..9d00e9460b0f 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -13,6 +13,8 @@
 
 #include <linux/types.h>
 
+#include <asm/sysreg.h>
+
 #ifdef CONFIG_KASAN_HW_TAGS
 
 /* Whether the MTE asynchronous mode is enabled. */
@@ -34,27 +36,23 @@ static inline bool system_uses_mte_async_or_asymm_mode(void)
 
 #ifdef CONFIG_ARM64_MTE
 
+/*
+ * These functions are meant to be only used from KASAN runtime through
+ * the arch_*() interface defined in asm/memory.h.
+ * These functions don't include system_supports_mte() checks,
+ * as KASAN only calls them when MTE is supported and enabled.
+ */
+
 static inline void mte_disable_tco(void)
 {
-	if (!IS_ENABLED(CONFIG_KASAN_HW_TAGS))
-		return;
-	asm volatile(ALTERNATIVE("nop", SET_PSTATE_TCO(0), ARM64_MTE));
+	set_pstate_tco(0);
 }
 
 static inline void mte_enable_tco(void)
 {
-	if (!IS_ENABLED(CONFIG_KASAN_HW_TAGS))
-		return;
-	asm volatile(ALTERNATIVE("nop", SET_PSTATE_TCO(1), ARM64_MTE));
+	set_pstate_tco(1);
 }
 
-/*
- * These functions are meant to be only used from KASAN runtime through
- * the arch_*() interface defined in asm/memory.h.
- * These functions don't include system_supports_mte() checks,
- * as KASAN only calls them when MTE is supported and enabled.
- */
-
 static inline u8 mte_get_ptr_tag(void *ptr)
 {
 	/* Note: The format of KASAN tags is 0xF<x> */
diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 614fff7dc35f..815843079eb9 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -271,7 +271,7 @@ static inline void mte_disable_tco_entry(struct task_struct *task)
 	 */
 	if (kasan_hw_tags_enabled() ||
 	    (task->thread.sctlr_user & (1UL << SCTLR_EL1_TCF0_SHIFT)))
-		asm volatile(SET_PSTATE_TCO(0));
+		set_pstate_tco(0);
 }
 
 #ifdef CONFIG_KASAN_HW_TAGS
diff --git a/arch/arm64/include/asm/sysreg.h b/arch/arm64/include/asm/sysreg.h
index e3b60ee03023..51c6a54d7146 100644
--- a/arch/arm64/include/asm/sysreg.h
+++ b/arch/arm64/include/asm/sysreg.h
@@ -105,6 +105,7 @@
 #define set_pstate_uao(x)		asm volatile(SET_PSTATE_UAO(x))
 #define set_pstate_ssbs(x)		asm volatile(SET_PSTATE_SSBS(x))
 #define set_pstate_dit(x)		asm volatile(SET_PSTATE_DIT(x))
+#define set_pstate_tco(x)		asm volatile(SET_PSTATE_TCO(x))
 
 /* Register-based PAN access, for save/restore purposes */
 #define SYS_PSTATE_PAN			sys_reg(3, 0, 4, 2, 3)
diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index 9e3957ee3642..aac7676da531 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -140,13 +140,13 @@ static inline void __uaccess_disable_mte(void)
 	 * such that a bad user address tag could trigger a kernel MTE fault.
 	 */
 	if (kasan_hw_tags_enabled())
-		asm volatile(SET_PSTATE_TCO(1));
+		set_pstate_tco(1);
 }
 
 static inline void __uaccess_enable_mte(void)
 {
 	if (kasan_hw_tags_enabled())
-		asm volatile(SET_PSTATE_TCO(0));
+		set_pstate_tco(0);
 }
 
 static inline void uaccess_disable_privileged(void)
@@ -249,12 +249,12 @@ do {										\
 do {									\
 	__label__ __rgm_err;						\
 									\
-	asm volatile(SET_PSTATE_TCO(1));				\
+	set_pstate_tco(1);						\
 	__raw_get_mem(ldr, x, ptr, __rgm_err, type);			\
-	asm volatile(SET_PSTATE_TCO(0));				\
+	set_pstate_tco(0);						\
 	break;								\
 __rgm_err:								\
-	asm volatile(SET_PSTATE_TCO(0));				\
+	set_pstate_tco(0);						\
 	goto label;							\
 } while (0)
 
@@ -355,12 +355,12 @@ do {										\
 do {									\
 	__label__ __rpm_err;						\
 									\
-	asm volatile(SET_PSTATE_TCO(1));				\
+	set_pstate_tco(1);						\
 	__raw_put_mem(str, x, ptr, __rpm_err, type);			\
-	asm volatile(SET_PSTATE_TCO(0));				\
+	set_pstate_tco(0);						\
 	break;								\
 __rpm_err:								\
-	asm volatile(SET_PSTATE_TCO(0));				\
+	set_pstate_tco(0);						\
 	goto label;							\
 } while (0)
 
diff --git a/arch/arm64/include/asm/word-at-a-time.h b/arch/arm64/include/asm/word-at-a-time.h
index ec72126210d1..430d5ab8f042 100644
--- a/arch/arm64/include/asm/word-at-a-time.h
+++ b/arch/arm64/include/asm/word-at-a-time.h
@@ -67,9 +67,9 @@ static inline unsigned long load_unaligned_zeropad(const void *addr)
 	if (!system_uses_mte_async_or_asymm_mode())
 		return __load_unaligned_zeropad(addr);
 
-	asm volatile(SET_PSTATE_TCO(1));
+	set_pstate_tco(1);
 	ret = __load_unaligned_zeropad(addr);
-	asm volatile(SET_PSTATE_TCO(0));
+	set_pstate_tco(0);
 
 	return ret;
 }

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 14/24] arm64: Add mte_{enter,exit}_*_{kernel,user}_mode()

[Pierre-Clément Tosi]

Abstract away and consolidate MTE switching code by aligning with the
concepts and naming from entry-common.c, similarly to the SME routines.

Define the helpers in entry-common.c, conveniently grouped for more
straightforward enter-exit logic pair evaluation.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/kernel/entry-common.c | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/arch/arm64/kernel/entry-common.c b/arch/arm64/kernel/entry-common.c
index 11d3f888fcd8..fbf9bafc6962 100644
--- a/arch/arm64/kernel/entry-common.c
+++ b/arch/arm64/kernel/entry-common.c
@@ -25,6 +25,7 @@
 #include <asm/irq_regs.h>
 #include <asm/kprobes.h>
 #include <asm/mmu.h>
+#include <asm/mte.h>
 #include <asm/percpu.h>
 #include <asm/processor.h>
 #include <asm/sdei.h>
@@ -68,6 +69,27 @@ static __always_inline void irqentry_exit_pcpu_adjust(struct pt_regs *regs)
 	pt_regs_write_reg(regs, reg_addr, addr);
 }
 
+static __always_inline void mte_enter_from_kernel_mode(void)
+{
+	mte_check_tfsr_entry();
+	mte_disable_tco_entry(current);
+}
+
+static __always_inline void mte_exit_to_kernel_mode(void)
+{
+	mte_check_tfsr_exit();
+}
+
+static __always_inline void mte_enter_from_user_mode(void)
+{
+	mte_disable_tco_entry(current);
+}
+
+static __always_inline void mte_exit_to_user_mode(void)
+{
+	mte_check_tfsr_exit();
+}
+
 /*
  * Handle IRQ/context state management when entering from kernel mode.
  * Before this function is called it is not safe to call regular kernel code,
@@ -78,8 +100,7 @@ static noinstr irqentry_state_t arm64_enter_from_kernel_mode(struct pt_regs *reg
 	irqentry_state_t state;
 
 	state = irqentry_enter_from_kernel_mode(regs);
-	mte_check_tfsr_entry();
-	mte_disable_tco_entry(current);
+	mte_enter_from_kernel_mode();
 
 	return state;
 }
@@ -94,7 +115,7 @@ static void noinstr __arm64_exit_to_kernel_mode(struct pt_regs *regs,
 {
 	local_daif_mask();
 	irqentry_exit_pcpu_adjust(regs);
-	mte_check_tfsr_exit();
+	mte_exit_to_kernel_mode();
 	irqentry_exit_to_kernel_mode_after_preempt(regs, state);
 }
 
@@ -123,7 +144,7 @@ static void noinstr arm64_exit_to_kernel_mode(struct pt_regs *regs,
 static __always_inline void arm64_syscall_enter_from_user_mode(struct pt_regs *regs)
 {
 	enter_from_user_mode(regs);
-	mte_disable_tco_entry(current);
+	mte_enter_from_user_mode();
 	sme_enter_from_user_mode();
 }
 
@@ -136,7 +157,7 @@ static __always_inline void arm64_enter_from_user_mode(struct pt_regs *regs)
 {
 	enter_from_user_mode(regs);
 	rseq_note_user_irq_entry();
-	mte_disable_tco_entry(current);
+	mte_enter_from_user_mode();
 	sme_enter_from_user_mode();
 }
 
@@ -146,7 +167,7 @@ static __always_inline void arm64_syscall_exit_to_user_mode(struct pt_regs *regs
 	syscall_exit_to_user_mode_prepare(regs);
 	local_daif_mask();
 	sme_exit_to_user_mode();
-	mte_check_tfsr_exit();
+	mte_exit_to_user_mode();
 	exit_to_user_mode();
 }
 
@@ -161,7 +182,7 @@ static __always_inline void arm64_exit_to_user_mode(struct pt_regs *regs)
 	irqentry_exit_to_user_mode_prepare(regs);
 	local_daif_mask();
 	sme_exit_to_user_mode();
-	mte_check_tfsr_exit();
+	mte_exit_to_user_mode();
 	exit_to_user_mode();
 }
 

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 16/24] arm64: Refactor KERNEL_GCR_EL1 out of <sysreg.h>

[Pierre-Clément Tosi]

Rework the constant to make the resulting sysreg value (and computation)
obvious at the usage site and clearly indicate that the GCR_EL1.Exclude
is specifically set by a KASAN policy.

As a result, <sysreg.h> is now KASAN-agnostic and writers of GCR_EL1
explicitly depend on KASAN headers, as they should.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h | 16 ++++++++++++++++
 arch/arm64/include/asm/sysreg.h    | 16 ----------------
 arch/arm64/kernel/entry.S          |  3 ++-
 arch/arm64/kernel/mte.c            |  4 +++-
 4 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 9d00e9460b0f..6f63c788470a 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -5,10 +5,26 @@
 #ifndef __ASM_MTE_KASAN_H
 #define __ASM_MTE_KASAN_H
 
+#include <linux/bits.h>
+#include <linux/const.h>
+#include <linux/kasan-tags.h>
+
 #include <asm/compiler.h>
 #include <asm/cputype.h>
 #include <asm/mte-def.h>
 
+/*
+ * KASAN always uses a whole byte for its tags. With CONFIG_KASAN_HW_TAGS it
+ * only uses tags in the range 0xF0-0xFF, which we map to MTE tags 0x0-0xF.
+ */
+#ifdef CONFIG_KASAN_HW_TAGS
+#define __MTE_TAG_MIN		(KASAN_TAG_MIN & 0xf)
+#define __MTE_TAG_MAX		(KASAN_TAG_MAX & 0xf)
+#define KASAN_GCR_INCL		GENMASK(__MTE_TAG_MAX, __MTE_TAG_MIN)
+#else
+#define KASAN_GCR_INCL		UL(0)
+#endif
+
 #ifndef __ASSEMBLER__
 
 #include <linux/types.h>
diff --git a/arch/arm64/include/asm/sysreg.h b/arch/arm64/include/asm/sysreg.h
index 51c6a54d7146..7d82c769ae69 100644
--- a/arch/arm64/include/asm/sysreg.h
+++ b/arch/arm64/include/asm/sysreg.h
@@ -11,7 +11,6 @@
 
 #include <linux/bits.h>
 #include <linux/stringify.h>
-#include <linux/kasan-tags.h>
 #include <linux/kconfig.h>
 
 #include <asm/gpr-num.h>
@@ -921,21 +920,6 @@
 #define SYS_GCR_EL1_RRND	(BIT(16))
 #define SYS_GCR_EL1_EXCL_MASK	0xffffUL
 
-#ifdef CONFIG_KASAN_HW_TAGS
-/*
- * KASAN always uses a whole byte for its tags. With CONFIG_KASAN_HW_TAGS it
- * only uses tags in the range 0xF0-0xFF, which we map to MTE tags 0x0-0xF.
- */
-#define __MTE_TAG_MIN		(KASAN_TAG_MIN & 0xf)
-#define __MTE_TAG_MAX		(KASAN_TAG_MAX & 0xf)
-#define __MTE_TAG_INCL		GENMASK(__MTE_TAG_MAX, __MTE_TAG_MIN)
-#define KERNEL_GCR_EL1_EXCL	(SYS_GCR_EL1_EXCL_MASK & ~__MTE_TAG_INCL)
-#else
-#define KERNEL_GCR_EL1_EXCL	SYS_GCR_EL1_EXCL_MASK
-#endif
-
-#define KERNEL_GCR_EL1		(SYS_GCR_EL1_RRND | KERNEL_GCR_EL1_EXCL)
-
 /* RGSR_EL1 Definitions */
 #define SYS_RGSR_EL1_TAG_MASK	0xfUL
 #define SYS_RGSR_EL1_SEED_SHIFT	8
diff --git a/arch/arm64/kernel/entry.S b/arch/arm64/kernel/entry.S
index 8d17f61d1b5b..d1d9a7a049a7 100644
--- a/arch/arm64/kernel/entry.S
+++ b/arch/arm64/kernel/entry.S
@@ -21,6 +21,7 @@
 #include <asm/esr.h>
 #include <asm/irq.h>
 #include <asm/memory.h>
+#include <asm/mte-kasan.h>
 #include <asm/mmu.h>
 #include <asm/processor.h>
 #include <asm/ptrace.h>
@@ -176,7 +177,7 @@ alternative_else_nop_endif
 alternative_cb	ARM64_ALWAYS_SYSTEM, kasan_hw_tags_enable
 	b	1f
 alternative_cb_end
-	mov	\tmp, KERNEL_GCR_EL1
+	mov	\tmp, #(SYS_GCR_EL1_RRND | (SYS_GCR_EL1_EXCL_MASK & ~KASAN_GCR_INCL))
 	msr_s	SYS_GCR_EL1, \tmp
 1:
 #endif
diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index 0d998289edcb..725f48faec3e 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -23,6 +23,7 @@
 #include <asm/barrier.h>
 #include <asm/cpufeature.h>
 #include <asm/mte.h>
+#include <asm/mte-kasan.h>
 #include <asm/ptrace.h>
 #include <asm/sysreg.h>
 
@@ -320,7 +321,8 @@ void mte_cpu_setup(void)
 			 MAIR_ATTRIDX(MAIR_ATTR_NORMAL_TAGGED,
 				      MT_NORMAL_TAGGED));
 
-	write_sysreg_s(KERNEL_GCR_EL1, SYS_GCR_EL1);
+	write_sysreg_s(SYS_GCR_EL1_RRND | (SYS_GCR_EL1_EXCL_MASK & ~KASAN_GCR_INCL),
+		       SYS_GCR_EL1);
 
 	/*
 	 * If GCR_EL1.RRND=1 is implemented the same way as RRND=0, then

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 17/24] arm64: Introduce MTE & KASAN tag handling helpers

[Pierre-Clément Tosi]

Introduce [__]mte_{tagged,tag_of}() helpers for MTE tags.

Use a naming convention which clearly indicates that the helpers are
dealing with the tag bits of the argument, as opposed to the tag value
actually associated in hardware with the memory the argument points to.
The get()/set() terminology may imply the latter, and actually is used
in that sense already e.g. mte_get_mem_tag() or mte_set_mem_tag_range().

Provide a const-compatible ("__"-prefixed) variant and type-checked
functions, these should provide compile-time validation (with zero
runtime overhead) of constant arguments.

Introduce [__]kasan_tag_{to,from}_mte() helpers for KASAN tags.

Here, the "__"-prefixed variants are asm-compatible.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-def.h   | 20 ++++++++++++++++++++
 arch/arm64/include/asm/mte-kasan.h | 29 +++++++++++++++++++++--------
 arch/arm64/mm/fault.c              |  2 +-
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/arch/arm64/include/asm/mte-def.h b/arch/arm64/include/asm/mte-def.h
index 14ee86b019c2..660948a3fa9c 100644
--- a/arch/arm64/include/asm/mte-def.h
+++ b/arch/arm64/include/asm/mte-def.h
@@ -15,4 +15,24 @@
 
 #define __MTE_PREAMBLE		ARM64_ASM_PREAMBLE ".arch_extension memtag\n"
 
+#ifndef __ASSEMBLER__
+
+#include <linux/bitfield.h>
+#include <linux/types.h>
+
+#define __mte_tag_of(addr)	FIELD_GET(MTE_TAG_MASK, addr)
+#define __mte_tagged(addr, tag)	u64_replace_bits(addr, tag, MTE_TAG_MASK)
+
+static __always_inline u8 mte_tag_of(const void *ptr)
+{
+	return __mte_tag_of((u64)ptr);
+}
+
+static __always_inline void *mte_tagged(const void *ptr, u8 tag)
+{
+	return (void *)__mte_tagged((u64)ptr, tag);
+}
+
+#endif /* __ASSEMBLER__ */
+
 #endif /* __ASM_MTE_DEF_H  */
diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 6f63c788470a..24d5d9d6d04b 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -13,24 +13,40 @@
 #include <asm/cputype.h>
 #include <asm/mte-def.h>
 
+#define __KASAN_TAG_MTE_MASK		GENMASK(MTE_TAG_SIZE - 1, 0)
 /*
  * KASAN always uses a whole byte for its tags. With CONFIG_KASAN_HW_TAGS it
  * only uses tags in the range 0xF0-0xFF, which we map to MTE tags 0x0-0xF.
  */
+#define __kasan_tag_to_mte(tag)		((tag) & __KASAN_TAG_MTE_MASK)
+#define __kasan_tag_from_mte(tag)	((KASAN_TAG_MIN & ~__KASAN_TAG_MTE_MASK) | (tag))
+
 #ifdef CONFIG_KASAN_HW_TAGS
-#define __MTE_TAG_MIN		(KASAN_TAG_MIN & 0xf)
-#define __MTE_TAG_MAX		(KASAN_TAG_MAX & 0xf)
-#define KASAN_GCR_INCL		GENMASK(__MTE_TAG_MAX, __MTE_TAG_MIN)
+#define KASAN_GCR_INCL \
+	GENMASK(__kasan_tag_to_mte(KASAN_TAG_MAX), __kasan_tag_to_mte(KASAN_TAG_MIN))
 #else
 #define KASAN_GCR_INCL		UL(0)
 #endif
 
 #ifndef __ASSEMBLER__
 
+#include <linux/bitfield.h>
 #include <linux/types.h>
 
 #include <asm/sysreg.h>
 
+#define KASAN_TAG_MTE_MASK		GENMASK_U8(MTE_TAG_SIZE - 1, 0)
+
+static __always_inline u8 kasan_tag_to_mte(u8 tag)
+{
+	return u8_get_bits(tag, KASAN_TAG_MTE_MASK);
+}
+
+static __always_inline u8 kasan_tag_from_mte(u8 tag)
+{
+	return u8_replace_bits(KASAN_TAG_MIN, tag, KASAN_TAG_MTE_MASK);
+}
+
 #ifdef CONFIG_KASAN_HW_TAGS
 
 /* Whether the MTE asynchronous mode is enabled. */
@@ -71,10 +87,7 @@ static inline void mte_enable_tco(void)
 
 static inline u8 mte_get_ptr_tag(void *ptr)
 {
-	/* Note: The format of KASAN tags is 0xF<x> */
-	u8 tag = 0xF0 | (u8)(((u64)(ptr)) >> MTE_TAG_SHIFT);
-
-	return tag;
+	return kasan_tag_from_mte(mte_tag_of(ptr));
 }
 
 /* Get allocation tag for the address. */
@@ -140,7 +153,7 @@ static inline void mte_set_mem_tag_range(void *addr, size_t size, u8 tag,
 	dczid_bs = 4ul << (dczid & 0xf);
 	dczid_dzp = (dczid >> 4) & 1;
 
-	curr = (u64)__tag_set(addr, tag);
+	curr = (u64)mte_tagged(addr, kasan_tag_to_mte(tag));
 	mask = dczid_bs - 1;
 	/* STG/STZG up to the end of the first block. */
 	end1 = curr | mask;
diff --git a/arch/arm64/mm/fault.c b/arch/arm64/mm/fault.c
index b22736853fcd..cf855aa1c30f 100644
--- a/arch/arm64/mm/fault.c
+++ b/arch/arm64/mm/fault.c
@@ -915,7 +915,7 @@ static int do_tag_check_fault(unsigned long far, unsigned long esr,
 	 * Otherwise, bits 63:60 of FAR_EL1 are not UNKNOWN.
 	 */
 	if (!cpus_have_cap(ARM64_MTE_FAR))
-		far = (__untagged_addr(far) & ~MTE_TAG_MASK) | (far & MTE_TAG_MASK);
+		far = __mte_tagged(__untagged_addr(far), __mte_tag_of(far));
 
 	do_bad_area(far, esr, regs);
 	return 0;

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 19/24] arm64: Rework static branches for kernel MTE

[Pierre-Clément Tosi]

Rename system_uses_mte_async_or_asymm_mode() to the more concise
system_uses_async_mte_faults(), which better reflects the intention, and
relocate it from <asm/mte-kasan.h> to <asm/mte.h> as it reflects the
state of MTE, not KASAN. This removes <asm/mte-kasan.h> from <asm/mte.h>
and further progresses our untangling of the features.

Introduce the all-encompassing system_uses_mte(), tracking when any mode
of MTE is used by the kernel (mimicking <asm/cpufeature.h>) and use it
in place of KASAN functions, when the intention was to check MTE state.

For now, leave the mte_enable_kernel_*() unchanged even though they
clearly don't support arbitrary successive transitions, as commented in
mte_enable_kernel_{async,asymm}(). Once we want to clean that up, note
that the TCF bits map cleanly to the static branches:

        TCF[0]  <-->  mte_sync_or_asymm_mode
        TCF[1]  <-->  mte_async_or_asymm_mode

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h      | 19 -------------------
 arch/arm64/include/asm/mte.h            | 32 +++++++++++++++++++++++++++++++-
 arch/arm64/include/asm/uaccess.h        |  9 ++++-----
 arch/arm64/include/asm/word-at-a-time.h |  4 ++--
 arch/arm64/kernel/entry-common.c        | 10 +++++-----
 arch/arm64/kernel/mte.c                 | 26 +++++++++++++++++---------
 6 files changed, 59 insertions(+), 41 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 24d5d9d6d04b..9d018d3a0b3b 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -47,25 +47,6 @@ static __always_inline u8 kasan_tag_from_mte(u8 tag)
 	return u8_replace_bits(KASAN_TAG_MIN, tag, KASAN_TAG_MTE_MASK);
 }
 
-#ifdef CONFIG_KASAN_HW_TAGS
-
-/* Whether the MTE asynchronous mode is enabled. */
-DECLARE_STATIC_KEY_FALSE(mte_async_or_asymm_mode);
-
-static inline bool system_uses_mte_async_or_asymm_mode(void)
-{
-	return static_branch_unlikely(&mte_async_or_asymm_mode);
-}
-
-#else /* CONFIG_KASAN_HW_TAGS */
-
-static inline bool system_uses_mte_async_or_asymm_mode(void)
-{
-	return false;
-}
-
-#endif /* CONFIG_KASAN_HW_TAGS */
-
 #ifdef CONFIG_ARM64_MTE
 
 /*
diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 1e20a43c0a67..8826d64c8e33 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -7,13 +7,13 @@
 
 #include <asm/compiler.h>
 #include <asm/mte-def.h>
-#include <asm/mte-kasan.h>
 
 #ifndef __ASSEMBLER__
 
 #include <linux/bitfield.h>
 #include <linux/page-flags.h>
 #include <linux/sched.h>
+#include <linux/static_key.h>
 #include <linux/types.h>
 
 #include <asm/barrier.h>
@@ -33,6 +33,36 @@ void mte_invalidate_tags_area(int type);
 void *mte_allocate_tag_storage(void);
 void mte_free_tag_storage(char *storage);
 
+#ifdef CONFIG_KASAN_HW_TAGS
+DECLARE_STATIC_KEY_FALSE(mte_sync_or_asymm_mode);
+DECLARE_STATIC_KEY_FALSE(mte_async_or_asymm_mode);
+
+static inline bool system_uses_sync_mte_faults(void)
+{
+	return static_branch_unlikely(&mte_sync_or_asymm_mode);
+}
+
+static inline bool system_uses_async_mte_faults(void)
+{
+	return static_branch_unlikely(&mte_async_or_asymm_mode);
+}
+#else
+static inline bool system_uses_sync_mte_faults(void)
+{
+	return false;
+}
+
+static inline bool system_uses_async_mte_faults(void)
+{
+	return false;
+}
+#endif
+
+static inline bool system_uses_mte(void)
+{
+	return system_uses_sync_mte_faults() || system_uses_async_mte_faults();
+}
+
 static __always_inline bool __mte_get_reset_tfsr(void)
 {
 	u64 tfsr_el1 = read_sysreg_s(SYS_TFSR_EL1);
diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index e99faabd5dfc..22a176ac8341 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -16,7 +16,6 @@
  */
 #include <linux/bitops.h>
 #include <linux/kasan-checks.h>
-#include <linux/kasan-enabled.h>
 #include <linux/string.h>
 
 #include <asm/asm-extable.h>
@@ -145,13 +144,13 @@ static inline void __uaccess_disable_mte(void)
 	 * Unlike get_user(), privileged uaccess (TTBR0) uses TCF (not TCF0)
 	 * such that a bad user address tag could trigger a kernel MTE fault.
 	 */
-	if (kasan_hw_tags_enabled())
+	if (system_uses_mte())
 		set_pstate_tco(1);
 }
 
 static inline void __uaccess_enable_mte(void)
 {
-	if (kasan_hw_tags_enabled())
+	if (system_uses_mte())
 		set_pstate_tco(0);
 }
 
@@ -316,7 +315,7 @@ do {									\
 	__typeof__(dst) __gkn_dst = (dst);				\
 	__typeof__(src) __gkn_src = (src);				\
 									\
-	if (system_uses_mte_async_or_asymm_mode())			\
+	if (system_uses_async_mte_faults())				\
 		__raw_get_mem_tco("ldr", *((type *)(__gkn_dst)),	\
 		      (__force type *)(__gkn_src), err_label, K);	\
 	else								\
@@ -422,7 +421,7 @@ do {									\
 	__typeof__(dst) __pkn_dst = (dst);				\
 	__typeof__(src) __pkn_src = (src);				\
 									\
-	if (system_uses_mte_async_or_asymm_mode())			\
+	if (system_uses_async_mte_faults())				\
 		__raw_put_mem_tco("str", *((type *)(__pkn_src)),	\
 			      (__force type *)(__pkn_dst), err_label, K); \
 	else								\
diff --git a/arch/arm64/include/asm/word-at-a-time.h b/arch/arm64/include/asm/word-at-a-time.h
index 430d5ab8f042..53fd58460235 100644
--- a/arch/arm64/include/asm/word-at-a-time.h
+++ b/arch/arm64/include/asm/word-at-a-time.h
@@ -10,7 +10,7 @@
 #include <linux/bitops.h>
 #include <linux/wordpart.h>
 
-#include <asm/mte-kasan.h>
+#include <asm/mte.h>
 #include <asm/sysreg.h>
 
 struct word_at_a_time {
@@ -64,7 +64,7 @@ static inline unsigned long load_unaligned_zeropad(const void *addr)
 	unsigned long ret;
 
 	/* We handle synchronous exceptions so only use TCO in non-sync MTE. */
-	if (!system_uses_mte_async_or_asymm_mode())
+	if (!system_uses_async_mte_faults())
 		return __load_unaligned_zeropad(addr);
 
 	set_pstate_tco(1);
diff --git a/arch/arm64/kernel/entry-common.c b/arch/arm64/kernel/entry-common.c
index db4ff003a5b6..c6a9cc4e99a2 100644
--- a/arch/arm64/kernel/entry-common.c
+++ b/arch/arm64/kernel/entry-common.c
@@ -77,14 +77,14 @@ static __always_inline void mte_enter_from_kernel_mode(void)
 	 * Taking an exception sets TCO but clearing it can be expensive on
 	 * some architectures so first check that MTE is actually required.
 	 */
-	if (kasan_hw_tags_enabled() || uaccess_requires_mte(current))
+	if (system_uses_mte() || uaccess_requires_mte(current))
 		set_pstate_tco(0);
 
 	if (!kasan_hw_tags_enabled())
 		return;
 
 	/* Asynchronous faults are sync'ed automatically on kernel entry. */
-	if (system_uses_mte_async_or_asymm_mode() && __mte_get_reset_tfsr())
+	if (system_uses_async_mte_faults() && __mte_get_reset_tfsr())
 		kasan_report_async();
 }
 
@@ -94,13 +94,13 @@ static __always_inline void mte_exit_to_kernel_mode(void)
 		return;
 
 	/* Asynchronous faults require explicit sync'ing on kernel exit. */
-	if (system_uses_mte_async_or_asymm_mode() && mte_get_reset_tfsr())
+	if (system_uses_async_mte_faults() && mte_get_reset_tfsr())
 		kasan_report_async();
 }
 
 static __always_inline void mte_enter_from_user_mode(void)
 {
-	if (kasan_hw_tags_enabled() || uaccess_requires_mte(current))
+	if (system_uses_mte() || uaccess_requires_mte(current))
 		set_pstate_tco(0);
 }
 
@@ -109,7 +109,7 @@ static __always_inline void mte_exit_to_user_mode(void)
 	if (!kasan_hw_tags_enabled())
 		return;
 
-	if (system_uses_mte_async_or_asymm_mode() && mte_get_reset_tfsr())
+	if (system_uses_async_mte_faults() && mte_get_reset_tfsr())
 		kasan_report_async();
 }
 
diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index b65e3aa3bd07..675ba6a803cf 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -13,6 +13,7 @@
 #include <linux/ptrace.h>
 #include <linux/sched.h>
 #include <linux/sched/mm.h>
+#include <linux/static_key.h>
 #include <linux/string.h>
 #include <linux/swap.h>
 #include <linux/swapops.h>
@@ -37,6 +38,8 @@ static DEFINE_PER_CPU_READ_MOSTLY(u64, mte_tcf_preferred);
  */
 DEFINE_STATIC_KEY_FALSE(mte_async_or_asymm_mode);
 EXPORT_SYMBOL_GPL(mte_async_or_asymm_mode);
+DEFINE_STATIC_KEY_FALSE(mte_sync_or_asymm_mode);
+EXPORT_SYMBOL_GPL(mte_sync_or_asymm_mode);
 #endif
 
 void mte_sync_tags(pte_t pte, unsigned int nr_pages)
@@ -114,10 +117,13 @@ void mte_enable_kernel_sync(void)
 	 * Make sure we enter this function when no PE has set
 	 * async mode previously.
 	 */
-	WARN_ONCE(system_uses_mte_async_or_asymm_mode(),
+	WARN_ONCE(system_uses_async_mte_faults(),
 			"MTE async mode enabled system wide!");
 
 	__mte_enable_kernel("synchronous", SCTLR_EL1_TCF_SYNC);
+
+	if (!system_uses_sync_mte_faults())
+		static_branch_enable(&mte_sync_or_asymm_mode);
 }
 
 void mte_enable_kernel_async(void)
@@ -132,7 +138,7 @@ void mte_enable_kernel_async(void)
 	 * mode in between sync and async, this strategy needs
 	 * to be reviewed.
 	 */
-	if (!system_uses_mte_async_or_asymm_mode())
+	if (!system_uses_async_mte_faults())
 		static_branch_enable(&mte_async_or_asymm_mode);
 }
 
@@ -150,7 +156,9 @@ void mte_enable_kernel_asymm(void)
 		 * mode in between sync and async, this strategy needs
 		 * to be reviewed.
 		 */
-		if (!system_uses_mte_async_or_asymm_mode())
+		if (!system_uses_sync_mte_faults())
+			static_branch_enable(&mte_sync_or_asymm_mode);
+		if (!system_uses_async_mte_faults())
 			static_branch_enable(&mte_async_or_asymm_mode);
 	} else {
 		/*
@@ -227,9 +235,9 @@ static void mte_update_gcr_excl(struct task_struct *task)
 {
 	/*
 	 * SYS_GCR_EL1 will be set to current->thread.mte_ctrl value by
-	 * mte_set_user_gcr() in kernel_exit, but only if KASAN is enabled.
+	 * mte_set_user_gcr() in kernel_exit, when the kernel uses MTE.
 	 */
-	if (kasan_hw_tags_enabled())
+	if (system_uses_mte())
 		return;
 
 	write_sysreg_s(
@@ -275,13 +283,13 @@ void mte_thread_switch(struct task_struct *next)
 	mte_update_gcr_excl(next);
 
 	/* TCO may not have been disabled on exception entry for the current task. */
-	if (!kasan_hw_tags_enabled() && uaccess_requires_mte(next))
+	if (!system_uses_mte() && uaccess_requires_mte(next))
 		set_pstate_tco(0);
 
 	if (!kasan_hw_tags_enabled())
 		return;
 
-	if (!system_uses_mte_async_or_asymm_mode())
+	if (!system_uses_async_mte_faults())
 		return;
 
 	/*
@@ -348,7 +356,7 @@ void mte_suspend_enter(void)
 	if (!kasan_hw_tags_enabled())
 		return;
 
-	if (system_uses_mte_async_or_asymm_mode() && mte_get_reset_tfsr())
+	if (system_uses_async_mte_faults() && mte_get_reset_tfsr())
 		kasan_report_async();
 }
 
@@ -393,7 +401,7 @@ long set_mte_ctrl(struct task_struct *task, unsigned long arg)
 		mte_update_sctlr_user(task);
 		mte_update_gcr_excl(task);
 		update_sctlr_el1(task->thread.sctlr_user);
-		if (!kasan_hw_tags_enabled() && uaccess_requires_mte(task))
+		if (!system_uses_mte() && uaccess_requires_mte(task))
 			set_pstate_tco(0);
 		preempt_enable();
 	}

-- 
2.56.0.455.g6c5a503b38-goog

[PATCH 20/24] arm64: Move MTE insn wrappers to <asm/mte-def.h>

[Pierre-Clément Tosi]

Relocate all assembly from <asm/mte-kasan.h> to a KASAN-agnostic header.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-def.h   | 45 ++++++++++++++++++++++++++++++++++
 arch/arm64/include/asm/mte-kasan.h | 49 +++-----------------------------------
 2 files changed, 48 insertions(+), 46 deletions(-)

diff --git a/arch/arm64/include/asm/mte-def.h b/arch/arm64/include/asm/mte-def.h
index 660948a3fa9c..b86d0e7e45ed 100644
--- a/arch/arm64/include/asm/mte-def.h
+++ b/arch/arm64/include/asm/mte-def.h
@@ -19,6 +19,7 @@
 
 #include <linux/bitfield.h>
 #include <linux/types.h>
+#include <asm/compiler.h>
 
 #define __mte_tag_of(addr)	FIELD_GET(MTE_TAG_MASK, addr)
 #define __mte_tagged(addr, tag)	u64_replace_bits(addr, tag, MTE_TAG_MASK)
@@ -33,6 +34,50 @@ static __always_inline void *mte_tagged(const void *ptr, u8 tag)
 	return (void *)__mte_tagged((u64)ptr, tag);
 }
 
+static __always_inline u64 __irg(u64 addr)
+{
+	asm volatile(__MTE_PREAMBLE "irg %0, %0"
+		     : "+r" (addr));
+
+	return addr;
+}
+
+static __always_inline u64 __ldg(u64 addr)
+{
+	asm(__MTE_PREAMBLE "ldg %0, [%0]"
+	    : "+r" (addr));
+
+	return addr;
+}
+
+static __always_inline u64 __stg_post(u64 addr)
+{
+	asm volatile(__MTE_PREAMBLE "stg %0, [%0], #16"
+		     : "+r" (addr)
+		     :
+		     : "memory");
+	return addr;
+}
+
+static __always_inline u64 __stzg_post(u64 addr)
+{
+	asm volatile(__MTE_PREAMBLE "stzg %0, [%0], #16"
+		     : "+r" (addr)
+		     :
+		     : "memory");
+	return addr;
+}
+
+static __always_inline void __dc_gva(u64 addr)
+{
+	asm volatile(__MTE_PREAMBLE "dc gva, %0" : : "r" (addr) : "memory");
+}
+
+static __always_inline void __dc_gzva(u64 addr)
+{
+	asm volatile(__MTE_PREAMBLE "dc gzva, %0" : : "r" (addr) : "memory");
+}
+
 #endif /* __ASSEMBLER__ */
 
 #endif /* __ASM_MTE_DEF_H  */
diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 9d018d3a0b3b..66d72d0d3c43 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -9,7 +9,6 @@
 #include <linux/const.h>
 #include <linux/kasan-tags.h>
 
-#include <asm/compiler.h>
 #include <asm/cputype.h>
 #include <asm/mte-def.h>
 
@@ -66,57 +65,20 @@ static inline void mte_enable_tco(void)
 	set_pstate_tco(1);
 }
 
-static inline u8 mte_get_ptr_tag(void *ptr)
-{
-	return kasan_tag_from_mte(mte_tag_of(ptr));
-}
-
 /* Get allocation tag for the address. */
 static inline u8 mte_get_mem_tag(void *addr)
 {
-	asm(__MTE_PREAMBLE "ldg %0, [%0]"
-		: "+r" (addr));
-
-	return mte_get_ptr_tag(addr);
+	return kasan_tag_from_mte(__mte_tag_of(__ldg((u64)addr)));
 }
 
 /* Generate a random tag. */
 static inline void *mte_set_random_tag(const void *addr)
 {
 	/* Writes the random MTE tag in bits [59:56]. */
-	asm volatile(__MTE_PREAMBLE "irg %0, %0"
-		: "+r" (addr));
+	addr = (const void *)__irg((u64)addr);
 
 	/* Sets bits [63:60] as KASAN expects. */
-	return (void *)__tag_set(addr, mte_get_ptr_tag((void *)addr));
-}
-
-static inline u64 __stg_post(u64 p)
-{
-	asm volatile(__MTE_PREAMBLE "stg %0, [%0], #16"
-		     : "+r"(p)
-		     :
-		     : "memory");
-	return p;
-}
-
-static inline u64 __stzg_post(u64 p)
-{
-	asm volatile(__MTE_PREAMBLE "stzg %0, [%0], #16"
-		     : "+r"(p)
-		     :
-		     : "memory");
-	return p;
-}
-
-static inline void __dc_gva(u64 p)
-{
-	asm volatile(__MTE_PREAMBLE "dc gva, %0" : : "r"(p) : "memory");
-}
-
-static inline void __dc_gzva(u64 p)
-{
-	asm volatile(__MTE_PREAMBLE "dc gzva, %0" : : "r"(p) : "memory");
+	return (void *)__tag_set(addr, kasan_tag_from_mte(mte_tag_of(addr)));
 }
 
 /*
@@ -188,11 +150,6 @@ static inline void mte_enable_tco(void)
 {
 }
 
-static inline u8 mte_get_ptr_tag(void *ptr)
-{
-	return 0xFF;
-}
-
 static inline u8 mte_get_mem_tag(void *addr)
 {
 	return 0xFF;

-- 
2.56.0.455.g6c5a503b38-goog