arm64: mm: Set stack size from the kernel command line

As now THREAD_SHIFT can be variable with ARCH_HAS_VARIABLE_STACK_SIZE,
select it for ARM64, and set it from the kernel command line. Update
the stack overflow detection path to use the dynamic thread_shift.

The max value that can be set is 16 (64K) and the min values is
MAX(13, PAGE_SHIFT) as the stack has to be page aligned.

One caveat:
- As init_task stack size is the max possible, that means overflow
  checks will fail for it. So we slide the stack bottom to match
  the configured stack size, that is done twice; once at boot to
  match the default stack size and once if the kernel command
  line was passed.

Signed-off-by: Mostafa Saleh <smostafa@google.com>
diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 1cc0f9f..c3149bb 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -151,6 +151,7 @@
 	select HAVE_ARCH_AUDITSYSCALL
 	select HAVE_ARCH_BITREVERSE if BITREVERSE
 	select HAVE_ARCH_COMPILER_H
+	select ARCH_HAS_VARIABLE_STACK_SIZE
 	select HAVE_ARCH_HUGE_VMALLOC
 	select HAVE_ARCH_HUGE_VMAP
 	select HAVE_ARCH_JUMP_LABEL
diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index 8041432..ea76239 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -112,28 +112,32 @@
 
 #define DIRECT_MAP_PHYSMEM_END	__pa(PAGE_END - 1)
 
-#define MIN_THREAD_SHIFT	(14 + KASAN_THREAD_SHIFT)
+#define MIN_THREAD_SHIFT	MAX(13, PAGE_SHIFT)
 
 /* Max possible stack size 64KB */
 #define MAX_THREAD_SHIFT	16
 #define MAX_THREAD_SIZE		(UL(1) << MAX_THREAD_SHIFT)
 
+#define __DEFAULT_THREAD_SHIFT (14 + KASAN_THREAD_SHIFT)
 /*
  * VMAP'd stacks are allocated at page granularity, so we must ensure that such
  * stacks are a multiple of page size.
  */
-#if (MIN_THREAD_SHIFT < PAGE_SHIFT)
+#if (__DEFAULT_THREAD_SHIFT < PAGE_SHIFT)
 #define DEFAULT_THREAD_SHIFT	PAGE_SHIFT
 #else
-#define DEFAULT_THREAD_SHIFT	MIN_THREAD_SHIFT
+#define DEFAULT_THREAD_SHIFT	__DEFAULT_THREAD_SHIFT
 #endif
 
 #define DEFAULT_THREAD_SIZE	(UL(1) << DEFAULT_THREAD_SHIFT)
 #define DEFAULT_THREAD_ALIGN	DEFAULT_THREAD_SIZE
 
-#define THREAD_SHIFT		DEFAULT_THREAD_SHIFT
-#define THREAD_SIZE_ORDER	(THREAD_SHIFT - PAGE_SHIFT)
-#define THREAD_SIZE		(UL(1) << THREAD_SHIFT)
+#ifndef __ASSEMBLY__
+extern unsigned long thread_shift;
+#define THREAD_SHIFT		thread_shift
+#define THREAD_SIZE		(UL(1) << thread_shift)
+#define THREAD_SIZE_ORDER	(thread_shift - PAGE_SHIFT)
+#endif
 
 #define THREAD_ALIGN		THREAD_SIZE
 
@@ -142,7 +146,7 @@
  * size, then later at boot it is truncated to use the system's
  * THREAD_SHIFT
  */
-#define INIT_THREAD_SIZE	(UL(1) << DEFAULT_THREAD_SHIFT)
+#define INIT_THREAD_SIZE	(UL(1) << MAX_THREAD_SHIFT)
 #define INIT_THREAD_ALIGN	INIT_THREAD_SIZE
 
 #define IRQ_STACK_SIZE		(UL(1) << DEFAULT_THREAD_SHIFT)
diff --git a/arch/arm64/kernel/entry.S b/arch/arm64/kernel/entry.S
index d2a4050..7463da0 100644
--- a/arch/arm64/kernel/entry.S
+++ b/arch/arm64/kernel/entry.S
@@ -75,7 +75,9 @@
 	msr	spsel, #0
 	sub	x1, x1, x0
 
