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

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

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