[RFC PATCH 14/15] arm64: mm: Relax kernel stack alignment

[Mostafa Saleh <smostafa@google.com>]

On the kernel entry it was not possible to use the stack until it
was checked for overflow which was done by a clever trick relying
on aligning the kernel stack to double it's size, so if the
bit at THREAD_SHIFT was set it means that the stack pointer has
overflowed.

Now, we can easily switch the sp to sp_el1 which is the overflow
stack, push some registers and execute more complex flow.

Rework the kernel entry code to eliminate the tbnz check and the
alignment requirement.

Signed-off-by: Mostafa Saleh <smostafa@google.com>
---
 arch/arm64/include/asm/memory.h |  7 +--
 arch/arm64/kernel/entry.S       | 87 ++++++++++++++++++++++++++++-----
 2 files changed, 77 insertions(+), 17 deletions(-)

diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index 93ce6ef65573..4545d9b39bd5 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -130,12 +130,7 @@
 
 #define THREAD_SIZE		(UL(1) << THREAD_SHIFT)
 
-/*
- * By aligning VMAP'd stacks to 2 * THREAD_SIZE, we can detect overflow by
- * checking sp & (1 << THREAD_SHIFT), which we can do cheaply in the entry
- * assembly.
- */
-#define THREAD_ALIGN		(2 * THREAD_SIZE)
+#define THREAD_ALIGN		THREAD_SIZE
 
 #define IRQ_STACK_SIZE		THREAD_SIZE
 
diff --git a/arch/arm64/kernel/entry.S b/arch/arm64/kernel/entry.S
index a31ef890a2ee..ea733b673970 100644
--- a/arch/arm64/kernel/entry.S
+++ b/arch/arm64/kernel/entry.S
@@ -62,15 +62,33 @@
 	sub	sp, sp, #PT_REGS_SIZE
 
 	/*
-	 * Test whether the SP has overflowed, without corrupting a GPR.
-	 * Task and IRQ stacks are aligned so that SP & (1 << THREAD_SHIFT)
-	 * should always be zero.
+	 * Test whether the SP has overflowed, using the overflow stack.
+	 * As we do not know if the CPU was in process or irq context, first
+	 * check against the task stack if that failed it means it might have
+	 * been an interrupt, so check against an interrupt stack, if both
+	 * failed it means that at least one of the stacks overflowed.
 	 */
-	add	sp, sp, x0			// sp' = sp + x0
-	sub	x0, sp, x0			// x0' = sp' - x0 = (sp + x0) - x0 = sp
-	tbnz	x0, #THREAD_SHIFT, __bad_stack
-	sub	x0, sp, x0			// x0'' = sp' - x0' = (sp + x0) - sp = x0
-	sub	sp, sp, x0			// sp'' = sp' - x0 = (sp + x0) - x0 = sp
+	msr	spsel, #1
+	stp	x0, x1, [sp, #-16]!
+
+	mrs	x0, tpidrro_el0
+	ldr	x0, [x0, #TSK_STACK]
+	msr	spsel, #0
+	mov	x1, sp
+	msr	spsel, #1
+	sub	x1, x1, x0
+	cmp	x1, #THREAD_SIZE
+	b.lo	.Lstack_ok\@
+
+	/* Check whether we are on the IRQ, SDEI or EFI stacks. */
+	stp	x2, x30, [sp, #-16]!
+	bl	__check_ext_stacks
+	ldp	x2, x30, [sp], #16
+
+.Lstack_ok\@:
+	ldp	x0, x1, [sp], #16
+	msr	spsel, #0
+
 	b	el1t_\regsize\()_\label
 	.endm
 
@@ -528,13 +546,60 @@ SYM_CODE_START(vectors)
 	kernel_ventry	0, t, 32, error		// Error 32-bit EL0
 SYM_CODE_END(vectors)
 
+
+	.macro check_stack_overflow type, ptr, size
+	.ifc \type, percpu
+	ldr_this_cpu x1, \ptr, x0
+	.else
+	adr_l	x1, \ptr
+	ldr	x1, [x1]
+	.endif
+	cbz	x1, 1f
+	.ifc \type, top
+	sub	x1, x1, #\size
+	.endif
+	msr	spsel, #0
+	mov	x0, sp
+	msr	spsel, #1
+	sub	x0, x0, x1
+	cmp	x0, #\size
+	b.lo	2f
+1:
+	.endm
+
+SYM_CODE_START_LOCAL(__check_ext_stacks)
+	/* IRQ stack check */
+	check_stack_overflow percpu, irq_stack_ptr, IRQ_STACK_SIZE
+#ifdef CONFIG_ARM_SDE_INTERFACE
+	/* SDEI normal stack check */
+	check_stack_overflow percpu, sdei_stack_normal_ptr, IRQ_STACK_SIZE
+	/* SDEI critical stack check */
+	check_stack_overflow percpu, sdei_stack_critical_ptr, IRQ_STACK_SIZE
+#endif
+
+#ifdef CONFIG_EFI
+	/* EFI runtime stack check */
+	check_stack_overflow top, efi_rt_stack_top, THREAD_SIZE
+#endif
+
+	/* All checks failed => it's a real overflow */
+	ldp	x2, x30, [sp], #16
+	b	__bad_stack
+
+2:
+	/* A check succeeded => return to kernel_ventry */
+	ret
+SYM_CODE_END(__check_ext_stacks)
+
 SYM_CODE_START_LOCAL(__bad_stack)
 	/*
 	 * We detected an overflow in kernel_ventry.
-	 * Restore SP and X0.
+	 * Restore X0 and X1, and pop the overflow stack.
 	 */
-	sub	x0, sp, x0
-	sub	sp, sp, x0
+	ldp	x0, x1, [sp], #16
+
+	/* Restore SP_EL0 */
+	msr	spsel, #0
 	add	sp, sp, #PT_REGS_SIZE
 
 	/* Switch to the overflow stack */
-- 
2.56.0.rc1.315.gc6ed9934b7-goog