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

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

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