-	lsr	x1, x1, #THREAD_SHIFT
+	adr_l	x0, thread_shift
+	ldr	x0, [x0]
+	lsr	x1, x1, x0
 	cbz	x1, .Lstack_ok\@
 
 	/*
@@ -585,7 +587,10 @@
 	mov	x0, sp
 	msr	spsel, #0
 	sub	x0, x1, x0
-	cmp	x0, #THREAD_SIZE
+	ldr_l	x2, thread_shift
+	mov	x1, #1
+	lsl	x1, x1, x2
+	cmp	x0, x1
 	b.ls	2f
 1:
 #endif
diff --git a/arch/arm64/kernel/head.S b/arch/arm64/kernel/head.S
index 86e8755..c13fb38 100644
--- a/arch/arm64/kernel/head.S
+++ b/arch/arm64/kernel/head.S
@@ -401,7 +401,10 @@
 	cbz	x2, __secondary_too_slow
 
 	ldr	x1, [x2, #TSK_STACK]
-	add	sp, x1, #THREAD_SIZE
+	mov	x3, #1
+	ldr_l	x0, thread_shift
+	lsl	x3, x3, x0
+	add	sp, x1, x3
 
 	init_cpu_task x2, x1, x3
 
diff --git a/arch/arm64/kernel/setup.c b/arch/arm64/kernel/setup.c
index 23c05dc..655f031 100644
--- a/arch/arm64/kernel/setup.c
+++ b/arch/arm64/kernel/setup.c
@@ -30,6 +30,54 @@
 #include <linux/efi.h>
 #include <linux/psci.h>
 #include <linux/sched/task.h>
+#include <linux/sched/task_stack.h>
+#include <linux/mm.h>
+
+unsigned long thread_shift __ro_after_init = MAX_THREAD_SHIFT;
+EXPORT_SYMBOL(thread_shift);
+
+static void __init arm64_stack_slide(void)
+{
+	unsigned long offset = (1UL << MAX_THREAD_SHIFT) - (1UL << thread_shift);
+	void *new_stack = (void *)&init_thread_union + offset;
+
+	if ((unsigned long)current_stack_pointer < (unsigned long)new_stack)
+		panic("init_task stack would be corrupted by sliding (sp: 0x%lx, new_stack: 0x%lx)\n",
+		      (unsigned long)current_stack_pointer, (unsigned long)new_stack);
+
+	init_task.stack = new_stack;
+	set_task_stack_end_magic(&init_task);
+}
+
+static int __init setup_stack_size(char *str)
+{
+	unsigned long _thread_shift;
+	int ret;
+
+	if (!str)
+		return -EINVAL;
+
+	ret = kstrtoul(str,10, &_thread_shift);
+	if (ret || (_thread_shift < MIN_THREAD_SHIFT) || (_thread_shift > MAX_THREAD_SHIFT)) {
+		_thread_shift = DEFAULT_THREAD_SHIFT;
+		pr_err("Invalid kernel_thread_shift=%s, should be between %d and %d\n",
+		       str, MIN_THREAD_SHIFT, MAX_THREAD_SHIFT);
+	}
+
+	/*
+	 * The init task, runs with the maximum possible stack, to avoid false,
+	 * positive stack overflow, slide the stack base to match the configured
+	 * one (as the stack grows downwards)
+	 */
+
+	thread_shift = _thread_shift;
+	arm64_stack_slide();
+
+	pr_err("Overriden kernel stack to %lu KB\n", THREAD_SIZE / 1024);
+	return 0;
+}
+early_param("kernel_thread_shift", setup_stack_size);
+
 #include <linux/scs.h>
 #include <linux/mm.h>
 
@@ -282,6 +330,16 @@ void __init __no_sanitize_address setup_arch(char **cmdline_p)
 {
 	setup_initial_init_mm(_text, _etext, _edata, _end);
 
+	/*
+	 * At this point init_task.stack is MAX_THREAD_SHIFT and thread_size is matching
+	 * that. As in case an early BUG was hit, it does not get mistaken for an overflow.
+	 * Now we set thread_shift to the default system value and truncate the init_task
+	 * stack to this value and it remains like this unless kernel_thread_shift was
+	 * set later.
+	 */
+	thread_shift = DEFAULT_THREAD_SHIFT;
+	arm64_stack_slide();
+
 	*cmdline_p = boot_command_line;
 
 	kaslr_init();