[PATCH v2 4/5] lib: sbi: Create a spot to place Smrnmi detection before traps and after DT is ready

[Evgeny Voevodin]

Since Smrnmi is detected from the device tree, move
sbi_platform_extensions_init() before trap-based feature detection so
the extension is known when the next commit acts on it.

The Zkr seed-CSR access can trap if Zkr is not implemented. On Smrnmi
platforms NMIE=0 by default after reset, which routes that trap to
NMEVEC rather than MTVEC. Until Smrnmi handlers are installed and
NMIE=1, no NMEVEC handler is in place, so an early-boot trap there
would crash. Refactor the inline seed loop in fw_base.S into a
callable __stack_chk_guard_init function and invoke it from
init_coldboot() after sbi_hart_init() returns; the next commit makes
sbi_hart_init() install Smrnmi handlers and set NMIE=1, so by the time
__stack_chk_guard_init runs the trap path is safe.

Signed-off-by: Evgeny Voevodin <evvoevod at tenstorrent.com>
---
 firmware/fw_base.S | 64 +++++++++++++++++++++++++++++-----------------
 lib/sbi/sbi_hart.c | 16 +++++++-----
 lib/sbi/sbi_init.c | 10 ++++++++
 3 files changed, 60 insertions(+), 30 deletions(-)

diff --git a/firmware/fw_base.S b/firmware/fw_base.S
index 043de7d7..ebecfecc 100644
--- a/firmware/fw_base.S
+++ b/firmware/fw_base.S
@@ -107,30 +107,6 @@ _bss_zero:
 	add	s4, s4, __SIZEOF_POINTER__
 	blt	s4, s5, _bss_zero
 
-	/* Trying to initialize the stack guard via the Zkr extension */
-	lla	t0, __stack_chk_guard_done
-	csrw	CSR_MTVEC, t0
-	li	t0, 0
-	li	t3, SEED_OPTS_ES16
-	li	t4, SEED_ENTROPY_MASK
-	li	t5, __SIZEOF_POINTER__
-__stack_chk_guard_loop:
-	csrrw	t1, CSR_SEED, x0
-	li	t2, SEED_OPTS_MASK
-	and	t2, t2, t1
-	bgtu	t2, t3, __stack_chk_guard_done
-	bltu	t2, t3, __stack_chk_guard_loop
-	and	t1, t1, t4
-	slli	t0, t0, 16
-	or	t0, t0, t1
-	addi	t5, t5, -2
-	bgtz	t5, __stack_chk_guard_loop
-	lla	t1, __stack_chk_guard
-	REG_S	t0, 0(t1)
-	j	__stack_chk_guard_done
-	.align 3
-__stack_chk_guard_done:
-
 	/* Setup temporary trap handler */
 	lla	s4, _start_hang
 	csrw	CSR_MTVEC, s4
@@ -895,6 +871,46 @@ __stack_chk_fail:
 	la	a0, .Lstack_corrupt_msg
 	call	sbi_panic
 
+	/*
+	 * Initialize __stack_chk_guard from the Zkr seed CSR. Called from C
+	 * after Smrnmi has been enabled (NMIE=1), so that a missing-Zkr
+	 * illegal-instruction trap is delivered to MTVEC normally and is
+	 * caught by our temporary MTVEC redirect below.
+	 */
+	.align 3
+	.globl __stack_chk_guard_init
+	.type __stack_chk_guard_init, %function
+__stack_chk_guard_init:
+	/* Save current MTVEC so we can restore it on return */
+	csrr	t6, CSR_MTVEC
+
+	/* Redirect MTVEC to local skip target for Zkr-not-supported case */
+	lla	t0, 1f
+	csrw	CSR_MTVEC, t0
+
+	li	t0, 0
+	li	t3, SEED_OPTS_ES16
+	li	t4, SEED_ENTROPY_MASK
+	li	t5, __SIZEOF_POINTER__
+2:
+	csrrw	t1, CSR_SEED, x0
+	li	t2, SEED_OPTS_MASK
+	and	t2, t2, t1
+	bgtu	t2, t3, 1f
+	bltu	t2, t3, 2b
+	and	t1, t1, t4
+	slli	t0, t0, 16
+	or	t0, t0, t1
+	addi	t5, t5, -2
+	bgtz	t5, 2b
+	lla	t1, __stack_chk_guard
+	REG_S	t0, 0(t1)
+	j	1f
+	.align 3
+1:
+	csrw	CSR_MTVEC, t6
+	ret
+
 	/* Initial value of the stack guard variable */
 	.section .data
 	.align 3
diff --git a/lib/sbi/sbi_hart.c b/lib/sbi/sbi_hart.c
index 4aefb759..781161e5 100644
--- a/lib/sbi/sbi_hart.c
+++ b/lib/sbi/sbi_hart.c
@@ -522,6 +522,16 @@ static int hart_detect_features(struct sbi_scratch *scratch)
 	hfeatures->mhpm_mask = 0;
 	hfeatures->priv_version = SBI_HART_PRIV_VER_UNKNOWN;
 
+	/*
+	 * Parse device tree extensions early, before any trap-based checks.
+	 * Needed to detect Smrnmi and install NMI handlers before CSR probes
+	 * that may trigger traps.
+	 */
+	rc = sbi_platform_extensions_init(sbi_platform_thishart_ptr(),
+					  hfeatures);
+	if (rc)
+		return rc;
+
 #define __check_hpm_csr(__csr, __mask) 					  \
 	oldval = csr_read_allowed(__csr, &trap);			  \
 	if (!trap.cause) {						  \
@@ -676,12 +686,6 @@ __pmp_skip:
 
 #undef __check_csr_existence
 
-	/* Let platform populate extensions */
-	rc = sbi_platform_extensions_init(sbi_platform_thishart_ptr(),
-					  hfeatures);
-	if (rc)
-		return rc;
-
 	/* Zicntr should only be detected using traps */
 	__sbi_hart_update_extension(hfeatures, SBI_HART_EXT_ZICNTR,
 			    sbi_hart_has_csr(scratch, SBI_HART_CSR_CYCLE) &&
diff --git a/lib/sbi/sbi_init.c b/lib/sbi/sbi_init.c
index aae035e1..644f62cb 100644
--- a/lib/sbi/sbi_init.c
+++ b/lib/sbi/sbi_init.c
@@ -36,6 +36,8 @@
 #include <sbi/sbi_version.h>
 #include <sbi/sbi_unit_test.h>
 
+extern void __stack_chk_guard_init(void);
+
 #define BANNER                                              \
 	"   ____                    _____ ____ _____\n"     \
 	"  / __ \\                  / ____|  _ \\_   _|\n"  \
@@ -269,6 +271,14 @@ static void __noreturn init_coldboot(struct sbi_scratch *scratch, u32 hartid)
 	if (rc)
 		sbi_hart_hang();
 
+	/*
+	 * Initialize stack-guard canary now that hart_detect_features()
+	 * has enabled Smrnmi (NMIE=1). Done from a __noreturn function so
+	 * the canary mutation doesn't trip our own exit check. The asm
+	 * helper saves and restores MTVEC around its operation.
+	 */
+	__stack_chk_guard_init();
+
 	rc = sbi_timer_init(scratch, true);
 	if (rc)
 		sbi_hart_hang();
-- 
2.43.0