[PATCH v0] RISC-V: Use Zkr to seed KASLR base address

[Jesse Taube]

Dectect the Zkr extension and use it to seed the kernel base address.

Detection of the extension can not be done in the typical fashion, as
this is very early in the boot process. Instead, add a trap handler
and run it to see if the extension is present.

Signed-off-by: Jesse Taube <jesse at rivosinc.com>
---
 arch/riscv/kernel/pi/Makefile           |  2 +-
 arch/riscv/kernel/pi/archrandom_early.c | 71 +++++++++++++++++++++++++
 arch/riscv/mm/init.c                    |  3 ++
 3 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 arch/riscv/kernel/pi/archrandom_early.c

diff --git a/arch/riscv/kernel/pi/Makefile b/arch/riscv/kernel/pi/Makefile
index 50bc5ef7dd2f..9025eb52945a 100644
--- a/arch/riscv/kernel/pi/Makefile
+++ b/arch/riscv/kernel/pi/Makefile
@@ -32,5 +32,5 @@ $(obj)/string.o: $(srctree)/lib/string.c FORCE
 $(obj)/ctype.o: $(srctree)/lib/ctype.c FORCE
 	$(call if_changed_rule,cc_o_c)
 
-obj-y		:= cmdline_early.pi.o fdt_early.pi.o string.pi.o ctype.pi.o lib-fdt.pi.o lib-fdt_ro.pi.o
+obj-y		:= cmdline_early.pi.o fdt_early.pi.o string.pi.o ctype.pi.o lib-fdt.pi.o lib-fdt_ro.pi.o archrandom_early.pi.o
 extra-y		:= $(patsubst %.pi.o,%.o,$(obj-y))
diff --git a/arch/riscv/kernel/pi/archrandom_early.c b/arch/riscv/kernel/pi/archrandom_early.c
new file mode 100644
index 000000000000..311be9388b5c
--- /dev/null
+++ b/arch/riscv/kernel/pi/archrandom_early.c
@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+/*
+ * To avoid rewriteing code include asm/archrandom.h and create macros
+ * for the functions that won't be included.
+ */
+
+#define riscv_has_extension_likely(...) false
+#define pr_err_once(...)
+
+#include <linux/types.h>
+#include <asm/hwcap.h>
+#include <asm/archrandom.h>
+
+/*
+ * Asm goto is needed so that the compiler does not remove the label.
+ */
+
+#define csr_goto_swap(csr, val)						\
+({									\
+	unsigned long __v;						\
+	__asm__ __volatile__ goto("csrrw %0, " __ASM_STR(csr) ", %1"	\
+				  : "=r" (__v) : "rK" (&&val)		\
+				  : "memory" : val);			\
+	__v;								\
+})
+
+/*
+ * Declare the functions that are exported (but prefixed) here so that LLVM
+ * does not complain it lacks the 'static' keyword (which, if added, makes
+ * LLVM complain because the function is actually unused in this file).
+ */
+
+u64 get_kaslr_seed_zkr(void);
+
+/*
+ * This function is called by setup_vm to check if the kernel has the ZKR.
+ * Traps haven't been set up yet, but save and restore the TVEC to avoid
+ * any side effects.
+ */
+
+static inline bool __must_check riscv_has_zkr(void)
+{
+	unsigned long tvec;
+
+	tvec = csr_goto_swap(CSR_TVEC, not_zkr);
+	csr_swap(CSR_SEED, 0);
+	csr_write(CSR_TVEC, tvec);
+	return true;
+not_zkr:
+	csr_write(CSR_TVEC, tvec);
+	return false;
+}
+
+u64 get_kaslr_seed_zkr(void)
+{
+	const int needed_seeds = sizeof(u64) / sizeof(long);
+	int i = 0;
+	u64 seed = 0;
+	long *entropy = (long *)(&seed);
+
+	if (!riscv_has_zkr())
+		return 0;
+
+	for (i = 0; i < needed_seeds; i++) {
+		if (!csr_seed_long(&entropy[i]))
+			return 0;
+	}
+
+	return seed;
+}
diff --git a/arch/riscv/mm/init.c b/arch/riscv/mm/init.c
index 9940171c79f0..8ef1edd2cddd 100644
--- a/arch/riscv/mm/init.c
+++ b/arch/riscv/mm/init.c
@@ -1025,6 +1025,7 @@ static void __init pt_ops_set_late(void)
 #ifdef CONFIG_RANDOMIZE_BASE
 extern bool __init __pi_set_nokaslr_from_cmdline(uintptr_t dtb_pa);
 extern u64 __init __pi_get_kaslr_seed(uintptr_t dtb_pa);
+extern u64 __init __pi_get_kaslr_seed_zkr(void);
 
 static int __init print_nokaslr(char *p)
 {
@@ -1049,6 +1050,8 @@ asmlinkage void __init setup_vm(uintptr_t dtb_pa)
 		u32 kernel_size = (uintptr_t)(&_end) - (uintptr_t)(&_start);
 		u32 nr_pos;
 
+		if (kaslr_seed == 0)
+			kaslr_seed = __pi_get_kaslr_seed_zkr();
 		/*
 		 * Compute the number of positions available: we are limited
 		 * by the early page table that only has one PUD and we must
-- 
2.43.0