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

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

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