With Hardware Tag-Based KASAN in asynchronous or asymmetric mode, arm64
sets and clears PSTATE.TCO around every kernel nofault load or store. A
loop pays that cost for every access even though tag checking can stay
disabled until the operation finishes.

Separate TCO management from the fault-tolerant access and add begin and
end hooks for callers that want to batch several accesses. Keep the
existing accessors self-contained, and provide aliases and no-op hooks
for architectures that do not need special handling.

A context switch re-enables tag checking, so a batched region must not
schedule. Continue to evaluate accessor arguments before overriding TCO,
as those expressions may block.

Signed-off-by: Muhammad Usama Anjum <usama.anjum@arm.com>
---
 arch/arm64/include/asm/uaccess.h | 71 ++++++++++++++++++++++++--------
 include/linux/uaccess.h          | 16 +++++++
 2 files changed, 69 insertions(+), 18 deletions(-)

diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index 9f5bd9c69c249..1a14eb2a51dce 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -270,28 +270,43 @@ 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.
+ * Nofault load without TCO management for use inside a
+ * __begin/__end_kernel_nofault_bare() region.
  */
-#define __get_kernel_nofault(dst, src, type, err_label)			\
+#define __get_kernel_nofault_bare(dst, src, type, err_label)		\
 do {									\
 	__typeof__(dst) __gkn_dst = (dst);				\
 	__typeof__(src) __gkn_src = (src);				\
 	do { 								\
 		__label__ __gkn_label;					\
-									\
-		__mte_enable_tco_async();				\
 		__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);							\
 } while (0)
 
+/*
+ * 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,
+ * evaluate them before overriding TCO.
+ */
+#define __get_kernel_nofault(dst, src, type, err_label)			\
+do {									\
+	__label__ __gkn_tco_err;					\
+	__typeof__(dst) __gkn_tco_dst = (dst);				\
+	__typeof__(src) __gkn_tco_src = (src);				\
+	__mte_enable_tco_async();					\
+	__get_kernel_nofault_bare(__gkn_tco_dst, __gkn_tco_src, type,	\
+				     __gkn_tco_err);			\
+	__mte_disable_tco_async();					\
+	break;								\
+__gkn_tco_err:								\
+	__mte_disable_tco_async();					\
+	goto err_label;							\
+} while (0)
+
 #define __put_mem_asm(store, reg, x, addr, label, type)			\
 	asm goto(							\
 	"1:	" store "	" reg "0, [%1]\n"			\
@@ -366,28 +381,48 @@ 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.
- */
-#define __put_kernel_nofault(dst, src, type, err_label)			\
+/* Nofault store without TCO management; see __get_kernel_nofault_bare. */
+#define __put_kernel_nofault_bare(dst, src, type, err_label)		\
 do {									\
 	__typeof__(dst) __pkn_dst = (dst);				\
 	__typeof__(src) __pkn_src = (src);				\
 									\
 	do {								\
 		__label__ __pkn_err;					\
-		__mte_enable_tco_async();				\
 		__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)
+} while (0)
+
+/*
+ * 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.
+ */
+#define __put_kernel_nofault(dst, src, type, err_label)			\
+do {									\
+	__label__ __pkn_tco_err;					\
+	__typeof__(dst) __pkn_tco_dst = (dst);				\
+	__typeof__(src) __pkn_tco_src = (src);				\
+	__mte_enable_tco_async();					\
+	__put_kernel_nofault_bare(__pkn_tco_dst, __pkn_tco_src, type,	\
+				     __pkn_tco_err);			\
+	__mte_disable_tco_async();					\
+	break;								\
+__pkn_tco_err:								\
+	__mte_disable_tco_async();					\
+	goto err_label;							\
+} while (0)
+
+/*
+ * A context switch re-enables tag checking, hence the no-scheduling
+ * requirement for a bare nofault region.
+ */
+#define __begin_kernel_nofault_bare()	__mte_enable_tco_async()
+#define __end_kernel_nofault_bare()	__mte_disable_tco_async()
 
 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/include/linux/uaccess.h b/include/linux/uaccess.h
index eddbbb65ccc4f..7ae1854673471 100644
--- a/include/linux/uaccess.h
+++ b/include/linux/uaccess.h
@@ -637,6 +637,22 @@ do {							\
 
 #endif  /* !__get_kernel_nofault */
 
+/*
+ * Architectures may use the begin/end hooks to establish state shared by a
+ * sequence of bare nofault accesses. Every path out of the region must call
+ * the end hook. The region, including expressions passed to the bare
+ * accessors, must not call into the scheduler.
+ */
+#ifndef __get_kernel_nofault_bare
+#define __get_kernel_nofault_bare	__get_kernel_nofault
+#define __put_kernel_nofault_bare	__put_kernel_nofault
+#endif
+
+#ifndef __begin_kernel_nofault_bare
+#define __begin_kernel_nofault_bare()	do {} while (0)
+#define __end_kernel_nofault_bare()	do {} while (0)
+#endif
+
 /**
  * get_kernel_nofault(): safely attempt to read from a location
  * @val: read into this variable
-- 
2.47.3