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

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

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